diff --git a/lib/python3.12/site-packages/peft/tuners/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..434de8c0704b974573310313fd78bc689ce8fea8 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/__pycache__/_buffer_dict.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/__pycache__/_buffer_dict.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..082da419415f2f74b74d7eac408ced49c7be520e Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/__pycache__/_buffer_dict.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/__pycache__/lycoris_utils.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/__pycache__/lycoris_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84526752508b5369bbaffec670de2d57befacdc2 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/__pycache__/lycoris_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/__pycache__/tuners_utils.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/__pycache__/tuners_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..982141f8e23cc2f48c961d6ce64d32f7de072ffe Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/__pycache__/tuners_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/__init__.py b/lib/python3.12/site-packages/peft/tuners/adalora/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..64d5f3e5ce6ba3c5873f01ee88f4c4766e1fde75 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adalora/__init__.py @@ -0,0 +1,43 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 peft.import_utils import is_bnb_4bit_available, is_bnb_available +from peft.utils import register_peft_method + +from .config import AdaLoraConfig +from .gptq import SVDQuantLinear +from .layer import AdaLoraLayer, RankAllocator, SVDLinear +from .model import AdaLoraModel + + +__all__ = ["AdaLoraConfig", "AdaLoraLayer", "AdaLoraModel", "RankAllocator", "SVDLinear", "SVDQuantLinear"] + + +register_peft_method( + name="adalora", config_cls=AdaLoraConfig, model_cls=AdaLoraModel, prefix="lora_", is_mixed_compatible=True +) + + +def __getattr__(name): + if (name == "SVDLinear8bitLt") and is_bnb_available(): + from .bnb import SVDLinear8bitLt + + return SVDLinear8bitLt + + if (name == "SVDLinear4bit") and is_bnb_4bit_available(): + from .bnb import SVDLinear4bit + + return SVDLinear4bit + + raise AttributeError(f"module {__name__} has no attribute {name}") diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d819fbb7433bd9dfbeba28fa388b6da82505503 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/bnb.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/bnb.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da0bf639c4667f933278d455020e91710c4daa7c Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/bnb.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0cd11333629132872b5e70cdd02fdb55cc60a040 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/gptq.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/gptq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..81d8cb63d92fd40abd43723ea19e30d3647d10e3 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/gptq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/layer.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4e4f039b6ea0d0f3d0f8dac57ea4f50c15cd65d2 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9a0e7aaf73328fa5f02fe744ec1e75b2ca33a6e8 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adalora/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/bnb.py b/lib/python3.12/site-packages/peft/tuners/adalora/bnb.py new file mode 100644 index 0000000000000000000000000000000000000000..fef3d25e65c459d20855e49832aabe56881734c0 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adalora/bnb.py @@ -0,0 +1,143 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 typing import Any + +import torch + +from peft.import_utils import is_bnb_4bit_available, is_bnb_available + +from .layer import AdaLoraLayer + + +if is_bnb_available(): + + class SVDLinear8bitLt(torch.nn.Module, AdaLoraLayer): + # Low-rank matrix for SVD-based adaptation + def __init__( + self, + base_layer: torch.nn.Module, + adapter_name: str, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + init_lora_weights: bool = True, + **kwargs, + ) -> None: + super().__init__() + AdaLoraLayer.__init__(self, base_layer) + # Freezing the pre-trained weight matrix + self.get_base_layer().weight.requires_grad = False + + self._active_adapter = adapter_name + self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # note: no check for self.merged because merging is not supported (yet) + result = self.base_layer(x) + + if self.disable_adapters: + return result + + for active_adapter in self.active_adapters: + if active_adapter not in self.lora_A.keys(): + continue + requires_conversion = not torch.is_autocast_enabled() + if requires_conversion: + expected_dtype = result.dtype + if x.dtype != torch.float32: + x = x.float() + + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + lora_E = self.lora_E[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + ranknum = self.ranknum[active_adapter] + 1e-5 + + output = dropout(x) @ (lora_A * lora_E).T @ lora_B.T + if requires_conversion: + output = output.to(expected_dtype) + output = output * scaling / ranknum + # inplace operation on view is forbidden for MatMul8bitLtBackward, so avoid it + result = result + output + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "adalora." + rep + + +if is_bnb_4bit_available(): + + class SVDLinear4bit(torch.nn.Module, AdaLoraLayer): + # Low-rank matrix for SVD-based adaptation + def __init__( + self, + base_layer: torch.nn.Module, + adapter_name: str, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + init_lora_weights: bool = True, + **kwargs, + ) -> None: + super().__init__() + AdaLoraLayer.__init__(self, base_layer) + # Freezing the pre-trained weight matrix + self.get_base_layer().weight.requires_grad = False + + self._active_adapter = adapter_name + self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights) + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + # note: no check for self.merged because merging is not supported (yet) + result = self.base_layer(x, *args, **kwargs) + + if self.disable_adapters: + return result + + # As per Tim Dettmers, for 4bit, we need to defensively clone here. + # The reason is that in some cases, an error can occur that backprop + # does not work on a manipulated view. This issue may be solved with + # newer PyTorch versions but this would need extensive testing to be + # sure. + result = result.clone() + + for active_adapter in self.active_adapters: + if active_adapter not in self.lora_A.keys(): + continue + + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + lora_E = self.lora_E[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + ranknum = self.ranknum[active_adapter] + 1e-5 + + requires_conversion = not torch.is_autocast_enabled() + if requires_conversion: + expected_dtype = result.dtype + x = self._cast_input_dtype(x, lora_A.dtype) + + output = dropout(x) @ (lora_A * lora_E).T @ lora_B.T + if requires_conversion: + output = output.to(expected_dtype) + output = output * scaling / ranknum + result += output + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "adalora." + rep diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/config.py b/lib/python3.12/site-packages/peft/tuners/adalora/config.py new file mode 100644 index 0000000000000000000000000000000000000000..004c26b0fe5ecf88212bd11de062baec2f10de83 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adalora/config.py @@ -0,0 +1,108 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +import warnings +from dataclasses import dataclass, field +from typing import Optional + +from peft.tuners.lora import LoraConfig +from peft.utils import PeftType + + +@dataclass +class AdaLoraConfig(LoraConfig): + """ + This is the configuration class to store the configuration of a [`~peft.AdaLora`]. + + AdaLoRA has three phases defined by `tinit`, `tfinal` and `total_step`. + + The initial phase can be understood as a step for pre-training the adapters so that when reducing their rank, there + is already some information encoded that can be reduced instead of random matrices. This phase is defined by + supplying `tinit`. + + After the initial phase is over (`tinit` steps have passed) and the final phase has not begun, AdaLoRA reduces the + budget of how much rank each layer is allowed to have with each step. This is where the reduction of rank is + happening. This goes on until `total_step - tfinal` steps are reached. + + The last phase, beginning once `total_step - tfinal` steps are reached, does not change the layer ranks anymore but + fine-tunes the reduced-rank layers that resulted from the previous phase. + + A practical example: `tinit` is 10, `tfinal` is 20, `total_step` is 100. We spend 10 steps doing pre-training + without rank reduction because our budget is constant (init phase), then we spend 80 (100-20) steps in the + reduction phase where our budget decreases step-wise and, finally, 20 steps in the final fine-tuning stage without + reduction. + + Args: + target_r (`int`): The target average rank of incremental matrix. + init_r (`int`): The initial rank for each incremental matrix. + tinit (`int`): The steps of initial fine-tuning warmup. + tfinal (`int`): The number of steps of final fine-tuning. + deltaT (`int`): The time internval between two budget allocations. + beta1 (`float`): The hyperparameter of EMA for sensitivity smoothing. + beta2 (`float`): The hyperparameter of EMA for undertainty quantification. + orth_reg_weight (`float`): The coefficient of orthogonal regularization. + total_step (`int`): The total training steps that should be specified before training. + rank_pattern (`list`): The allocated rank for each weight matrix by RankAllocator. + """ + + target_r: int = field(default=8, metadata={"help": "Target Lora matrix dimension."}) + init_r: int = field(default=12, metadata={"help": "Initial Lora matrix dimension."}) + tinit: int = field(default=0, metadata={"help": "The steps of initial warmup."}) + tfinal: int = field(default=0, metadata={"help": "The steps of final warmup."}) + deltaT: int = field(default=1, metadata={"help": "Step interval of rank allocation."}) + beta1: float = field(default=0.85, metadata={"help": "Hyperparameter of EMA."}) + beta2: float = field(default=0.85, metadata={"help": "Hyperparameter of EMA."}) + orth_reg_weight: float = field(default=0.5, metadata={"help": "The orthogonal regularization coefficient."}) + total_step: Optional[int] = field(default=None, metadata={"help": "The total training steps."}) + rank_pattern: Optional[dict] = field(default=None, metadata={"help": "The saved rank pattern."}) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.ADALORA + + if self.use_dora: + raise ValueError(f"{self.peft_type} does not support DoRA.") + + if self.loftq_config: + raise ValueError(f"{self.peft_type} does not support LOFTQ.") + + self.target_modules = ( + set(self.target_modules) if isinstance(self.target_modules, list) else self.target_modules + ) + self.exclude_modules = ( + set(self.exclude_modules) if isinstance(self.exclude_modules, list) else self.exclude_modules + ) + # if target_modules is a regex expression, then layers_to_transform should be None + if isinstance(self.target_modules, str) and self.layers_to_transform is not None: + raise ValueError("`layers_to_transform` cannot be used when `target_modules` is a str.") + + # check for layers_to_transform and layers_pattern + if self.layers_pattern and not self.layers_to_transform: + raise ValueError("When `layers_pattern` is specified, `layers_to_transform` must also be specified. ") + + # Check if 'r' has been set to a non-default value + if self.r != 8: # 8 is the default value for 'r' in LoraConfig + warnings.warn( + "Note that `r` is not used in AdaLora and will be ignored." + "If you intended to set the initial rank, use `init_r` instead." + ) + + if self.total_step is None or self.total_step <= 0: + raise ValueError("AdaLoRA does not work when `total_step` is None, supply a value > 0.") + + if self.tinit >= (self.total_step - self.tfinal): + raise ValueError( + "The supplied schedule values don't allow for a budgeting phase. Decrease `tfinal`/`tinit` or " + "increase `total_step`." + ) diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/gptq.py b/lib/python3.12/site-packages/peft/tuners/adalora/gptq.py new file mode 100644 index 0000000000000000000000000000000000000000..bed1a0a7ca8dabb9d068bcf2470a97e34ec348fe --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adalora/gptq.py @@ -0,0 +1,71 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. +import torch + +from .layer import AdaLoraLayer + + +class SVDQuantLinear(torch.nn.Module, AdaLoraLayer): + def __init__( + self, + base_layer, + adapter_name, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + init_lora_weights: bool = True, + **kwargs, + ) -> None: + super().__init__() + AdaLoraLayer.__init__(self, base_layer) + + # self.base_layer and self.quant_linear_module are the same; we need the former for consistency and the latter + # for backwards compatibility + self.quant_linear_module = base_layer + self._active_adapter = adapter_name + self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + result = self.quant_linear_module(x) + + if self.disable_adapters: + return result + + for active_adapter in self.active_adapters: + if active_adapter not in self.lora_A.keys(): + continue + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + lora_E = self.lora_E[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + ranknum = self.ranknum[active_adapter] + 1e-5 + + requires_conversion = not torch.is_autocast_enabled() + if requires_conversion: + expected_dtype = result.dtype + x = self._cast_input_dtype(x, torch.float32) + + output = (dropout(x) @ (lora_A * lora_E).T @ lora_B.T) * scaling / ranknum + # TODO: here, the dtype conversion is applied on the *whole expression*, + # not the intermediate result, unlike for SVDLinear8bitLT and + # SVDLinear4bit, is that correct? + if requires_conversion: + output = output.to(expected_dtype) + result += output + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "adalora." + rep diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/layer.py b/lib/python3.12/site-packages/peft/tuners/adalora/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..ea142c512f45073de282aeea52531ef08a82f248 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adalora/layer.py @@ -0,0 +1,358 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +import warnings +from typing import Any, List, Optional + +import packaging +import torch +import transformers +from torch import nn + +from peft.tuners.lora import LoraLayer +from peft.tuners.tuners_utils import check_adapters_to_merge +from peft.utils import transpose + + +if packaging.version.parse(transformers.__version__) >= packaging.version.parse("4.33.0"): + from transformers.integrations import deepspeed_config +else: + from transformers.deepspeed import deepspeed_config + + +class AdaLoraLayer(LoraLayer): + # List all names of layers that may contain adapter weights + # Note: ranknum doesn't need to be included as it is not an nn.Module + adapter_layer_names = ("lora_A", "lora_B", "lora_E", "lora_embedding_A", "lora_embedding_B") + # All names of other parameters that may contain adapter-related parameters + other_param_names = ("r", "lora_alpha", "scaling", "lora_dropout", "ranknum") + + def __init__(self, base_layer: nn.Module) -> None: + super().__init__(base_layer) + self.lora_E = nn.ParameterDict({}) + self.lora_A = nn.ParameterDict({}) + self.lora_B = nn.ParameterDict({}) + self.ranknum = nn.ParameterDict({}) + + def update_layer(self, adapter_name, r, lora_alpha, lora_dropout, init_lora_weights): + if r < 0: + # note: r == 0 is allowed for AdaLora, see #1539 + raise ValueError(f"`r` should be a positive integer or 0, but the value passed is {r}") + + self.r[adapter_name] = r + self.lora_alpha[adapter_name] = lora_alpha + if lora_dropout > 0.0: + lora_dropout_layer = nn.Dropout(p=lora_dropout) + else: + lora_dropout_layer = nn.Identity() + + self.lora_dropout[adapter_name] = lora_dropout_layer + # Actual trainable parameters + # Right singular vectors + self.lora_A[adapter_name] = nn.Parameter(torch.randn(r, self.in_features)) + # Singular values + self.lora_E[adapter_name] = nn.Parameter(torch.randn(r, 1)) + # Left singular vectors + self.lora_B[adapter_name] = nn.Parameter(torch.randn(self.out_features, r)) + # The current rank + self.ranknum[adapter_name] = nn.Parameter(torch.randn(1), requires_grad=False) + self.ranknum[adapter_name].data.fill_(float(r)) + self.ranknum[adapter_name].requires_grad = False + self.scaling[adapter_name] = lora_alpha if lora_alpha > 0 else float(r) + if init_lora_weights: + self.reset_lora_parameters(adapter_name) + + self._move_adapter_to_device_of_base_layer(adapter_name) + self.set_adapter(self.active_adapters) + + def reset_lora_parameters(self, adapter_name): + if adapter_name in self.lora_A.keys(): + nn.init.zeros_(self.lora_E[adapter_name]) + nn.init.normal_(self.lora_A[adapter_name], mean=0.0, std=0.02) + nn.init.normal_(self.lora_B[adapter_name], mean=0.0, std=0.02) + + +class SVDLinear(nn.Module, AdaLoraLayer): + # SVD-based adaptation by a dense layer + def __init__( + self, + base_layer: nn.Module, + adapter_name: str, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + fan_in_fan_out: bool = False, + init_lora_weights: bool = True, + **kwargs, + ) -> None: + super().__init__() + AdaLoraLayer.__init__(self, base_layer) + # Freezing the pre-trained weight matrix + self.get_base_layer().weight.requires_grad = False + + self.fan_in_fan_out = fan_in_fan_out + self._active_adapter = adapter_name + self.update_layer(adapter_name, r, lora_alpha, lora_dropout, init_lora_weights) + + def merge(self, safe_merge: bool = False, adapter_names: Optional[List[str]] = None) -> None: + """ + Merge the active adapter weights into the base weights + + Args: + safe_merge (`bool`, *optional*): + If True, the merge operation will be performed in a copy of the original weights and check for NaNs + before merging the weights. This is useful if you want to check if the merge operation will produce + NaNs. Defaults to `False`. + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + """ + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + # no adapter to merge + return + + for active_adapter in adapter_names: + base_layer = self.get_base_layer() + if active_adapter in self.lora_A.keys(): + if safe_merge: + # Note that safe_merge will be slower than the normal merge + # because of the copy operation. + orig_weights = base_layer.weight.data.clone() + orig_weights += self.get_delta_weight(active_adapter) + + if not torch.isfinite(orig_weights).all(): + raise ValueError( + f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" + ) + + base_layer.weight.data = orig_weights + else: + base_layer.weight.data += self.get_delta_weight(active_adapter) + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + """ + This method unmerges all merged adapter layers from the base weights. + """ + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter in self.lora_A.keys(): + self.get_base_layer().weight.data -= self.get_delta_weight(active_adapter) + + def get_delta_weight(self, adapter) -> torch.Tensor: + return ( + transpose(self.lora_B[adapter] @ (self.lora_A[adapter] * self.lora_E[adapter]), self.fan_in_fan_out) + * self.scaling[adapter] + / (self.ranknum[adapter] + 1e-5) + ) + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + result = self.base_layer(x, *args, **kwargs) + for active_adapter in self.active_adapters: + if active_adapter not in self.lora_A.keys(): + continue + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + lora_E = self.lora_E[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + ranknum = self.ranknum[active_adapter] + 1e-5 + + x = self._cast_input_dtype(x, lora_A.dtype) + result += (dropout(x) @ (lora_A * lora_E).T @ lora_B.T) * scaling / ranknum + + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "adalora." + rep + + +class RankAllocator: + """ + The RankAllocator for AdaLoraModel. Paper: https://openreview.net/pdf?id=lq62uWRJjiY + + Args: + config ([`AdaLoraConfig`]): The configuration of the AdaLora model. + model: the model that we apply AdaLoRA to. + + """ + + def __init__(self, model, peft_config, adapter_name): + self.peft_config = peft_config + self.adapter_name = adapter_name + self.beta1 = peft_config.beta1 + self.beta2 = peft_config.beta2 + assert self.beta1 > 0 and self.beta1 < 1 + assert self.beta2 > 0 and self.beta2 < 1 + + self.reset_ipt() + self._set_budget_scheduler(model) + + def set_total_step(self, total_step): + self.peft_config.total_step = total_step + + def reset_ipt(self): + self.ipt = {} + self.exp_avg_ipt = {} + self.exp_avg_unc = {} + + def _set_budget_scheduler(self, model): + self.init_bgt = 0 + self.name_set = set() + for n, p in model.named_parameters(): + if f"lora_A.{self.adapter_name}" in n: + self.init_bgt += p.size(0) + self.name_set.add(n.replace("lora_A", "%s")) + self.name_set = sorted(self.name_set) + # The total final rank budget + self.target_bgt = self.peft_config.target_r * len(self.name_set) + + def budget_schedule(self, step: int): + tinit = self.peft_config.tinit + tfinal = self.peft_config.tfinal + total_step = self.peft_config.total_step + # Initial warmup + if step <= tinit: + budget = self.init_bgt + mask_ind = False + # Final fine-tuning + elif step > total_step - tfinal: + budget = self.target_bgt + mask_ind = True + else: + # Budget decreasing with a cubic scheduler + mul_coeff = 1 - (step - tinit) / (total_step - tfinal - tinit) + budget = int((self.init_bgt - self.target_bgt) * (mul_coeff**3) + self.target_bgt) + mask_ind = True if step % self.peft_config.deltaT == 0 else False + return budget, mask_ind + + def update_ipt(self, model): + # Update the sensitivity and uncertainty for every weight + for n, p in model.named_parameters(): + if "lora_" in n and self.adapter_name in n: + if n not in self.ipt: + self.ipt[n] = torch.zeros_like(p) + self.exp_avg_ipt[n] = torch.zeros_like(p) + self.exp_avg_unc[n] = torch.zeros_like(p) + with torch.no_grad(): + if deepspeed_config() is not None: + import deepspeed + + grad = deepspeed.utils.safe_get_full_grad(p) + self.ipt[n] = (p * grad).abs().detach() + else: + self.ipt[n] = (p * p.grad).abs().detach() + # Sensitivity smoothing + self.exp_avg_ipt[n] = self.beta1 * self.exp_avg_ipt[n] + (1 - self.beta1) * self.ipt[n] + # Uncertainty quantification + self.exp_avg_unc[n] = ( + self.beta2 * self.exp_avg_unc[n] + (1 - self.beta2) * (self.ipt[n] - self.exp_avg_ipt[n]).abs() + ) + + def _element_score(self, n): + return self.exp_avg_ipt[n] * self.exp_avg_unc[n] + + def _combine_ipt(self, ipt_E, ipt_AB): + ipt_AB = ipt_AB.sum(dim=1, keepdim=False) + sum_ipt = ipt_E.view(-1) + ipt_AB.view(-1) + return sum_ipt + + def mask_to_budget(self, model, budget): + value_ipt = {} + vector_ipt = {} + triplet_ipt = {} + # Get the importance score for A, E, B + for n, p in model.named_parameters(): + if f"lora_A.{self.adapter_name}" in n: + entry_ipt = self._element_score(n) + comb_ipt = torch.mean(entry_ipt, dim=1, keepdim=True) + name_m = n.replace("lora_A", "%s") + if name_m not in vector_ipt: + vector_ipt[name_m] = [comb_ipt] + else: + vector_ipt[name_m].append(comb_ipt) + if f"lora_B.{self.adapter_name}" in n: + entry_ipt = self._element_score(n) + comb_ipt = torch.mean(entry_ipt, dim=0, keepdim=False).view(-1, 1) + name_m = n.replace("lora_B", "%s") + if name_m not in vector_ipt: + vector_ipt[name_m] = [comb_ipt] + else: + vector_ipt[name_m].append(comb_ipt) + if f"lora_E.{self.adapter_name}" in n: + entry_ipt = self._element_score(n) + name_m = n.replace("lora_E", "%s") + value_ipt[name_m] = entry_ipt + + all_score = [] + # Calculate the score for each triplet + for name_m in vector_ipt: + ipt_E = value_ipt[name_m] + ipt_AB = torch.cat(vector_ipt[name_m], dim=1) + sum_ipt = self._combine_ipt(ipt_E, ipt_AB) + name_E = name_m % "lora_E" + triplet_ipt[name_E] = sum_ipt.view(-1, 1) + all_score.append(sum_ipt.view(-1)) + + # Get the threshold by ranking ipt + mask_threshold = torch.kthvalue( + torch.cat(all_score), + k=self.init_bgt - budget, + )[0].item() + + rank_pattern = {} + # Mask the unimportant triplets + with torch.no_grad(): + for n, p in model.named_parameters(): + if f"lora_E.{self.adapter_name}" in n: + p.masked_fill_(triplet_ipt[n] <= mask_threshold, 0.0) + rank_pattern[n] = (~(triplet_ipt[n] <= mask_threshold)).view(-1).tolist() + return rank_pattern + + def update_and_allocate(self, model, global_step, force_mask=False): + # # Update the importance score and allocate the budget + if global_step < self.peft_config.total_step - self.peft_config.tfinal: + self.update_ipt(model) + budget, mask_ind = self.budget_schedule(global_step) + # Allocate the budget according to importance scores + if mask_ind or force_mask: + rank_pattern = self.mask_to_budget(model, budget) + else: + rank_pattern = None + return budget, rank_pattern + + def mask_using_rank_pattern(self, model, rank_pattern): + # Mask the unimportant triplets + is_adapter_name_truncated = False + if self.adapter_name not in next(iter(rank_pattern.keys())): + is_adapter_name_truncated = True + + with torch.no_grad(): + for n, p in model.named_parameters(): + if f"lora_E.{self.adapter_name}" in n: + key = n if not is_adapter_name_truncated else n.replace(f".{self.adapter_name}", "") + mask = torch.Tensor(rank_pattern[key]).unsqueeze(-1).to(p.device) + p.masked_fill_(~mask.bool(), 0.0) diff --git a/lib/python3.12/site-packages/peft/tuners/adalora/model.py b/lib/python3.12/site-packages/peft/tuners/adalora/model.py new file mode 100644 index 0000000000000000000000000000000000000000..3c52ecdf2f9117d073998d58f4b3510e0cfbf58a --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adalora/model.py @@ -0,0 +1,364 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +import warnings + +import torch +from transformers.pytorch_utils import Conv1D + +from peft.import_utils import is_bnb_4bit_available, is_bnb_available, is_gptqmodel_available +from peft.tuners.lora import LoraConfig, LoraModel +from peft.tuners.tuners_utils import BaseTunerLayer +from peft.utils import ( + TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING, + _freeze_adapter, + _get_submodules, + get_auto_gptq_quant_linear, + get_gptqmodel_quant_linear, + get_quantization_config, +) +from peft.utils.integrations import gather_params_ctx + +from .gptq import SVDQuantLinear +from .layer import AdaLoraLayer, RankAllocator, SVDLinear + + +class AdaLoraModel(LoraModel): + """ + Creates AdaLoRA (Adaptive LoRA) model from a pretrained transformers model. Paper: + https://openreview.net/forum?id=lq62uWRJjiY + + Args: + model ([`transformers.PreTrainedModel`]): The model to be adapted. + config ([`AdaLoraConfig`]): The configuration of the AdaLora model. + adapter_name (`str`): The name of the adapter, defaults to `"default"`. + low_cpu_mem_usage (`bool`, `optional`, defaults to `False`): + Create empty adapter weights on meta device. Useful to speed up the loading process. + + Returns: + `torch.nn.Module`: The AdaLora model. + + Example:: + + >>> from transformers import AutoModelForSeq2SeqLM >>> from peft import LoraConfig, AdaLoraModel, AdaLoraConfig + >>> config = AdaLoraConfig( + peft_type="ADALORA", task_type="SEQ_2_SEQ_LM", init_r=12, lora_alpha=32, target_modules=["q", "v"], + lora_dropout=0.01, + ) + >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> model = AdaLoraModel(model, config, "default") + + **Attributes**: + - **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted. + - **peft_config** ([`AdaLoraConfig`]): The configuration of the AdaLora model. + """ + + # Note: don't redefine prefix here, it should be inherited from LoraModel + + def __init__(self, model, config, adapter_name): + super().__init__(model, config, adapter_name) + + traininable_mode_counter = 0 + for config in self.peft_config.values(): + if not config.inference_mode: + traininable_mode_counter += 1 + + if traininable_mode_counter > 1: + raise ValueError( + "AdaLoraModel supports only 1 trainable adapter. " + "When using multiple adapters, set inference_mode to True for all adapters except the one you want to train." + ) + + if self.peft_config[adapter_name].inference_mode: + _freeze_adapter(self.model, adapter_name) + else: + self.trainable_adapter_name = adapter_name + self.rankallocator = RankAllocator(self.model, self.peft_config[adapter_name], self.trainable_adapter_name) + + def _check_new_adapter_config(self, config: LoraConfig) -> None: + """ + A helper method to check the config when a new adapter is being added. + + Raise a ValueError if there is something wrong with the config or if it conflicts with existing adapters. + + """ + super()._check_new_adapter_config(config) + + traininable_mode_counter = 0 + for config_ in self.peft_config.values(): + if not config_.inference_mode: + traininable_mode_counter += 1 + + if traininable_mode_counter > 1: + raise ValueError( + f"{self.__class__.__name__} supports only 1 trainable adapter. " + "When using multiple adapters, set inference_mode to True for all adapters except the one " + "you want to train." + ) + + def _create_and_replace( + self, + lora_config, + adapter_name, + target, + target_name, + parent, + current_key, + ): + kwargs = { + "r": lora_config.init_r, + "lora_alpha": lora_config.lora_alpha, + "lora_dropout": lora_config.lora_dropout, + "fan_in_fan_out": lora_config.fan_in_fan_out, + "init_lora_weights": lora_config.init_lora_weights, + "loaded_in_8bit": getattr(self.model, "is_loaded_in_8bit", False), + "loaded_in_4bit": getattr(self.model, "is_loaded_in_4bit", False), + } + if (kwargs["loaded_in_8bit"] or kwargs["loaded_in_4bit"]) and not is_bnb_available(): + raise ImportError( + "To use AdaLora with 8-bit quantization, please install the `bitsandbytes` package. " + "You can install it with `pip install bitsandbytes`." + ) + + quantization_config = get_quantization_config(self.model, method="gptq") + if quantization_config is not None: + kwargs["gptq_quantization_config"] = quantization_config + + # If it is not an AdaLoraLayer, create a new module, else update it with new adapters + if not isinstance(target, AdaLoraLayer): + device_map = self.model.hf_device_map if hasattr(self.model, "hf_device_map") else None + new_module = self._create_new_module(lora_config, adapter_name, target, device_map=device_map, **kwargs) + if adapter_name not in self.active_adapters: + # adding an additional adapter: it is not automatically trainable + new_module.requires_grad_(False) + self._replace_module(parent, target_name, new_module, target) + else: + target.update_layer( + adapter_name, + lora_config.init_r, + lora_config.lora_alpha, + lora_config.lora_dropout, + lora_config.init_lora_weights, + ) + + @staticmethod + def _create_new_module(lora_config, adapter_name, target, device_map=None, **kwargs): + # avoid eager bnb import + if is_bnb_available(): + import bitsandbytes as bnb + + from .bnb import SVDLinear8bitLt + if is_bnb_4bit_available(): + from .bnb import SVDLinear4bit + + gptq_quantization_config = kwargs.get("gptq_quantization_config", None) + + if is_gptqmodel_available(): + QuantLinear = get_gptqmodel_quant_linear(gptq_quantization_config, device_map=device_map) + else: + QuantLinear = get_auto_gptq_quant_linear(gptq_quantization_config) + + loaded_in_8bit = kwargs.pop("loaded_in_8bit", False) + loaded_in_4bit = kwargs.pop("loaded_in_4bit", False) + + if isinstance(target, BaseTunerLayer): + target_base_layer = target.get_base_layer() + else: + target_base_layer = target + + if loaded_in_8bit and isinstance(target_base_layer, bnb.nn.Linear8bitLt): + kwargs.update( + { + "has_fp16_weights": target_base_layer.state.has_fp16_weights, + "threshold": target_base_layer.state.threshold, + "index": target_base_layer.index, + } + ) + new_module = SVDLinear8bitLt(target, adapter_name, **kwargs) + elif loaded_in_4bit and is_bnb_4bit_available() and isinstance(target_base_layer, bnb.nn.Linear4bit): + fourbit_kwargs = kwargs.copy() + fourbit_kwargs.update( + { + "compute_dtype": target_base_layer.compute_dtype, + "compress_statistics": target_base_layer.weight.compress_statistics, + "quant_type": target_base_layer.weight.quant_type, + } + ) + new_module = SVDLinear4bit(target, adapter_name, **fourbit_kwargs) + elif QuantLinear is not None and isinstance(target, QuantLinear): + new_module = SVDQuantLinear(target, adapter_name, **kwargs) + else: + if isinstance(target_base_layer, torch.nn.Linear): + if kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. " + "Setting fan_in_fan_out to False." + ) + kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = False + elif isinstance(target_base_layer, Conv1D): + if not kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to False but the target module is `Conv1D`. " + "Setting fan_in_fan_out to True." + ) + kwargs["fan_in_fan_out"] = lora_config.fan_in_fan_out = True + else: + raise ValueError( + f"Target module {target} is not supported. " + f"Currently, only `torch.nn.Linear` and `Conv1D` are supported." + ) + new_module = SVDLinear(target, adapter_name, **kwargs) + + return new_module + + @staticmethod + def _prepare_adapter_config(peft_config, model_config): + if peft_config.target_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING: + raise ValueError("Please specify `target_modules` in `peft_config`") + peft_config.target_modules = TRANSFORMERS_MODELS_TO_ADALORA_TARGET_MODULES_MAPPING[ + model_config["model_type"] + ] + return peft_config + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + if name == "model": # see #1892: prevent infinite recursion if class is not initialized + raise + return getattr(self.model, name) + + def forward(self, *args, **kwargs): + outputs = self.model.forward(*args, **kwargs) + + if (getattr(outputs, "loss", None) is not None) and isinstance(outputs.loss, torch.Tensor): + # Calculate the orthogonal regularization + orth_reg_weight = self.peft_config[self.trainable_adapter_name].orth_reg_weight + + if orth_reg_weight <= 0: + raise ValueError("orth_reg_weight should be greater than 0. ") + + regu_loss = 0 + num_param = 0 + for n, p in self.model.named_parameters(): + if ("lora_A" in n or "lora_B" in n) and self.trainable_adapter_name in n: + if p.shape == torch.Size([0]): + with gather_params_ctx(p, fwd_module=self): + para_cov = p @ p.T if "lora_A" in n else p.T @ p + else: + para_cov = p @ p.T if "lora_A" in n else p.T @ p + I = torch.eye(*para_cov.size(), out=torch.empty_like(para_cov)) # noqa: E741 + I.requires_grad = False + num_param += 1 + regu_loss += torch.norm(para_cov - I, p="fro") + if num_param > 0: + regu_loss = regu_loss / num_param + else: + regu_loss = 0 + outputs.loss += orth_reg_weight * regu_loss + return outputs + + def resize_modules_by_rank_pattern(self, rank_pattern, adapter_name): + lora_config = self.peft_config[adapter_name] + for name, rank_idx in rank_pattern.items(): + if isinstance(rank_idx, list): + rank = sum(rank_idx) + elif isinstance(rank_idx, torch.Tensor): + rank_idx = rank_idx.view(-1) + rank = rank_idx.sum().item() + else: + raise ValueError("Unexpected type of rank_idx") + key = ".".join(name.split(".")[0:-2]) if adapter_name in name else ".".join(name.split(".")[0:-1]) + _, target, _ = _get_submodules(self.model, key) + lora_E_weights = target.lora_E[adapter_name][rank_idx] + lora_A_weights = target.lora_A[adapter_name][rank_idx] + lora_B_weights = target.lora_B[adapter_name][:, rank_idx] + ranknum = target.ranknum[adapter_name] + target.update_layer( + adapter_name, + rank, + lora_config.lora_alpha, + lora_config.lora_dropout, + lora_config.init_lora_weights, + ) + with torch.no_grad(): + if rank > 0: + target.lora_E[adapter_name].copy_(lora_E_weights) + target.lora_A[adapter_name].copy_(lora_A_weights) + target.lora_B[adapter_name].copy_(lora_B_weights) + # The scaling is exactly as the previous + target.ranknum[adapter_name].copy_(ranknum) + + def resize_state_dict_by_rank_pattern(self, rank_pattern, state_dict, adapter_name): + for name, rank_idx in rank_pattern.items(): + rank = sum(rank_idx) + prefix = ".".join(name.split(".")[0:-2]) if adapter_name in name else ".".join(name.split(".")[0:-1]) + for layer in ["lora_E", "lora_A", "lora_B"]: + key = f"base_model.model.{prefix}.{layer}.{adapter_name}" + if layer != "lora_B": + state_dict[key] = ( + state_dict[key][rank_idx] if rank != state_dict[key].shape[0] else state_dict[key] + ) + else: + state_dict[key] = ( + state_dict[key][:, rank_idx] if rank != state_dict[key].shape[1] else state_dict[key] + ) + return state_dict + + def update_and_allocate(self, global_step): + """ + This method updates Adalora budget and mask. + + This should be called in every training step after `loss.backward()` and before `zero_grad()`. + + `tinit`, `tfinal` and `deltaT` are handled with in the method. + + Args: + global_step (`int`): The current training step, it is used to calculate adalora budget. + + Example: + + ```python + >>> loss = model(**input).loss + >>> loss.backward() + >>> optimizer.step() + >>> model.base_model.update_and_allocate(i_step) + >>> optimizer.zero_grad() + ``` + """ + lora_config = self.peft_config[self.trainable_adapter_name] + # Update the importance score and allocate the budget + if global_step < lora_config.total_step - lora_config.tfinal: + _, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step) + if rank_pattern: + lora_config.rank_pattern = rank_pattern + # Finalize the budget allocation + elif global_step == lora_config.total_step - lora_config.tfinal: + _, rank_pattern = self.rankallocator.update_and_allocate(self.model, global_step, force_mask=True) + # for some reason, this freezes the trainable parameters and nothing gets updates + # self.resize_modules_by_rank_pattern(rank_pattern, self.trainable_adapter_name) + lora_config.rank_pattern = rank_pattern + self.rankallocator.reset_ipt() + # Currently using inefficient way to mask the unimportant weights using the rank pattern + # due to problem mentioned above + elif global_step > lora_config.total_step - lora_config.tfinal: + self.rankallocator.mask_using_rank_pattern(self.model, lora_config.rank_pattern) + # Pass the function and do forward propagation + else: + return None + + def add_weighted_adapter(self, *args, **kwargs): + """This method is not supported for AdaLoRA, use LoRA instead.""" + raise TypeError(f"{self.__class__.__name__} does not support add_weighted_adapter method.") diff --git a/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__init__.py b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..68882a222625eedc770d28816a3a0646964739bb --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__init__.py @@ -0,0 +1,23 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 peft.utils import register_peft_method + +from .config import AdaptionPromptConfig +from .layer import AdaptedAttention +from .model import AdaptionPromptModel + + +__all__ = ["AdaptedAttention", "AdaptionPromptConfig", "AdaptionPromptModel"] + +register_peft_method(name="adaption_prompt", config_cls=AdaptionPromptConfig, model_cls=AdaptionPromptModel) diff --git a/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cefab2e7c50e66dd50764154d77956b7569cd9c7 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d73da71f99e09fb1dfcd2e4c3a336f2040d86ff0 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/layer.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..975617aec476931705754bb4f3c6ad8765bb290f Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c02468bfff9596a946947940ce26b97c3213fe1 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f04fb58086e2b121420c9dbefe3a50372775b624 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/adaption_prompt/config.py b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/config.py new file mode 100644 index 0000000000000000000000000000000000000000..52a405cebe821d68ba446d1a3e7a844c07b135de --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/config.py @@ -0,0 +1,81 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 collections import namedtuple +from dataclasses import dataclass, field + +from peft.config import PeftConfig +from peft.utils import PeftType + +from .utils import llama_compute_query_states + + +@dataclass +class AdaptionPromptConfig(PeftConfig): + """Stores the configuration of an [`AdaptionPromptModel`].""" + + target_modules: str = field( + default=None, metadata={"help": "Name of the attention submodules to insert adaption prompts into."} + ) + adapter_len: int = field(default=None, metadata={"help": "Number of adapter tokens to insert"}) + adapter_layers: int = field(default=None, metadata={"help": "Number of adapter layers (from the top)"}) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.ADAPTION_PROMPT + + @property + def is_adaption_prompt(self) -> bool: + """Return True if this is an adaption prompt config.""" + return True + + +# Contains the config that is specific to a transformers model type. +ModelTypeConfig = namedtuple( + "ModelTypeConfig", ["compute_query_states", "target_modules", "k_proj_layer", "v_proj_layer", "o_proj_layer"] +) + +# Mapping of transformers model types to their specific configuration. +TRANSFORMERS_MODEL_CONFIG = { + "llama": ModelTypeConfig( + compute_query_states=llama_compute_query_states, + target_modules="self_attn", + k_proj_layer="k_proj", + v_proj_layer="v_proj", + o_proj_layer="o_proj", + ), + "mistral": ModelTypeConfig( # same as llama, + compute_query_states=llama_compute_query_states, + target_modules="self_attn", + k_proj_layer="k_proj", + v_proj_layer="v_proj", + o_proj_layer="o_proj", + ), +} + + +def prepare_config( + peft_config: AdaptionPromptConfig, + model, +) -> AdaptionPromptConfig: + """Prepare the config based on the llama model type.""" + if model.config.model_type not in TRANSFORMERS_MODEL_CONFIG: + raise ValueError("Unsupported model type for adaption prompt: '{model.config.model_type}'.") + + model_config = TRANSFORMERS_MODEL_CONFIG[model.config.model_type] + + if peft_config.target_modules is None: + peft_config.target_modules = model_config.target_modules + + return peft_config diff --git a/lib/python3.12/site-packages/peft/tuners/adaption_prompt/layer.py b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..e15439de0507db5d1fc00189260eb97906bb4759 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/layer.py @@ -0,0 +1,138 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .config import TRANSFORMERS_MODEL_CONFIG + + +class AdaptedAttention(nn.Module): + """This module wraps a LLamaAttention module and injects adaption prompts.""" + + def __init__(self, model_type: str, adapter_len: int, model): + """ + Initialize object. + + Args: + model_type: The transformer model type. This is used to retrieve the right method to + compute query states. + adapter_len: The length of the adaption prompt to insert. + model: The original transformer attention module that is being wrapped. + """ + assert not isinstance(model, AdaptedAttention) + super().__init__() + self.model_type = model_type + self.model = model + self.adapter_len = adapter_len + # Assume all parameters of the attention model we are wrapping are on the same device. + device = next(model.parameters()).device + # Don't think this was specified in the paper, but we follow the official repo which used an Embedding + # which initializes the tokens with standard normal values. + # https://github.com/ZrrSkywalker/LLaMA-Adapter/blob/41c3546fe1997ab8a65809dc8d8f9252b19d9faf/llama/model.py#L234 + # (bsz, adapter_len, hidden_size) + target_dtype = ( + model.q_proj.weight.dtype if model.q_proj.weight.dtype not in [torch.int8, torch.uint8] else torch.float32 + ) + if hasattr(self.model, "hidden_size"): + # TODO: remove this clause after 2026-01-01 + hidden_size = self.model.hidden_size + else: # changed in https://github.com/huggingface/transformers/pull/35235 + hidden_size = self.model.config.hidden_size + self.adaption_prompt = nn.Parameter( + torch.empty(1, adapter_len, hidden_size, device=device, dtype=target_dtype).normal_() + ) + # Initialize the gate to 0 as this is "zero-init". + self.adaption_gate = nn.Parameter(torch.zeros(1, device=device, dtype=target_dtype)) + + def forward(self, **kwargs): + """ + Forward pass for the adapter which wraps the original LlamaAttention module. + + "Official" paper implementation: + https://github.com/ZrrSkywalker/LLaMA-Adapter/blob/41c3546fe1997ab8a65809dc8d8f9252b19d9faf/llama/model.py#L141 + + Args: + kwargs: See the original LlamaAttention module. + """ + if kwargs.get("output_attention", False): + raise NotImplementedError("output_attention is not currently supported.") + + output, *_ = self.model(**kwargs) + bsz = output.shape[0] + q_len = output.shape[1] + embed_dim = output.shape[2] + k_proj_layer = TRANSFORMERS_MODEL_CONFIG[self.model_type].k_proj_layer + v_proj_layer = TRANSFORMERS_MODEL_CONFIG[self.model_type].v_proj_layer + o_proj_layer = TRANSFORMERS_MODEL_CONFIG[self.model_type].o_proj_layer + factor = ( + self.model.k_proj.in_features // self.model.k_proj.out_features + ) # Mistral has different input and output dimension for k_proj and v_proj layers + + if k_proj_layer == v_proj_layer: + _, key, value = getattr(self.model, k_proj_layer)(self.adaption_prompt).split(embed_dim, dim=2) + else: + key = getattr(self.model, k_proj_layer)(self.adaption_prompt) + value = getattr(self.model, v_proj_layer)(self.adaption_prompt) + + if hasattr(self.model, "num_heads"): + # TODO: remove this clause after 2026-01-01 + num_heads = self.model.num_heads + else: # changed in https://github.com/huggingface/transformers/pull/35235 + num_heads = self.model.config.num_attention_heads + # (bsz, num_key_value_heads, adapter_len, head_dim) + adapter_k = ( + key.view(1, self.adapter_len, (num_heads // factor), self.model.head_dim) + .repeat(bsz, 1, 1, 1) + .transpose(1, 2) + ) + adapter_v = ( + value.view(1, self.adapter_len, (num_heads // factor), self.model.head_dim) + .repeat(bsz, 1, 1, 1) + .transpose(1, 2) + ) + # Below is taken from https://github.com/huggingface/transformers/blob/e547458c43dfdbbb8f6a7757237e234c44e20a8f/src/transformers/models/mistral/modeling_mistral.py#L181 + # (bsz, num_heads, adapter_len, head_dim) + adapter_k = torch.repeat_interleave(adapter_k, repeats=factor, dim=1) + adapter_v = torch.repeat_interleave(adapter_v, repeats=factor, dim=1) + # Recompute query states. + compute_query_states = TRANSFORMERS_MODEL_CONFIG[self.model_type].compute_query_states + # (bsz, num_heads, q_len, head_dim) + query_states = compute_query_states(model=self.model, **kwargs) + + previous_dtype = query_states.dtype + + # (bsz, num_heads, q_len, adapter_len) + scores = torch.matmul(query_states, adapter_k.transpose(2, 3).to(previous_dtype)) / math.sqrt( + self.model.head_dim + ) + # Upcast attention to fp32 + # (bsz, num_heads, q_len, adapter_len) + scores = self.adaption_gate * F.softmax(scores, dim=-1, dtype=torch.float32).to(previous_dtype) + # (bsz, q_len, num_heads * head_dim) + adapter_output = torch.matmul(scores, adapter_v).transpose(1, 2).reshape(bsz, q_len, -1) + + # (bsz, q_len, hidden_size) + if o_proj_layer is not None: + adapter_output = getattr(self.model, o_proj_layer)(adapter_output) + + # Add adaption prompt output to original output. + output = output + adapter_output + + # Restore original dtype. + output = output.to(previous_dtype) + return output, *_ diff --git a/lib/python3.12/site-packages/peft/tuners/adaption_prompt/model.py b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/model.py new file mode 100644 index 0000000000000000000000000000000000000000..1ee98cd8fbdb28f19a0a382199941255b602309f --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/model.py @@ -0,0 +1,162 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 typing import Dict, List + +import torch.nn as nn + +from peft.utils import _freeze_adapter, _get_submodules + +from .config import AdaptionPromptConfig, prepare_config +from .layer import AdaptedAttention +from .utils import is_adaption_prompt_trainable + + +class AdaptionPromptModel(nn.Module): + """ + Implements adaption prompts as described in https://arxiv.org/pdf/2303.16199.pdf. + + The top L attention modules are replaced with AdaptedAttention modules that wrap the original ones, but insert + trainable prompts with gates (for zero init). + + Notes on the multi-adapter pattern: + - We store the states of different adapters by keeping a dictionary of AdaptedAttention modules indexed by adapter + name. + - Every time we switch adapters, we remove the modules of the currently active adapter from the model, store them + in the dictionary, and replace them with the modules of the new adapter. + - To avoid duplicated and potentially inconsistent state, the currently active adapter is always removed from the + dictionary. + - Disabling the adapter would also result in the modules being removed from the model. + """ + + def __init__(self, model, configs: Dict, adapter_name: str): + super().__init__() + self.model = model + # Store adapter configs by name. + self.peft_config: Dict[str, AdaptionPromptConfig] = {} + # Store lists of the parents of the affected attention modules by adapter name. + # We keep references to the parents so we can swap the adapters in-and-out of the model. + self._parents: Dict[str, List[nn.Module]] = {} + # Store lists of cached AdaptedAttention modules by name. + self._cached_adapters: Dict[str, List] = {} + # The name of the currently active adapter. + self._active_adapter = None + # Whether the adapter is enabled. + self._enabled = True + self.forward = self.model.forward + self.add_adapter(adapter_name, configs[adapter_name]) + self._mark_only_adaption_prompts_as_trainable(self.model) + + def add_adapter(self, adapter_name: str, config: AdaptionPromptConfig) -> None: + """Add an adapter with the given name and config.""" + config = prepare_config(config, self.model) + if adapter_name in self.peft_config: + raise ValueError(f"Adapter with name '{adapter_name}' already exists.") + + parents = [] + for name, _ in self.model.named_modules(): + if name.endswith(config.target_modules): + par, _, _ = _get_submodules(self.model, name) + parents.append(par) + if len(parents) < config.adapter_layers: + raise ValueError( + f"Config specifies more adapter layers '{config.adapter_layers}' than the model has '{len(parents)}'." + ) + # Note that if the target modules are not in Sequential, ModuleList, or + # some other PyTorch ordered container, the behavior is undefined as we + # assume here that the order of the modules is the same as the order of + # the transformer decoder layers. + parents = parents[-config.adapter_layers :] + self._parents[adapter_name] = parents + + # It is only None during initialization. + # If it is disabled, we don't have to remove the modules. + if self._active_adapter is not None and self._enabled: + self._remove_adapted_attentions(self._active_adapter) + self._active_adapter = adapter_name + self.peft_config[adapter_name] = config + self._create_adapted_attentions(config, parents) + if not self._enabled: + self._remove_adapted_attentions(self._active_adapter) + + if config.inference_mode: + _freeze_adapter(self.model, adapter_name) + + def set_adapter(self, adapter_name: str) -> None: + """Set the model to use the adapter with the given name.""" + if self._active_adapter == adapter_name: + return + if adapter_name not in self.peft_config: + raise ValueError(f"Adapter with name '{adapter_name}' does not exist.") + + if self._enabled: + self._remove_adapted_attentions(self._active_adapter) + self._set_adapted_attentions(adapter_name) + + self._active_adapter = adapter_name + + def enable_adapter_layers(self): + """Enable adapter layers by swapping in cached AdaptedAttention modules.""" + self._enabled = True + self._set_adapted_attentions(self._active_adapter) + + def disable_adapter_layers(self): + """Disable adapter layers by swapping out AdaptedAttention modules.""" + self._enabled = False + self._remove_adapted_attentions(self._active_adapter) + + def _create_adapted_attentions(self, config: AdaptionPromptConfig, parents: List[nn.Module]) -> None: + """Wrap LlamaAttention modules with newly created AdaptedAttention modules.""" + for par in parents: + attn = AdaptedAttention( + model_type=self.model.config.model_type, + adapter_len=config.adapter_len, + model=getattr(par, config.target_modules), + ) + setattr(par, config.target_modules, attn) + + def _set_adapted_attentions(self, adapter_name: str) -> None: + """Replace LlamaAttention modules with cached AdaptedAttention modules.""" + cached = self._cached_adapters[adapter_name] + del self._cached_adapters[adapter_name] + config = self.peft_config[adapter_name] + for i, par in enumerate(self._parents[adapter_name]): + setattr(par, config.target_modules, cached[i]) + + def _remove_adapted_attentions(self, adapter_name: str) -> None: + """Remove AdaptedAttention modules from the model and store them in the cache.""" + config = self.peft_config[adapter_name] + adapted_attentions = [] + for par in self._parents[adapter_name]: + attn = getattr(par, config.target_modules) + adapted_attentions.append(attn) + setattr(par, config.target_modules, attn.model) + self._cached_adapters[adapter_name] = adapted_attentions + + def _mark_only_adaption_prompts_as_trainable(self, model: nn.Module) -> None: + """Freeze all parameters of the model except the adaption prompts.""" + for n, p in model.named_parameters(): + if not is_adaption_prompt_trainable(n): + p.requires_grad = False + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + # This is necessary as e.g. causal models have various methods that we + # don't want to re-implement here. + if name == "model": # see #1892: prevent infinite recursion if class is not initialized + raise + return getattr(self.model, name) diff --git a/lib/python3.12/site-packages/peft/tuners/adaption_prompt/utils.py b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3058c4ef2fd5350e33e87b1746aaf3b9b6d5ef01 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/adaption_prompt/utils.py @@ -0,0 +1,132 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. +import inspect + +import torch +import torch.nn as nn + + +def llama_rotate_half(x: torch.Tensor) -> torch.Tensor: + """ + Rotate half the hidden dims of the input. + + This function was duplicated verbatim from: + https://github.com/huggingface/transformers/blob/1de8ce9ee1191ba761a593ac15d9ccbf5851bfc5/src/transformers/models/llama/modeling_llama.py#L126 + + This was done to eliminate the Llama transformers implementation as a dependency of this file. Note that some other + functions were also adapted from the transformers implementation but were modified. + """ + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def llama_apply_rotary_pos_emb(q, cos, sin, position_ids): + """ + Apply rotary position embedding to query states in the Llama model. + + This function was adapted from: + https://github.com/huggingface/transformers/blob/1de8ce9ee1191ba761a593ac15d9ccbf5851bfc5/src/transformers/models/llama/modeling_llama.py#L133 + + It was modified to remove unnecessary processing of key states. The method is compatible with transformers <= + 4.34.2 and also with the latest version (>=4.35). + """ + # In previous transformers version cos/sin cached had a shape of 4D + if len(cos.shape) == 4: + gather_indices = position_ids[:, None, :, None] # [bs, 1, seq_len, 1] + gather_indices = gather_indices.repeat(1, cos.shape[1], 1, cos.shape[3]) + cos = torch.gather(cos.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices) + sin = torch.gather(sin.repeat(gather_indices.shape[0], 1, 1, 1), 2, gather_indices) + # In the new version, it is 2D so we fall back to the new implementation + # https://github.com/huggingface/transformers/blame/eef7ea98c31a333bacdc7ae7a2372bde772be8e4/src/transformers/models/llama/modeling_llama.py#L222-L226 + else: + cos = cos[position_ids].unsqueeze(1) + sin = sin[position_ids].unsqueeze(1) + q_embed = (q * cos) + (llama_rotate_half(q) * sin) + return q_embed + + +def llama_compute_query_states(model: nn.Module, **kwargs) -> torch.Tensor: + """ + Compute query states for Llama models specifically. They need to be recomputed as the forward() method of the + original LlamaModel in the transformers library does not return them. See the related discussion in the PR: + https://github.com/huggingface/peft/pull/268 + """ + hidden_states = kwargs.get("hidden_states") + position_ids = kwargs.get("position_ids") + past_key_value = kwargs.get("past_key_value") + bsz, q_len, _ = hidden_states.size() + if hasattr(model, "num_heads"): + # TODO: remove this clause after 2026-01-01 + num_heads = model.num_heads + else: # changed in https://github.com/huggingface/transformers/pull/35235 + num_heads = model.config.num_attention_heads + query_states = model.q_proj(hidden_states).view(bsz, q_len, num_heads, model.head_dim).transpose(1, 2) + + factor = model.k_proj.in_features // model.k_proj.out_features + value_states = model.v_proj(hidden_states).view(bsz, q_len, (num_heads // factor), model.head_dim).transpose(1, 2) + + seq_len = q_len + + if past_key_value is not None: + if isinstance(past_key_value, tuple): + # for transformers <= 4.35 + seq_len += past_key_value[0].shape[-2] + else: + # since transformers 4.36, this is a DynamicCache instance + seq_len += past_key_value.get_seq_length(model.layer_idx) + + # model.rotary_emb is deprecated and will be removed in transformers > 4.47.0. Instead, the position embeddings are + # passed via the kwargs + if "position_embeddings" in kwargs: + cos, sin = kwargs["position_embeddings"] + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return (query_states * cos) + (llama_rotate_half(query_states) * sin) + + # For transformers > 4.37.2 `position_ids` became a required arguments in the rotary embedding's forward pass. + if "position_ids" not in inspect.signature(model.rotary_emb.forward).parameters: + # TODO we assume that position_ids is not None here, not sure if that is safe but the old code also did that + cos, sin = model.rotary_emb(value_states, seq_len=seq_len) + return llama_apply_rotary_pos_emb(query_states, cos, sin, position_ids) + + past_seen_tokens = 0 + if position_ids is None: + # Compute position_ids, since they are required for transformers > 4.37.2 + if past_key_value is None: + new_cache_positions = torch.arange(q_len, q_len + q_len, device=value_states.device) + else: + past_seen_tokens = past_key_value.get_usable_length(q_len, model.layer_idx) + new_cache_positions = torch.arange(past_seen_tokens, past_seen_tokens + q_len, device=value_states.device) + position_ids = new_cache_positions.unsqueeze(0) + + rotary_emb_kwargs = {"position_ids": position_ids} + # The `seq_len` argument has been officially removed in transformers >= 4.39.0 + if "seq_len" in inspect.signature(model.rotary_emb.forward).parameters: + rotary_emb_kwargs["seq_len"] = q_len + past_seen_tokens + + cos, sin = model.rotary_emb(value_states, **rotary_emb_kwargs) + + # For batched inference unsqueeze it on the correct dim + # since: https://github.com/huggingface/transformers/pull/29109 + if len(cos.shape) == 3: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + + return (query_states * cos) + (llama_rotate_half(query_states) * sin) + + +def is_adaption_prompt_trainable(params: str) -> bool: + """Return True if module is trainable under adaption prompt fine-tuning.""" + return params.split(".")[-1].startswith("adaption_") diff --git a/lib/python3.12/site-packages/peft/tuners/boft/__init__.py b/lib/python3.12/site-packages/peft/tuners/boft/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c84b8358da8fc52d3f332c195f22b8b9620c665f --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/boft/__init__.py @@ -0,0 +1,24 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 peft.utils import register_peft_method + +from .config import BOFTConfig +from .layer import BOFTLayer +from .model import BOFTModel + + +__all__ = ["BOFTConfig", "BOFTLayer", "BOFTModel"] + +register_peft_method(name="boft", config_cls=BOFTConfig, model_cls=BOFTModel) diff --git a/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0ec6e1e90e8e5fa28ec7b12926a3b27c66eb95f Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..afa256547cc6e97230102814b337b89f1f4cb2a6 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/layer.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..186a5c8cd2b9b7eeed3dc23ce6d49757337ad1b2 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4a103f6e3496e02fb982ba579418cebdb6cd149f Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/boft/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/boft/config.py b/lib/python3.12/site-packages/peft/tuners/boft/config.py new file mode 100644 index 0000000000000000000000000000000000000000..bedd915664c9d9970cd1da91978415e32b173008 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/boft/config.py @@ -0,0 +1,160 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +# The implementation is based on "Parameter-Efficient Orthogonal Finetuning +# via Butterfly Factorization" (https://arxiv.org/abs/2311.06243) in ICLR 2024. + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Optional, Union + +from peft.config import PeftConfig +from peft.utils import PeftType + + +@dataclass +class BOFTConfig(PeftConfig): + """ + This is the configuration class to store the configuration of a [`BOFTModel`]. + + Args: + boft_block_size (`int`): BOFT block size across different layers. + boft_block_num (`int`): Number of BOFT blocks per injected layer. + boft_n_butterfly_factor (`int`): Number of butterfly factors across different layers. + target_modules (`Union[List[str],str]`): The names of the modules to apply the adapter to. + exclude_modules (`Optional[Union[List[str], str]]`): + The names of the modules to not apply the adapter. When passing a string, a regex match will be performed. + When passing a list of strings, either an exact match will be performed or it is checked if the name of the + module ends with any of the passed strings. + boft_dropout (`float`): + The multiplicative dropout probability, by setting OFT blocks to identity during training, similar to the + dropout layer in LoRA. + fan_in_fan_out (`bool`): Set this to True if the layer to replace stores weight like (fan_in, fan_out). + For example, gpt-2 uses `Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set + to `True`. + bias (`str`): Bias type for BOFT. Can be 'none', 'all' or 'boft_only'. If 'all' or 'boft_only', the + corresponding biases will be updated during training. Be aware that this means that, even when disabling + the adapters, the model will not produce the same output as the base model would have without adaptation. + modules_to_save (`List[str]`):List of modules apart from BOFT layers to be set as trainable + and saved in the final checkpoint. + layers_to_transform (`Union[List[int],int]`): + The layer indexes to transform, if this argument is specified, it will apply the BOFT transformations on + the layer indexes that are specified in this list. If a single integer is passed, it will apply the BOFT + transformations on the layer at this index. + layers_pattern (`Optional[Union[List[str], str]]`): + The layer pattern name, used only if `layers_to_transform` is different from `None` and if the layer + pattern is not in the common layers pattern. This should target the `nn.ModuleList` of the model, which is + often called `'layers'` or `'h'`. + """ + + boft_block_size: int = field( + default=4, + metadata={ + "help": "BOFT block size across different layers.", + "note": "You can only specify either boft_block_size or boft_block_num, but not both simultaneously, because boft_block_size x boft_block_num = layer dimension.", + }, + ) + boft_block_num: int = field( + default=0, + metadata={ + "help": "Number of BOFT blocks per injected layer.", + "note": "You can only specify either boft_block_size or boft_block_num, but not both simultaneously, because boft_block_size x boft_block_num = layer dimension.", + }, + ) + boft_n_butterfly_factor: int = field( + default=1, + metadata={ + "help": "Number of butterfly factors.", + "note": ( + "for example, boft_n_butterfly_factor=2, the effective block size of OFT becomes twice as big and the number of blocks become half.", + "note: for boft_n_butterfly_factor=1, BOFT is the same as vanilla OFT.", + ), + }, + ) + target_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": "List of module names or regex expression of the module names to replace with BOFT.", + "example": "For example, ['q', 'v'] or '.*decoder.*(SelfAttention|EncDecAttention).*(q|v)$' ", + }, + ) + exclude_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={"help": "List of module names or regex expression of the module names to exclude from BOFT."}, + ) + boft_dropout: float = field( + default=0.0, + metadata={ + "help": "BOFT multiplicative dropout, randomly setting blocks of OFT to be identity matrix, similar to the dropout layer in LoRA." + }, + ) + fan_in_fan_out: bool = field( + default=False, + metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"}, + ) + bias: str = field(default="none", metadata={"help": "Bias type for BOFT. Can be 'none', 'all' or 'boft_only'"}) + modules_to_save: Optional[list[str]] = field( + default=None, + metadata={ + "help": "List of modules apart from BOFT layers to be set as trainable and saved in the final checkpoint. ", + "note": ( + "For example, in Sequence Classification or Token Classification tasks, ", + "the final layer `classifier/score` are randomly initialized and as such need to be trainable and saved.", + ), + }, + ) + init_weights: bool = field( + default=True, + metadata={ + "help": ( + "Whether to initialize the weights of the BOFT layers with their default initialization. Don't change ", + "this setting, except if you know exactly what you're doing.", + ), + }, + ) + layers_to_transform: Optional[Union[list[int], int]] = field( + default=None, + metadata={ + "help": "The layer indexes to transform, is this argument is specified, PEFT will transform only the layers indexes that are specified inside this list. If a single integer is passed, PEFT will transform only the layer at this index." + }, + ) + layers_pattern: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": "The layer pattern name, used only if `layers_to_transform` is different to None and if the layer pattern is not in the common layers pattern. " + "This should target the `nn.ModuleList` of the model, which is often called `'layers'` or `'h'`." + }, + ) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.BOFT + self.target_modules = ( + set(self.target_modules) if isinstance(self.target_modules, list) else self.target_modules + ) + self.exclude_modules = ( + set(self.exclude_modules) if isinstance(self.exclude_modules, list) else self.exclude_modules + ) + # check for layers_to_transform and layers_pattern + if self.layers_pattern and not self.layers_to_transform: + raise ValueError("When `layers_pattern` is specified, `layers_to_transform` must also be specified. ") + if self.boft_block_size == 0 and self.boft_block_num == 0: + raise ValueError( + f"Either `boft_block_size` or `boft_block_num` must be non-zero. Currently, boft_block_size = {self.boft_block_size} and boft_block_num = {self.boft_block_num}." + ) + if not (self.boft_block_size != 0) ^ (self.boft_block_num != 0): + raise ValueError( + f"You can only specify either boft_block_size ({self.boft_block_size}) or boft_block_num ({self.boft_block_num}), but not both simultaneously, because boft_block_size x boft_block_num == in_features." + ) diff --git a/lib/python3.12/site-packages/peft/tuners/boft/fbd/__init__.py b/lib/python3.12/site-packages/peft/tuners/boft/fbd/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/peft/tuners/boft/fbd/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/boft/fbd/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e3298aac6e9dc8a96964780462a4500b455dd84 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/boft/fbd/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/boft/fbd/fbd_cuda.cpp b/lib/python3.12/site-packages/peft/tuners/boft/fbd/fbd_cuda.cpp new file mode 100644 index 0000000000000000000000000000000000000000..d63111b04083f3b1359aba0a9de9656060a4d515 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/boft/fbd/fbd_cuda.cpp @@ -0,0 +1,28 @@ +#include +#include +#include +#include + +std::vector forward_fast_block_diag_cuda( + at::Tensor input); + +std::vector forward_fast_block_diag( + at::Tensor input + ) { + return forward_fast_block_diag_cuda(input); +} + +std::vector backward_fast_block_diag_cuda( + at::Tensor grad_output, + at::Tensor input); +std::vector backward_fast_block_diag( + at::Tensor grad_output, + at::Tensor input + ) { + return backward_fast_block_diag_cuda(grad_output, input); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &forward_fast_block_diag, "FAST BLOCK DIAG (CUDA)"); + m.def("backward", &backward_fast_block_diag, "FAST BLOCK DIAG backward (CUDA)"); +} diff --git a/lib/python3.12/site-packages/peft/tuners/boft/fbd/fbd_cuda_kernel.cu b/lib/python3.12/site-packages/peft/tuners/boft/fbd/fbd_cuda_kernel.cu new file mode 100644 index 0000000000000000000000000000000000000000..9f307455349e6ea8067d47ebf8a39b98ba452d09 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/boft/fbd/fbd_cuda_kernel.cu @@ -0,0 +1,109 @@ +// Author: Yao Feng +// Date: 2023/08 +// Description: cuda kernel for fast block diag + +#include + +#include +#include +#include + +namespace{ +template +__global__ void forward_fast_block_diag_cuda_kernel( + const scalar_t* __restrict__ input, //[z, N, b, b] + scalar_t* output, //[z, Nxb, Nxb] + int z, int N, int b + ) { + + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= z*N*b*b) { + return; + } + const int zi = i/(N*b*b); + const int Ni = (i%(N*b*b))/(b*b); + const int x = ((i%(N*b*b))%(b*b))/b; + const int y = ((i%(N*b*b))%(b*b))%b; + + output[zi*N*b*N*b + (Ni*b+x)*N*b + Ni*b + y] = input[zi*N*b*b + Ni*b*b + x*b + y]; + +} + +template +__global__ void backward_fast_block_diag_cuda_kernel( + const scalar_t* __restrict__ grad_output, + scalar_t* grad_input, + int z, int N, int b + ) { + + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= z*N*b*b) { + return; + } + const int zi = i/(N*b*b); + const int Ni = (i%(N*b*b))/(b*b); + const int x = ((i%(N*b*b))%(b*b))/b; + const int y = ((i%(N*b*b))%(b*b))%b; + + grad_input[zi*N*b*b + Ni*b*b + x*b + y] = grad_output[zi*N*b*N*b + (Ni*b+x)*N*b + Ni*b + y]; + +} // namespace +} + +std::vector forward_fast_block_diag_cuda( + at::Tensor input + ){ + const auto z = input.size(0); + const auto N = input.size(1); + const auto b = input.size(2); + + // print(channel_size) + const int threads = 512; + const dim3 blocks_1 ((z*N*b*b - 1) / threads +1); + // initlaize output + auto output = at::zeros({z, N*b, N*b}, input.options()); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "forward_fast_block_diag1", ([&] { + forward_fast_block_diag_cuda_kernel<<>>( + input.data_ptr(), + output.data_ptr(), + z, N, b); + })); + + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + printf("Error in forward_fast_block_diag_cuda_kernel: %s\n", cudaGetErrorString(err)); + + return {output}; +} + +std::vector backward_fast_block_diag_cuda( + at::Tensor grad_output, + at::Tensor input + ){ + + const auto z = input.size(0); + const auto N = input.size(1); + const auto b = input.size(2); + + // print(channel_size) + const int threads = 512; + const dim3 blocks_1 ((z*N*b*b - 1) / threads +1); + + // initialize grad input + auto grad_input = at::zeros_like(input); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_output.type(), "backward_fast_block_diag", ([&] { + backward_fast_block_diag_cuda_kernel<<>>( + grad_output.data_ptr(), + grad_input.data_ptr(), + z, N, b); + })); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + printf("Error in backward_fast_block_diag_cuda_kernel: %s\n", cudaGetErrorString(err)); + + return {grad_input}; +} diff --git a/lib/python3.12/site-packages/peft/tuners/boft/layer.py b/lib/python3.12/site-packages/peft/tuners/boft/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..06f372659c4ce5c04c87ae6d18a98d1407ac2bbb --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/boft/layer.py @@ -0,0 +1,984 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +# The implementation is based on "Parameter-Efficient Orthogonal Finetuning +# via Butterfly Factorization" (https://arxiv.org/abs/2311.06243) in ICLR 2024. + +from __future__ import annotations + +import math +import os +import warnings +from contextlib import contextmanager +from typing import Any, Optional, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Function + +from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge + + +_FBD_CUDA = None + + +# this function is a 1:1 copy from accelerate +@contextmanager +def patch_environment(**kwargs): + """ + A context manager that will add each keyword argument passed to `os.environ` and remove them when exiting. + + Will convert the values in `kwargs` to strings and upper-case all the keys. + + Example: + + ```python + >>> import os + >>> from accelerate.utils import patch_environment + + >>> with patch_environment(FOO="bar"): + ... print(os.environ["FOO"]) # prints "bar" + >>> print(os.environ["FOO"]) # raises KeyError + ``` + """ + existing_vars = {} + for key, value in kwargs.items(): + key = key.upper() + if key in os.environ: + existing_vars[key] = os.environ[key] + os.environ[key] = str(value) + + yield + + for key in kwargs: + key = key.upper() + if key in existing_vars: + # restore previous value + os.environ[key] = existing_vars[key] + else: + os.environ.pop(key, None) + + +def get_fbd_cuda(): + global _FBD_CUDA + + if _FBD_CUDA is not None: + return _FBD_CUDA + + # This import initializes cuda context and should thus be local, see issue 1877 + from torch.utils.cpp_extension import load + + curr_dir = os.path.dirname(__file__) + # need ninja to build the extension + try: + with patch_environment(CC="gcc", CXX="gcc"): + fbd_cuda = load( + name="fbd_cuda", + sources=[f"{curr_dir}/fbd/fbd_cuda.cpp", f"{curr_dir}/fbd/fbd_cuda_kernel.cu"], + verbose=True, + # build_directory='/tmp/' # for debugging + ) + # extra_cuda_cflags = ['-std=c++14', '-ccbin=$$(which gcc-7)']) # cuda10.2 is not compatible with gcc9. Specify gcc 7 + except Exception as e: + warnings.warn(f"Failed to load the CUDA extension: {e}, check if ninja is available.") + warnings.warn("Setting boft_n_butterfly_factor to 1 to speed up the finetuning process.") + fbd_cuda = None + + _FBD_CUDA = fbd_cuda + return _FBD_CUDA + + +class FastBlockDiag(Function): + """ + Implements a custom autograd Function for a fast block diagonal operation using CUDA. + + This function is optimized for 4D tensors where the last two dimensions are equal, representing block diagonal + matrices for efficient computation on CUDA devices. + """ + + @staticmethod + def forward(ctx, input): + """ + The forward method for FastBlockDiag. + + Computes the block diagonal operation on the input tensor using a CUDA-optimized function. This method assumes + that the input is a 4D tensor where the last two dimensions are equal, which represent the blocks to be + diagonalized. + + Parameters: + ctx: A context object that can be used to stash information for backward computation. + input (Tensor): The input tensor of shape (N, D, H, H), where `N` is the batch size, + `D` represents one additional dimension (In BOFT, the number of BOFT blocks), and `H` is the + size of the square blocks along the last two dimensions (In BOFT, the block size). + + Returns: + Tensor: The resulting tensor after applying the block diagonal operation, + will have the shape (N, DxH, DxH). + """ + output = get_fbd_cuda().forward(input)[0] + ctx.save_for_backward(input) + return output + + @staticmethod + def backward(ctx, grad_output): + (input,) = ctx.saved_tensors + grad_input = get_fbd_cuda().backward(grad_output, input)[0] + return grad_input + + +class MultiplicativeDropoutLayer(nn.Module): + """ + Implements the multiplicative dropout layer for BOFT. + """ + + def __init__(self, p=0.0): + """ + Initializes the multiplicative dropout layer. + + Parameters: + p (float): The probability of dropping out a block. Defaults to 0.0. + """ + super().__init__() + self.p = p + + def forward(self, x): + """ + Applies multiplicative dropout to the input tensor. + + Parameters: + x (Tensor): The input tensor of shape (N, D, H, H), where `N` is the batch size, `D` represents + one additional dimension (In BOFT, the number of BOFT blocks), and `H` is the size of the square + blocks along the last two dimensions (In BOFT, the block size). + """ + if self.training: + # Ensure the last two dimensions are the same + if x.shape[-1] != x.shape[-2]: + raise ValueError("The last two dimensions of input should be the same!") + + N, D, H, _ = x.shape + + # Randomly select one from N + n_random = torch.randint(0, N, (1,)).item() + + # Create a mask with 1s for matrices to be replaced with identity and 0s otherwise + num_to_replace = int(self.p * D) + num_zeros = D - num_to_replace + + # Generate a flat tensor with desired number of 1s and 0s + mask = torch.cat([torch.ones(num_to_replace, device=x.device), torch.zeros(num_zeros, device=x.device)]) + + # Shuffle and reshape the mask + mask = mask[torch.randperm(D)].view(1, D, 1, 1) + + full_mask = torch.zeros(N, D, 1, 1, device=x.device) + full_mask[n_random] = mask + + # Use the mask to combine original matrices and identity matrices + eye_matrix = torch.eye(H, device=x.device).repeat(N, D, 1, 1) + x = (1 - full_mask) * x + full_mask * eye_matrix + return x + + +class BOFTLayer(BaseTunerLayer): + """ + Implements the BOFT layer. + """ + + # All names of layers that may contain (trainable) adapter weights + adapter_layer_names = ("boft_R", "boft_s") + # All names of other parameters that may contain adapter-related parameters + other_param_names = ("boft_block_size", "boft_block_num", "boft_dropout") + + def __init__(self, base_layer: nn.Module, **kwargs) -> None: + """ + Initializes the BOFT layer. + + Note, currently only support linear layer and convolutional layer, with further support for other layers to be + added soon. + + Parameters: + base_layer: the pretrained model layer + """ + self.base_layer = base_layer + self.boft_block_size = {} + self.boft_block_num = {} + self.boft_dropout = nn.ModuleDict({}) + self.boft_R = nn.ParameterDict({}) + self.boft_s = nn.ParameterDict({}) + # Mark the weight as unmerged + self._disable_adapters = False + self.merged_adapters = [] + self.kwargs = kwargs + + base_layer = self.get_base_layer() + + if isinstance(base_layer, nn.Linear): + in_features, out_features = base_layer.in_features, base_layer.out_features + elif isinstance(base_layer, nn.Conv2d): + in_features, out_features = base_layer.in_channels, base_layer.out_channels + else: + raise ValueError(f"Unsupported layer type {type(base_layer)}") + + self.in_features = in_features + self.out_features = out_features + + def set_scale(self, adapter, scale): + if adapter not in self.scaling: + # Ignore the case where the adapter is not in the layer + return + + warnings.warn("Scaling operation for BOFT not supported! Automatically set scale to 1.") + + def scale_layer(self, scale: float) -> None: + if scale == 1: + return + + for active_adapter in self.active_adapters: + if active_adapter not in self.boft_R.keys(): + continue + + warnings.warn("Scaling operation for BOFT not supported! Automatically set scale to 1.") + + def unscale_layer(self, scale=None) -> None: + for active_adapter in self.active_adapters: + if active_adapter not in self.boft_R.keys(): + continue + + warnings.warn("Unscaling operation for BOFT not supported! Keeping scale to 1.") + + def update_layer( + self, adapter_name, boft_block_size, boft_block_num, boft_n_butterfly_factor, boft_dropout, init_weights + ): + """ + Update the linear layer with trainable BOFT weights. Override for other layer types. + """ + # Attempt to load the CUDA extension during model initialization + if not get_fbd_cuda(): + self.fbd_cuda_available = False + # If the CUDA extension is not available, set the butterfly factor to 1 to speed up the finetuning process + boft_n_butterfly_factor = 1 + else: + self.fbd_cuda_available = True + + # to be consistent with the paper notation + boft_n_butterfly_factor = boft_n_butterfly_factor - 1 + if boft_n_butterfly_factor < 0: + raise ValueError( + f"You can only specify boft_n_butterfly_factor {boft_n_butterfly_factor + 1} to be a positive integer number." + ) + + # Initialize the MultiplicativeDropoutLayer for boft_dropout > 0.0. + if boft_dropout > 0.0: + boft_dropout_layer = MultiplicativeDropoutLayer(p=boft_dropout) + else: + boft_dropout_layer = nn.Identity() + self.boft_dropout.update(nn.ModuleDict({adapter_name: boft_dropout_layer})) + + if boft_block_size == 0 and boft_block_num != 0: + if self.in_features % boft_block_num != 0: + raise ValueError( + f"in_features ({self.in_features}) must be divisible by boft_block_num ({boft_block_num})!" + ) + + if boft_n_butterfly_factor != 0: + if boft_n_butterfly_factor > int(math.log2(boft_block_num)): + raise ValueError( + f"Invalid combination of boft_n_butterfly_factor ({boft_n_butterfly_factor + 1}) and boft_block_num ({boft_block_num})!" + ) + if boft_block_num % (2**boft_n_butterfly_factor) != 0: + raise ValueError( + f"boft_block_num ({boft_block_num}) must be a multiple of 2 raised to the power of boft_n_butterfly_factor ({boft_n_butterfly_factor + 1})!" + ) + + boft_block_size = int(self.in_features // boft_block_num) + + elif boft_block_size != 0 and boft_block_num == 0: + if self.in_features % boft_block_size != 0: + raise ValueError( + f"in_features ({self.in_features}) must be divisible by boft_block_size ({boft_block_size})!" + ) + + if boft_n_butterfly_factor != 0: + if self.in_features < (boft_block_size * (2**boft_n_butterfly_factor)): + raise ValueError( + f"Invalid combination of in_features ({self.in_features}), boft_n_butterfly_factor ({boft_n_butterfly_factor + 1}) and boft_block_size ({boft_block_size})!" + ) + if self.in_features % (boft_block_size * (2**boft_n_butterfly_factor)) != 0: + raise ValueError( + f"Invalid combination of in_features ({self.in_features}), boft_n_butterfly_factor ({boft_n_butterfly_factor + 1}) and boft_block_size ({boft_block_size})!" + ) + + boft_block_num = int(self.in_features // boft_block_size) + + else: + raise ValueError( + "Something went wrong, please report this error: https://github.com/huggingface/peft/issues" + ) + + # In OFT you can specify the number of blocks to be 1 + if boft_n_butterfly_factor != 0: + if boft_block_num % 2 != 0: + raise ValueError(f"boft_block_num ({boft_block_num}) must be an even number!") + + if boft_block_size % 2 != 0: + raise ValueError(f"boft_block_size ({boft_block_size}) must be an even number!") + + # If there is no butterfly factor, then permutation matrix P will be an identity matrix. + P = torch.empty((boft_n_butterfly_factor + 1, self.in_features, self.in_features)) + for i in range(boft_n_butterfly_factor + 1): + perm = self.block_butterfly_perm( + self.in_features, int(boft_block_num / (2 ** (i))), int(boft_block_size / 2), boft_n_butterfly_factor + ) + perm_mat = self.perm2mat(perm) + P[i] = perm_mat + + self.register_buffer("boft_P", P, persistent=False) + + self.boft_R[adapter_name] = nn.Parameter( + torch.zeros(boft_n_butterfly_factor + 1, boft_block_num, boft_block_size, boft_block_size) + ) + self.boft_s[adapter_name] = nn.Parameter(torch.ones(int(self.out_features), 1)) + + self.reset_boft_parameters(adapter_name, init_weights) + + # set the boft block size and number + self.boft_block_size[adapter_name] = boft_block_size + self.boft_block_num[adapter_name] = boft_block_num + + self._move_adapter_to_device_of_base_layer(adapter_name) + self.set_adapter(self.active_adapters) + + def reset_boft_parameters(self, adapter_name, init_weights): + """ + Reset the BOFT parameters. + """ + if init_weights is False: + nn.init.normal_(self.boft_R[adapter_name], mean=0.0, std=0.1) + nn.init.normal_(self.boft_s[adapter_name], mean=1.0, std=0.1) + return + + if adapter_name in self.boft_R.keys(): + if init_weights is True: + # initialize R to zero + nn.init.zeros_(self.boft_R[adapter_name]) + nn.init.ones_(self.boft_s[adapter_name]) + else: + raise ValueError(f"Unknown initialization {init_weights=}") + + def perm2mat(self, indices): + """ + Convert permutation indices to permutation matrix. + + Args: + indices: A list of indices representing the permutation. + """ + # Number of indices determines the size of the square matrix + n = len(indices) + + # Initialize a matrix of zeros + perm_mat = torch.zeros((n, n)) + + # Set the 1s according to the indices + for i, idx in enumerate(indices): + perm_mat[i, idx] = 1 + + return perm_mat + + def block_butterfly_perm(self, n, b, r=3, n_butterfly_factor=1): + """ + Define the permutation matrix for the block butterfly permutation. + + Args: + n: size of the permutation matrix + b: desired number of blocks after multiplying with the permutation matrix + r: base block size of the block diagonal matrix, e.g. 2x2, 3x3, 5x5 etc. + """ + + if n_butterfly_factor == 0: + return torch.arange(n) + + if b * r * 2 > n: + raise ValueError("Invalid number of blocks!") + + block_size = int(n // b) + indices = torch.arange(n) + + def sort_block(b, r): + step = b / r + initial_order = torch.arange(b) + sorted_order = torch.empty(b, dtype=torch.long) + + evens = torch.arange(0, step, 2) + odds = torch.arange(1, step, 2) + sorted_seq = torch.cat((evens, odds), dim=0) + for i, pos in enumerate(sorted_seq): + sorted_order[int(i * r) : int(i * r + r)] = initial_order[int(pos * r) : int(pos * r + r)] + return sorted_order + + sorted_order = sort_block(block_size, r) + + for i in range(0, n, block_size): + block_end = i + block_size + tmp_indices = indices[i:block_end] + indices[i:block_end] = tmp_indices[sorted_order] + return indices + + def cayley_batch(self, data): + """ + Perform the Cayley parametrization on a batch of skew-symmetric matrices. + + Args: + data: A batch of skew-symmetric matrices of shape (b, r, c). + """ + b, r, c = data.shape + # Ensure the input matrix is skew-symmetric + skew_mat = 0.5 * (data - data.transpose(1, 2)) + id_mat = torch.eye(r, device=data.device).unsqueeze(0).expand(b, r, c) + + # Perform the Cayley parametrization + Q = torch.linalg.solve(id_mat + skew_mat, id_mat - skew_mat, left=False) + + return Q + + +class Linear(nn.Module, BOFTLayer): + """ + BOFT implemented in a dense layer. + """ + + def __init__( + self, + base_layer, + adapter_name: str, + boft_block_size: int = 8, + boft_block_num: int = 0, + boft_n_butterfly_factor: int = 0, + boft_dropout: float = 0.1, + fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) + init_weights: Union[bool, str] = True, + is_target_conv_1d_layer: bool = False, + **kwargs, + ) -> None: + super().__init__() + BOFTLayer.__init__(self, base_layer, **kwargs) + self.fan_in_fan_out = fan_in_fan_out + + self._active_adapter = adapter_name + + self.update_layer( + adapter_name, boft_block_size, boft_block_num, boft_n_butterfly_factor, boft_dropout, init_weights + ) + self.is_target_conv_1d_layer = is_target_conv_1d_layer + + def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None: + """ + Merge the active adapter weights into the base weights + + Args: + safe_merge (`bool`, *optional*): + If True, the merge operation will be performed in a copy of the original weights and check for NaNs + before merging the weights. This is useful if you want to check if the merge operation will produce + NaNs. Defaults to `False`. + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + """ + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + # no adapter to merge + return + + for active_adapter in adapter_names: + if active_adapter in self.boft_R.keys(): + base_layer = self.get_base_layer() + if safe_merge: + # Note that safe_merge will be slower than the normal merge + # because of the copy operation. + orig_weight = base_layer.weight.data.clone() + butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter) + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = torch.mm(butterfly_oft_mat, orig_weight) + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = orig_weight * boft_s + + if not torch.isfinite(orig_weight).all(): + raise ValueError( + f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" + ) + + self.base_layer.weight.data = orig_weight.contiguous() + else: + butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter) + orig_weight = base_layer.weight.data.clone() + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = torch.mm(butterfly_oft_mat, orig_weight) + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = orig_weight * boft_s + + self.base_layer.weight.data = orig_weight.contiguous() + + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + """ + This method unmerges all merged adapter layers from the base weights. + """ + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter in self.boft_R.keys(): + butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter) + + orig_weight = self.get_base_layer().weight.data.clone() + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = torch.mm(butterfly_oft_mat.t(), orig_weight) + orig_weight = torch.transpose(orig_weight, 0, 1) + + self.get_base_layer().weight.data = orig_weight * (1 / boft_s) + + def get_delta_weight(self, adapter) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute the delta weight for the given adapter. + + Args: + adapter (str): + The name of the adapter for which the delta weight should be computed. + """ + boft_R = self.boft_R[adapter] + boft_s = self.boft_s[adapter] + + N, D, H, _ = boft_R.shape + boft_R = boft_R.view(N * D, H, H) + orth_rotate_butterfly = self.cayley_batch(boft_R) + orth_rotate_butterfly = orth_rotate_butterfly.view(N, D, H, H) + if self.fbd_cuda_available: + block_diagonal_butterfly = FastBlockDiag.apply(orth_rotate_butterfly) + else: + orth_rotate_butterfly = orth_rotate_butterfly.squeeze(0) + block_diagonal_butterfly = torch.block_diag(*torch.unbind(orth_rotate_butterfly)) + block_diagonal_butterfly = block_diagonal_butterfly.unsqueeze(0) + + boft_P = self.boft_P.to(block_diagonal_butterfly.device) + butterfly_oft_mat_batch = torch.bmm(block_diagonal_butterfly, boft_P.permute(0, 2, 1)) + butterfly_oft_mat_batch = torch.bmm(boft_P, butterfly_oft_mat_batch) + butterfly_oft_mat = butterfly_oft_mat_batch[0] + + for i in range(1, butterfly_oft_mat_batch.shape[0]): + butterfly_oft_mat = butterfly_oft_mat_batch[i] @ butterfly_oft_mat + + return butterfly_oft_mat, boft_s + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + previous_dtype = x.dtype + + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + boft_rotation = torch.eye(self.in_features, device=x.device, dtype=previous_dtype) + boft_scale = torch.ones((int(self.out_features), 1), device=x.device, dtype=previous_dtype) + + for active_adapter in self.active_adapters: + if active_adapter not in self.boft_R.keys(): + continue + boft_R = self.boft_R[active_adapter] + boft_s = self.boft_s[active_adapter] + dropout = self.boft_dropout[active_adapter] + + N, D, H, _ = boft_R.shape + boft_R = boft_R.view(N * D, H, H) + orth_rotate_butterfly = self.cayley_batch(boft_R) + orth_rotate_butterfly = orth_rotate_butterfly.view(N, D, H, H) + orth_rotate_butterfly = dropout(orth_rotate_butterfly) + if self.fbd_cuda_available: + block_diagonal_butterfly = FastBlockDiag.apply(orth_rotate_butterfly) + else: + orth_rotate_butterfly = orth_rotate_butterfly.squeeze(0) + block_diagonal_butterfly = torch.block_diag(*torch.unbind(orth_rotate_butterfly)) + block_diagonal_butterfly = block_diagonal_butterfly.unsqueeze(0) + + # The BOFT author's cayley_batch, dropout and FastBlockDiag ONLY return fp32 outputs. + boft_P = self.boft_P.to(x) + block_diagonal_butterfly = block_diagonal_butterfly.to(x) + butterfly_oft_mat_batch = torch.bmm(block_diagonal_butterfly, boft_P.permute(0, 2, 1)) + butterfly_oft_mat_batch = torch.bmm(boft_P, butterfly_oft_mat_batch) + butterfly_oft_mat = butterfly_oft_mat_batch[0] + + for i in range(1, butterfly_oft_mat_batch.shape[0]): + butterfly_oft_mat = butterfly_oft_mat_batch[i] @ butterfly_oft_mat + + boft_rotation = butterfly_oft_mat @ boft_rotation + boft_scale = boft_s * boft_scale + + x = x.to(self.get_base_layer().weight.data.dtype) + + orig_weight = self.get_base_layer().weight.data + orig_weight = torch.transpose(orig_weight, 0, 1) + boft_rotation = boft_rotation.to(previous_dtype) + orig_weight = orig_weight.to(previous_dtype) + rotated_weight = torch.mm(boft_rotation, orig_weight) + rotated_weight = torch.transpose(rotated_weight, 0, 1) + + scaled_rotated_weight = rotated_weight * boft_scale + + scaled_rotated_weight = scaled_rotated_weight.to(previous_dtype) + if self.base_layer.bias is not None: + self.base_layer.bias = self.base_layer.bias.to(previous_dtype) + result = F.linear(input=x, weight=scaled_rotated_weight, bias=self.base_layer.bias) + + result = result.to(previous_dtype) + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "boft." + rep + + +class Conv2d(nn.Module, BOFTLayer): + """ + BOFT implemented in a Conv2d layer. + """ + + def __init__( + self, + base_layer: nn.Module, + adapter_name: str, + boft_block_size: int = 8, + boft_block_num: int = 0, + boft_n_butterfly_factor: int = 0, + boft_dropout: float = 0.1, + init_weights: Union[bool, str] = True, + **kwargs, + ) -> None: + super().__init__() + BOFTLayer.__init__(self, base_layer) + + self._active_adapter = adapter_name + self.update_layer( + adapter_name, boft_block_size, boft_block_num, boft_n_butterfly_factor, boft_dropout, init_weights + ) + + def update_layer( + self, adapter_name, boft_block_size, boft_block_num, boft_n_butterfly_factor, boft_dropout, init_weights + ): + """ + Update the conv2d layer with trainable BOFT weights. + """ + + # Attempt to load the CUDA extension during model initialization + if not get_fbd_cuda(): + self.fbd_cuda_available = False + # If the CUDA extension is not available, set the butterfly factor to 1 to speed up the finetuning process + boft_n_butterfly_factor = 1 + else: + self.fbd_cuda_available = True + + # to be consistent with the paper notation + boft_n_butterfly_factor = boft_n_butterfly_factor - 1 + if boft_n_butterfly_factor < 0: + raise ValueError( + f"You can only specify boft_n_butterfly_factor {boft_n_butterfly_factor + 1} to be a positive integer number." + ) + + # Initialize the MultiplicativeDropoutLayer for boft_dropout > 0.0. + if boft_dropout > 0.0: + boft_dropout_layer = MultiplicativeDropoutLayer(p=boft_dropout) + else: + boft_dropout_layer = nn.Identity() + self.boft_dropout.update(nn.ModuleDict({adapter_name: boft_dropout_layer})) + + # layer information from the base layer + base_layer = self.get_base_layer() + conv_filter_dim = self.in_features * base_layer.kernel_size[0] * base_layer.kernel_size[0] + + # Initialize the BOFT parameters. + if boft_block_size == 0 and boft_block_num != 0: + if conv_filter_dim % boft_block_num != 0: + raise ValueError( + f"Convolutional kernel dimension ({conv_filter_dim}) must be divisible by boft_block_num ({boft_block_num})!" + ) + + if boft_n_butterfly_factor != 0: + if boft_n_butterfly_factor > int(math.log2(boft_block_num)): + raise ValueError( + f"Invalid combination of boft_n_butterfly_factor ({boft_n_butterfly_factor + 1}) and boft_block_num ({boft_block_num})!" + ) + if boft_block_num % (2**boft_n_butterfly_factor) != 0: + raise ValueError( + f"boft_block_num ({boft_block_num}) must be a multiple of 2 raised to the power of boft_n_butterfly_factor ({boft_n_butterfly_factor + 1})!" + ) + + boft_block_size = int(conv_filter_dim // boft_block_num) + + elif boft_block_size != 0 and boft_block_num == 0: + if conv_filter_dim % boft_block_size != 0: + raise ValueError( + f"Convolutional kernel dimension ({conv_filter_dim}) must be divisible by boft_block_size ({boft_block_size})!" + ) + + if boft_n_butterfly_factor != 0: + if conv_filter_dim < (boft_block_size * (2**boft_n_butterfly_factor)): + raise ValueError( + f"Invalid combination of convolutional kernel dimension ({conv_filter_dim}), boft_n_butterfly_factor ({boft_n_butterfly_factor + 1}) and boft_block_size ({boft_block_size})!" + ) + if conv_filter_dim % (boft_block_size * (2**boft_n_butterfly_factor)) != 0: + raise ValueError( + f"Invalid combination of convolutional kernel dimension ({conv_filter_dim}), boft_n_butterfly_factor ({boft_n_butterfly_factor + 1}) and boft_block_size ({boft_block_size})!" + ) + + boft_block_num = int(conv_filter_dim // boft_block_size) + + else: + raise ValueError( + "Something went wrong, please report this error: https://github.com/huggingface/peft/issues" + ) + + # In OFT you can specify the number of blocks to be 1 + if boft_n_butterfly_factor != 0: + if boft_block_num % 2 != 0: + raise ValueError(f"boft_block_num ({boft_block_num}) must be an even number!") + + if boft_block_size % 2 != 0: + raise ValueError(f"boft_block_size ({boft_block_size}) must be an even number!") + + # If there is no butterfly factor, then permutation matrix P will be an identity matrix. + P = torch.empty((boft_n_butterfly_factor + 1, conv_filter_dim, conv_filter_dim)) + for i in range(boft_n_butterfly_factor + 1): + perm = self.block_butterfly_perm( + conv_filter_dim, int(boft_block_num / (2 ** (i))), int(boft_block_size / 2), boft_n_butterfly_factor + ) + perm_mat = self.perm2mat(perm) + P[i] = perm_mat + + self.register_buffer("boft_P", P, persistent=False) + + self.boft_R[adapter_name] = nn.Parameter( + torch.zeros(boft_n_butterfly_factor + 1, boft_block_num, boft_block_size, boft_block_size) + ) + self.boft_s[adapter_name] = nn.Parameter(torch.ones(1, int(self.out_features))) + + self.reset_boft_parameters(adapter_name, init_weights) + + # set the boft block size and number + self.boft_block_size[adapter_name] = boft_block_size + self.boft_block_num[adapter_name] = boft_block_num + + self._move_adapter_to_device_of_base_layer(adapter_name) + self.set_adapter(self.active_adapters) + + def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None: + """ + Merge the active adapter weights into the base weights + + Args: + safe_merge (`bool`, *optional*): + If True, the merge operation will be performed in a copy of the original weights and check for NaNs + before merging the weights. This is useful if you want to check if the merge operation will produce + NaNs. Defaults to `False`. + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + """ + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + # no adapter to merge + return + + for active_adapter in adapter_names: + if active_adapter in self.boft_R.keys(): + base_layer = self.get_base_layer() + if safe_merge: + # Note that safe_merge will be slower than the normal merge + # because of the copy operation. + orig_weight = base_layer.weight.data.clone() + butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter) + + orig_weight = orig_weight.view( + self.out_features, self.in_features * base_layer.kernel_size[0] * base_layer.kernel_size[0] + ) + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = torch.mm(butterfly_oft_mat, orig_weight) + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = orig_weight * boft_s + orig_weight = orig_weight.view( + self.out_features, self.in_features, base_layer.kernel_size[0], base_layer.kernel_size[0] + ) + + self.base_layer.weight.data = orig_weight.contiguous() + else: + butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter) + + orig_weight = base_layer.weight.data.clone() + orig_weight = orig_weight.view( + self.out_features, self.in_features * base_layer.kernel_size[0] * base_layer.kernel_size[0] + ) + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = torch.mm(butterfly_oft_mat, orig_weight) + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = orig_weight * boft_s + orig_weight = orig_weight.view( + self.out_features, self.in_features, base_layer.kernel_size[0], base_layer.kernel_size[0] + ) + + self.base_layer.weight.data = orig_weight.contiguous() + + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + """ + This method unmerges all merged adapter layers from the base weights. + """ + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter in self.boft_R.keys(): + butterfly_oft_mat, boft_s = self.get_delta_weight(active_adapter) + + orig_weight = self.get_base_layer().weight.data.clone() + orig_weight = orig_weight.view( + self.out_features, + self.in_features * self.get_base_layer().kernel_size[0] * self.get_base_layer().kernel_size[0], + ) + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = torch.mm(butterfly_oft_mat.t(), orig_weight) + orig_weight = torch.transpose(orig_weight, 0, 1) + orig_weight = orig_weight * (1 / boft_s) + orig_weight = orig_weight.view( + self.out_features, + self.in_features, + self.get_base_layer().kernel_size[0], + self.get_base_layer().kernel_size[0], + ) + + self.get_base_layer().weight.data = orig_weight + + def get_delta_weight(self, adapter) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute the delta weight for the given adapter. + + Args: + adapter (str): + The name of the adapter for which the delta weight should be computed. + """ + + boft_R = self.boft_R[adapter] + boft_s = self.boft_s[adapter].transpose(0, 1) + + N, D, H, _ = boft_R.shape + boft_R = boft_R.view(N * D, H, H) + orth_rotate_butterfly = self.cayley_batch(boft_R) + orth_rotate_butterfly = orth_rotate_butterfly.view(N, D, H, H) + if self.fbd_cuda_available: + block_diagonal_butterfly = FastBlockDiag.apply(orth_rotate_butterfly) + else: + orth_rotate_butterfly = orth_rotate_butterfly.squeeze(0) + block_diagonal_butterfly = torch.block_diag(*torch.unbind(orth_rotate_butterfly)) + block_diagonal_butterfly = block_diagonal_butterfly.unsqueeze(0) + + boft_P = self.boft_P.to(block_diagonal_butterfly.device) + butterfly_oft_mat_batch = torch.bmm(block_diagonal_butterfly, boft_P.permute(0, 2, 1)) + butterfly_oft_mat_batch = torch.bmm(boft_P, butterfly_oft_mat_batch) + butterfly_oft_mat = butterfly_oft_mat_batch[0] + + for i in range(1, butterfly_oft_mat_batch.shape[0]): + butterfly_oft_mat = butterfly_oft_mat_batch[i] @ butterfly_oft_mat + + return butterfly_oft_mat, boft_s + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + previous_dtype = x.dtype + + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + boft_rotation = torch.eye( + self.in_features * self.base_layer.kernel_size[0] * self.base_layer.kernel_size[0], + device=x.device, + dtype=x.dtype, + ) + boft_scale = torch.ones((int(self.out_features), 1), device=x.device, dtype=x.dtype) + + for active_adapter in self.active_adapters: + if active_adapter not in self.boft_R.keys(): + continue + boft_R = self.boft_R[active_adapter] + boft_s = self.boft_s[active_adapter].transpose(0, 1) + dropout = self.boft_dropout[active_adapter] + + N, D, H, _ = boft_R.shape + boft_R = boft_R.view(N * D, H, H) + orth_rotate_butterfly = self.cayley_batch(boft_R) + orth_rotate_butterfly = orth_rotate_butterfly.view(N, D, H, H) + orth_rotate_butterfly = dropout(orth_rotate_butterfly) + if self.fbd_cuda_available: + block_diagonal_butterfly = FastBlockDiag.apply(orth_rotate_butterfly) + else: + orth_rotate_butterfly = orth_rotate_butterfly.squeeze(0) + block_diagonal_butterfly = torch.block_diag(*torch.unbind(orth_rotate_butterfly)) + block_diagonal_butterfly = block_diagonal_butterfly.unsqueeze(0) + + boft_P = self.boft_P.to(x) + block_diagonal_butterfly = block_diagonal_butterfly.to(x) + butterfly_oft_mat_batch = torch.bmm(block_diagonal_butterfly, boft_P.permute(0, 2, 1)) + butterfly_oft_mat_batch = torch.bmm(boft_P, butterfly_oft_mat_batch) + butterfly_oft_mat = butterfly_oft_mat_batch[0] + + for i in range(1, butterfly_oft_mat_batch.shape[0]): + butterfly_oft_mat = butterfly_oft_mat_batch[i] @ butterfly_oft_mat + + boft_rotation = butterfly_oft_mat @ boft_rotation + boft_scale = boft_s * boft_scale + + x = x.to(self.base_layer.weight.data.dtype) + + orig_weight = self.base_layer.weight.data + orig_weight = orig_weight.view( + self.out_features, + self.in_features * self.base_layer.kernel_size[0] * self.base_layer.kernel_size[0], + ) + orig_weight = torch.transpose(orig_weight, 0, 1) + rotated_weight = torch.mm(boft_rotation, orig_weight) + rotated_weight = torch.transpose(rotated_weight, 0, 1) + + scaled_rotated_weight = rotated_weight * boft_scale + + scaled_rotated_weight = scaled_rotated_weight.view( + self.out_features, self.in_features, self.base_layer.kernel_size[0], self.base_layer.kernel_size[0] + ) + result = F.conv2d( + input=x, + weight=scaled_rotated_weight, + bias=self.base_layer.bias, + padding=self.base_layer.padding[0], + stride=self.base_layer.stride[0], + ) + + result = result.to(previous_dtype) + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "boft." + rep diff --git a/lib/python3.12/site-packages/peft/tuners/boft/model.py b/lib/python3.12/site-packages/peft/tuners/boft/model.py new file mode 100644 index 0000000000000000000000000000000000000000..4a4bb8158fbf8f30508548d7c0d66161b57ded33 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/boft/model.py @@ -0,0 +1,353 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +# The implementation is based on "Parameter-Efficient Orthogonal Finetuning +# via Butterfly Factorization" (https://arxiv.org/abs/2311.06243) in ICLR 2024. + +import warnings +from dataclasses import asdict +from enum import Enum +from typing import List, Optional + +import torch +from torch import nn +from tqdm import tqdm + +from peft.tuners.tuners_utils import ( + BaseTuner, + BaseTunerLayer, + check_target_module_exists, + onload_layer, +) +from peft.utils import ( + TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, + ModulesToSaveWrapper, + _get_submodules, +) + +from .config import BOFTConfig +from .layer import BOFTLayer, Conv2d, Linear + + +class BOFTModel(BaseTuner): + """ + Creates BOFT and OFT model from a pretrained transformers model. Paper: https://arxiv.org/abs/2311.06243 + https://arxiv.org/abs/2306.07280 + + Args: + model ([`transformers.PreTrainedModel`]): The model to be adapted. + config ([`BOFTConfig`]): The configuration of the BOFT model. + adapter_name (`str`): The name of the adapter, defaults to `"default"`. + low_cpu_mem_usage (`bool`, `optional`, defaults to `False`): + Create empty adapter weights on meta device. Useful to speed up the loading process. + + Returns: + `torch.nn.Module`: The BOFT model. + + Example:: + + >>> import transformers >>> from transformers import AutoModelForSeq2SeqLM, BOFTConfig >>> from peft import + BOFTConfig, get_peft_model + + >>> config = BOFTConfig( ... boft_block_size=8, ... boft_n_butterfly_factor=1, ... target_modules=["query", + "value", "key", "output.dense", "mlp.fc1", "mlp.fc2"], ... boft_dropout=0.1, ... bias="boft_only", ... + modules_to_save=["classifier"], ... ) + + >>> model = transformers.Dinov2ForImageClassification.from_pretrained( ... "facebook/dinov2-large", ... + num_labels=100, ... ) >>> boft_model = get_peft_model(model, config) + + **Attributes**: + - **model** ([`transformers.PreTrainedModel`]) -- The model to be adapted. + - **peft_config** ([`BOFTConfig`]): The configuration of the BOFT model. + """ + + prefix: str = "boft_" + + def __init__(self, model, config, adapter_name, low_cpu_mem_usage: bool = False) -> None: + super().__init__(model, config, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage) + + def _check_new_adapter_config(self, config: BOFTConfig) -> None: + """ + A helper method to check the config when a new adapter is being added. + + Raise a ValueError if there is something wrong with the config or if it conflicts with existing adapters. + + """ + # TODO: there should be a check if any of the existing adapters actually has bias != "none", or else the check + # does not fully correspond to the error message. + if (len(self.peft_config) > 1) and (config.bias != "none"): + raise ValueError( + f"{self.__class__.__name__} supports only 1 adapter with bias. When using multiple adapters, " + "set bias to 'none' for all adapters." + ) + + @staticmethod + def _check_target_module_exists(boft_config, key): + return check_target_module_exists(boft_config, key) + + def _create_and_replace( + self, + boft_config, + adapter_name, + target, + target_name, + parent, + current_key, + **optional_kwargs, + ): + if current_key is None: + raise ValueError("Current Key shouldn't be `None`") + + bias = hasattr(target, "bias") and target.bias is not None + kwargs = { + "boft_block_size": boft_config.boft_block_size, + "boft_block_num": boft_config.boft_block_num, + "boft_n_butterfly_factor": boft_config.boft_n_butterfly_factor, + "boft_dropout": boft_config.boft_dropout, + "fan_in_fan_out": boft_config.fan_in_fan_out, + "init_weights": boft_config.init_weights, + } + kwargs["bias"] = bias + + # If it is not a BOFTLayer, create a new module, else update it with new adapters + if not isinstance(target, BOFTLayer): + new_module = self._create_new_module(boft_config, adapter_name, target, **kwargs) + if adapter_name not in self.active_adapters: + # adding an additional adapter: it is not automatically trainable + new_module.requires_grad_(False) + self._replace_module(parent, target_name, new_module, target) + else: + target.update_layer( + adapter_name, + boft_block_size=boft_config.boft_block_size, + boft_block_num=boft_config.boft_block_num, + boft_n_butterfly_factor=boft_config.boft_n_butterfly_factor, + boft_dropout=boft_config.boft_dropout, + init_weights=boft_config.init_weights, + ) + + def _replace_module(self, parent, child_name, new_module, child): + setattr(parent, child_name, new_module) + # It's not necessary to set requires_grad here, as that is handled by + # _mark_only_adapters_as_trainable + + # child layer wraps the original module, unpack it + if hasattr(child, "base_layer"): + child = child.base_layer + + if not hasattr(new_module, "base_layer"): + new_module.weight = child.weight + if hasattr(child, "bias"): + new_module.bias = child.bias + + if getattr(child, "state", None) is not None: + if hasattr(new_module, "base_layer"): + new_module.base_layer.state = child.state + else: + new_module.state = child.state + new_module.to(child.weight.device) + + meta = torch.device("meta") + # dispatch to correct device + for name, module in new_module.named_modules(): + if self.prefix in name: + if not any(p.device == meta for p in module.parameters()): + module.to(child.weight.device) + + def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None: + for n, p in model.named_parameters(): + if self.prefix not in n: + p.requires_grad = False + + for active_adapter in self.active_adapters: + bias = self.peft_config[active_adapter].bias + if bias == "none": + continue + + if bias == "all": + for n, p in model.named_parameters(): + if "bias" in n: + p.requires_grad = True + elif bias == "boft_only": + for name, m in model.named_modules(): + if isinstance(m, BOFTLayer) and hasattr(m, "bias") and m.bias is not None: + m.bias.requires_grad = True + else: + raise NotImplementedError(f"Requested bias: {bias}, is not implemented.") + + @staticmethod + def _create_new_module(boft_config, adapter_name, target, **kwargs): + if isinstance(target, BaseTunerLayer): + target_base_layer = target.get_base_layer() + else: + target_base_layer = target + + if isinstance(target_base_layer, torch.nn.Linear): + if kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. " + "Setting fan_in_fan_out to False." + ) + kwargs["fan_in_fan_out"] = boft_config.fan_in_fan_out = False + new_module = Linear(target, adapter_name, **kwargs) + elif isinstance(target_base_layer, torch.nn.Conv2d): + new_module = Conv2d(target, adapter_name, **kwargs) + else: + raise ValueError( + f"Target module {target} is not supported. " + "Currently, only `torch.nn.Linear` and `torch.nn.Conv2d` are supported." + ) + + return new_module + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + if name == "model": # see #1892: prevent infinite recursion if class is not initialized + raise + return getattr(self.model, name) + + def get_peft_config_as_dict(self, inference: bool = False): + config_dict = {} + for key, value in self.peft_config.items(): + config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} + if inference: + config["inference_mode"] = True + config_dict[key] = config + return config + + def _set_adapter_layers(self, enabled=True): + for module in self.model.modules(): + if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)): + module.enable_adapters(enabled) + + def enable_adapter_layers(self): + self._set_adapter_layers(enabled=True) + + def disable_adapter_layers(self): + for active_adapter in self.active_adapters: + val = self.peft_config[active_adapter].bias + if val != "none": + msg = ( + f"Careful, disabling adapter layers with bias configured to be '{val}' does not produce the same " + "output as the the base model would without adaption." + ) + warnings.warn(msg) + self._set_adapter_layers(enabled=False) + + def set_adapter(self, adapter_name): + for module in self.model.modules(): + if isinstance(module, BOFTLayer): + if module.merged: + warnings.warn("Adapter cannot be set when the model is merged. Unmerging the model first.") + module.unmerge() + module.set_adapter(adapter_name) + self.active_adapter = adapter_name + + @staticmethod + def _prepare_adapter_config(peft_config, model_config): + if peft_config.target_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING: + raise ValueError("Please specify `target_modules` in `peft_config`") + peft_config.target_modules = set( + TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]] + ) + return peft_config + + def _unload_and_optionally_merge( + self, + merge=True, + progressbar: bool = False, + safe_merge: bool = False, + adapter_names: Optional[List[str]] = None, + ): + if merge: + self._check_merge_allowed() + + key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] + desc = "Unloading " + ("and merging " if merge else "") + "model" + for key in tqdm(key_list, disable=not progressbar, desc=desc): + try: + parent, target, target_name = _get_submodules(self.model, key) + except AttributeError: + continue + with onload_layer(target): + if hasattr(target, "base_layer"): + if merge: + target.merge(safe_merge=safe_merge, adapter_names=adapter_names) + self._replace_module(parent, target_name, target.get_base_layer(), target) + elif isinstance(target, ModulesToSaveWrapper): + # save any additional trainable modules part of `modules_to_save` + new_module = target.modules_to_save[target.active_adapter] + if hasattr(new_module, "base_layer"): + # check if the module is itself a tuner layer + if merge: + new_module.merge(safe_merge=safe_merge, adapter_names=adapter_names) + new_module = new_module.get_base_layer() + setattr(parent, target_name, new_module) + + return self.model + + def delete_adapter(self, adapter_name: str) -> None: + """ + Deletes an existing adapter. + + Args: + adapter_name (str): Name of the adapter to be deleted. + """ + if adapter_name not in list(self.peft_config.keys()): + raise ValueError(f"Adapter {adapter_name} does not exist") + del self.peft_config[adapter_name] + + key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] + new_adapter = None + for key in key_list: + _, target, _ = _get_submodules(self.model, key) + if isinstance(target, BOFTLayer): + target.delete_adapter(adapter_name) + if new_adapter is None: + new_adapter = target.active_adapters[:] + + self.active_adapter = new_adapter or [] + + def merge_and_unload( + self, progressbar: bool = False, safe_merge: bool = False, adapter_names: Optional[List[str]] = None + ) -> torch.nn.Module: + r""" + This method merges the BOFT layers into the base model. This is needed if someone wants to use the base model + as a standalone model. + + Args: + progressbar (`bool`): + whether to show a progressbar indicating the unload and merge process + safe_merge (`bool`): + whether to activate the safe merging check to check if there is any potential Nan in the adapter + weights + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + + """ + return self._unload_and_optionally_merge( + progressbar=progressbar, safe_merge=safe_merge, adapter_names=adapter_names + ) + + def unload(self) -> torch.nn.Module: + """ + Gets back the base model by removing all the boft modules without merging. This gives back the original base + model. + """ + return self._unload_and_optionally_merge(merge=False) diff --git a/lib/python3.12/site-packages/peft/tuners/cpt/__init__.py b/lib/python3.12/site-packages/peft/tuners/cpt/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fcd4de85989168ad6834712679c805cba39c1df1 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/cpt/__init__.py @@ -0,0 +1,24 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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 peft.utils import register_peft_method + +from .config import CPTConfig +from .model import CPTEmbedding + + +__all__ = ["CPTConfig", "CPTEmbedding"] + +register_peft_method(name="cpt", config_cls=CPTConfig, model_cls=CPTEmbedding) diff --git a/lib/python3.12/site-packages/peft/tuners/cpt/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/cpt/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..233f9dc625053181040eb4a227ff58dc982bcc16 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/cpt/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/cpt/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/cpt/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b2fa274a78ba21c08d29dcbb358804a363e1704 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/cpt/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/cpt/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/cpt/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50d4484fc78492e8fe7c5a70949a304c3f714811 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/cpt/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/cpt/config.py b/lib/python3.12/site-packages/peft/tuners/cpt/config.py new file mode 100644 index 0000000000000000000000000000000000000000..971c30983f20ebabe24f2fe88a24df8f662d6a95 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/cpt/config.py @@ -0,0 +1,98 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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 dataclasses import dataclass, field +from typing import Literal, Optional + +from peft.config import PromptLearningConfig +from peft.utils import PeftType + + +@dataclass +class CPTConfig(PromptLearningConfig): + """ + CPT Configuration class extending PeftConfig for Context-aware Prompt Tuning (CPT). + + This class introduces additional parameters required for CPT, such as: + - Token type masks + - Prompt tuning initialization + - Loss weighting + - Projection settings + + For more details, see the paper: https://arxiv.org/abs/2410.17222 + """ + + # Token-related configurations + cpt_token_ids: Optional[list[int]] = field( + default=None, metadata={"help": "Tensor of token IDs used for CPT prompts."} + ) + cpt_mask: Optional[list[int]] = field(default=None, metadata={"help": "Tensor mask applied to CPT tokens."}) + cpt_tokens_type_mask: Optional[list[int]] = field( + default=None, metadata={"help": "Mask indicating the type of each CPT token."} + ) + + # Loss-related configurations + opt_weighted_loss_type: Optional[Literal["none", "decay"]] = field( + default="none", metadata={"help": "Type of weighted loss: 'none' or 'decay'."} + ) + opt_loss_decay_factor: Optional[float] = field( + default=1.0, metadata={"help": "Factor for exponential decay in loss weighting."} + ) + + # Projection-related configurations + opt_projection_epsilon: Optional[float] = field( + default=0.1, metadata={"help": "Epsilon value for input projection."} + ) + opt_projection_format_epsilon: Optional[float] = field( + default=0.1, metadata={"help": "Epsilon value for format projection."} + ) + + # Tokenizer configuration + tokenizer_name_or_path: Optional[str] = field( + default=None, + metadata={ + "help": "The tokenizer to use for prompt tuning initialization. Only used if prompt_tuning_init is `TEXT`" + }, + ) + # Neet to define CPT-specific static attributes + is_prompt_learning = True # Indicates that CPT is a prompt-learning method. + + def __post_init__(self): + """ + Post-initialization hook to set additional attributes after the config is initialized. + """ + # CPT-specific static attributes + self.is_prompt_learning = True # Indicates that CPT is a prompt-learning method. + self.num_layers = None # Number of layers (optional, not always required). + self.token_dim = None # Dimension of token embeddings. + self.num_attention_heads = None # Number of attention heads (if applicable). + self.num_transformer_submodules = 1 # Number of transformer submodules used. + self.peft_type = PeftType.CPT # Specifies that the PEFT type is CPT. + self.task_type = "CAUSAL_LM" # Ensures task type is causal language modeling. + + if self.cpt_token_ids is None: + self.cpt_token_ids = [0] + + self.num_virtual_tokens = len(self.cpt_token_ids) + + if self.cpt_mask is None: + self.cpt_mask = [1 for _ in self.cpt_token_ids] + + if self.cpt_tokens_type_mask is None: + self.cpt_tokens_type_mask = [1 for _ in self.cpt_token_ids] + + if not ( + len(self.cpt_token_ids) == len(self.cpt_mask) == len(self.cpt_tokens_type_mask) == self.num_virtual_tokens + ): + raise ValueError("cpt_token_ids, cpt_mask and cpt_tokens_type_mask must have the same length.") diff --git a/lib/python3.12/site-packages/peft/tuners/cpt/model.py b/lib/python3.12/site-packages/peft/tuners/cpt/model.py new file mode 100644 index 0000000000000000000000000000000000000000..934a3b7928c125bd433441f36cf4b26de15b535f --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/cpt/model.py @@ -0,0 +1,200 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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. + +import copy + +import torch +from torch.nn import CrossEntropyLoss + +from peft.utils.integrations import gather_params_ctx + + +class CPTEmbedding(torch.nn.Module): + """ + CPTEmbedding is a custom embedding layer designed for Context-aware Prompt Tuning (CPT) in PEFT. It initializes + embeddings, applies prompt-specific projections, and computes loss using label masks. + """ + + def __init__(self, config, word_embeddings): + """ + Initializes the CPTEmbedding module. + + Args: + config (Namespace): + Configuration object containing model hyperparameters and CPT-specific settings. + word_embeddings (torch.nn.Embedding): + The base word embedding layer used to initialize CPT embeddings. + """ + super().__init__() + self.config = copy.deepcopy(config) + num_virtual_tokens = config.num_virtual_tokens + + # Initialize embeddings with virtual token dimensions + self.embedding = torch.nn.Embedding(num_virtual_tokens, config.token_dim) + + # Initialize embeddings using text-based prompt tuning, if configured + if not config.inference_mode: + assert config.num_virtual_tokens == len(config.cpt_token_ids) + + init_token_ids = torch.LongTensor(config.cpt_token_ids).to(word_embeddings.weight.device) + with gather_params_ctx(word_embeddings.parameters()): + word_embedding_weights = word_embeddings(init_token_ids).detach().clone() + word_embedding_weights = word_embedding_weights.to(torch.float32) + self.embedding.weight = torch.nn.Parameter(word_embedding_weights) + + # Initialize delta embedding with zero weights + self.delta_embedding = torch.nn.Embedding(num_virtual_tokens, config.token_dim) + self.delta_embedding.weight.data = torch.zeros_like(self.delta_embedding.weight).to(torch.float32) + + # Apply hook for backward gradient updates + self.set_updated_tokens() + + def forward(self, indices): + """ + Computes the prompt embeddings and applies delta adjustments. + + Args: + indices (torch.Tensor): + Indices of the tokens to be embedded. + + Returns: + torch.Tensor: + Sum of prompt embeddings and delta embeddings. + """ + with torch.no_grad(): + prompt_embeddings = self.embedding(indices) + + self.delta_embedding.weight.data = self.get_projection() # Apply epsilon-based projection + + delta_prompt_embeddings = self.delta_embedding(indices) + + return prompt_embeddings + delta_prompt_embeddings + + def set_updated_tokens(self): + """ + Sets up a backward hook to selectively update token gradients based on the CPT token type mask. + """ + tensor_ICL_mask = torch.Tensor(self.config.cpt_tokens_type_mask).long() + mask_input_template = torch.remainder(tensor_ICL_mask, 4) == 1 + mask_input = torch.remainder(tensor_ICL_mask, 4) == 2 + mask_output_template = torch.remainder(tensor_ICL_mask, 4) == 3 + mask = mask_input_template | mask_input | mask_output_template + mask = mask.view(-1, 1) + + def backward_hook(grad): + grad = grad * mask.to(grad.device) # Apply mask to gradients + return grad + + self.delta_embedding.weight.register_hook(backward_hook) + + def get_epsilon(self): + cpt_tokens_type_mask = self.config.cpt_tokens_type_mask + + MIN_VALUE = 1e-10 + + # Calculate normalized epsilon values for input, output, and format tokens + normalized_format_eps = self.config.opt_projection_format_epsilon * torch.sqrt( + torch.Tensor([self.config.token_dim / 2048]) + ) + normalized_input_eps = self.config.opt_projection_epsilon * torch.sqrt( + torch.Tensor([self.config.token_dim / 2048]) + ) + + epsilon = torch.ones_like(torch.Tensor(cpt_tokens_type_mask)).to(torch.float32) * MIN_VALUE + cpt_tokens_type_mask = torch.Tensor(cpt_tokens_type_mask).long() + + epsilon[(cpt_tokens_type_mask > 0) & (torch.remainder(cpt_tokens_type_mask, 4) == 1)] = normalized_format_eps + epsilon[(cpt_tokens_type_mask > 0) & (torch.remainder(cpt_tokens_type_mask, 4) == 3)] = normalized_format_eps + epsilon[(cpt_tokens_type_mask > 0) & (torch.remainder(cpt_tokens_type_mask, 4) == 2)] = normalized_input_eps + + return epsilon + + def get_projection(self): + """ + Applies epsilon-based projection to the delta embeddings to control their norm. + """ + + # Apply projection to control delta embedding norm + with torch.no_grad(): + new_embeddings_weights = self.delta_embedding.weight.clone().to(self.delta_embedding.weight.device) + token_norm = torch.norm(new_embeddings_weights, p=2, dim=1) + + projection_mask = token_norm > 0 + if torch.any(projection_mask): + epsilon = self.get_epsilon().to(self.delta_embedding.weight.device) + new_embeddings_weights[projection_mask] *= ( + epsilon[projection_mask] / (token_norm[projection_mask].clamp(min=epsilon[projection_mask])) + ).view(-1, 1) + return new_embeddings_weights + + @staticmethod + def calculate_loss(base_model_output, labels, cpt_type_mask, config): + """ + Computes the loss for CPT models with optional exponential decay. + + Args: + base_model_output (ModelOutput): + Output from the base model containing logits. + labels (torch.Tensor): + Ground-truth labels for the input tokens. + cpt_type_mask (torch.Tensor): + Token type mask used for filtering valid loss terms. + config (Namespace): + Configuration object containing loss-related hyperparameters. + + Returns: + ModelOutput: + The base model output with computed loss. + """ + + device = base_model_output.logits.device + + lm_logits = base_model_output.logits + labels = labels.to(device) + + # Shift logits and labels for token prediction + shift_logits = lm_logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + shift_cpt_type_mask = cpt_type_mask[..., 1:].contiguous() + + shift_labels_bool = (shift_labels.clone().detach() != -100).bool() + batch_size, seq_length, vocab_size = shift_logits.shape + + # Compute cross-entropy loss + loss_fct = CrossEntropyLoss(reduction="none", ignore_index=-100) + loss = loss_fct( + shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length) + ) + loss = loss.view(batch_size, seq_length) + # Apply exponential decay weights to the loss + shift_labels_weights = shift_labels_bool.clone().detach().float() + + for i in range(batch_size): + idx_labels = (shift_cpt_type_mask[i] > 0) & (shift_cpt_type_mask[i] % 4 == 0) + labels_ids = shift_cpt_type_mask[i][idx_labels].unique() + + exponential_decay = torch.ones_like(shift_cpt_type_mask[i]).to(device=device).float() + decay_value = 1 + for label_mask_idx in torch.flip(labels_ids, [0]): + exponential_decay[shift_cpt_type_mask[i] == label_mask_idx] = decay_value + decay_value *= config.opt_loss_decay_factor + if config.opt_weighted_loss_type == "decay": + shift_labels_weights[i] *= exponential_decay + + # Compute the weighted mean loss + loss = (loss[shift_labels_bool] * shift_labels_weights[shift_labels_bool]).mean() + + base_model_output.loss = loss + + return base_model_output diff --git a/lib/python3.12/site-packages/peft/tuners/fourierft/__init__.py b/lib/python3.12/site-packages/peft/tuners/fourierft/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dfe3f5d89e36123eefef7fc6e675c6b80bd8d44c --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/fourierft/__init__.py @@ -0,0 +1,24 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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 peft.utils import register_peft_method + +from .config import FourierFTConfig +from .layer import FourierFTLayer, FourierFTLinear +from .model import FourierFTModel + + +__all__ = ["FourierFTConfig", "FourierFTLayer", "FourierFTLinear", "FourierFTModel"] + +register_peft_method(name="fourierft", model_cls=FourierFTModel, config_cls=FourierFTConfig) diff --git a/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..88661d7e248dff9b8443e5898d77ff290bb19cfd Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4215644b72ba33c52c62e9664f7f8c234f7883e4 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/layer.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ae3e6361bf0d817e1165f8f95da6b039cd0d56f0 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3ff149fd685c97fd6d02647ffd180496502b79de Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/fourierft/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/fourierft/config.py b/lib/python3.12/site-packages/peft/tuners/fourierft/config.py new file mode 100644 index 0000000000000000000000000000000000000000..b8e354a207f4b2af6fad1c89017d6c889f47bf39 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/fourierft/config.py @@ -0,0 +1,205 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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 + +from dataclasses import dataclass, field +from typing import Optional, Union + +from peft.config import PeftConfig +from peft.utils import PeftType + + +@dataclass +class FourierFTConfig(PeftConfig): + """ + This is the configuration class to store the configuration of a [`FourierFTModel`]. + + Args: + n_frequency (`int`): + Num of learnable frequencies for the Discrete Fourier Transform. 'n_frequency' is an integer that is + greater than 0 and less than or equal to d^2 (assuming the weight W has dimensions of d by d). + Additionally, it is the number of trainable parameters required to update each delta W weight. + 'n_frequency' will affect the performance and efficiency for PEFT. Specifically, it has little impact on + training speed, but higher values of it (typically) result in larger GPU memory costs and better accuracy. + With the same `target_modules`, the number of parameters of LoRA is (2*d*r/n_frequency) times that of + FourierFT. The following examples of settings regarding 'n_frequency' can be used as reference for users. + For NLU tasks with the RoBERTa-large model, adopting 'n_frequency': 1000 can almost achieve similar results + as 'r': 8 in LoRA. At this time, the number of parameters of LoRA is about 16 times that of FourierFT. For + image classification tasks with Vit-large models, adopting 'n_frequency': 3000 can almost achieve similar + results as 'r': 16 in LoRA, where the number of parameters of LoRA is about 11 times that of FourierFT. + scaling (`float`): + The scaling value for the delta W matrix. This is an important hyperparameter used for scaling, similar to + the 'lora_alpha' parameter in the LoRA method. 'scaling' can be determined during the hyperparameter search + process. However, if users want to skip this process, one can refer to the settings in the following + scenarios. This parameter can be set to 100.0 or 150.0 for both RoBERTa-base and RoBERTa-large models + across all NLU (GLUE) tasks. This parameter can be set to 300.0 for both LLaMA family models for all + instruction tuning. This parameter can be set to 300.0 for both ViT-base and ViT-large models across all + image classification tasks. + random_loc_seed (`int`): + Seed for the random location of the frequencies, i.e., the spectral entry matrix. + target_modules (`Union[list[str],str]`): + List of module names or regex expression of the module names to replace with FourierFT. For example, ['q', + 'v'] or '.*decoder.*(SelfAttention|EncDecAttention).*(q|v)$'. Only linear layers are supported. + exclude_modules (`Optional[Union[List[str], str]]`): + The names of the modules to not apply the adapter. When passing a string, a regex match will be performed. + When passing a list of strings, either an exact match will be performed or it is checked if the name of the + module ends with any of the passed strings. + fan_in_fan_out (`bool`): + Set this to True if the layer to replace stores weight like (fan_in, fan_out). + bias (`str`): + Bias type for FourierFT. Can be 'none', 'all' or 'fourier_only'. + modules_to_save (`list[str]`): + List of modules apart from FourierFT layers to be set as trainable and saved in the final checkpoint. For + example, in Sequence Classification or Token Classification tasks, the final layer `classifier/score` are + randomly initialized and as such need to be trainable and saved. + layers_to_transform (`Union[list[int],int]`): + The layer indexes to transform, is this argument is specified, PEFT will transform only the layers indexes + that are specified inside this list. If a single integer is passed, PEFT will transform only the layer at + this index. + layers_pattern (`Optional[Union[List[str], str]]`): + The layer pattern name, used only if `layers_to_transform` is different to None and if the layer pattern is + not in the common layers pattern. This should target the `nn.ModuleList` of the model, which is often + called `'layers'` or `'h'`. + n_frequency_pattern (`dict`): + The mapping from layer names or regexp expression to n_frequency which are different from the default + specified. For example, `{model.decoder.layers.0.encoder_attn.k_proj: 1000`}. + init_weights (`bool`): + The initialization of the Fourier weights. Set this to False if the spectrum are initialized to a standard + normal distribution. Set this to True if the spectrum are initialized to zeros. + """ + + n_frequency: int = field( + default=1000, + metadata={ + "help": ( + "Num of learnable frequencies for the Discrete Fourier Transform. 'n_frequency' is an integer that is" + "greater than 0 and less than or equal to d^2 (assuming the weight W has dimensions of d by d)." + "Additionally, it is the number of trainable parameters required to update each delta W weight." + "'n_frequency' will affect the performance and efficiency for PEFT. Specifically, it has little impact on" + "training speed, but higher values of it (typically) result in larger GPU memory costs and better accuracy." + "With the same `target_modules`, the number of parameters of LoRA is (2*d*r/n_frequency) times that of FourierFT." + "The following examples of settings regarding 'n_frequency' can be used as reference for users. For NLU" + "tasks with the RoBERTa-large model, adopting 'n_frequency': 1000 can almost achieve similar results as" + "'r': 8 in LoRA. At this time, the number of parameters of LoRA is about 16 times that of FourierFT." + "For image classification tasks with Vit-large models, adopting 'n_frequency': 3000 can almost achieve" + "similar results as 'r': 16 in LoRA, where the number of parameters of LoRA is about 11 times that of FourierFT." + ) + }, + ) + scaling: float = field( + default=150.0, + metadata={ + "help": ( + "The scaling value for the delta W matrix. This is an important hyperparameter used for scaling, similar to the" + "'lora_alpha' parameter in the LoRA method. 'scaling' can be determined during the hyperparameter search process." + "However, if users want to skip this process, one can refer to the settings in the following scenarios." + "This parameter can be set to 100.0 or 150.0 for both RoBERTa-base and RoBERTa-large models across all NLU (GLUE) tasks." + "This parameter can be set to 300.0 for both LLaMA family models for all instruction tuning." + "This parameter can be set to 300.0 for both ViT-base and ViT-large models across all image classification tasks." + ) + }, + ) + random_loc_seed: Optional[int] = field( + default=777, metadata={"help": "Seed for the random location of the frequencies."} + ) + fan_in_fan_out: bool = field( + default=False, + metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"}, + ) + target_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": ( + "List of module names or regex expression of the module names to replace with FourierFT." + "For example, ['q', 'v'] or '.*decoder.*(SelfAttention|EncDecAttention).*(q|v)$'. " + "Only linear layers are supported." + ) + }, + ) + exclude_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={"help": "List of module names or regex expression of the module names to exclude from fourierft."}, + ) + bias: str = field( + default="none", metadata={"help": "Bias type for FourierFT. Can be 'none', 'all' or 'fourier_only'."} + ) + modules_to_save: Optional[list[str]] = field( + default=None, + metadata={ + "help": ( + "List of modules apart from FourierFT layers to be set as trainable and saved in the final checkpoint. For" + " example, in Sequence Classification or Token Classification tasks, the final layer" + " `classifier/score` are randomly initialized and as such need to be trainable and saved." + ) + }, + ) + layers_to_transform: Optional[Union[list[int], int]] = field( + default=None, + metadata={ + "help": ( + "The layer indexes to transform, is this argument is specified, PEFT will transform only the layers" + " indexes that are specified inside this list. If a single integer is passed, PEFT will transform only" + " the layer at this index." + ) + }, + ) + layers_pattern: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": ( + "The layer pattern name, used only if `layers_to_transform` is different to None and if the layer" + " pattern is not in the common layers pattern. This should target the `nn.ModuleList` of the " + "model, which is often called `'layers'` or `'h'`." + ) + }, + ) + n_frequency_pattern: Optional[dict] = field( + default_factory=dict, + metadata={ + "help": ( + "The mapping from layer names or regexp expression to n_frequency which are different from the default specified." + "For example, `{model.decoder.layers.0.encoder_attn.k_proj: 500`}." + ) + }, + ) + init_weights: bool = field( + default=False, + metadata={ + "help": ( + "The initialization of the Fourier weights. Set this to False if the spectrum should be initialized to a standard normal distribution." + "Set this to True if the spectrum should be initialized to zeros." + ) + }, + ) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.FOURIERFT + self.target_modules = ( + set(self.target_modules) if isinstance(self.target_modules, list) else self.target_modules + ) + self.exclude_modules = ( + set(self.exclude_modules) if isinstance(self.exclude_modules, list) else self.exclude_modules + ) + # if target_modules is a regex expression, then layers_to_transform should be None + if isinstance(self.target_modules, str) and self.layers_to_transform is not None: + raise ValueError("`layers_to_transform` cannot be used when `target_modules` is a str.") + + # if target_modules is a regex expression, then layers_pattern should be None + if isinstance(self.target_modules, str) and self.layers_pattern is not None: + raise ValueError("`layers_pattern` cannot be used when `target_modules` is a str.") + # check for layers_to_transform and layers_pattern + if self.layers_pattern and not self.layers_to_transform: + raise ValueError("When `layers_pattern` is specified, `layers_to_transform` must also be specified. ") diff --git a/lib/python3.12/site-packages/peft/tuners/fourierft/layer.py b/lib/python3.12/site-packages/peft/tuners/fourierft/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..a8f615a989de08fd71d69d65dbb71913026166f4 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/fourierft/layer.py @@ -0,0 +1,190 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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. + +import warnings +from typing import Any, List, Optional, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.pytorch_utils import Conv1D + +from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge + + +class FourierFTLayer(BaseTunerLayer): + # All names of layers that may contain (trainable) adapter weights + adapter_layer_names = ("fourierft_spectrum",) + # All names of other parameters that may contain adapter-related parameters + other_param_names = ("fourierft_n_frequency", "fourierft_scaling", "fourierft_random_loc_seed") + + def __init__(self, base_layer: nn.Module, **kwargs) -> None: + self.base_layer = base_layer + self.fourierft_n_frequency = {} + self.fourierft_scaling = {} + self.fourierft_spectrum = nn.ParameterDict({}) + self.indices = {} + self.fourierft_random_loc_seed = {} + # Mark the weight as unmerged + self._disable_adapters = False + self.merged_adapters = [] + self.kwargs = kwargs + + base_layer = self.get_base_layer() + if isinstance(base_layer, nn.Linear): + self.in_features, self.out_features = base_layer.in_features, base_layer.out_features + elif isinstance(base_layer, Conv1D): + self.in_features, self.out_features = ( + base_layer.weight.ds_shape if hasattr(base_layer.weight, "ds_shape") else base_layer.weight.shape + ) + else: + raise ValueError(f"Unsupported layer type {type(base_layer)}") + + def update_layer(self, adapter_name, n_frequency, scaling, init_weights, random_loc_seed): + if n_frequency <= 0: + raise ValueError(f"`n_frequency` should be a positive integer value but the value passed is {n_frequency}") + if n_frequency > self.in_features * self.out_features: + raise ValueError( + f"`n_frequency` should be less than or equal to the product of the input and output dimensions " + f"but the value passed is {n_frequency} and the product is {self.in_features * self.out_features}" + ) + self.fourierft_n_frequency[adapter_name] = n_frequency + self.fourierft_random_loc_seed[adapter_name] = random_loc_seed + self.indices[adapter_name] = torch.randperm( + self.out_features * self.in_features, + generator=torch.Generator().manual_seed(self.fourierft_random_loc_seed[adapter_name]), + )[:n_frequency] + self.indices[adapter_name] = torch.stack( + [self.indices[adapter_name] // self.in_features, self.indices[adapter_name] % self.in_features], dim=0 + ) + self.fourierft_scaling[adapter_name] = scaling + # Actual trainable parameters + self.fourierft_spectrum[adapter_name] = nn.Parameter(torch.randn(n_frequency), requires_grad=True) + + if init_weights: + self.reset_fourier_parameters(adapter_name) + + self._move_adapter_to_device_of_base_layer(adapter_name) + self.set_adapter(self.active_adapters) + + @torch.no_grad() + def reset_fourier_parameters(self, adapter_name): + if adapter_name in self.fourierft_spectrum.keys(): + nn.init.zeros_(self.fourierft_spectrum[adapter_name]) + + def get_delta_weight(self, adapter) -> torch.Tensor: + spectrum = self.fourierft_spectrum[adapter] + indices = self.indices[adapter].to(spectrum.device) + dense_spectrum = torch.zeros(self.out_features, self.in_features, device=spectrum.device, dtype=spectrum.dtype) + dense_spectrum[indices[0, :], indices[1, :]] = spectrum + delta_weight = torch.fft.ifft2(dense_spectrum).real * self.fourierft_scaling[adapter] + return delta_weight + + +class FourierFTLinear(nn.Module, FourierFTLayer): + # FourierFT implemented in a dense layer + def __init__( + self, + base_layer, + adapter_name: str, + n_frequency: int = 1000, + scaling: float = 150.0, + fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) + init_weights: Union[bool, str] = False, + random_loc_seed: int = 777, + **kwargs, + ) -> None: + super().__init__() + FourierFTLayer.__init__(self, base_layer, **kwargs) + self.fan_in_fan_out = fan_in_fan_out + self._active_adapter = adapter_name + self.update_layer(adapter_name, n_frequency, scaling, init_weights, random_loc_seed) + + def merge(self, safe_merge: bool = False, adapter_names: Optional[List[str]] = None) -> None: + """ + Merge the active adapter weights into the base weights + + Args: + safe_merge (`bool`, *optional*): + If True, the merge operation will be performed in a copy of the original weights and check for NaNs + before merging the weights. This is useful if you want to check if the merge operation will produce + NaNs. Defaults to `False`. + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + """ + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + # no adapter to merge + return + + for active_adapter in adapter_names: + if active_adapter in self.fourierft_spectrum.keys(): + base_layer = self.get_base_layer() + if safe_merge: + # Note that safe_merge will be slower than the normal merge + # because of the copy operation. + orig_weights = base_layer.weight.data.clone() + orig_weights += self.get_delta_weight(active_adapter) + + if not torch.isfinite(orig_weights).all(): + raise ValueError( + f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" + ) + + base_layer.weight.data = orig_weights + else: + base_layer.weight.data += self.get_delta_weight(active_adapter) + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + """ + This method unmerges all merged adapter layers from the base weights. + """ + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter in self.fourierft_spectrum.keys(): + self.get_base_layer().weight.data -= self.get_delta_weight(active_adapter) + + def get_delta_weight(self, adapter) -> torch.Tensor: + return super().get_delta_weight(adapter) + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + previous_dtype = x.dtype + + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + result = self.base_layer(x, *args, **kwargs) + for active_adapter in self.active_adapters: + if active_adapter not in self.fourierft_spectrum.keys(): + continue + + delta_w = self.get_delta_weight(active_adapter) + x = x.to(delta_w.dtype) + result = result + F.linear(x, delta_w) + + result = result.to(previous_dtype) + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "fourierft." + rep diff --git a/lib/python3.12/site-packages/peft/tuners/fourierft/model.py b/lib/python3.12/site-packages/peft/tuners/fourierft/model.py new file mode 100644 index 0000000000000000000000000000000000000000..711c8ffaa2cc4db8d720c618f613814d0fa7714d --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/fourierft/model.py @@ -0,0 +1,349 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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 re +import warnings +from dataclasses import asdict +from enum import Enum +from itertools import chain +from typing import Optional + +import torch +from tqdm import tqdm +from transformers.pytorch_utils import Conv1D + +from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists +from peft.utils import ( + TRANSFORMERS_MODELS_TO_FOURIERFT_TARGET_MODULES_MAPPING, + ModulesToSaveWrapper, + _get_submodules, +) + +from .config import FourierFTConfig +from .layer import FourierFTLayer, FourierFTLinear + + +class FourierFTModel(BaseTuner): + """ + Creates FourierFT model from a pretrained transformers model. + + The method is described in detail in https://arxiv.org/abs/2405.03003. + + Args: + model ([`torch.nn.Module`]): The model to be adapted. + config ([`FourierFTConfig`]): The configuration of the FourierFT model. + adapter_name (`str`): The name of the adapter, defaults to `"default"`. + low_cpu_mem_usage (`bool`, `optional`, defaults to `False`): + Create empty adapter weights on meta device. Useful to speed up the loading process. + + Returns: + `torch.nn.Module`: The FourierFT model. + + **Attributes**: + - **model** ([`~transformers.PreTrainedModel`]) -- The model to be adapted. + - **peft_config** ([`FourierFTConfig`]): The configuration of the Fourier model. + """ + + prefix: str = "fourierft_" + + def __init__(self, model, config, adapter_name, low_cpu_mem_usage: bool = False) -> None: + super().__init__(model, config, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage) + + def _check_new_adapter_config(self, config: FourierFTConfig) -> None: + """ + A helper method to check the config when a new adapter is being added. + + Raise a ValueError if there is something wrong with the config or if it conflicts with existing adapters. + + """ + # TODO: there should be a check if any of the existing adapters actually has bias != "none", or else the check + # does not fully correspond to the error message. + if (len(self.peft_config) > 1) and (config.bias != "none"): + raise ValueError( + f"{self.__class__.__name__} supports only 1 adapter with bias. When using multiple adapters, " + "set bias to 'none' for all adapters." + ) + + @staticmethod + def _check_target_module_exists(fourierft_config, key): + return check_target_module_exists(fourierft_config, key) + + def _create_and_replace( + self, + fourierft_config, + adapter_name, + target, + target_name, + parent, + current_key, + **optional_kwargs, + ): + if current_key is None: + raise ValueError("Current Key shouldn't be `None`") + # Regexp matching - Find key which matches current target_name in patterns provided + pattern_keys = list(chain(fourierft_config.n_frequency_pattern.keys())) + target_name_key = next(filter(lambda key: re.match(rf".*\.{key}$", current_key), pattern_keys), current_key) + + n_frequency = fourierft_config.n_frequency_pattern.get(target_name_key, fourierft_config.n_frequency) + scaling = fourierft_config.scaling + random_loc_seed = fourierft_config.random_loc_seed + bias = hasattr(target, "bias") and target.bias is not None + kwargs = { + "n_frequency": n_frequency, + "scaling": scaling, + "fan_in_fan_out": fourierft_config.fan_in_fan_out, + "init_weights": fourierft_config.init_weights, + "random_loc_seed": fourierft_config.random_loc_seed, + } + kwargs["bias"] = bias + if isinstance(target, FourierFTLayer): + target.update_layer( + adapter_name, + n_frequency, + scaling, + fourierft_config.init_weights, + random_loc_seed, + ) + else: + new_module = self._create_new_module(fourierft_config, adapter_name, target, **kwargs) + if adapter_name != self.active_adapter: + # adding an additional adapter: it is not automatically trainable + new_module.requires_grad_(False) + self._replace_module(parent, target_name, new_module, target) + + def _replace_module(self, parent, child_name, new_module, child): + setattr(parent, child_name, new_module) + # It's not necessary to set requires_grad here, as that is handled by + # _mark_only_adapters_as_trainable + + # child layer wraps the original module, unpack it + if hasattr(child, "base_layer"): + child = child.base_layer + + if not hasattr(new_module, "base_layer"): + new_module.weight = child.weight + if hasattr(child, "bias"): + new_module.bias = child.bias + + if getattr(child, "state", None) is not None: + if hasattr(new_module, "base_layer"): + new_module.base_layer.state = child.state + else: + new_module.state = child.state + new_module.to(child.weight.device) + + meta = torch.device("meta") + # dispatch to correct device + for name, module in new_module.named_modules(): + if "fourierft_" in name: + if not any(p.device == meta for p in module.parameters()): + module.to(child.weight.device) + + def _mark_only_adapters_as_trainable(self, model: torch.nn.Module) -> None: + for n, p in model.named_parameters(): + if self.prefix not in n: + p.requires_grad = False + + for active_adapter in self.active_adapters: + bias = self.peft_config[active_adapter].bias + if bias == "none": + continue + + if bias == "all": + for n, p in model.named_parameters(): + if "bias" in n: + p.requires_grad = True + elif bias == "fourier_only": + for m in model.modules(): + if isinstance(m, FourierFTLayer) and hasattr(m, "bias") and m.bias is not None: + m.bias.requires_grad = True + else: + raise NotImplementedError(f"Requested bias: {bias}, is not implemented.") + + @staticmethod + def _create_new_module(fourierft_config, adapter_name, target, **kwargs): + if isinstance(target, BaseTunerLayer): + target_base_layer = target.get_base_layer() + else: + target_base_layer = target + + if isinstance(target_base_layer, torch.nn.Linear): + if kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. " + "Setting fan_in_fan_out to False." + ) + kwargs["fan_in_fan_out"] = fourierft_config.fan_in_fan_out = False + elif isinstance(target_base_layer, Conv1D): + kwargs["is_target_conv_1d_layer"] = True + if not kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to False but the target module is `Conv1D`. Setting fan_in_fan_out to True." + ) + kwargs["fan_in_fan_out"] = fourierft_config.fan_in_fan_out = True + else: + raise ValueError( + f"Target module {target} is not supported. Currently, only the following modules are supported: " + "`torch.nn.Linear`." + ) + + new_module = FourierFTLinear(target, adapter_name, **kwargs) + + return new_module + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + if name == "model": + raise + return getattr(self.model, name) + + def get_peft_config_as_dict(self, inference: bool = False): + config_dict = {} + for key, value in self.peft_config.items(): + config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} + if inference: + config["inference_mode"] = True + config_dict[key] = config + return config + + def _set_adapter_layers(self, enabled: bool = True) -> None: + for module in self.model.modules(): + if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)): + module.enable_adapters(enabled) + + def enable_adapter_layers(self) -> None: + """Enable all adapters. + + Call this if you have previously disabled all adapters and want to re-enable them. + """ + self._set_adapter_layers(enabled=True) + + def disable_adapter_layers(self) -> None: + """Disable all adapters. + + When disabling all adapters, the model output corresponds to the output of the base model. + """ + for active_adapter in self.active_adapters: + val = self.peft_config[active_adapter].bias + if val != "none": + msg = ( + f"Careful, disabling adapter layers with bias configured to be '{val}' does not produce the same " + "output as the the base model would without adaption." + ) + warnings.warn(msg) + self._set_adapter_layers(enabled=False) + + def set_adapter(self, adapter_name: str | list[str]) -> None: + """Set the active adapter(s). + + Args: + adapter_name (`str` or `list[str]`): Name of the adapter(s) to be activated. + """ + for module in self.model.modules(): + if isinstance(module, FourierFTLayer): + if module.merged: + warnings.warn("Adapter cannot be set when the model is merged. Unmerging the model first.") + module.unmerge() + module.set_adapter(adapter_name) + self.active_adapter = adapter_name + + @staticmethod + def _prepare_adapter_config(peft_config, model_config): + if peft_config.target_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_FOURIERFT_TARGET_MODULES_MAPPING: + raise ValueError("Please specify `target_modules` in `peft_config`") + peft_config.target_modules = set( + TRANSFORMERS_MODELS_TO_FOURIERFT_TARGET_MODULES_MAPPING[model_config["model_type"]] + ) + return peft_config + + def _unload_and_optionally_merge( + self, + merge=True, + progressbar: bool = False, + safe_merge: bool = False, + adapter_names: Optional[list[str]] = None, + ): + key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] + desc = "Unloading " + ("and merging " if merge else "") + "model" + for key in tqdm(key_list, disable=not progressbar, desc=desc): + try: + parent, target, target_name = _get_submodules(self.model, key) + except AttributeError: + continue + + if hasattr(target, "base_layer"): + if merge: + target.merge(safe_merge=safe_merge, adapter_names=adapter_names) + self._replace_module(parent, target_name, target.get_base_layer(), target) + elif isinstance(target, ModulesToSaveWrapper): + # save any additional trainable modules part of `modules_to_save` + setattr(parent, target_name, target.modules_to_save[target.active_adapter]) + + return self.model + + def delete_adapter(self, adapter_name: str): + """ + Deletes an existing adapter. + + Args: + adapter_name (str): Name of the adapter to be deleted. + """ + if adapter_name not in list(self.peft_config.keys()): + raise ValueError(f"Adapter {adapter_name} does not exist") + del self.peft_config[adapter_name] + + # we cannot use self.prefix as we want to include non-trainable fourierft parameters + key_list = [key for key, _ in self.model.named_modules() if "fourierft" not in key] + new_adapter = None + for key in key_list: + _, target, _ = _get_submodules(self.model, key) + if isinstance(target, FourierFTLayer): + target.delete_adapter(adapter_name) + if new_adapter is None: + new_adapter = target.active_adapter[:] + + self.active_adapter = new_adapter or [] + + def merge_and_unload( + self, progressbar: bool = False, safe_merge: bool = False, adapter_names: Optional[list[str]] = None + ) -> torch.nn.Module: + r""" + This method merges the Fourier layers into the base model. This is needed if someone wants to use the base + model as a standalone model. + + Args: + progressbar (`bool`): + whether to show a progressbar indicating the unload and merge process + safe_merge (`bool`): + whether to activate the safe merging check to check if there is any potential Nan in the adapter + weights + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + """ + return self._unload_and_optionally_merge( + progressbar=progressbar, safe_merge=safe_merge, adapter_names=adapter_names + ) + + def unload(self) -> torch.nn.Module: + """ + Gets back the base model by removing all the Fourier modules without merging. This gives back the original base + model. + """ + return self._unload_and_optionally_merge(merge=False) diff --git a/lib/python3.12/site-packages/peft/tuners/hra/__init__.py b/lib/python3.12/site-packages/peft/tuners/hra/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8f5f6a54435d0fcaa1bd275c623168b5491f2d9e --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/hra/__init__.py @@ -0,0 +1,24 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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 peft.utils import register_peft_method + +from .config import HRAConfig +from .layer import HRAConv2d, HRALayer, HRALinear +from .model import HRAModel + + +__all__ = ["HRAConfig", "HRAConv2d", "HRALayer", "HRALinear", "HRAModel"] + +register_peft_method(name="hra", config_cls=HRAConfig, model_cls=HRAModel) diff --git a/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32926a82c5fcde1ed43e0a3f0498ee55c423e4a6 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c27ab6827309332b3cf059383e1e979f3dc9ff79 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/layer.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2cb880f13ba2b9ea980c61bf9915367c28a2d57 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..55eba30783c5397a1b29cc0af2eed0a087d57824 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/hra/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/hra/config.py b/lib/python3.12/site-packages/peft/tuners/hra/config.py new file mode 100644 index 0000000000000000000000000000000000000000..9e6d87cc17c3ba003eac963467bb1253725ebe10 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/hra/config.py @@ -0,0 +1,133 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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 + +from dataclasses import dataclass, field +from typing import Optional, Union + +from peft.config import PeftConfig +from peft.utils import PeftType + + +@dataclass +class HRAConfig(PeftConfig): + """ + This is the configuration class to store the configuration of a [`HRAModel`]. + + Args: + r (`int`): + The rank of HRA across different layers. It is best to set 'r' to an even number; otherwise, the default + initialization method will not work. + apply_GS (`bool`): + Whether to apply Gram-Schmidt orthogonalization. + target_modules (`Optional[Union[List[str], str]]`): + The names of the modules to apply the adapter to. If this is specified, only the modules with the specified + names will be replaced. When passing a string, a regex match will be performed. When passing a list of + strings, either an exact match will be performed or it is checked if the name of the module ends with any + of the passed strings. If this is specified as 'all-linear', then all linear modules are chosen, excluding + the output layer. If this is not specified, modules will be chosen according to the model architecture. If + the architecture is not known, an error will be raised -- in this case, you should specify the target + modules manually. + exclude_modules (`Optional[Union[List[str], str]]`): + The names of the modules to not apply the adapter. When passing a string, a regex match will be performed. + When passing a list of strings, either an exact match will be performed or it is checked if the name of the + module ends with any of the passed strings. + init_weights (`bool`): + Whether to perform initialization of HRA weights. + layers_to_transform (`Union[List[int], int]`): + The layer indices to transform. If a list of ints is passed, it will apply the adapter to the layer indices + that are specified in this list. If a single integer is passed, it will apply the transformations on the + layer at this index. + layers_pattern (`Optional[Union[List[str], str]]`): + The layer pattern name, used only if `layers_to_transform` is different from `None`. This should target the + `nn.ModuleList` of the model, which is often called `'layers'` or `'h'`. + modules_to_save (`List[str]`): + List of modules apart from adapter layers to be set as trainable and saved in the final checkpoint. + """ + + r: int = field( + default=8, + metadata={ + "help": "The rank of HRA across different layers.", + "note": "It is best to set 'r' to an even number; otherwise, the default initialization method will not work.", + }, + ) + apply_GS: bool = field( + default=False, + metadata={"help": "Whether to apply Gram-Schmidt orthogonalization or not."}, + ) + target_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": "List of module names or regex expression of the module names to replace with HRA.", + "example": "For example, ['q', 'v'] or '.*decoder.*(SelfAttention|EncDecAttention).*(q|v)$' ", + }, + ) + exclude_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={"help": "List of module names or regex expression of the module names to exclude from HRA."}, + ) + init_weights: bool = field( + default=True, + metadata={ + "help": ( + "Whether to initialize the weights of the HRA layers with their default initialization. Don't change " + "this setting, except if you know exactly what you're doing." + ), + }, + ) + layers_to_transform: Optional[Union[list[int], int]] = field( + default=None, + metadata={ + "help": "The layer indexes to transform, is this argument is specified, PEFT will transform only the layers indexes that are specified inside this list. If a single integer is passed, PEFT will transform only the layer at this index." + }, + ) + layers_pattern: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": "The layer pattern name, used only if `layers_to_transform` is different to None and if the layer pattern is not in the common layers pattern. " + "This should target the `nn.ModuleList` of the model, which is often called `'layers'` or `'h'`." + }, + ) + bias: str = field(default="none", metadata={"help": "Bias type for HRA. Can be 'none', 'all' or 'hra_only'"}) + modules_to_save: Optional[list[str]] = field( + default=None, + metadata={ + "help": "List of modules apart from HRA layers to be set as trainable and saved in the final checkpoint. " + "For example, in Sequence Classification or Token Classification tasks, " + "the final layer `classifier/score` are randomly initialized and as such need to be trainable and saved." + }, + ) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.HRA + self.target_modules = ( + set(self.target_modules) if isinstance(self.target_modules, list) else self.target_modules + ) + self.exclude_modules = ( + set(self.exclude_modules) if isinstance(self.exclude_modules, list) else self.exclude_modules + ) + # if target_modules is a regex expression, then layers_to_transform should be None + if isinstance(self.target_modules, str) and self.layers_to_transform is not None: + raise ValueError("`layers_to_transform` cannot be used when `target_modules` is a str.") + + # if target_modules is a regex expression, then layers_pattern should be None + if isinstance(self.target_modules, str) and self.layers_pattern is not None: + raise ValueError("`layers_pattern` cannot be used when `target_modules` is a str.") + + # check for layers_to_transform and layers_pattern + if self.layers_pattern and not self.layers_to_transform: + raise ValueError("When `layers_pattern` is specified, `layers_to_transform` must also be specified. ") diff --git a/lib/python3.12/site-packages/peft/tuners/hra/layer.py b/lib/python3.12/site-packages/peft/tuners/hra/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..2494510c60ad7d5491993612ede1bd6e55f442c7 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/hra/layer.py @@ -0,0 +1,440 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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. + +import math +import warnings +from typing import Any, List, Optional, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge + + +class HRALayer(BaseTunerLayer): + # All names of layers that may contain (trainable) adapter weights + adapter_layer_names = ("hra_u",) + # All names of other parameters that may contain adapter-related parameters + other_param_names = ("hra_r", "hra_apply_GS") + + def __init__(self, base_layer: nn.Module, **kwargs) -> None: + self.base_layer = base_layer + self.hra_r = {} + self.hra_apply_GS = {} + self.hra_u = nn.ParameterDict({}) + # Mark the weight as unmerged + self._disable_adapters = False + self.merged_adapters = [] + self.kwargs = kwargs + + base_layer = self.get_base_layer() + if isinstance(base_layer, nn.Linear): + self.in_features, self.out_features = base_layer.in_features, base_layer.out_features + elif isinstance(base_layer, nn.Conv2d): + self.in_features, self.out_features = base_layer.in_channels, base_layer.out_channels + else: + raise ValueError(f"Unsupported layer type {type(base_layer)}") + + def update_layer( + self, + adapter_name: str, + r: int, + apply_GS: bool, + init_weights: bool, + **kwargs, + ) -> None: + """Internal function to create hra adapter + + Args: + adapter_name (`str`): Name for the adapter to add. + r (`int`): Rank for the added adapter. + init_weights (`bool`): Whether to initialize weights. + apply_GS (`bool`): Whether to apply Gram-Schmidt orthogonalization or not. + """ + if r <= 0: + raise ValueError(f"`r` should be a positive integer value but the value passed is {r}") + + self.hra_r[adapter_name] = r + self.hra_apply_GS[adapter_name] = apply_GS + + # Determine shape of HRA weights + base_layer = self.get_base_layer() + if isinstance(base_layer, nn.Linear): + self.hra_u[adapter_name] = nn.Parameter(torch.empty(self.in_features, r), requires_grad=True) + elif isinstance(base_layer, nn.Conv2d): + self.hra_u[adapter_name] = nn.Parameter( + torch.empty(self.in_features * base_layer.kernel_size[0] * base_layer.kernel_size[0], r), + requires_grad=True, + ) + else: + raise TypeError(f"HRA is not implemented for base layers of type {type(base_layer).__name__}") + + # Initialize weights + if init_weights: + self.reset_hra_parameters(adapter_name) + else: + self.reset_hra_parameters_random(adapter_name) + + # Move new weights to device + self._move_adapter_to_device_of_base_layer(adapter_name) + self.set_adapter(self.active_adapters) + + def reset_hra_parameters(self, adapter_name: str): + if self.hra_r[adapter_name] % 2 != 0: + warnings.warn("The symmetric initialization can NOT be performed when r is odd!") + nn.init.kaiming_uniform_(self.hra_u[adapter_name], a=math.sqrt(5)) + else: + shape = self.hra_u[adapter_name].shape + half_u = torch.zeros(shape[0], shape[1] // 2) + nn.init.kaiming_uniform_(half_u, a=math.sqrt(5)) + self.hra_u[adapter_name] = nn.Parameter(torch.repeat_interleave(half_u, 2, dim=1)) + + def reset_hra_parameters_random(self, adapter_name: str): + nn.init.kaiming_uniform_(self.hra_u[adapter_name], a=math.sqrt(5)) + + def scale_layer(self, scale: float) -> None: + if scale == 1: + return + + for active_adapter in self.active_adapters: + if active_adapter not in self.hra_u.keys(): + continue + + warnings.warn("Scaling operation for HRA not supported! Automatically set scale to 1.") + + def unscale_layer(self, scale=None) -> None: + for active_adapter in self.active_adapters: + if active_adapter not in self.hra_u.keys(): + continue + + warnings.warn("Unscaling operation for HRA not supported! Keeping scale at 1.") + + +class HRALinear(nn.Module, HRALayer): + """ + HRA implemented in a dense layer. + """ + + def __init__( + self, + base_layer, + adapter_name: str, + r: int = 0, + apply_GS: bool = False, + init_weights: Union[bool, str] = True, + **kwargs, + ) -> None: + super().__init__() + HRALayer.__init__(self, base_layer, **kwargs) + self._active_adapter = adapter_name + self.update_layer(adapter_name, r, apply_GS, init_weights, **kwargs) + + def merge(self, safe_merge: bool = False, adapter_names: Optional[List[str]] = None) -> None: + """ + Merge the active adapter weights into the base weights + + Args: + safe_merge (`bool`, *optional*): + If `True`, the merge operation will be performed in a copy of the original weights and check for NaNs + before merging the weights. This is useful if you want to check if the merge operation will produce + NaNs. Defaults to `False`. + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If `None`, all active adapters will be merged. + Defaults to `None`. + """ + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + # no adapter to merge + return + + for active_adapter in adapter_names: + if active_adapter in self.hra_u.keys(): + base_layer = self.get_base_layer() + if safe_merge: + # Note that safe_merge will be slower than the normal merge + # because of the copy operation. + orig_weight = base_layer.weight.data.clone() + delta_weight = self.get_delta_weight(active_adapter) + orig_weight = torch.mm(orig_weight, delta_weight) + + if not torch.isfinite(orig_weight).all(): + raise ValueError( + f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" + ) + + self.base_layer.weight.data = orig_weight + else: + delta_weight = self.get_delta_weight(active_adapter) + self.base_layer.weight.data = torch.mm(self.base_layer.weight.data, delta_weight) + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + """ + This method unmerges all merged adapter layers from the base weights. + """ + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter in self.hra_u.keys(): + orig_weight = self.get_base_layer().weight.data.clone() + delta_weight = self.get_delta_weight(active_adapter, reverse=True) + self.get_base_layer().weight.data = torch.mm(orig_weight, delta_weight) + + def get_delta_weight(self, adapter_name: str, reverse: bool = False) -> torch.Tensor: + rank = self.hra_r[adapter_name] + apply_GS = self.hra_apply_GS[adapter_name] + opt_u = self.hra_u[adapter_name] + shape = opt_u.shape + + if apply_GS: + weight = [(opt_u[:, 0] / opt_u[:, 0].norm()).view(-1, 1)] + for i in range(1, rank): + ui = opt_u[:, i].view(-1, 1) + for j in range(i): + ui = ui - (weight[j].t() @ ui) * weight[j] + weight.append((ui / ui.norm()).view(-1, 1)) + weight = torch.cat(weight, dim=1) + weight = torch.eye(shape[0], device=opt_u.device, dtype=opt_u.dtype) - 2 * weight @ weight.t() + + else: + opt_u = opt_u / opt_u.norm(dim=0) + weight = torch.eye(shape[0], device=opt_u.device, dtype=opt_u.dtype) + if reverse: + indices = range(rank - 1, -1, -1) + else: + indices = range(rank) + + for i in indices: + ui = opt_u[:, i].view(-1, 1) + weight = weight - 2 * weight @ ui @ ui.t() + + return weight + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + previous_dtype = x.dtype + + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + new_weight = torch.eye(self.in_features, device=x.device) + + for active_adapter in self.active_adapters: + if active_adapter not in self.hra_u.keys(): + continue + delta_weight = self.get_delta_weight(active_adapter) + new_weight = torch.mm(new_weight, delta_weight) + + x = x.to(self.get_base_layer().weight.data.dtype) + orig_weight = self.get_base_layer().weight.data + new_weight = torch.mm(orig_weight, new_weight) + + result = F.linear(input=x, weight=new_weight, bias=self.base_layer.bias) + + result = result.to(previous_dtype) + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "hra." + rep + + +class HRAConv2d(nn.Module, HRALayer): + """HRA implemented in Conv2d layer""" + + def __init__( + self, + base_layer, + adapter_name: str, + r: int = 0, + apply_GS: bool = False, + init_weights: Union[bool, str] = True, + **kwargs, + ): + super().__init__() + HRALayer.__init__(self, base_layer) + self._active_adapter = adapter_name + self.update_layer(adapter_name, r, apply_GS, init_weights, **kwargs) + + def merge(self, safe_merge: bool = False, adapter_names: Optional[List[str]] = None) -> None: + """ + Merge the active adapter weights into the base weights + + Args: + safe_merge (`bool`, *optional*): + If `True`, the merge operation will be performed in a copy of the original weights and check for NaNs + before merging the weights. This is useful if you want to check if the merge operation will produce + NaNs. Defaults to `False`. + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If `None`, all active adapters will be merged. + Defaults to `None`. + """ + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + # no adapter to merge + return + + for active_adapter in adapter_names: + if active_adapter in self.hra_u.keys(): + base_layer = self.get_base_layer() + if safe_merge: + # Note that safe_merge will be slower than the normal merge + # because of the copy operation. + orig_weight = base_layer.weight.data.clone() + orig_weight = orig_weight.view( + self.out_features, + self.in_features * self.base_layer.kernel_size[0] * self.base_layer.kernel_size[0], + ) + delta_weight = self.get_delta_weight(active_adapter) + orig_weight = torch.mm(orig_weight, delta_weight) + orig_weight = orig_weight.view( + self.out_features, + self.in_features, + self.base_layer.kernel_size[0], + self.base_layer.kernel_size[0], + ) + + if not torch.isfinite(orig_weight).all(): + raise ValueError( + f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" + ) + + self.base_layer.weight.data = orig_weight + else: + orig_weight = base_layer.weight.data + orig_weight = orig_weight.view( + self.out_features, + self.in_features * self.base_layer.kernel_size[0] * self.base_layer.kernel_size[0], + ) + delta_weight = self.get_delta_weight(active_adapter) + orig_weight = torch.mm(orig_weight, delta_weight) + orig_weight = orig_weight.view( + self.out_features, + self.in_features, + self.base_layer.kernel_size[0], + self.base_layer.kernel_size[0], + ) + + self.base_layer.weight.data = orig_weight + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + """ + This method unmerges all merged adapter layers from the base weights. + """ + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter in self.hra_u.keys(): + orig_weight = self.get_base_layer().weight.data.clone() + orig_weight = orig_weight.view( + self.out_features, + self.in_features * self.base_layer.kernel_size[0] * self.base_layer.kernel_size[0], + ) + delta_weight = self.get_delta_weight(active_adapter, reverse=True) + orig_weight = torch.mm(orig_weight, delta_weight) + orig_weight = orig_weight.view( + self.out_features, self.in_features, self.base_layer.kernel_size[0], self.base_layer.kernel_size[0] + ) + + self.get_base_layer().weight.data = orig_weight + + def get_delta_weight(self, adapter_name: str, reverse: bool = False) -> torch.Tensor: + rank = self.hra_r[adapter_name] + apply_GS = self.hra_apply_GS[adapter_name] + opt_u = self.hra_u[adapter_name] + shape = opt_u.shape + + if apply_GS: + weight = [(opt_u[:, 0] / opt_u[:, 0].norm()).view(-1, 1)] + for i in range(1, rank): + ui = opt_u[:, i].view(-1, 1) + for j in range(i): + ui = ui - (weight[j].t() @ ui) * weight[j] + weight.append((ui / ui.norm()).view(-1, 1)) + weight = torch.cat(weight, dim=1) + weight = torch.eye(shape[0], device=opt_u.device, dtype=opt_u.dtype) - 2 * weight @ weight.t() + + else: + opt_u = opt_u / opt_u.norm(dim=0) + weight = torch.eye(shape[0], device=opt_u.device, dtype=opt_u.dtype) + if reverse: + indices = range(rank - 1, -1, -1) + else: + indices = range(rank) + + for i in indices: + ui = opt_u[:, i].view(-1, 1) + weight = weight - 2 * weight @ ui @ ui.t() + + return weight + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + previous_dtype = x.dtype + + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + new_weight = torch.eye( + self.in_features * self.base_layer.kernel_size[0] * self.base_layer.kernel_size[0], + device=x.device, + dtype=previous_dtype, + ) + for active_adapter in self.active_adapters: + if active_adapter not in self.hra_u.keys(): + continue + delta_weight = self.get_delta_weight(active_adapter) + new_weight = torch.mm(new_weight, delta_weight) + + x = x.to(self.base_layer.weight.data.dtype) + + orig_weight = self.base_layer.weight.data + orig_weight = orig_weight.view( + self.out_features, + self.in_features * self.base_layer.kernel_size[0] * self.base_layer.kernel_size[0], + ) + new_weight = torch.mm(orig_weight, new_weight) + new_weight = new_weight.view( + self.out_features, + self.in_features, + self.base_layer.kernel_size[0], + self.base_layer.kernel_size[0], + ) + + result = F.conv2d( + input=x, + weight=new_weight, + bias=self.base_layer.bias, + padding=self.base_layer.padding[0], + stride=self.base_layer.stride[0], + ) + + result = result.to(previous_dtype) + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "hra." + rep diff --git a/lib/python3.12/site-packages/peft/tuners/hra/model.py b/lib/python3.12/site-packages/peft/tuners/hra/model.py new file mode 100644 index 0000000000000000000000000000000000000000..1bfa0f12060542da5801b95cb14268eaaa91105c --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/hra/model.py @@ -0,0 +1,341 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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. + +import warnings +from dataclasses import asdict +from enum import Enum +from typing import List, Optional + +import torch +from torch import nn +from tqdm import tqdm + +from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists +from peft.utils import ( + TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, + ModulesToSaveWrapper, + _get_submodules, +) + +from .config import HRAConfig +from .layer import HRAConv2d, HRALayer, HRALinear + + +class HRAModel(BaseTuner): + """ + Creates Householder reflection adaptation (HRA) model from a pretrained model. The method is described in + https://arxiv.org/abs/2405.17484 + + Args: + model (`torch.nn.Module`): The model to which the adapter tuner layers will be attached. + config ([`HRAConfig`]): The configuration of the HRA model. + adapter_name (`str`): The name of the adapter, defaults to `"default"`. + low_cpu_mem_usage (`bool`, `optional`, defaults to `False`): + Create empty adapter weights on meta device. Useful to speed up the loading process. + + Returns: + `torch.nn.Module`: The HRA model. + + Example: + ```py + >>> from diffusers import StableDiffusionPipeline + >>> from peft import HRAModel, HRAConfig + + >>> config_te = HRAConfig( + ... r=8, + ... target_modules=["k_proj", "q_proj", "v_proj", "out_proj", "fc1", "fc2"], + ... init_weights=True, + ... ) + >>> config_unet = HRAConfig( + ... r=8, + ... target_modules=[ + ... "proj_in", + ... "proj_out", + ... "to_k", + ... "to_q", + ... "to_v", + ... "to_out.0", + ... "ff.net.0.proj", + ... "ff.net.2", + ... ], + ... init_weights=True, + ... ) + + >>> model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5") + >>> model.text_encoder = HRAModel(model.text_encoder, config_te, "default") + >>> model.unet = HRAModel(model.unet, config_unet, "default") + ``` + + **Attributes**: + - **model** ([`~torch.nn.Module`]) -- The model to be adapted. + - **peft_config** ([`HRAConfig`]): The configuration of the HRA model. + """ + + prefix: str = "hra_" + + def _check_new_adapter_config(self, config: HRAConfig) -> None: + """ + A helper method to check the config when a new adapter is being added. + + Raise a ValueError if there is something wrong with the config or if it conflicts with existing adapters. + + """ + # TODO: there should be a check if any of the existing adapters actually has bias != "none", or else the check + # does not fully correspond to the error message. + if (len(self.peft_config) > 1) and (config.bias != "none"): + raise ValueError( + f"{self.__class__.__name__} supports only 1 adapter with bias. When using multiple adapters, " + "set bias to 'none' for all adapters." + ) + + @staticmethod + def _check_target_module_exists(hra_config, key): + return check_target_module_exists(hra_config, key) + + def _create_and_replace( + self, + hra_config, + adapter_name, + target, + target_name, + parent, + current_key, + **optional_kwargs, + ): + if current_key is None: + raise ValueError("Current Key shouldn't be `None`") + + bias = hasattr(target, "bias") and target.bias is not None + kwargs = { + "r": hra_config.r, + "apply_GS": hra_config.apply_GS, + "init_weights": hra_config.init_weights, + } + kwargs["bias"] = bias + + # If it is not a HRALayer, create a new module, else update it with new adapters + if not isinstance(target, HRALayer): + new_module = self._create_new_module(hra_config, adapter_name, target, **kwargs) + if adapter_name not in self.active_adapters: + # adding an additional adapter: it is not automatically trainable + new_module.requires_grad_(False) + self._replace_module(parent, target_name, new_module, target) + else: + target.update_layer( + adapter_name, + r=hra_config.r, + apply_GS=hra_config.apply_GS, + init_weights=hra_config.init_weights, + ) + + def _replace_module(self, parent, child_name, new_module, child): + setattr(parent, child_name, new_module) + # It's not necessary to set requires_grad here, as that is handled by + # _mark_only_adapters_as_trainable + + # child layer wraps the original module, unpack it + if hasattr(child, "base_layer"): + child = child.base_layer + + if not hasattr(new_module, "base_layer"): + new_module.weight = child.weight + if hasattr(child, "bias"): + new_module.bias = child.bias + + if getattr(child, "state", None) is not None: + if hasattr(new_module, "base_layer"): + new_module.base_layer.state = child.state + else: + new_module.state = child.state + new_module.to(child.weight.device) + + meta = torch.device("meta") + # dispatch to correct device + for name, module in new_module.named_modules(): + if self.prefix in name: + if not any(p.device == meta for p in module.parameters()): + module.to(child.weight.device) + + def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None: + for n, p in model.named_parameters(): + if self.prefix not in n: + p.requires_grad = False + + for active_adapter in self.active_adapters: + bias = self.peft_config[active_adapter].bias + if bias == "none": + continue + + if bias == "all": + for n, p in model.named_parameters(): + if "bias" in n: + p.requires_grad = True + elif bias == "hra_only": + for name, m in model.named_modules(): + if isinstance(m, HRALayer) and hasattr(m, "bias") and m.bias is not None: + m.bias.requires_grad = True + else: + raise NotImplementedError(f"Requested bias: {bias}, is not implemented.") + + @staticmethod + def _create_new_module(hra_config, adapter_name, target, **kwargs): + if isinstance(target, BaseTunerLayer): + target_base_layer = target.get_base_layer() + else: + target_base_layer = target + + if isinstance(target_base_layer, torch.nn.Linear): + new_module = HRALinear(target, adapter_name, **kwargs) + elif isinstance(target_base_layer, torch.nn.Conv2d): + new_module = HRAConv2d(target, adapter_name, **kwargs) + else: + raise ValueError( + f"Target module {target} is not supported. " + "Currently, only `torch.nn.Linear` and `torch.nn.Conv2d` are supported." + ) + + return new_module + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + if name == "base_model": + raise + return getattr(self.model, name) + + def get_peft_config_as_dict(self, inference: bool = False): + config_dict = {} + for key, value in self.peft_config.items(): + config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} + if inference: + config["inference_mode"] = True + config_dict[key] = config + return config + + def _set_adapter_layers(self, enabled=True): + for module in self.model.modules(): + if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)): + module.enable_adapters(enabled) + + def enable_adapter_layers(self): + self._set_adapter_layers(enabled=True) + + def disable_adapter_layers(self): + for active_adapter in self.active_adapters: + val = self.peft_config[active_adapter].bias + if val != "none": + msg = ( + f"Careful, disabling adapter layers with bias configured to be '{val}' does not produce the same " + "output as the the base model would without adaption." + ) + warnings.warn(msg) + self._set_adapter_layers(enabled=False) + + def set_adapter(self, adapter_name): + for module in self.model.modules(): + if isinstance(module, HRALayer): + if module.merged: + warnings.warn("Adapter cannot be set when the model is merged. Unmerging the model first.") + module.unmerge() + module.set_adapter(adapter_name) + self.active_adapter = adapter_name + + @staticmethod + def _prepare_adapter_config(peft_config, model_config): + if peft_config.target_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING: + raise ValueError("Please specify `target_modules` in `peft_config`") + peft_config.target_modules = set( + TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]] + ) + return peft_config + + def _unload_and_optionally_merge( + self, + merge=True, + progressbar: bool = False, + safe_merge: bool = False, + adapter_names: Optional[List[str]] = None, + ): + self._unloading_checks(adapter_names) + key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] + desc = "Unloading " + ("and merging " if merge else "") + "model" + for key in tqdm(key_list, disable=not progressbar, desc=desc): + try: + parent, target, target_name = _get_submodules(self.model, key) + except AttributeError: + continue + + if hasattr(target, "base_layer"): + if merge: + target.merge(safe_merge=safe_merge, adapter_names=adapter_names) + self._replace_module(parent, target_name, target.get_base_layer(), target) + elif isinstance(target, ModulesToSaveWrapper): + # save any additional trainable modules part of `modules_to_save` + setattr(parent, target_name, target.modules_to_save[target.active_adapter]) + + return self.model + + def delete_adapter(self, adapter_name: str) -> None: + """ + Deletes an existing adapter. + + Args: + adapter_name (str): Name of the adapter to be deleted. + """ + if adapter_name not in list(self.peft_config.keys()): + raise ValueError(f"Adapter {adapter_name} does not exist") + del self.peft_config[adapter_name] + + key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] + new_adapter = None + for key in key_list: + _, target, _ = _get_submodules(self.model, key) + if isinstance(target, HRALayer): + target.delete_adapter(adapter_name) + if new_adapter is None: + new_adapter = target.active_adapters[:] + + self.active_adapter = new_adapter or [] + + def merge_and_unload( + self, progressbar: bool = False, safe_merge: bool = False, adapter_names: Optional[List[str]] = None + ) -> torch.nn.Module: + r""" + This method merges the HRA layers into the base model. This is needed if someone wants to use the base model as + a standalone model. + + Args: + progressbar (`bool`): + whether to show a progressbar indicating the unload and merge process + safe_merge (`bool`): + whether to activate the safe merging check to check if there is any potential Nan in the adapter + weights + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + + """ + return self._unload_and_optionally_merge( + progressbar=progressbar, safe_merge=safe_merge, adapter_names=adapter_names + ) + + def unload(self) -> torch.nn.Module: + """ + Gets back the base model by removing all the hra modules without merging. This gives back the original base + model. + """ + return self._unload_and_optionally_merge(merge=False) diff --git a/lib/python3.12/site-packages/peft/tuners/ia3/__init__.py b/lib/python3.12/site-packages/peft/tuners/ia3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..21cab4d6d8a6437766f395e90079ce6ecd9e4f26 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/ia3/__init__.py @@ -0,0 +1,39 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 peft.import_utils import is_bnb_4bit_available, is_bnb_available +from peft.utils import register_peft_method + +from .config import IA3Config +from .layer import Conv2d, Conv3d, IA3Layer, Linear +from .model import IA3Model + + +__all__ = ["Conv2d", "Conv3d", "IA3Config", "IA3Layer", "IA3Model", "Linear"] + +register_peft_method(name="ia3", config_cls=IA3Config, model_cls=IA3Model, is_mixed_compatible=True) + + +def __getattr__(name): + if (name == "Linear8bitLt") and is_bnb_available(): + from .bnb import Linear8bitLt + + return Linear8bitLt + + if (name == "Linear4bit") and is_bnb_4bit_available(): + from .bnb import Linear4bit + + return Linear4bit + + raise AttributeError(f"module {__name__} has no attribute {name}") diff --git a/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46d1ab1e8d15f86bdf8c94518a28d1584ca61bd9 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/bnb.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/bnb.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ed9479208a2ff31a0c09feef56de4f95db23950c Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/bnb.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76e169b15652e108ee8ff74dcb102633f3836500 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/layer.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d388155b867e5850f76452f38660a5436df9c79a Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f74b3a6153d35afb0876a0c0df8e5f2f11c1e75 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/ia3/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/ia3/bnb.py b/lib/python3.12/site-packages/peft/tuners/ia3/bnb.py new file mode 100644 index 0000000000000000000000000000000000000000..628e3ce7229528a0b3157da349b2b34153573c51 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/ia3/bnb.py @@ -0,0 +1,129 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 typing import Any + +import torch + +from peft.import_utils import is_bnb_4bit_available, is_bnb_available + +from .layer import IA3Layer + + +if is_bnb_available(): + + class Linear8bitLt(torch.nn.Module, IA3Layer): + # (IA)^3 implemented in a dense layer + def __init__( + self, + base_layer: torch.nn.Module, + adapter_name: str, + is_feedforward: bool, + init_ia3_weights: bool = True, + **kwargs, + ) -> None: + super().__init__() + IA3Layer.__init__(self, base_layer, is_feedforward=is_feedforward) + + # Freezing the pre-trained weight matrix + self.get_base_layer().weight.requires_grad = False + self._active_adapter = adapter_name + self.update_layer(adapter_name, init_ia3_weights) + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + # note: no check for self.merged because merging is not supported (yet) + if self.disable_adapters: + return self.base_layer(x) + + ia3_scaling = 1 + for active_adapter in self.active_adapters: + if active_adapter not in self.ia3_l.keys(): + continue + ia3_scaling *= self.ia3_l[active_adapter].flatten() + + requires_conversion = (not torch.is_autocast_enabled()) and (x.dtype != torch.float32) + if requires_conversion: + x = x.float() + if self.is_feedforward: + result = self.base_layer(x * ia3_scaling) + expected_dtype = result.dtype + else: + result = self.base_layer(x) + expected_dtype = result.dtype + result = result * ia3_scaling + + if requires_conversion: + result = result.to(expected_dtype) + + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "ia3." + rep + + +if is_bnb_4bit_available(): + + class Linear4bit(torch.nn.Module, IA3Layer): + # IA3 implemented in a dense layer + def __init__( + self, + base_layer: torch.nn.Module, + adapter_name: str, + is_feedforward: bool, + init_ia3_weights: bool = True, + **kwargs, + ) -> None: + super().__init__() + IA3Layer.__init__(self, base_layer, is_feedforward=is_feedforward) + + # Freezing the pre-trained weight matrix + self.get_base_layer().weight.requires_grad = False + self._active_adapter = adapter_name + self.update_layer(adapter_name, init_ia3_weights) + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + # note: no check for self.merged because merging is not supported (yet) + if self.disable_adapters: + return self.base_layer(x) + + ia3_scaling = 1 + for active_adapter in self.active_adapters: + if active_adapter not in self.ia3_l.keys(): + continue + ia3_scaling *= self.ia3_l[active_adapter].flatten() + + requires_conversion = (not torch.is_autocast_enabled()) and (x.dtype != torch.float32) + if requires_conversion: + x = x.float() + if self.is_feedforward: + result = self.base_layer(x * ia3_scaling) + expected_dtype = result.dtype + else: + result = self.base_layer(x) + expected_dtype = result.dtype + result = result * ia3_scaling + + result = result.clone() + # adalora.py and lora.py both suggest that this is necessary for 4-bit training on older versions of Pytorch. + # This has been duplicated here. + + if requires_conversion: + result = result.to(expected_dtype) + + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "ia3." + rep diff --git a/lib/python3.12/site-packages/peft/tuners/ia3/config.py b/lib/python3.12/site-packages/peft/tuners/ia3/config.py new file mode 100644 index 0000000000000000000000000000000000000000..1c4161e5ed97998017b800535d4315c66c9d14cd --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/ia3/config.py @@ -0,0 +1,112 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 + +from dataclasses import dataclass, field +from typing import Optional, Union + +from peft.config import PeftConfig +from peft.utils import PeftType + + +@dataclass +class IA3Config(PeftConfig): + """ + This is the configuration class to store the configuration of a [`IA3Model`]. + + Args: + target_modules (`Optional[Union[List[str], str]]`): + The names of the modules to apply the adapter to. If this is specified, only the modules with the specified + names will be replaced. When passing a string, a regex match will be performed. When passing a list of + strings, either an exact match will be performed or it is checked if the name of the module ends with any + of the passed strings. If this is specified as 'all-linear', then all linear/Conv1D modules are chosen, + excluding the output layer. If this is not specified, modules will be chosen according to the model + architecture. If the architecture is not known, an error will be raised -- in this case, you should specify + the target modules manually. + exclude_modules (`Optional[Union[List[str], str]]`): + The names of the modules to not apply the adapter. When passing a string, a regex match will be performed. + When passing a list of strings, either an exact match will be performed or it is checked if the name of the + module ends with any of the passed strings. + feedforward_modules (`Optional[Union[List[str], str]]`): + The names of the modules to be treated as feedforward modules, as in the original paper. These modules will + have (IA)³ vectors multiplied to the input, instead of the output. `feedforward_modules` must be a name or + a subset of names present in `target_modules`. + fan_in_fan_out (`bool`): + Set this to True if the layer to replace stores weight like (fan_in, fan_out). For example, gpt-2 uses + `Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set to `True`. + modules_to_save (`Optional[List[str]]`): + List of modules apart from (IA)³ layers to be set as trainable and saved in the final checkpoint. + init_ia3_weights (`bool`): + Whether to initialize the vectors in the (IA)³ layers, defaults to `True`. Setting this to `False` is + discouraged. + """ + + target_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": ( + "List of module names or regex expression of the module names to replace with (IA)³." + "For example, ['q', 'v'] or '.*decoder.*(SelfAttention|EncDecAttention).*(q|v)$'." + "This can also be a wildcard 'all-linear' which matches all linear/Conv1D layers except the output layer." + "If not specified, modules will be chosen according to the model architecture, If the architecture is " + "not known, an error will be raised -- in this case, you should specify the target modules manually." + ), + }, + ) + exclude_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={"help": "List of module names or regex expression of the module names to exclude from (IA)³."}, + ) + feedforward_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": "List of module names or a regex expression of module names which are feedforward" + "For example, ['output.dense']" + }, + ) + fan_in_fan_out: bool = field( + default=False, + metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"}, + ) + modules_to_save: Optional[list[str]] = field( + default=None, + metadata={ + "help": "List of modules apart from (IA)^3 layers to be set as trainable and saved in the final checkpoint. " + "For example, in Sequence Classification or Token Classification tasks, " + "the final layer `classifier/score` are randomly initialized and as such need to be trainable and saved." + }, + ) + init_ia3_weights: bool = field( + default=True, + metadata={"help": "Whether to initialize the vectors in the (IA)^3 layers."}, + ) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.IA3 + self.target_modules = ( + set(self.target_modules) if isinstance(self.target_modules, list) else self.target_modules + ) + self.exclude_modules = ( + set(self.exclude_modules) if isinstance(self.exclude_modules, list) else self.exclude_modules + ) + self.feedforward_modules = ( + set(self.feedforward_modules) if isinstance(self.feedforward_modules, list) else self.feedforward_modules + ) + + # check if feedforward_modules is a subset of target_modules. run the check only if both are sets + if isinstance(self.feedforward_modules, set) and isinstance(self.target_modules, set): + if not self.feedforward_modules.issubset(self.target_modules): + raise ValueError("`feedforward_modules` should be a subset of `target_modules`") diff --git a/lib/python3.12/site-packages/peft/tuners/ia3/layer.py b/lib/python3.12/site-packages/peft/tuners/ia3/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..e941089315e3f01f259fa8366791120db9a9daba --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/ia3/layer.py @@ -0,0 +1,327 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +import warnings +from typing import Any, List, Optional + +import torch +import torch.nn as nn +from transformers.pytorch_utils import Conv1D + +from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge +from peft.utils import transpose + + +class IA3Layer(BaseTunerLayer): + # All names of layers that may contain adapter weights + adapter_layer_names = ("ia3_l",) + + def __init__(self, base_layer: nn.Module, is_feedforward: bool, **kwargs) -> None: + self.base_layer = base_layer + self.ia3_l = nn.ParameterDict({}) + # Mark the weight as unmerged + self._disable_adapters = False + self.merged_adapters = [] + self.is_feedforward = is_feedforward + + base_layer = self.get_base_layer() + if isinstance(base_layer, nn.Linear): + in_features, out_features = base_layer.in_features, base_layer.out_features + elif isinstance(base_layer, (nn.Conv2d, nn.Conv3d)): + in_features, out_features = base_layer.in_channels, base_layer.out_channels + elif isinstance(base_layer, nn.Embedding): + in_features, out_features = base_layer.num_embeddings, base_layer.embedding_dim + elif isinstance(base_layer, Conv1D): + in_features, out_features = ( + base_layer.weight.ds_shape if hasattr(base_layer.weight, "ds_shape") else base_layer.weight.shape + ) + else: + raise ValueError(f"Unsupported layer type {type(base_layer)}") + self.in_features = in_features + self.out_features = out_features + + def update_layer(self, adapter_name, init_ia3_weights): + # This code works for linear layers, override for other layer types + # Actual trainable parameters + if self.is_feedforward: + weight = torch.randn((1, self.in_features)) + else: + weight = torch.randn((self.out_features, 1)) + self.ia3_l[adapter_name] = nn.Parameter(weight) + if init_ia3_weights: + self.reset_ia3_parameters(adapter_name) + self._move_adapter_to_device_of_base_layer(adapter_name) + self.set_adapter(self.active_adapters) + + def reset_ia3_parameters(self, adapter_name): + if adapter_name in self.ia3_l.keys(): + # initialize learned vector with torch.ones + nn.init.constant_(self.ia3_l[adapter_name], 1.0) + + +class Linear(nn.Module, IA3Layer): + # (IA)^3 implemented in a dense layer + def __init__( + self, + base_layer: nn.Module, + adapter_name: str, + fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) + is_feedforward: bool = False, # Set to True if the layer is treated as a feedforward layer + is_target_conv_1d_layer: bool = False, # whether target module is a conv1d layer. useful while unloading later + init_ia3_weights: bool = True, # whether to initialize IA3 weights + **kwargs, + ) -> None: + super().__init__() + IA3Layer.__init__(self, base_layer, is_feedforward=is_feedforward) + self.fan_in_fan_out = fan_in_fan_out + self.is_target_conv_1d_layer = is_target_conv_1d_layer + self._active_adapter = adapter_name + self.update_layer(adapter_name, init_ia3_weights) + + def merge(self, safe_merge: bool = False, adapter_names: Optional[List[str]] = None) -> None: + """ + Merge the active adapter weights into the base weights + + Args: + safe_merge (`bool`, *optional*): + If True, the merge operation will be performed in a copy of the original weights and check for NaNs + before merging the weights. This is useful if you want to check if the merge operation will produce + NaNs. Defaults to `False`. + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + """ + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + # no adapter to merge + return + + for active_adapter in adapter_names: + if active_adapter in self.ia3_l.keys(): + base_layer = self.get_base_layer() + ia3_l = transpose(self.ia3_l[active_adapter].data, self.fan_in_fan_out) + orig_dtype = base_layer.weight.data.dtype + if safe_merge: + orig_weights = base_layer.weight.data + orig_weights = torch.mul(orig_weights, ia3_l) + + if not torch.isfinite(orig_weights).all(): + raise ValueError( + f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" + ) + base_layer.weight.data = orig_weights.to(orig_dtype) + else: + base_layer.weight.data = torch.mul(base_layer.weight.data, ia3_l).to(orig_dtype) + + if not self.is_feedforward and (base_layer.bias is not None): + scaling = self.ia3_l[active_adapter].reshape(base_layer.bias.shape) + orig_dtype = base_layer.bias.data.dtype + base_layer.bias.data = torch.mul(base_layer.bias.data, scaling.data).to(orig_dtype) + + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + """ + This method unmerges all merged adapter layers from the base weights. + """ + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + + warnings.warn("Unmerge result can be inaccurate for (IA)^3.") + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter in self.ia3_l.keys(): + base_layer = self.get_base_layer() + # Add tolerace to avoid division by zero + ia3_l = transpose(self.ia3_l[active_adapter].data, self.fan_in_fan_out) + 1e-8 + orig_dtype = base_layer.weight.data.dtype + base_layer.weight.data = torch.div(base_layer.weight.data, ia3_l).to(orig_dtype) + + if not self.is_feedforward and (base_layer.bias is not None): + scaling = self.ia3_l[active_adapter].reshape(base_layer.bias.shape) + orig_dtype = base_layer.bias.data.dtype + base_layer.bias.data = torch.div(base_layer.bias.data, scaling.data + 1e-8).to(orig_dtype) + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + dtype = previous_dtype = x.dtype + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + ia3_scaling = 1 + for active_adapter in self.active_adapters: + if active_adapter not in self.ia3_l.keys(): + continue + dtype = self.ia3_l[active_adapter].dtype + ia3_scaling *= self.ia3_l[active_adapter].flatten() + + if self.is_feedforward: + x = x.to(dtype) + # TODO: weight.dtype can be != self.ia3_l[self.active_adapters].dtype + # e.g. bf16 vs fp32. Is that okay? + interm = (x * ia3_scaling).to(previous_dtype) + result = self.base_layer(interm, *args, **kwargs) + else: + result = self.base_layer(x, *args, **kwargs) + result_dtype = result.dtype + result = (result * ia3_scaling).to(result_dtype) + + return result + + +class _ConvNd(nn.Module, IA3Layer): + def __init__( + self, + base_layer: nn.Module, + adapter_name: str, + fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) + is_feedforward: bool = False, # Set to True if the layer is treated as a feedforward layer + init_ia3_weights: bool = True, + **kwargs, + ) -> None: + super().__init__() + IA3Layer.__init__(self, base_layer, is_feedforward=is_feedforward) + self.fan_in_fan_out = fan_in_fan_out + self._active_adapter = adapter_name + self._kernel_dim = base_layer.weight.dim() + + self.update_layer(adapter_name, init_ia3_weights) + + def update_layer(self, adapter_name, init_ia3_weights): + # Actual trainable parameters + num_features = self.in_features if self.is_feedforward else self.out_features + weights_size = (1, num_features) + (1,) * (self._kernel_dim - 2) + weight = torch.randn(weights_size) + self.ia3_l[adapter_name] = nn.Parameter(weight) + if init_ia3_weights: + self.reset_ia3_parameters(adapter_name) + self._move_adapter_to_device_of_base_layer(adapter_name) + self.set_adapter(self.active_adapters) + + def merge(self, safe_merge: bool = False, adapter_names: Optional[List[str]] = None) -> None: + """ + Merge the active adapter weights into the base weights + + Args: + safe_merge (`bool`, *optional*): + If True, the merge operation will be performed in a copy of the original weights and check for NaNs + before merging the weights. This is useful if you want to check if the merge operation will produce + NaNs. Defaults to `False`. + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + """ + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + # no adapter to merge + return + + for active_adapter in adapter_names: + if active_adapter in self.ia3_l.keys(): + base_layer = self.get_base_layer() + ia3_scaling = self.ia3_l[active_adapter].data + if not self.is_feedforward: + ia3_scaling = ia3_scaling.transpose(0, 1) + + if safe_merge: + output_weight = torch.mul(base_layer.weight.data, ia3_scaling).clone() + + if not torch.isfinite(output_weight).all(): + raise ValueError( + f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" + ) + + base_layer.weight.data = output_weight + else: + base_layer.weight.data = torch.mul(base_layer.weight.data, ia3_scaling) + + if not self.is_feedforward and (base_layer.bias is not None): + scaling = self.ia3_l[active_adapter].reshape(base_layer.bias.shape) + base_layer.bias.data = torch.mul(base_layer.bias.data, scaling.data) + + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + """ + This method unmerges all merged adapter layers from the base weights. + """ + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + + warnings.warn("Unmerge result can be inaccurate for (IA)^3.") + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter in self.ia3_l.keys(): + base_layer = self.get_base_layer() + # divide by (IA)^3 vector. Add tolerace to avoid division by zero + ia3_scaling = self.ia3_l[active_adapter].data + if not self.is_feedforward: + ia3_scaling = ia3_scaling.transpose(0, 1) + base_layer.weight.data = torch.div(base_layer.weight.data, ia3_scaling + 1e-8) + + if not self.is_feedforward and (base_layer.bias is not None): + scaling = self.ia3_l[active_adapter].reshape(base_layer.bias.shape) + base_layer.bias.data = torch.mul(base_layer.bias.data, scaling.data) + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + dtype = previous_dtype = x.dtype + + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + ia3_scaling = 1 + for active_adapter in self.active_adapters: + if active_adapter not in self.ia3_l.keys(): + continue + dtype = self.ia3_l[active_adapter].dtype + ia3_scaling *= self.ia3_l[active_adapter] + + if self.is_feedforward: + x = x.to(dtype) + # TODO: weight.dtype can be != self.ia3_l[self.active_adapters].dtype + # e.g. bf16 vs fp32. Is that okay? + interm = (x * ia3_scaling).to(self.get_base_layer().weight.dtype) + result = self.base_layer(interm, *args, **kwargs) + else: + result = self.base_layer(x, *args, **kwargs) + result = result.to(dtype) * ia3_scaling + + result = result.to(previous_dtype) + return result + + +class Conv2d(_ConvNd): + # IA3 implemented in a 2D convolutional layer + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if not self._kernel_dim == 4: + raise ValueError(f"Conv2d layer kernel must have 4 dimensions, not {self._kernel_dim}") + + +class Conv3d(_ConvNd): + # IA3 implemented in a 3D convolutional layer + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if not self._kernel_dim == 5: + raise ValueError(f"Conv2d layer kernel must have 5 dimensions, not {self._kernel_dim}") diff --git a/lib/python3.12/site-packages/peft/tuners/ia3/model.py b/lib/python3.12/site-packages/peft/tuners/ia3/model.py new file mode 100644 index 0000000000000000000000000000000000000000..453b62f3a19d51526465869ecedecc05800ff40b --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/ia3/model.py @@ -0,0 +1,496 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 re +import warnings +from dataclasses import asdict, replace +from enum import Enum +from typing import Optional + +import torch +from torch import nn +from transformers.pytorch_utils import Conv1D + +from peft.import_utils import is_bnb_4bit_available, is_bnb_available +from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists +from peft.utils import ( + TRANSFORMERS_MODELS_TO_IA3_FEEDFORWARD_MODULES_MAPPING, + TRANSFORMERS_MODELS_TO_IA3_TARGET_MODULES_MAPPING, + ModulesToSaveWrapper, + _freeze_adapter, + _get_submodules, +) + +from .layer import Conv2d, Conv3d, IA3Layer, Linear + + +class IA3Model(BaseTuner): + """ + Creates a Infused Adapter by Inhibiting and Amplifying Inner Activations ((IA)^3) model from a pretrained + transformers model. The method is described in detail in https://arxiv.org/abs/2205.05638 + + Args: + model ([`~transformers.PreTrainedModel`]): The model to be adapted. + config ([`IA3Config`]): The configuration of the (IA)^3 model. + adapter_name (`str`): The name of the adapter, defaults to `"default"`. + low_cpu_mem_usage (`bool`, `optional`, defaults to `False`): + Create empty adapter weights on meta device. Useful to speed up the loading process. + + Returns: + `torch.nn.Module`: The (IA)^3 model. + + Example: + + ```py + >>> from transformers import AutoModelForSeq2SeqLM, ia3Config + >>> from peft import IA3Model, IA3Config + + >>> config = IA3Config( + ... peft_type="IA3", + ... task_type="SEQ_2_SEQ_LM", + ... target_modules=["k", "v", "w0"], + ... feedforward_modules=["w0"], + ... ) + + >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") + >>> ia3_model = IA3Model(config, model) + ``` + + **Attributes**: + - **model** ([`~transformers.PreTrainedModel`]) -- The model to be adapted. + - **peft_config** ([`ia3Config`]): The configuration of the (IA)^3 model. + """ + + prefix: str = "ia3_" + + def __init__(self, model, config, adapter_name, low_cpu_mem_usage: bool = False): + super().__init__(model, config, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage) + + @staticmethod + def _create_new_module(ia3_config, adapter_name, target, **kwargs): + # avoid eager bnb import + if is_bnb_available(): + import bitsandbytes as bnb + + from .bnb import Linear8bitLt + + if is_bnb_4bit_available(): + from .bnb import Linear4bit + + loaded_in_8bit = kwargs.pop("loaded_in_8bit", False) + loaded_in_4bit = kwargs.pop("loaded_in_4bit", False) + is_feedforward = kwargs.pop("is_feedforward", False) + + if isinstance(target, BaseTunerLayer): + target_base_layer = target.get_base_layer() + else: + target_base_layer = target + + if loaded_in_8bit and isinstance(target_base_layer, bnb.nn.Linear8bitLt): + eightbit_kwargs = kwargs.copy() + eightbit_kwargs.update( + { + "has_fp16_weights": target_base_layer.state.has_fp16_weights, + "threshold": target_base_layer.state.threshold, + "index": target_base_layer.index, + } + ) + new_module = Linear8bitLt(target, adapter_name, is_feedforward=is_feedforward, **eightbit_kwargs) + elif loaded_in_4bit and isinstance(target_base_layer, bnb.nn.Linear4bit): + fourbit_kwargs = kwargs.copy() + fourbit_kwargs.update( + { + "compute_dtype": target_base_layer.compute_dtype, + "compress_statistics": target_base_layer.weight.compress_statistics, + "quant_type": target_base_layer.weight.quant_type, + } + ) + new_module = Linear4bit(target, adapter_name, is_feedforward=is_feedforward, **fourbit_kwargs) + elif isinstance(target, torch.nn.Conv2d): + new_module = Conv2d(target, adapter_name, is_feedforward=is_feedforward, **kwargs) + elif isinstance(target, torch.nn.Conv3d): + new_module = Conv3d(target, adapter_name, is_feedforward=is_feedforward, **kwargs) + elif isinstance(target_base_layer, torch.nn.Linear): + if kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. " + "Setting fan_in_fan_out to False." + ) + kwargs["fan_in_fan_out"] = ia3_config.fan_in_fan_out = False + new_module = Linear(target, adapter_name, is_feedforward=is_feedforward, **kwargs) + elif isinstance(target_base_layer, Conv1D): + if not kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to False but the target module is `Conv1D`. Setting fan_in_fan_out to True." + ) + kwargs["fan_in_fan_out"] = ia3_config.fan_in_fan_out = True + new_module = Linear( + target, adapter_name, is_feedforward=is_feedforward, is_target_conv_1d_layer=True, **kwargs + ) + else: + raise ValueError( + f"Target module {target} is not supported. " + f"Currently, only `torch.nn.Linear`, `torch.nn.Conv2d`, and `Conv1D` are supported." + ) + return new_module + + @staticmethod + def _check_target_module_exists(ia3_config, key): + return check_target_module_exists(ia3_config, key) + + def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None: + for n, p in model.named_parameters(): + if self.prefix not in n: + p.requires_grad = False + + def _create_and_replace( + self, + ia3_config, + adapter_name, + target, + target_name, + parent, + current_key, + ): + # check if target module is in feedforward_modules + is_feedforward = self._check_target_module_feedforward(ia3_config, current_key) + + kwargs = { + "fan_in_fan_out": ia3_config.fan_in_fan_out, + "init_ia3_weights": ia3_config.init_ia3_weights, + "is_feedforward": is_feedforward, + "loaded_in_8bit": getattr(self.model, "is_loaded_in_8bit", False), + "loaded_in_4bit": getattr(self.model, "is_loaded_in_4bit", False), + } + + if isinstance(target, IA3Layer): + target.update_layer( + adapter_name, + ia3_config.init_ia3_weights, + ) + else: + new_module = self._create_new_module(ia3_config, adapter_name, target, **kwargs) + if adapter_name not in self.active_adapters: + # adding an additional adapter: it is not automatically trainable + new_module.requires_grad_(False) + self._replace_module(parent, target_name, new_module, target) + + @staticmethod + def _check_target_module_feedforward(ia3_config, key) -> bool: + """ + A helper private method that checks if the target module `key` matches with a feedforward module specified in + `ia3_config` + """ + if isinstance(ia3_config.feedforward_modules, str): + is_feedforward = bool(re.fullmatch(ia3_config.feedforward_modules, key)) + else: + is_feedforward = any(key.endswith(target_key) for target_key in ia3_config.feedforward_modules) + return is_feedforward + + def _replace_module(self, parent, child_name, new_module, child): + setattr(parent, child_name, new_module) + + # child layer wraps the original module, unpack it + if hasattr(child, "base_layer"): + child = child.base_layer + + # layers with base_layer don't need the weight to be copied, as they have a reference already + if not hasattr(new_module, "base_layer"): + new_module.weight = child.weight + if hasattr(child, "bias"): + new_module.bias = child.bias + + if getattr(child, "state", None) is not None: + if hasattr(new_module, "base_layer"): + new_module.base_layer.state = child.state + else: + new_module.state = child.state + new_module.to(child.weight.device) + + meta = torch.device("meta") + # dispatch to correct device + for name, module in new_module.named_modules(): + if self.prefix in name: + if not any(p.device == meta for p in module.parameters()): + module.to(child.weight.device) + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + if name == "model": # see #1892: prevent infinite recursion if class is not initialized + raise + return getattr(self.model, name) + + def get_peft_config_as_dict(self, inference: bool = False): + config_dict = {} + for key, value in self.peft_config.items(): + config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} + if inference: + config["inference_mode"] = True + config_dict[key] = config + return config + + def _set_adapter_layers(self, enabled=True): + for module in self.model.modules(): + if isinstance(module, (IA3Layer, ModulesToSaveWrapper)): + module.enable_adapters(enabled) + + def enable_adapter_layers(self) -> None: + """Enable all adapters. + + Call this if you have previously disabled all adapters and want to re-enable them. + """ + self._set_adapter_layers(enabled=True) + + def disable_adapter_layers(self) -> None: + """Disable all adapters. + + When disabling all adapters, the model output corresponds to the output of the base model. + """ + self._set_adapter_layers(enabled=False) + + def set_adapter(self, adapter_name: str | list[str]) -> None: + """Set the active adapter(s). + + Additionally, this function will set the specified adapters to trainable (i.e., requires_grad=True). If this is + not desired, use the following code. + + ```py + >>> for name, param in model_peft.named_parameters(): + ... if ...: # some check on name (ex. if 'lora' in name) + ... param.requires_grad = False + ``` + + Args: + adapter_name (`str` or `list[str]`): Name of the adapter(s) to be activated. + """ + for module in self.model.modules(): + if isinstance(module, IA3Layer): + if module.merged: + warnings.warn("Adapter cannot be set when the model is merged. Unmerging the model first.") + module.unmerge() + module.set_adapter(adapter_name) + self.active_adapter = adapter_name + + @staticmethod + def _prepare_adapter_config(peft_config, model_config): + if peft_config.target_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_IA3_TARGET_MODULES_MAPPING: + raise ValueError("Please specify `target_modules` in `peft_config`") + peft_config.target_modules = set( + TRANSFORMERS_MODELS_TO_IA3_TARGET_MODULES_MAPPING[model_config["model_type"]] + ) + if peft_config.feedforward_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_IA3_FEEDFORWARD_MODULES_MAPPING: + raise ValueError("Please specify `feedforward_modules` in `peft_config`") + peft_config.feedforward_modules = set( + TRANSFORMERS_MODELS_TO_IA3_FEEDFORWARD_MODULES_MAPPING[model_config["model_type"]] + ) + return peft_config + + def _unload_and_optionally_merge( + self, merge: bool = True, safe_merge: bool = False, adapter_names: Optional[list[str]] = None + ): + r""" + This method merges the (IA)^3 layers into the base model. This is needed if someone wants to use the base model + as a standalone model. + + Args: + safe_merge (`bool`, `optional`, defaults to `False`): + If True, the merge operation will be performed in a copy of the original weights and check for NaNs + before merging the weights. This is useful if you want to check if the merge operation will produce + NaNs. Defaults to `False`. + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + """ + if getattr(self.model, "is_loaded_in_8bit", False): + raise ValueError("Cannot merge ia3 layers when the model is loaded in 8-bit mode") + + if getattr(self.model, "is_loaded_in_4bit", False): + raise ValueError("Cannot merge ia3 layers when the model is loaded in 4-bit mode") + + self._unloading_checks(adapter_names) + key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] + for key in key_list: + try: + parent, target, target_name = _get_submodules(self.model, key) + except AttributeError: + continue + + if hasattr(target, "base_layer"): + if merge: + target.merge(safe_merge=safe_merge, adapter_names=adapter_names) + self._replace_module(parent, target_name, target.get_base_layer(), target) + elif isinstance(target, ModulesToSaveWrapper): + # save any additional trainable modules part of `modules_to_save` + new_module = target.modules_to_save[target.active_adapter] + if hasattr(new_module, "base_layer"): + # check if the module is itself a tuner layer + if merge: + new_module.merge(safe_merge=safe_merge, adapter_names=adapter_names) + new_module = new_module.get_base_layer() + setattr(parent, target_name, new_module) + + return self.model + + def merge_and_unload(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> torch.nn.Module: + r""" + This method merges the IA³ layers into the base model. This is needed if someone wants to use the base model as + a standalone model. + + Args: + safe_merge (`bool`): + whether to activate the safe merging check to check if there is any potential Nan in the adapter + weights + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + + Example: + + ```py + >>> from transformers import AutoModelForCausalLM + >>> from peft import PeftModel + + >>> base_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-40b") + >>> peft_model_id = "smangrul/falcon-40B-int4-peft-lora-sfttrainer-sample" + >>> model = PeftModel.from_pretrained(base_model, peft_model_id) + >>> merged_model = model.merge_and_unload() + ``` + """ + return self._unload_and_optionally_merge(safe_merge=safe_merge, adapter_names=adapter_names) + + def unload(self) -> torch.nn.Module: + """ + Gets back the base model by removing all the IA³ modules without merging. This gives back the original base + model. + """ + return self._unload_and_optionally_merge(merge=False) + + def delete_adapter(self, adapter_name: str) -> None: + """ + Deletes an existing adapter. + + Args: + adapter_name (str): Name of the adapter to be deleted. + """ + if adapter_name not in self.peft_config: + raise ValueError(f"Adapter {adapter_name} does not exist") + del self.peft_config[adapter_name] + + key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] + new_adapter = None + for key in key_list: + _, target, _ = _get_submodules(self.model, key) + if isinstance(target, IA3Layer): + target.delete_adapter(adapter_name) + if new_adapter is None: + new_adapter = target.active_adapters[:] + + self.active_adapter = new_adapter or [] + + def _check_add_weighted_adapter(self, adapters: list[str]) -> tuple[str, str]: + """ + Helper function to check if the arguments to add_weighted_adapter are valid and compatible with the underlying + model. + """ + # Validate existence of adapters + for adapter in adapters: + if adapter not in self.peft_config: + raise ValueError(f"Adapter {adapter} does not exist") + + # Check for conflicting modules_to_save + modules_to_save_wrappers = [module for module in self.modules() if isinstance(module, ModulesToSaveWrapper)] + if any( + sum(adapter in wrapper.modules_to_save for adapter in adapters) > 1 for wrapper in modules_to_save_wrappers + ): + raise ValueError("Cannot add weighted adapters targeting the same module with modules_to_save.") + + # Ensure all adapters have compatible target and feedforward module types + target_module_types = {type(self.peft_config[adapter].target_modules) for adapter in adapters} + feedforward_module_types = {type(self.peft_config[adapter].feedforward_modules) for adapter in adapters} + if len(target_module_types) > 1 or len(feedforward_module_types) > 1: + raise ValueError("All adapter configs should have the same type for target and feedforward modules.") + + # Combine target and feedforward modules + if str in target_module_types: + new_target_modules = "|".join(f"({self.peft_config[adapter].target_modules})" for adapter in adapters) + else: + new_target_modules = set.union(*(self.peft_config[adapter].target_modules for adapter in adapters)) + + if str in feedforward_module_types: + new_feedforward_modules = "|".join( + f"({self.peft_config[adapter].feedforward_modules})" for adapter in adapters + ) + else: + new_feedforward_modules = set.union( + *(self.peft_config[adapter].feedforward_modules for adapter in adapters) + ) + + return new_target_modules, new_feedforward_modules + + def add_weighted_adapter( + self, + adapters: list[str], + weights: list[float], + adapter_name: str, + ) -> None: + """ + This method adds a new adapter by merging the given adapters with the given weights. + + Args: + adapters (`list`): + List of adapter names to be merged. + weights (`list`): + List of weights for each adapter. + adapter_name (`str`): + Name of the new adapter. + """ + if adapter_name in list(self.peft_config.keys()): + return + + new_target_modules, new_feedforward_modules = self._check_add_weighted_adapter( + adapters=adapters, + ) + + self.peft_config[adapter_name] = replace( + self.peft_config[adapters[0]], + target_modules=new_target_modules, + feedforward_modules=new_feedforward_modules, + ) + self.inject_adapter(self.model, adapter_name) + + # Do we really need that? + _freeze_adapter(self.model, adapter_name) + + key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] + for key in key_list: + _, target, _ = _get_submodules(self.model, key) + if isinstance(target, IA3Layer): + if adapter_name in target.ia3_l: + target_ia3_l = target.ia3_l[adapter_name] + else: + continue + + target_ia3_l.data = target_ia3_l.data.zero_() + for adapter, weight in zip(adapters, weights): + if adapter in target.ia3_l: + current_adapter_ia3_l = target.ia3_l[adapter] + else: + continue + target_ia3_l.data += current_adapter_ia3_l.data * weight diff --git a/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__init__.py b/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fe692a93378254601753cc9b17218fdd5a92b8de --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__init__.py @@ -0,0 +1,25 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 peft.utils import register_peft_method + +from .config import MultitaskPromptTuningConfig, MultitaskPromptTuningInit +from .model import MultitaskPromptEmbedding + + +__all__ = ["MultitaskPromptEmbedding", "MultitaskPromptTuningConfig", "MultitaskPromptTuningInit"] + +register_peft_method( + name="multitask_prompt_tuning", config_cls=MultitaskPromptTuningConfig, model_cls=MultitaskPromptEmbedding +) diff --git a/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e66c5d4da5a830d1ebd9ce66fac997142d4903f7 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..53c176bfb7798bb05f8cc86fbbc0f66ce579a776 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a099569b5dd69c60d247edd18242996a63e8c5b Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/config.py b/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/config.py new file mode 100644 index 0000000000000000000000000000000000000000..6cb279573559e4fa33a973aa35e052647a6ab8f9 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/multitask_prompt_tuning/config.py @@ -0,0 +1,62 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +import enum +from dataclasses import dataclass, field +from typing import Optional, Union + +from peft.tuners.prompt_tuning import PromptTuningConfig +from peft.utils import PeftType + + +class MultitaskPromptTuningInit(str, enum.Enum): + # initialize prompt with text + TEXT = "TEXT" + # initialize prompt with random matrix + RANDOM = "RANDOM" + # average the prefix and column matrices obtained during source training + AVERAGE_SOURCE_TASKS = "AVERAGE_SOURCE_TASKS" + # pick prefix and column matrices for a particular task obtained during source training + EXACT_SOURCE_TASK = "EXACT_SOURCE_TASK" + # only use the prompt embeddings trained during source training + ONLY_SOURCE_SHARED = "ONLY_SOURCE_SHARED" + + +@dataclass +class MultitaskPromptTuningConfig(PromptTuningConfig): + prompt_tuning_init: Union[MultitaskPromptTuningInit, str] = field( + default=MultitaskPromptTuningInit.RANDOM, + metadata={ + "help": ( + "How to initialize the prompt tuning parameters. Can be one of TEXT, RANDOM, AVERAGE_SOURCE_TASKS, " + "EXACT_SOURCE_TASK, ONLY_SOURCE_SHARED." + ), + }, + ) + prompt_tuning_init_state_dict_path: Optional[str] = field( + default=None, + metadata={ + "help": ( + "The path of source state dict. This is required when training the downstream target prompt from " + "the pretrained source prompt" + ), + }, + ) + prompt_tuning_init_task: Optional[int] = field(default=0, metadata={"help": "source task id for initialization"}) + num_ranks: Optional[int] = field(default=1, metadata={"help": "ranks"}) + num_tasks: Optional[int] = field(default=1, metadata={"help": "number of tasks"}) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.MULTITASK_PROMPT_TUNING diff --git a/lib/python3.12/site-packages/peft/tuners/p_tuning/__init__.py b/lib/python3.12/site-packages/peft/tuners/p_tuning/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9195c0d75d3d11e752d0477b64edd79599bdaa01 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/p_tuning/__init__.py @@ -0,0 +1,23 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 peft.utils import register_peft_method + +from .config import PromptEncoderConfig, PromptEncoderReparameterizationType +from .model import PromptEncoder + + +__all__ = ["PromptEncoder", "PromptEncoderConfig", "PromptEncoderReparameterizationType"] + +register_peft_method(name="p_tuning", config_cls=PromptEncoderConfig, model_cls=PromptEncoder) diff --git a/lib/python3.12/site-packages/peft/tuners/p_tuning/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/p_tuning/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..796f4f6f583cb13f748313b128f61e55c4f254f3 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/p_tuning/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/p_tuning/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/p_tuning/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2727d386b876009f4ecebaa513d384ee93a1eef1 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/p_tuning/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/p_tuning/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/p_tuning/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1302928c3c765fba2b0de9a2501c5a1d5d1547f2 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/p_tuning/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/p_tuning/config.py b/lib/python3.12/site-packages/peft/tuners/p_tuning/config.py new file mode 100644 index 0000000000000000000000000000000000000000..a69c13db9c8a0f57a7daa7d312472625251fb6c8 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/p_tuning/config.py @@ -0,0 +1,60 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +import enum +from dataclasses import dataclass, field +from typing import Union + +from peft.config import PromptLearningConfig +from peft.utils import PeftType + + +class PromptEncoderReparameterizationType(str, enum.Enum): + MLP = "MLP" + LSTM = "LSTM" + + +@dataclass +class PromptEncoderConfig(PromptLearningConfig): + """ + This is the configuration class to store the configuration of a [`PromptEncoder`]. + + Args: + encoder_reparameterization_type (Union[[`PromptEncoderReparameterizationType`], `str`]): + The type of reparameterization to use. + encoder_hidden_size (`int`): The hidden size of the prompt encoder. + encoder_num_layers (`int`): The number of layers of the prompt encoder. + encoder_dropout (`float`): The dropout probability of the prompt encoder. + """ + + encoder_reparameterization_type: Union[str, PromptEncoderReparameterizationType] = field( + default=PromptEncoderReparameterizationType.MLP, + metadata={"help": "How to reparameterize the prompt encoder"}, + ) + encoder_hidden_size: int = field( + default=None, + metadata={"help": "The hidden size of the prompt encoder"}, + ) + encoder_num_layers: int = field( + default=2, + metadata={"help": "The number of layers of the prompt encoder"}, + ) + encoder_dropout: float = field( + default=0.0, + metadata={"help": "The dropout of the prompt encoder"}, + ) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.P_TUNING diff --git a/lib/python3.12/site-packages/peft/tuners/p_tuning/model.py b/lib/python3.12/site-packages/peft/tuners/p_tuning/model.py new file mode 100644 index 0000000000000000000000000000000000000000..ade2b1128158376c134441687803b85d444cfb96 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/p_tuning/model.py @@ -0,0 +1,130 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +# Based on https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/prompt_encoder.py +# with some refactor +import warnings + +import torch + +from .config import PromptEncoderConfig, PromptEncoderReparameterizationType + + +class PromptEncoder(torch.nn.Module): + """ + The prompt encoder network that is used to generate the virtual token embeddings for p-tuning. + + Args: + config ([`PromptEncoderConfig`]): The configuration of the prompt encoder. + + Example: + + ```py + >>> from peft import PromptEncoder, PromptEncoderConfig + + >>> config = PromptEncoderConfig( + ... peft_type="P_TUNING", + ... task_type="SEQ_2_SEQ_LM", + ... num_virtual_tokens=20, + ... token_dim=768, + ... num_transformer_submodules=1, + ... num_attention_heads=12, + ... num_layers=12, + ... encoder_reparameterization_type="MLP", + ... encoder_hidden_size=768, + ... ) + + >>> prompt_encoder = PromptEncoder(config) + ``` + + **Attributes**: + - **embedding** (`torch.nn.Embedding`) -- The embedding layer of the prompt encoder. + - **mlp_head** (`torch.nn.Sequential`) -- The MLP head of the prompt encoder if `inference_mode=False`. + - **lstm_head** (`torch.nn.LSTM`) -- The LSTM head of the prompt encoder if `inference_mode=False` and + `encoder_reparameterization_type="LSTM"`. + - **token_dim** (`int`) -- The hidden embedding dimension of the base transformer model. + - **input_size** (`int`) -- The input size of the prompt encoder. + - **output_size** (`int`) -- The output size of the prompt encoder. + - **hidden_size** (`int`) -- The hidden size of the prompt encoder. + - **total_virtual_tokens** (`int`): The total number of virtual tokens of the + prompt encoder. + - **encoder_type** (Union[[`PromptEncoderReparameterizationType`], `str`]): The encoder type of the prompt + encoder. + + + Input shape: (`batch_size`, `total_virtual_tokens`) + + Output shape: (`batch_size`, `total_virtual_tokens`, `token_dim`) + """ + + def __init__(self, config): + super().__init__() + self.token_dim = config.token_dim + self.input_size = self.token_dim + self.output_size = self.token_dim + self.hidden_size = config.encoder_hidden_size + self.total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules + self.encoder_type = config.encoder_reparameterization_type + + # embedding + self.embedding = torch.nn.Embedding(self.total_virtual_tokens, self.token_dim) + if not config.inference_mode: + if self.encoder_type == PromptEncoderReparameterizationType.LSTM: + lstm_dropout = config.encoder_dropout + num_layers = config.encoder_num_layers + # LSTM + self.lstm_head = torch.nn.LSTM( + input_size=self.input_size, + hidden_size=self.hidden_size, + num_layers=num_layers, + dropout=lstm_dropout, + bidirectional=True, + batch_first=True, + ) + + self.mlp_head = torch.nn.Sequential( + torch.nn.Linear(self.hidden_size * 2, self.hidden_size * 2), + torch.nn.ReLU(), + torch.nn.Linear(self.hidden_size * 2, self.output_size), + ) + + elif self.encoder_type == PromptEncoderReparameterizationType.MLP: + encoder_num_layers_default = PromptEncoderConfig.encoder_num_layers + if config.encoder_num_layers != encoder_num_layers_default: + warnings.warn( + f"for {self.encoder_type.value}, the argument `encoder_num_layers` is ignored. " + f"Exactly {encoder_num_layers_default} MLP layers are used." + ) + layers = [ + torch.nn.Linear(self.input_size, self.hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(self.hidden_size, self.hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(self.hidden_size, self.output_size), + ] + self.mlp_head = torch.nn.Sequential(*layers) + + else: + raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.") + + def forward(self, indices): + input_embeds = self.embedding(indices) + if self.encoder_type == PromptEncoderReparameterizationType.LSTM: + output_embeds = self.mlp_head(self.lstm_head(input_embeds)[0]) + elif self.encoder_type == PromptEncoderReparameterizationType.MLP: + output_embeds = self.mlp_head(input_embeds) + else: + raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.") + + return output_embeds diff --git a/lib/python3.12/site-packages/peft/tuners/poly/__init__.py b/lib/python3.12/site-packages/peft/tuners/poly/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1c18933eba3fa44106ba9fa89ba34ecd12a2bed4 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/poly/__init__.py @@ -0,0 +1,24 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 peft.utils import register_peft_method + +from .config import PolyConfig +from .layer import Linear, PolyLayer +from .model import PolyModel + + +__all__ = ["Linear", "PolyConfig", "PolyLayer", "PolyModel"] + +register_peft_method(name="poly", config_cls=PolyConfig, model_cls=PolyModel) diff --git a/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cce9e6b754b3aec68cc9db7e8e7d0d0d0afd1023 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6211823bf31e2a5ea7aae5b4e22baa03ddd38516 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/layer.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f09d1b5cb73774febb5f2508ec87cb80914dd9c1 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1421bd42e581774dbd11aa845c877f45053e6cd8 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/router.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/router.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ccda1636cf1268738223a1d3d389a1c65255eb3 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/poly/__pycache__/router.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/poly/config.py b/lib/python3.12/site-packages/peft/tuners/poly/config.py new file mode 100644 index 0000000000000000000000000000000000000000..040c1647a005fa59bc2046640dacdd4783aa11fe --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/poly/config.py @@ -0,0 +1,103 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 + +from dataclasses import dataclass, field +from typing import Literal, Optional, Union + +from peft.config import PeftConfig +from peft.utils import PeftType + + +@dataclass +class PolyConfig(PeftConfig): + """ + This is the configuration class to store the configuration of a [`PolyModel`]. + - [Polytropon (Poly)](https://arxiv.org/abs/2202.13914) + - [Multi-Head Routing (MHR)](https://arxiv.org/abs/2211.03831) + + Args: + r (`int`): Attention dimension of each Lora in Poly. + target_modules (`Union[List[str],str]`): The names of the modules to apply Poly to. + exclude_modules (`Optional[Union[List[str], str]]`): + The names of the modules to not apply the adapter. When passing a string, a regex match will be performed. + When passing a list of strings, either an exact match will be performed or it is checked if the name of the + module ends with any of the passed strings. + modules_to_save (`List[str]`): List of modules apart from Poly layers to be set as trainable + and saved in the final checkpoint. + init_weights (bool): Whether to perform initialization of Poly weights. + poly_type (`Literal["poly"]`): The variant of the Poly module to use. Currently, only "poly" + is supported. + n_tasks (`int`): The number of tasks in a multitasking scenario. + n_skills (`int`): The number of skills (LoRA) in each Poly layer. + n_splits (`int`): The number of splits within each LoRA of a Poly layer. A value greater + than 1 indicates the use of Multi-Head Routing (MHR). + """ + + r: int = field(default=8, metadata={"help": "Lora attention dimension"}) + target_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": "List of module names or regex expression of the module names to replace with Poly." + "For example, ['q', 'v'] or '.*decoder.*(SelfAttention|EncDecAttention).*(q|v)$' " + }, + ) + exclude_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={"help": "List of module names or regex expression of the module names to exclude from Poly."}, + ) + modules_to_save: Optional[list[str]] = field( + default=None, + metadata={ + "help": "List of modules apart from Poly layers to be set as trainable and saved in the final checkpoint. " + "For example, in Sequence Classification or Token Classification tasks, " + "the final layer `classifier/score` are randomly initialized and as such need to be trainable and saved." + }, + ) + init_weights: bool = field( + default=True, + metadata={ + "help": ( + "Whether to initialize the weights of the Poly layers with their default initialization. Don't change " + "this setting, except if you know exactly what you're doing." + ), + }, + ) + poly_type: Literal["poly"] = field( + default="poly", + metadata={"help": 'Type of Poly modules to be used. Currently only "poly" is supported.'}, + ) + n_tasks: int = field( + default=1, + metadata={"help": "Number of tasks in multitasking scenario."}, + ) + n_skills: int = field( + default=4, + metadata={"help": "Number of skills (LoRA) in each Poly layer."}, + ) + n_splits: int = field( + default=1, + metadata={"help": "Number of splits within each LoRA of a Poly layer."}, + ) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.POLY + self.target_modules = ( + set(self.target_modules) if isinstance(self.target_modules, list) else self.target_modules + ) + self.exclude_modules = ( + set(self.exclude_modules) if isinstance(self.exclude_modules, list) else self.exclude_modules + ) diff --git a/lib/python3.12/site-packages/peft/tuners/poly/layer.py b/lib/python3.12/site-packages/peft/tuners/poly/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..5169d9ba72904898afdec7e032650c9fabbeec82 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/poly/layer.py @@ -0,0 +1,165 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +import math +from typing import Any + +import torch +import torch.nn as nn + +from peft.tuners.tuners_utils import BaseTunerLayer + +from .config import PolyConfig +from .router import get_router + + +class PolyLayer(BaseTunerLayer): + # All names of layers that may contain (trainable) adapter weights + adapter_layer_names = ("poly_lora_A", "poly_lora_B", "poly_router") + # All names of other parameters that may contain adapter-related parameters + other_param_names = ("r", "n_tasks", "n_skills", "n_splits") + + def __init__(self, base_layer: nn.Module, **kwargs): + self.base_layer = base_layer + self.r = {} + self.n_tasks = {} + self.n_skills = {} + self.n_splits = {} + self.poly_type = {} + self.poly_router = nn.ModuleDict() + self.poly_lora_A = nn.ParameterDict() + self.poly_lora_B = nn.ParameterDict() + self.kwargs = kwargs + + base_layer = self.get_base_layer() + if isinstance(base_layer, nn.Linear): + in_features, out_features = base_layer.in_features, base_layer.out_features + else: + raise ValueError(f"Unsupported layer type {type(base_layer)}") + + self.in_features = in_features + self.out_features = out_features + + def update_layer(self, adapter_name, poly_config): + if poly_config.r <= 0: + raise ValueError(f"`r` should be a positive integer value but the value passed is {poly_config.r}") + + self.r[adapter_name] = poly_config.r + self.n_tasks[adapter_name] = poly_config.n_tasks + self.n_skills[adapter_name] = poly_config.n_skills + self.n_splits[adapter_name] = poly_config.n_splits + self.poly_type[adapter_name] = poly_config.poly_type + + self.poly_lora_A[adapter_name] = nn.Parameter( + torch.empty( + poly_config.n_splits, + poly_config.n_skills, + self.in_features // poly_config.n_splits, + poly_config.r, + ) + ) + self.poly_lora_B[adapter_name] = nn.Parameter( + torch.empty( + poly_config.n_splits, + poly_config.n_skills, + poly_config.r, + self.out_features // poly_config.n_splits, + ) + ) + self.poly_router[adapter_name] = get_router(poly_config) + + self.reset_poly_parameters(adapter_name, init_weights=poly_config.init_weights) + + self._move_adapter_to_device_of_base_layer(adapter_name) + self.set_adapter(self.active_adapters) + + def reset_poly_parameters(self, adapter_name, init_weights): + if adapter_name in self.poly_lora_A.keys(): + # initialize A the same way as the default for nn.Linear + # https://github.com/microsoft/mttl/blob/ce4ca51dbca73be656feb9b3e5233633e3c5dec7/mttl/models/poly.py#L269 + n_splits, n_skills, d, r = self.poly_lora_A[adapter_name].shape + for skill in range(n_skills): + for split in range(n_splits): + param = torch.empty((r, d)) + torch.nn.init.kaiming_uniform_(param, a=math.sqrt(5)) + self.poly_lora_A[adapter_name].data[split, skill, :, :] = param.T + + if init_weights: + # initialize B to zero + torch.nn.init.zeros_(self.poly_lora_B[adapter_name]) + else: + # initialize B the same way as the default for nn.Linear + n_splits, n_skills, r, d = self.poly_lora_B[adapter_name].shape + for skill in range(n_skills): + for split in range(n_splits): + param = torch.empty((d, r)) + torch.nn.init.kaiming_uniform_(param, a=math.sqrt(5)) + self.poly_lora_B[adapter_name].data[split, skill, :, :] = param.T + + # initialized router + self.poly_router[adapter_name].reset() + + +class Linear(nn.Module, PolyLayer): + # Lora implemented in a dense layer + def __init__( + self, + base_layer, + adapter_name: str, + poly_config: PolyConfig, + **kwargs, + ) -> None: + super().__init__() + PolyLayer.__init__(self, base_layer, **kwargs) + + self._active_adapter = adapter_name + self.update_layer(adapter_name, poly_config) + + def forward(self, x: torch.Tensor, *args: Any, task_ids: torch.Tensor = None, **kwargs: Any) -> torch.Tensor: + previous_dtype = x.dtype + if self.disable_adapters: + result = self.base_layer(x, *args, **kwargs) + else: + result = self.base_layer(x, *args, **kwargs) + for active_adapter in self.active_adapters: + if active_adapter not in self.poly_lora_A.keys(): + continue + + r = self.r[active_adapter] + poly_router = self.poly_router[active_adapter] + poly_lora_A = self.poly_lora_A[active_adapter] + poly_lora_B = self.poly_lora_B[active_adapter] + + # Combine the output of LoRAs + # https://github.com/microsoft/mttl/blob/ce4ca51dbca73be656feb9b3e5233633e3c5dec7/mttl/models/poly.py#L293 + mixing_weights = poly_router(task_ids=task_ids, input_ids=x) + bs, n_splits, n_skills = mixing_weights.size() + + # A is n_splits, n_skills, D // n_splits, rank + # we want bs, n_splits, D // n_splits, rank + A = torch.einsum("bqs,qsdr->bqdr", (mixing_weights, poly_lora_A)) + B = torch.einsum("bqs,qsrd->bqrd", (mixing_weights, poly_lora_B)) + + A = A.reshape(bs, self.in_features, r) + B = B.transpose(1, 2).reshape(bs, r, self.out_features) + + x = x.to(A.dtype) + result += x.bmm(A).bmm(B) / r + + result = result.to(previous_dtype) + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "poly." + rep diff --git a/lib/python3.12/site-packages/peft/tuners/poly/model.py b/lib/python3.12/site-packages/peft/tuners/poly/model.py new file mode 100644 index 0000000000000000000000000000000000000000..30313cdabbc2aaac9621ccd418fa7720db5b8ce9 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/poly/model.py @@ -0,0 +1,189 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 contextlib import contextmanager +from dataclasses import asdict +from enum import Enum +from typing import Any + +import torch +from torch import nn + +from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists +from peft.utils import ( + TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING, + ModulesToSaveWrapper, +) + +from .config import PolyConfig +from .layer import Linear, PolyLayer + + +class PolyModel(BaseTuner): + prefix: str = "poly_" + + def __init__(self, model, config, adapter_name) -> None: + super().__init__(model, config, adapter_name) + + @staticmethod + def _check_target_module_exists(poly_config, key): + return check_target_module_exists(poly_config, key) + + def _create_and_replace( + self, + poly_config: PolyConfig, + adapter_name: str, + target: nn.Module, + target_name: str, + parent: nn.Module, + **optional_kwargs: Any, + ): + if isinstance(target, PolyLayer): + target.update_layer(adapter_name, poly_config) + else: + new_module = self._create_new_module( + poly_config, + adapter_name, + target, + ) + if adapter_name not in self.active_adapters: + # adding an additional adapter: it is not automatically trainable + new_module.requires_grad_(False) + self._replace_module(parent, target_name, new_module, target) + + def _replace_module(self, parent, child_name, new_module, child): + setattr(parent, child_name, new_module) + # It's not necessary to set requires_grad here, as that is handled by + # _mark_only_adapters_as_trainable + + # child layer wraps the original module, unpack it + if hasattr(child, "base_layer"): + child = child.base_layer + + if not hasattr(new_module, "base_layer"): + new_module.weight = child.weight + if hasattr(child, "bias"): + new_module.bias = child.bias + + if getattr(child, "state", None) is not None: + if hasattr(new_module, "base_layer"): + new_module.base_layer.state = child.state + else: + new_module.state = child.state + new_module.to(child.weight.device) + + # dispatch to correct device + for name, module in new_module.named_modules(): + if (self.prefix in name) or ("ranknum" in name): + weight = child.qweight if hasattr(child, "qweight") else child.weight + module.to(weight.device) + + def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None: + for n, p in model.named_parameters(): + if self.prefix not in n: + p.requires_grad = False + + @staticmethod + def _create_new_module(poly_config, adapter_name, target, **kwargs): + if isinstance(target, BaseTunerLayer): + target_base_layer = target.get_base_layer() + else: + target_base_layer = target + + if isinstance(target_base_layer, torch.nn.Linear): + return Linear(target, adapter_name, poly_config, **kwargs) + else: + raise ValueError( + f"Target module {target} is not supported. Currently, only the following modules are supported: " + "`torch.nn.Linear`." + ) + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + if name == "model": # see #1892: prevent infinite recursion if class is not initialized + raise + return getattr(self.model, name) + + def get_peft_config_as_dict(self, inference: bool = False): + config_dict = {} + for key, value in self.peft_config.items(): + config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} + if inference: + config["inference_mode"] = True + config_dict[key] = config + return config + + def _set_adapter_layers(self, enabled=True): + for module in self.model.modules(): + if isinstance(module, (PolyLayer, ModulesToSaveWrapper)): + module.enable_adapters(enabled) + + def enable_adapter_layers(self): + self._set_adapter_layers(enabled=True) + + def disable_adapter_layers(self): + self._set_adapter_layers(enabled=False) + + def set_adapter(self, adapter_name): + for module in self.model.modules(): + if isinstance(module, PolyLayer): + module.set_adapter(adapter_name) + + def _prepare_adapter_config(self, peft_config, model_config): + if peft_config.target_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING: + raise ValueError("Please specify `target_modules` in `peft_config`") + peft_config.target_modules = set( + TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]] + ) + return peft_config + + def _register_pre_hooks(self, task_ids): + """Helper method to register pre hooks.""" + if task_ids is None: + return [] + + def pre_hook(_, args, kwargs): + kwargs["task_ids"] = task_ids + return args, kwargs + + handles = [] + + for module in self.model.modules(): + if isinstance(module, Linear): + handle = module.register_forward_pre_hook(pre_hook, with_kwargs=True) + handles.append(handle) + + return handles + + @contextmanager + def _manage_pre_hooks(self, task_ids): + """Context manager to handle the lifecycle of pre hooks.""" + handles = self._register_pre_hooks(task_ids) + try: + yield + finally: + for handle in handles: + handle.remove() + + def forward(self, *args, task_ids=None, **kwargs): + with self._manage_pre_hooks(task_ids): + return self.model(*args, **kwargs) + + def generate(self, *args, task_ids=None, **kwargs): + with self._manage_pre_hooks(task_ids): + return self.model.generate(*args, **kwargs) diff --git a/lib/python3.12/site-packages/peft/tuners/poly/router.py b/lib/python3.12/site-packages/peft/tuners/poly/router.py new file mode 100644 index 0000000000000000000000000000000000000000..3dda3e75e35b6a9fbd5a2412815a0f05421f2ef4 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/poly/router.py @@ -0,0 +1,81 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 abc import ABC, abstractmethod + +import torch +from torch import nn +from torch.distributions.relaxed_bernoulli import RelaxedBernoulli + +from .config import PolyConfig + + +EPS = 1e-12 + + +def get_router(poly_config: PolyConfig) -> nn.Module: + if poly_config.poly_type == "poly": + return PolyRouter(poly_config) + else: + raise ValueError( + f"Unsupported poly_type: {poly_config.poly_type}. " + "Currently, only the following types are supported: " + "`poly`." + ) + + +class Router(nn.Module, ABC): + @abstractmethod + def reset(self): ... + + @abstractmethod + def forward(self, task_ids: torch.Tensor, input_ids: torch.Tensor): ... + + +class PolyRouter(Router): + # It's a simplified implementation of + # https://github.com/microsoft/mttl/blob/ce4ca51dbca73be656feb9b3e5233633e3c5dec7/mttl/models/poly.py#L138 + def __init__(self, poly_config: PolyConfig): + super().__init__() + + self.poly_type = poly_config.poly_type + self.n_tasks = poly_config.n_tasks + self.n_skills = poly_config.n_skills + self.n_splits = poly_config.n_splits + + self.module_logits = nn.Parameter(torch.empty((self.n_tasks, self.n_splits * self.n_skills))) + + def reset(self): + torch.nn.init.uniform_(self.module_logits, -1e-3, 1e-3) + + def forward(self, task_ids: torch.Tensor, input_ids: torch.Tensor): + if task_ids is None: + raise ValueError("task_ids should not be None.") + if task_ids.max().item() >= self.n_tasks: + raise ValueError(f"Only {self.n_tasks} tasks available. Found task id = {task_ids.max().item()}") + + # move task id to input's device + task_ids = task_ids.to(self.module_logits.device) + + module_logits = self.module_logits[task_ids] + module_logits = module_logits.view(-1, self.n_splits, self.n_skills) + + if self.training: + module_logits = RelaxedBernoulli(temperature=1.0, logits=module_logits).rsample() + else: + module_logits = torch.sigmoid(module_logits) + + module_weights = module_logits / (module_logits.sum(dim=-1, keepdim=True) + EPS) + + return module_weights diff --git a/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__init__.py b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..939f74d3f689f400dfdcb0139f4a2cf04cce52fc --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__init__.py @@ -0,0 +1,23 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 peft.utils import register_peft_method + +from .config import PrefixTuningConfig +from .model import PrefixEncoder + + +__all__ = ["PrefixEncoder", "PrefixTuningConfig"] + +register_peft_method(name="prefix_tuning", config_cls=PrefixTuningConfig, model_cls=PrefixEncoder) diff --git a/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da3940718e8c8196786f3a9fd38980a7a3f8e25e Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9e5f06af790d0fea2011fc8ac2292f0219c7080 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87d1c129db4d7858ff1e4af7d489ba1249a57e33 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/prefix_tuning/config.py b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/config.py new file mode 100644 index 0000000000000000000000000000000000000000..6eed77167a6e0b928e59e7e07fbf842c2d7a2d83 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/config.py @@ -0,0 +1,42 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 dataclasses import dataclass, field + +from peft.config import PromptLearningConfig +from peft.utils import PeftType + + +@dataclass +class PrefixTuningConfig(PromptLearningConfig): + """ + This is the configuration class to store the configuration of a [`PrefixEncoder`]. + + Args: + encoder_hidden_size (`int`): The hidden size of the prompt encoder. + prefix_projection (`bool`): Whether to project the prefix embeddings. + """ + + encoder_hidden_size: int = field( + default=None, + metadata={"help": "The hidden size of the encoder"}, + ) + prefix_projection: bool = field( + default=False, + metadata={"help": "Whether to project the prefix tokens"}, + ) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.PREFIX_TUNING diff --git a/lib/python3.12/site-packages/peft/tuners/prefix_tuning/model.py b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/model.py new file mode 100644 index 0000000000000000000000000000000000000000..ffd51892a3cc074406791f6bc7d1b088d25148e3 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/prefix_tuning/model.py @@ -0,0 +1,80 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +# Based on https://github.com/THUDM/P-tuning-v2/blob/main/model/prefix_encoder.py +# with some refactor +import torch + + +class PrefixEncoder(torch.nn.Module): + r""" + The `torch.nn` model to encode the prefix. + + Args: + config ([`PrefixTuningConfig`]): The configuration of the prefix encoder. + + Example: + + ```py + >>> from peft import PrefixEncoder, PrefixTuningConfig + + >>> config = PrefixTuningConfig( + ... peft_type="PREFIX_TUNING", + ... task_type="SEQ_2_SEQ_LM", + ... num_virtual_tokens=20, + ... token_dim=768, + ... num_transformer_submodules=1, + ... num_attention_heads=12, + ... num_layers=12, + ... encoder_hidden_size=768, + ... ) + >>> prefix_encoder = PrefixEncoder(config) + ``` + + **Attributes**: + - **embedding** (`torch.nn.Embedding`) -- The embedding layer of the prefix encoder. + - **transform** (`torch.nn.Sequential`) -- The two-layer MLP to transform the prefix embeddings if + `prefix_projection` is `True`. + - **prefix_projection** (`bool`) -- Whether to project the prefix embeddings. + + Input shape: (`batch_size`, `num_virtual_tokens`) + + Output shape: (`batch_size`, `num_virtual_tokens`, `2*layers*hidden`) + """ + + def __init__(self, config): + super().__init__() + self.prefix_projection = config.prefix_projection + token_dim = config.token_dim + num_layers = config.num_layers + encoder_hidden_size = config.encoder_hidden_size + num_virtual_tokens = config.num_virtual_tokens + if self.prefix_projection and not config.inference_mode: + # Use a two-layer MLP to encode the prefix + self.embedding = torch.nn.Embedding(num_virtual_tokens, token_dim) + self.transform = torch.nn.Sequential( + torch.nn.Linear(token_dim, encoder_hidden_size), + torch.nn.Tanh(), + torch.nn.Linear(encoder_hidden_size, num_layers * 2 * token_dim), + ) + else: + self.embedding = torch.nn.Embedding(num_virtual_tokens, num_layers * 2 * token_dim) + + def forward(self, prefix: torch.Tensor): + if self.prefix_projection: + prefix_tokens = self.embedding(prefix) + past_key_values = self.transform(prefix_tokens) + else: + past_key_values = self.embedding(prefix) + return past_key_values diff --git a/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__init__.py b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4aa7bf8e5fc4705cb7b190cee0de53ac8db89573 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__init__.py @@ -0,0 +1,33 @@ +# Copyright 2025-present the HuggingFace Inc. team. +# +# 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 peft.utils import register_peft_method + +from .config import TrainableTokensConfig +from .layer import TrainableTokensLayer +from .model import TrainableTokensModel + + +__all__ = [ + "TrainableTokensConfig", + "TrainableTokensLayer", + "TrainableTokensModel", +] + +register_peft_method( + name="trainable_tokens", + config_cls=TrainableTokensConfig, + model_cls=TrainableTokensModel, + is_mixed_compatible=False, +) diff --git a/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb41c212747d5e05365ade1fed486d3aa144ebd3 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d86f7cfce59a3972c1bf28dfd09b8a0edb3e0b5b Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/layer.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..494c178e3deada19d9fb36a145f7db0d7dc2100e Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71fc6a47c06b1cfe5b9bd89aa6c32d250ed8f1c8 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/trainable_tokens/config.py b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/config.py new file mode 100644 index 0000000000000000000000000000000000000000..c4a93880db5345de93b98978e7fb02ef43cf14fa --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/config.py @@ -0,0 +1,88 @@ +# Copyright 2025-present the HuggingFace Inc. team. +# +# 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 + +from dataclasses import dataclass, field +from typing import Optional, Union + +from peft.config import PeftConfig +from peft.utils import PeftType + + +@dataclass +class TrainableTokensConfig(PeftConfig): + """ + Configuration for the `TrainableTokens` method. + + Allows for training new tokens (and re-training existing ones) without training the full embedding matrix. By + marking a few select tokens (identified by their indices) trainable and leaving the rest untouched, this method can + be used to add new tokens or changing the embedding of existing tokens while saving on memory. Both storage as well + as working memory usage are reduced in contrast to training the embedding matrix fully. + + Note that training with FSDP/DeepSpeed might not yet be fully supported. + + Args: + token_indices (`list[int]`): + List of integers, signifying the indices of the tokens you want to be trainable. To find the index of a + token with a tokenizer, you can tokenize the string and look at the returned `input_ids`. The closer the + amount of indices is to the total amount of tokens, the less efficient this method gets. + target_modules (`Optional[Union[list[str], str]]`): + List of module names or regex expression of the module names to replace with our `TrainableTokensLayer`. If + not defined, it will attempt to get the model's input embedding layer if the model has a + `get_input_embeddings` method (transformer models usually do), if that fails the default is 'embed_tokens'. + Other example targets are `embedding`, `encoder.embeddings` or `decoder.embeddings`. + init_weights (`bool`): + By default the new token weights are initialized to be the same as the respective token embeddings. This + makes TrainableTokens a no-op when not trained. If set to `False` the weights will be random values. Do not + change this setting unless you know exactly what you're doing. + """ + + token_indices: list[int] = field( + default_factory=list, + metadata={ + "help": ( + "List of integers, signifying the indices of the tokens you want to be trainable. " + "To find the index of a token with a tokenizer, you can tokenize the string and " + "look at the returned `input_ids`. The closer the amount of indices is to the total amount of " + "tokens, the less efficient this method gets." + ) + }, + ) + target_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": ( + "List of module names or regex expression of the module names to replace with our " + "`TrainableTokensLayer`. If not defined, it will default to the model's input embedding layer if " + "the model has a `get_input_embeddings` method (transformer models usually do), if that fails the " + "default is 'embed_tokens'. Other example targets could be `embedding`, `encoder.embeddings` or " + "`decoder.embeddings`." + ), + }, + ) + + init_weights: bool = field( + default=True, + metadata={ + "help": ( + "By default the new token weights are initialized to be the same as the respective token embeddings. " + "This makes TrainableTokens a no-op when not trained. If set to `False` the weights will be random " + "values. Do not change this setting unless you know exactly what you're doing. " + ) + }, + ) + + def __post_init__(self): + self.peft_type = PeftType.TRAINABLE_TOKENS diff --git a/lib/python3.12/site-packages/peft/tuners/trainable_tokens/layer.py b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..73d7d49d554dd7d0655b7ac7efc20aa3b28c0f72 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/layer.py @@ -0,0 +1,208 @@ +# Copyright 2025-present the HuggingFace Inc. team. +# +# 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 warnings +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from peft.tuners._buffer_dict import BufferDict +from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge + + +class TrainableTokensLayer(nn.Module, BaseTunerLayer): + # All names of layers that may contain (trainable) adapter weights + adapter_layer_names = ("trainable_tokens_delta",) + + # All names of other parameters that may contain adapter-related parameters + other_param_names = ("token_indices", "trainable_tokens_original") + + def __init__( + self, + base_layer: nn.Module, + adapter_name: str, + token_indices: list[int], + tied_adapter: Optional[TrainableTokensLayer] = None, + **kwargs, + ) -> None: + super().__init__() + + self.base_layer = base_layer + self._active_adapter = adapter_name + self.kwargs = kwargs + + # wrap the tied adapter in a list so that it is excluded from .(named_)modules() and, therefore, + # not included in the state dict since it would be a copy of the tied adapter anyway. + self._tied_adapter = [tied_adapter] if tied_adapter else [] + + # we store the updated weights of particular tokens and their originals. we assume + # that the count of new tokens is far smaller than the number of total tokens. + # + # In case we have weight tying with another token adapter, we'll have no actual + # references on our own but use everything from the tied adapter. + if not self.tied_adapter: + self.trainable_tokens_delta = nn.ParameterDict({}) + self.trainable_tokens_original = BufferDict({}) + self.token_indices = {} + else: + self.trainable_tokens_delta = self.tied_adapter.trainable_tokens_delta + self.trainable_tokens_original = self.tied_adapter.trainable_tokens_original + self.token_indices = self.tied_adapter.token_indices + + # Mark the weight as unmerged + self.merged_adapters = [] + + @property + def tied_adapter(self): + if self._tied_adapter: + return self._tied_adapter[0] + return None + + def update_layer(self, adapter_name, **kwargs): + if kwargs.get("tied_adapter", None): + # as a tied adapter, we're just following whatever the adpater we're tied to does, we don't update anything. + return + + self.token_indices[adapter_name] = kwargs["token_indices"] + init_weights = kwargs.get("init_weights", True) + + # we initialize the delta embedding weights from the base embedding matrix and replace values instead of + # adding/subtracting deltas. we do it this way and use `embedding.weight.index_copy()` to write the updated + # values during `forward()` to avoid that the user resizing the embedding matrix, effectively filling the new + # token space with random values, training the model with TrainableTokensLayer, initializing the model anew - + # thus re-initializing the new embeddings again with new random variables. If we would add/subtract deltas + # onto the new values, we would get undefined behavior. By replacing the specific token values we always + # get defined behavior. + # + if init_weights: + values = self.get_base_layer().weight[self.token_indices[adapter_name]] + else: + values = torch.rand_like(self.get_base_layer().weight[self.token_indices[adapter_name]]) + + self.trainable_tokens_delta[adapter_name] = nn.Parameter(values.clone(), requires_grad=True) + self.trainable_tokens_original[adapter_name] = values.clone() + + self._move_adapter_to_device_of_base_layer(adapter_name) + + def _check_overlapping_tokens(self, adapter_names): + """Raises an error if the token indices of the given adapter names are overlapping. + This is currently not supported and can lead to undefined behavior of the model if no specific merging between + the overlapping indices' values is applied. + """ + if len(adapter_names) <= 1: + return + + indices = set() + + # we take already merged adapters into account as well since they can be overriden by new adapters as well. + for adapter_name in set(adapter_names + self.merged_adapters): + index_set = set(self.token_indices[adapter_name]) + if len(indices.intersection(index_set)): + raise ValueError( + f"Token indices of adapter {adapter_name} are already defined and would result in " + "undefined merging behavior. Only disjunct token indices are currently supported." + ) + indices.update(index_set) + + def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None: + adapter_names = check_adapters_to_merge(self, adapter_names) + + if not adapter_names: + # no adapter to merge + return + + self._check_overlapping_tokens(adapter_names) + + merged = self.base_layer.weight.data + + for adapter_name in adapter_names: + index = torch.tensor(self.token_indices[adapter_name]).to(merged.device) + deltas = self.trainable_tokens_delta[adapter_name].to(merged) + merged = merged.index_copy(dim=0, index=index, source=deltas) + + if safe_merge and not torch.isfinite(merged).all(): + raise ValueError(f"NaNs detected in the merged weights. The adapter {adapter_name} seems to be broken") + + self.base_layer.weight.data = merged + self.merged_adapters.extend(adapter_names) + + def unmerge(self) -> None: + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + + while len(self.merged_adapters) > 0: + adapter_name = self.merged_adapters.pop() + + index = torch.tensor(self.token_indices[adapter_name]).to(self.base_layer.weight.device) + originals = self.trainable_tokens_original[adapter_name].to(self.base_layer.weight) + self.base_layer.weight.data.index_copy_(dim=0, index=index, source=originals) + + def get_merged_weights(self, active_adapters): + W = self.base_layer.weight + + for adapter_name in active_adapters: + index = torch.tensor(self.token_indices[adapter_name]).to(W.device) + deltas = self.trainable_tokens_delta[adapter_name].to(W) + W = W.index_copy(dim=0, index=index, source=deltas) + + return W + + def forward_adapters(self, x: torch.Tensor, active_adapters, *args, **kwargs) -> torch.Tensor: + if self.disable_adapters or not active_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + self._check_overlapping_tokens(active_adapters) + + W = self.get_merged_weights(active_adapters) + + # Normally it should be very clear that we're wrapping Embedding layers but there are cases, such as + # tying weights with an LM head where the layer we wrap is a Linear layer. Therefore we must choose + # accordingly. + # + # TODO: the isinstance checks, especially the one for nn.Linear, may not hold for quantized layers; + # TODO: we may need to find a better way to detect quantized layers. + if isinstance(self.base_layer, torch.nn.Embedding): + result = F.embedding( + input=x, + weight=W, + padding_idx=self.base_layer.padding_idx, + max_norm=self.base_layer.max_norm, + norm_type=self.base_layer.norm_type, + scale_grad_by_freq=self.base_layer.scale_grad_by_freq, + sparse=self.base_layer.sparse, + ) + elif isinstance(self.base_layer, torch.nn.Linear): + # Probably a tied adapter that wraps an LM head. + result = F.linear( + input=x, + weight=W, + ) + else: + raise ValueError( + "TrainableTokensLayer wraps an unknown layer type, maybe you are targeting the wrong layer?" + ) + + return result + + def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + return self.forward_adapters(x, self.active_adapters, *args, **kwargs) diff --git a/lib/python3.12/site-packages/peft/tuners/trainable_tokens/model.py b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/model.py new file mode 100644 index 0000000000000000000000000000000000000000..b6b797ebaf9f9aa7555d10983985b18b18535e57 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/trainable_tokens/model.py @@ -0,0 +1,280 @@ +# Copyright 2025-present the HuggingFace Inc. team. +# +# 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 warnings +from typing import Optional + +import torch +import torch.nn as nn +from tqdm import tqdm + +from peft.config import PeftConfig +from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists, onload_layer +from peft.utils import AuxiliaryTrainingWrapper, _get_input_embeddings_name, _get_submodules + +from .layer import TrainableTokensLayer + + +class TrainableTokensModel(BaseTuner): + prefix: str = "trainable_tokens_" + + def __init__(self, model, config, adapter_name, low_cpu_mem_usage: bool = False): + super().__init__(model, config, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage) + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + return getattr(self.model, name) + + def _prepare_adapter_config(self, peft_config, model_config): + # target_modules can be none which prompts us to infer the embedding layer name ourselves. + if peft_config.target_modules is None: + peft_config.target_modules = _get_input_embeddings_name(self.model, "embed_tokens") + + return peft_config + + def inject_adapter( + self, model: nn.Module, adapter_name: str, autocast_adapter_dtype: bool = True, low_cpu_mem_usage: bool = False + ) -> None: + super().inject_adapter( + model=model, + adapter_name=adapter_name, + autocast_adapter_dtype=autocast_adapter_dtype, + low_cpu_mem_usage=low_cpu_mem_usage, + ) + + model_config = self.get_model_config(self) + + # In case of weight-tying we need to adapt the tied weights as well and use tie the embedding adapter. + # + # The TrainableTokensLayer supports being tied to another TrainableTokensLayer meaning that the layer will + # not do any changes on its own but solely rely on the weights from the tied adapter. We will search for the + # tied weights and put tied TrainableTokensLayer adapters on them, all tied to the adapter of the embedding + # matrix. + if ( + model_config.get("tie_word_embeddings", False) + # some models may be misconfigured to have weight tying enabled but don't define tied weights keys + and self.model._tied_weights_keys is not None + and isinstance(self.model.get_input_embeddings(), TrainableTokensLayer) + ): + module_keys = [".".join(n.split(".")[:-1]) for n in self.model._tied_weights_keys] + # disable removing of duplicates since we're essentially only dealing with duplicates (i.e. tied weights) + for name, module in self.model.named_modules(remove_duplicate=False): + matched_keys = [target_key for target_key in module_keys if name.endswith(target_key)] + if matched_keys: + parent, target, target_name = _get_submodules(model, name) + + peft_config = self.peft_config[adapter_name].to_dict() + peft_config["tied_adapter"] = self.model.get_input_embeddings() + + self._create_and_replace_dict( + peft_config, + adapter_name, + target, + target_name, + parent, + matched_keys[0], + ) + + def _get_tied_target_modules(self, *args, **kwargs): + # Normally this method would return the layers that target tied layers. + # + # We override this method since we explicitly support tied weights tied to the embedding layer. + # Therefore, we don't need the warning issued by returning the modules here. + return [] + + def _create_and_replace_dict( + self, + peft_config: dict, + adapter_name: str, + target: nn.Module, + target_name: str, + parent: nn.Module, + current_key: str, + ) -> None: + """ + The same as `_create_and_replace` but takes a dictionary instead of a peft config so that we can add keys that + are not present in the config, such as `tied_adapter`. + """ + kwargs = peft_config + + if isinstance(target, TrainableTokensLayer): + target.update_layer(adapter_name, **kwargs) + else: + new_module = self._create_new_module(peft_config, adapter_name, target, **kwargs) + self._replace_module(parent, target_name, new_module, target) + + def _create_and_replace( + self, + peft_config: PeftConfig, + adapter_name: str, + target: nn.Module, + target_name: str, + parent: nn.Module, + current_key: str, + ) -> None: + """ + A private method to create and replace the target module with the adapter module. + """ + kwargs = peft_config.to_dict() + self._create_and_replace_dict(kwargs, adapter_name, target, target_name, parent, current_key) + + def _check_target_module_exists(self, peft_config: PeftConfig, key: str) -> bool: + return check_target_module_exists(peft_config, key) + + @staticmethod + def _create_new_module(peft_config, adapter_name, target, **kwargs): + new_module = TrainableTokensLayer(target, adapter_name, **kwargs) + new_module.update_layer( + adapter_name, + init_weights=kwargs["init_weights"], + token_indices=kwargs["token_indices"], + tied_adapter=kwargs.get("tied_adapter", None), + ) + + return new_module + + def _replace_module(self, parent, child_name, new_module, child): + setattr(parent, child_name, new_module) + # It's not necessary to set requires_grad here, as that is handled by + # _mark_only_adapters_as_trainable + + # child layer wraps the original module, unpack it + if hasattr(child, "base_layer"): + child = child.base_layer + + if not hasattr(new_module, "base_layer"): + new_module.weight = child.weight + if hasattr(child, "bias"): + new_module.bias = child.bias + + if getattr(child, "state", None) is not None: + if hasattr(new_module, "base_layer"): + new_module.base_layer.state = child.state + else: + new_module.state = child.state + new_module.to(child.weight.device) + + meta = torch.device("meta") + # dispatch to correct device + for name, module in new_module.named_modules(): + if self.prefix in name: + if not any(p.device == meta for p in module.parameters()): + module.to(child.weight.device) + + def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None: + for n, p in model.named_parameters(): + if self.prefix not in n: + p.requires_grad = False + + def _set_adapter_layers(self, enabled: bool = True) -> None: + for module in self.model.modules(): + if isinstance(module, (BaseTunerLayer, AuxiliaryTrainingWrapper)): + module.enable_adapters(enabled) + + def enable_adapter_layers(self) -> None: + """Enable all adapters. + + Call this if you have previously disabled all adapters and want to re-enable them. + """ + self._set_adapter_layers(enabled=True) + + def disable_adapter_layers(self) -> None: + """Disable all adapters. + + When disabling all adapters, the model output corresponds to the output of the base model. + """ + self._set_adapter_layers(enabled=False) + + def set_adapter(self, adapter_name: str | list[str]) -> None: + """Set the active adapter(s). + + Additionally, this function will set the specified adapters to trainable (i.e., requires_grad=True). If this is + not desired, use the following code. + + ```py + >>> for name, param in model_peft.named_parameters(): + ... if ...: # some check on name (ex. if 'lora' in name) + ... param.requires_grad = False + ``` + + Args: + adapter_name (`str` or `list[str]`): Name of the adapter(s) to be activated. + """ + for module in self.model.modules(): + if isinstance(module, TrainableTokensLayer): + if module.merged: + warnings.warn("Adapter cannot be set when the model is merged. Unmerging the model first.") + module.unmerge() + module.set_adapter(adapter_name) + self.active_adapter = adapter_name + + def unload(self) -> torch.nn.Module: + """ + Gets back the base model by removing all the trainable tokens modules without merging. + """ + return self._unload_and_optionally_merge(merge=False) + + def merge_and_unload( + self, progressbar: bool = False, safe_merge: bool = False, adapter_names: Optional[list[str]] = None + ) -> torch.nn.Module: + r""" + This method merges the trained tokens into the targeted embedding layer(s) of the base model. This is needed if + someone wants to use the base model as a standalone model. + + Args: + progressbar (`bool`): + whether to show a progressbar indicating the unload and merge process + safe_merge (`bool`): + whether to activate the safe merging check to check if there is any potential Nan in the adapter + weights + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + """ + return self._unload_and_optionally_merge( + progressbar=progressbar, safe_merge=safe_merge, adapter_names=adapter_names + ) + + def _unload_and_optionally_merge( + self, + merge=True, + progressbar: bool = False, + safe_merge: bool = False, + adapter_names: Optional[list[str]] = None, + ): + key_list = [key for key, _ in self.model.named_modules() if self.prefix not in key] + desc = "Unloading " + ("and merging " if merge else "") + "model" + for key in tqdm(key_list, disable=not progressbar, desc=desc): + try: + parent, target, target_name = _get_submodules(self.model, key) + except AttributeError: + continue + with onload_layer(target): + if hasattr(target, "unload_and_optionally_merge_module"): + # if layers have special unloading method, like MultiheadAttention, use that + unloaded_module = target.unload_and_optionally_merge_module( + merge=merge, safe_merge=safe_merge, adapter_names=adapter_names + ) + self._replace_module(parent, target_name, unloaded_module, target) + elif hasattr(target, "base_layer"): + if merge: + target.merge(safe_merge=safe_merge, adapter_names=adapter_names) + self._replace_module(parent, target_name, target.get_base_layer(), target) + + return self.model diff --git a/lib/python3.12/site-packages/peft/tuners/vera/__init__.py b/lib/python3.12/site-packages/peft/tuners/vera/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..25c4a96619524bfdcd41a8f7df331533ba370782 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/vera/__init__.py @@ -0,0 +1,40 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 peft.import_utils import is_bnb_4bit_available, is_bnb_available +from peft.utils import register_peft_method + +from .config import VeraConfig +from .layer import Linear, VeraLayer +from .model import VeraModel + + +__all__ = ["Linear", "VeraConfig", "VeraLayer", "VeraModel"] + + +register_peft_method(name="vera", config_cls=VeraConfig, model_cls=VeraModel, prefix="vera_lambda_") + + +def __getattr__(name): + if (name == "Linear8bitLt") and is_bnb_available(): + from .bnb import Linear8bitLt + + return Linear8bitLt + + if (name == "Linear4bit") and is_bnb_4bit_available(): + from .bnb import Linear4bit + + return Linear4bit + + raise AttributeError(f"module {__name__} has no attribute {name}") diff --git a/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f664bdb1c7703f9619bac8d79006c948bdf43a25 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/bnb.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/bnb.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8985e3486f8fe351bfacb980bd8f08e294e1c17 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/bnb.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..19ee5286801a72bfb1d84d015566b49da4c2f60b Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/layer.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d4829031548eabcc0ec16e27f04a0eb93685e03 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6307265d41cb46e032193ff03745fb43da6df193 Binary files /dev/null and b/lib/python3.12/site-packages/peft/tuners/vera/__pycache__/model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/peft/tuners/vera/bnb.py b/lib/python3.12/site-packages/peft/tuners/vera/bnb.py new file mode 100644 index 0000000000000000000000000000000000000000..71d20e4b1163ab9cc93b7b33c2aa022ddd3eb11f --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/vera/bnb.py @@ -0,0 +1,411 @@ +# Copyright 2024-present the HuggingFace Inc. team. +# +# 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 warnings +from typing import Optional + +import bitsandbytes as bnb +import torch + +from peft.import_utils import is_bnb_4bit_available, is_bnb_available +from peft.tuners.tuners_utils import check_adapters_to_merge +from peft.utils.integrations import dequantize_bnb_weight +from peft.utils.other import transpose + +from .layer import VeraLayer + + +if is_bnb_available(): + + class Linear8bitLt(torch.nn.Module, VeraLayer): + def __init__( + self, + base_layer: torch.nn.Module, + adapter_name: str, + vera_A, + vera_B, + r: int = 0, + vera_dropout: float = 0.0, + fan_in_fan_out: bool = False, + init_weights: bool = True, + d_initial: float = 0.1, + **kwargs, + ) -> None: + super().__init__() + VeraLayer.__init__(self, base_layer) + self.fan_in_fan_out = fan_in_fan_out + + self._active_adapter = adapter_name + self.update_layer( + adapter_name, + vera_A, + vera_B, + r, + vera_dropout=vera_dropout, + init_weights=init_weights, + d_initial=d_initial, + ) + + def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None: + if self.merged: + warnings.warn( + f"Already following adapters were merged {','.join(self.merged_adapters)}. " + f"You are now additionally merging {','.join(self.active_adapters)}." + ) + + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + return + + for active_adapter in adapter_names: + if active_adapter not in self.vera_lambda_d.keys(): + continue + + warnings.warn( + "Merge vera module to 8-bit linear may get different generations due to rounding errors." + ) + vera_data = self.get_delta_weight(active_adapter) + + weight = self.get_base_layer().weight + state = self.get_base_layer().state + if state.SCB is None: + state.SCB = weight.SCB + + output = dequantize_bnb_weight(weight, state) + w_data = output.to(vera_data.dtype).to(vera_data.device) + vera_data + + if safe_merge and not torch.isfinite(w_data).all(): + raise ValueError( + f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" + ) + + self.get_base_layer().weight = bnb.nn.Int8Params( + w_data.to("cpu"), requires_grad=False, has_fp16_weights=weight.has_fp16_weights + ).to(weight.device) + state.reset_grads() + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + if not self.merged: + warnings.warn("Already unmerged. Nothing to do") + return + + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter not in self.vera_lambda_d.keys(): + continue + warnings.warn( + "Unmerge vera module to 8-bit linear may get different generations due to rounding errors." + ) + vera_data = self.get_delta_weight(active_adapter) + + weight = self.get_base_layer().weight + state = self.get_base_layer().state + if state.SCB is None: + state.SCB = weight.SCB + output = dequantize_bnb_weight(weight, state=state) + + w_data = output.to(vera_data.dtype).to(vera_data.device) - vera_data + + self.get_base_layer().weight = bnb.nn.Int8Params( + w_data.to("cpu"), requires_grad=False, has_fp16_weights=weight.has_fp16_weights + ).to(weight.device) + state.reset_grads() + + def get_delta_weight(self, adapter) -> torch.Tensor: + """ + Compute the delta weight for the given adapter. + + Args: + adapter (str): The name of the adapter for which the delta weight should be computed. + + Returns: + torch.Tensor: The computed delta weight for the VeRA adapter. + + Note: + This method implements the VeRA-specific weight update. Unlike LoRA, VeRA uses shared projection + matrices (vera_A and vera_B) across all layers, along with per-layer trainable parameters (lambda_d and + lambda_b). + """ + # Retrieve shared projection matrices + vera_A = self.vera_A[adapter] + vera_B = self.vera_B[adapter] + + # Retrieve per-layer trainable parameters + device = vera_B.device + dtype = vera_B.dtype + + # In case users wants to merge the adapter weights that are in + # (b)float16 while being on CPU, we need to cast the weights to float32, perform the merge and then cast back to + # (b)float16 because some CPUs have slow bf16/fp16 matmuls. + cast_to_fp32 = device.type == "cpu" and (dtype == torch.float16 or dtype == torch.bfloat16) + + lambda_d = self.vera_lambda_d[adapter] + lambda_b = self.vera_lambda_b[adapter] + + if cast_to_fp32: + vera_A = vera_A.float() + vera_B = vera_B.float() + lambda_d = lambda_d.float() + lambda_b = lambda_b.float() + + sliced_A = vera_A[:, : self.in_features].to(lambda_d.device) + sliced_B = vera_B[: self.out_features, :].to(lambda_d.device) + lambda_b = lambda_b.unsqueeze(-1) + lambda_d = lambda_d.unsqueeze(-1) + + # VeRA-specific computation: + # 1. Apply lambda_d to the input projection (vera_A) + # 2. Apply lambda_b to the output projection (vera_B) + # 3. Compute the outer product of the scaled projections + output_tensor = transpose((lambda_b * sliced_B) @ (lambda_d * sliced_A), self.fan_in_fan_out) + + if cast_to_fp32: + output_tensor = output_tensor.to(dtype=dtype) + + return output_tensor + + def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + """ + Perform the forward pass using the VeRA adapter. + + Args: + x (torch.Tensor): Input tensor. + + Returns: + torch.Tensor: Output tensor after applying the VeRA adaptation. + + Note: + This method implements the VeRA-specific forward pass. It applies the shared projections (vera_A and + vera_B) along with the per-layer trainable parameters (lambda_d and lambda_b) to compute the adapter + output. + """ + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + result = self.base_layer(x, *args, **kwargs) + for active_adapter in self.active_adapters: + if active_adapter not in self.vera_lambda_d.keys(): + continue + + lambda_d = self.vera_lambda_d[active_adapter] + lambda_b = self.vera_lambda_b[active_adapter] + + vera_A = self.vera_A[active_adapter] + vera_B = self.vera_B[active_adapter] + + dropout = self.vera_dropout[active_adapter] + + requires_conversion = not torch.is_autocast_enabled() + if requires_conversion: + expected_dtype = result.dtype + compute_dtype = lambda_d.dtype + if x.dtype != compute_dtype: + x = x.to(compute_dtype) + + sliced_A = vera_A[:, : self.in_features].to(x.device) + sliced_B = vera_B[: self.out_features, :].to(x.device) + + x_temp = dropout(x.to(lambda_d.dtype)) + + adapter_output = lambda_b * torch.nn.functional.linear( + lambda_d * torch.nn.functional.linear(x_temp, sliced_A), sliced_B + ) + + if requires_conversion: + adapter_output = adapter_output.to(expected_dtype) + + result = result + adapter_output + + # Ensure the output tensor has the same dtype as the input tensor + return result.to(x.dtype) + + def __repr__(self) -> str: + rep = super().__repr__() + return "vera." + rep + + +if is_bnb_4bit_available(): + + class Linear4bit(torch.nn.Module, VeraLayer): + def __init__( + self, + base_layer: torch.nn.Module, + adapter_name: str, + vera_A, + vera_B, + r: int = 0, + vera_dropout: float = 0.0, + fan_in_fan_out: bool = False, + init_weights: bool = True, + d_initial: float = 0.1, + **kwargs, + ) -> None: + super().__init__() + VeraLayer.__init__(self, base_layer) + self.fan_in_fan_out = fan_in_fan_out + + self._active_adapter = adapter_name + self.update_layer( + adapter_name, + vera_A, + vera_B, + r, + vera_dropout=vera_dropout, + init_weights=init_weights, + d_initial=d_initial, + ) + + def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None: + if self.merged: + warnings.warn( + f"Already following adapters were merged {','.join(self.merged_adapters)}. " + f"You are now additionally merging {','.join(self.active_adapters)}." + ) + + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + return + + for active_adapter in adapter_names: + if active_adapter not in self.vera_lambda_d.keys(): + continue + + warnings.warn( + "Merge vera module to 4-bit linear may get different generations due to rounding errors." + ) + vera_data = self.get_delta_weight(active_adapter) + + weight = self.get_base_layer().weight + kwargs = weight.__dict__ + # torch.compile can introduce attributes preceded by '_', remove them + kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")} + w_data = bnb.functional.dequantize_4bit(weight.data, weight.quant_state) + vera_data + + if safe_merge and not torch.isfinite(w_data).all(): + raise ValueError( + f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" + ) + + self.get_base_layer().weight = bnb.nn.Params4bit(w_data.to("cpu"), requires_grad=False, **kwargs).to( + weight.device + ) + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + if not self.merged: + warnings.warn("Already unmerged. Nothing to do") + return + + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter not in self.vera_lambda_d.keys(): + continue + warnings.warn( + "Unmerge vera module to 4-bit linear may get different generations due to rounding errors." + ) + vera_data = self.get_delta_weight(active_adapter) + + weight = self.get_base_layer().weight + kwargs = weight.__dict__ + w_data = bnb.functional.dequantize_4bit(weight.data, weight.quant_state) - vera_data + + self.get_base_layer().weight = bnb.nn.Params4bit(w_data.to("cpu"), requires_grad=False, **kwargs).to( + weight.device + ) + + def get_delta_weight(self, adapter) -> torch.Tensor: + vera_A = self.vera_A[adapter] + vera_B = self.vera_B[adapter] + + device = vera_B.device + dtype = vera_B.dtype + + cast_to_fp32 = device.type == "cpu" and (dtype == torch.float16 or dtype == torch.bfloat16) + + lambda_d = self.vera_lambda_d[adapter] + lambda_b = self.vera_lambda_b[adapter] + + if cast_to_fp32: + vera_A = vera_A.float() + vera_B = vera_B.float() + lambda_d = lambda_d.float() + lambda_b = lambda_b.float() + + sliced_A = vera_A[:, : self.in_features].to(lambda_d.device) + sliced_B = vera_B[: self.out_features, :].to(lambda_d.device) + lambda_b = lambda_b.unsqueeze(-1) + lambda_d = lambda_d.unsqueeze(-1) + + output_tensor = transpose((lambda_b * sliced_B) @ (lambda_d * sliced_A), self.fan_in_fan_out) + + if cast_to_fp32: + output_tensor = output_tensor.to(dtype=dtype) + + return output_tensor + + def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + result = self.base_layer(x, *args, **kwargs) + result = result.clone() + for active_adapter in self.active_adapters: + if active_adapter not in self.vera_lambda_d.keys(): + continue + + lambda_d = self.vera_lambda_d[active_adapter] + lambda_b = self.vera_lambda_b[active_adapter] + + vera_A = self.vera_A[active_adapter] + vera_B = self.vera_B[active_adapter] + + dropout = self.vera_dropout[active_adapter] + + requires_conversion = not torch.is_autocast_enabled() + if requires_conversion: + expected_dtype = result.dtype + compute_dtype = lambda_d.dtype + if x.dtype != compute_dtype: + x = x.to(compute_dtype) + + sliced_A = vera_A[:, : self.in_features].to(x.device) + sliced_B = vera_B[: self.out_features, :].to(x.device) + + x_temp = dropout(x.to(lambda_d.dtype)) + + adapter_output = lambda_b * torch.nn.functional.linear( + lambda_d * torch.nn.functional.linear(x_temp, sliced_A), sliced_B + ) + + if requires_conversion: + adapter_output = adapter_output.to(expected_dtype) + + result = result + adapter_output + + # Ensure the output tensor has the same dtype as the input tensor + return result.to(x.dtype) + + def __repr__(self) -> str: + rep = super().__repr__() + return "vera." + rep diff --git a/lib/python3.12/site-packages/peft/tuners/vera/config.py b/lib/python3.12/site-packages/peft/tuners/vera/config.py new file mode 100644 index 0000000000000000000000000000000000000000..5a8eaa4cdc27505bf3b3d6b9404f61eaa76f9494 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/vera/config.py @@ -0,0 +1,162 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 warnings +from dataclasses import dataclass, field +from typing import Optional, Union + +from peft.config import PeftConfig +from peft.utils import PeftType + + +@dataclass +class VeraConfig(PeftConfig): + """ + This is the configuration class to store the configuration of a [`VeraModel`]. + + Paper: https://arxiv.org/abs/2310.11454. + + Args: + r (`int`, *optional*, defaults to `256`): + VeRA parameter dimension ("rank"). Choose higher values than LoRA ranks here, since VeRA uses far fewer + parameters than LoRA (see Table 1). + target_modules (`Union[List[str], str]`): + The names of the modules to apply Vera to. Only linear layers are supported. + projection_prng_key (`int`): + Vera PRNG init key. Used for initialising vera_A and vera_B for new models or when loading a checkpoint + that did not include these projections. Defaults to `0`. + save_projection (`bool`): + Whether to save the vera_A / vera_B projections in the state dict alongside per layer lambda_b / lambda_d + weights. This will increase the size of the checkpoint, but guarantee that we can reload the checkpoint on + all system configurations. Defaults to `True`. + vera_dropout (`float`): + The dropout probability for Vera layers. + d_initial (`float`, *optional*, defaults to `0.1`): + Initial init value for `vera_lambda_d` vector used when initializing the VeRA parameters. Small values + (<=0.1) are recommended (see Table 6c in the paper). + fan_in_fan_out (`bool`): + Set this to True if the layer to replace stores weight like (fan_in, fan_out). For example, gpt-2 uses + `Conv1D` which stores weights like (fan_in, fan_out) and hence this should be set to `True`. + bias (`str`): + Bias type for Vera. Can be 'none', 'all' or 'vera_only'. If 'all' or 'vera_only', the corresponding biases + will be updated during training. Be aware that this means that, even when disabling the adapters, the model + will not produce the same output as the base model would have without adaptation. + modules_to_save (`List[str]`): + List of modules apart from Vera layers to be set as trainable and saved in the final checkpoint. + init_weights (`bool`): + Whether to initialize the weights of the Vera layers with their default initialization. Don't change this + setting, except if you know exactly what you're doing. + layers_to_transform (`Union[List[int],int]`): + The layer indexes to transform, if this argument is specified, it will apply the Vera transformations on + the layer indexes that are specified in this list. If a single integer is passed, it will apply the Vera + transformations on the layer at this index. + layers_pattern (`Optional[Union[List[str], str]]`): + The layer pattern name, used only if `layers_to_transform` is different from `None`. This should target the + `nn.ModuleList` of the model, which is often called `'layers'` or `'h'`. + """ + + r: int = field(default=256, metadata={"help": "Vera attention dimension"}) + + target_modules: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": ( + "List of module names or regex expression of the module names to replace with Vera." + "For example, ['q', 'v'] or '.*decoder.*(SelfAttention|EncDecAttention).*(q|v)$'. " + "Only linear layers are supported." + ) + }, + ) + projection_prng_key: int = field( + default=0, + metadata={ + "help": ( + "Vera PRNG init key. Used for initialising vera_A and vera_B for new models or when loading a " + "checkpoint that did not include these projections." + ) + }, + ) + save_projection: bool = field( + default=True, + metadata={ + "help": ( + "Whether to save the vera_A / vera_B projections in the state dict alongside per layer lambda_b / " + "lambda_d weights. This will increase the size of the checkpoint, but guarantee that we can reload " + "the checkpoint on all system configurations." + ) + }, + ) + vera_dropout: float = field(default=0.0, metadata={"help": "Vera dropout"}) + d_initial: float = field(default=0.1, metadata={"help": "Initial init value for d vector."}) + fan_in_fan_out: bool = field( + default=False, + metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"}, + ) + bias: str = field(default="none", metadata={"help": "Bias type for Vera. Can be 'none', 'all' or 'vera_only'"}) + modules_to_save: Optional[list[str]] = field( + default=None, + metadata={ + "help": ( + "List of modules apart from Vera layers to be set as trainable and saved in the final checkpoint. For" + " example, in Sequence Classification or Token Classification tasks, the final layer" + " `classifier/score` are randomly initialized and as such need to be trainable and saved." + ) + }, + ) + init_weights: bool = field( + default=True, + metadata={ + "help": ( + "Whether to initialize the weights of the Vera layers with their default initialization. Don't change " + "this setting, except if you know exactly what you're doing." + ), + }, + ) + layers_to_transform: Optional[Union[list[int], int]] = field( + default=None, + metadata={ + "help": ( + "The layer indexes to transform, is this argument is specified, PEFT will transform only the layers" + " indexes that are specified inside this list. If a single integer is passed, PEFT will transform only" + " the layer at this index." + ) + }, + ) + layers_pattern: Optional[Union[list[str], str]] = field( + default=None, + metadata={ + "help": ( + "The layer pattern name, used only if `layers_to_transform` is different to None and if the layer " + "pattern is not in the common layers pattern. This should target the `nn.ModuleList` of the " + "model, which is often called `'layers'` or `'h'`." + ) + }, + ) + + def __post_init__(self): + super().__post_init__() + self.peft_type = PeftType.VERA + self.target_modules = ( + set(self.target_modules) if isinstance(self.target_modules, list) else self.target_modules + ) + # check for layers_to_transform and layers_pattern + if self.layers_pattern and not self.layers_to_transform: + raise ValueError("When `layers_pattern` is specified, `layers_to_transform` must also be specified. ") + if not self.save_projection: + warnings.warn( + "Specified to not save vera_A and vera_B within the state dictionary, instead they will be restored " + "using the PRNG key store in `config.projection_prng_key`. Consider setting `config.save_projection` " + "to `True` to guarantee restoring the checkpoint correctly on all system configurations." + ) diff --git a/lib/python3.12/site-packages/peft/tuners/vera/layer.py b/lib/python3.12/site-packages/peft/tuners/vera/layer.py new file mode 100644 index 0000000000000000000000000000000000000000..25789e8f59cdf70de58a280b620cd6d2bd356c75 --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/vera/layer.py @@ -0,0 +1,294 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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. + +import warnings +from typing import List, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.pytorch_utils import Conv1D + +from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge +from peft.utils.other import transpose + +from .._buffer_dict import BufferDict + + +class VeraLayer(BaseTunerLayer): + # List all names of layers that may contain adapter weights + adapter_layer_names = ("vera_lambda_b", "vera_lambda_d") + other_param_names = ("vera_A", "vera_B") + + def __init__(self, base_layer: nn.Module, **kwargs): + self.base_layer = base_layer + self.r = {} + self.vera_dropout = nn.ModuleDict({}) + + # For storing vector scale + self.vera_lambda_b = nn.ParameterDict({}) + self.vera_lambda_d = nn.ParameterDict({}) + + # Stores a reference to the vera_A/B BufferDict. + # Set to `None` otherwise to avoid computation with random weights + self.vera_A: Optional[BufferDict] = None + self.vera_B: Optional[BufferDict] = None + + # Mark the weight as unmerged + self._disable_adapters = False + self.merged_adapters = [] + + base_layer = self.get_base_layer() + if isinstance(base_layer, nn.Linear): + in_features, out_features = base_layer.in_features, base_layer.out_features + elif isinstance(base_layer, Conv1D): + in_features, out_features = ( + base_layer.weight.ds_shape if hasattr(base_layer.weight, "ds_shape") else base_layer.weight.shape + ) + + self.in_features = in_features + self.out_features = out_features + self.kwargs = kwargs + + @property + def merged(self) -> bool: + return bool(self.merged_adapters) + + def update_layer( + self, + adapter_name, + vera_A: BufferDict, + vera_B: BufferDict, + r, + vera_dropout, + init_weights, + d_initial: float = 0.1, + ): + if r <= 0: + raise ValueError(f"`r` should be a positive integer value but the value passed is {r}") + self.r[adapter_name] = r + if vera_dropout > 0.0: + vera_dropout_layer = nn.Dropout(p=vera_dropout) + else: + vera_dropout_layer = nn.Identity() + + self.vera_dropout.update(nn.ModuleDict({adapter_name: vera_dropout_layer})) + # Actual trainable parameters + self.vera_lambda_b[adapter_name] = nn.Parameter(torch.ones(self.out_features), requires_grad=True) + self.vera_lambda_d[adapter_name] = nn.Parameter(torch.randn(r), requires_grad=True) + + # non trainable references to vera_A/B buffers + self.vera_A = vera_A + self.vera_B = vera_B + if adapter_name not in vera_A: + # This means that this is not the first VeRA adapter. We have to add an entry in the dict for this adapter. + if len(self.vera_A) < 1: + raise ValueError( + "The `vera_A` and `vera_B` buffers are empty. This should not happen. Please report this issue." + ) + # we can take any of the existing adapter's parameters, as they should all be identical + vera_A_param = list(self.vera_A.values())[0] + vera_B_param = list(self.vera_B.values())[0] + + error_tmpl = ( + "{} has a size of {} but {} or greater is required; this probably happened because an additional VeRA " + "adapter was added after the first one with incompatible shapes." + ) + # check input size + if vera_A_param.shape[1] < self.in_features: + raise ValueError(error_tmpl.format("vera_A", vera_A_param.shape[1], self.in_features)) + # check output size + if vera_B_param.shape[0] < self.out_features: + raise ValueError(error_tmpl.format("vera_B", vera_B_param.shape[0], self.out_features)) + # check r + error_tmpl = ( + "{} has a size of {} but {} or greater is required; this probably happened because an additional VeRA " + "adapter with a lower rank was added after the first one; loading the adapters " + "in reverse order may solve this." + ) + if vera_A_param.shape[0] < self.r[adapter_name]: + raise ValueError(error_tmpl.format("vera_A", vera_A_param.shape[0], self.r[adapter_name])) + if vera_B_param.shape[1] < self.r[adapter_name]: + raise ValueError(error_tmpl.format("vera_B", vera_B_param.shape[1], self.r[adapter_name])) + + self.vera_A[adapter_name] = vera_A_param + self.vera_B[adapter_name] = vera_B_param + + if init_weights: + self.reset_vera_parameters(adapter_name, d_initial=d_initial) + + self._move_adapter_to_device_of_base_layer(adapter_name) + self.set_adapter(self.active_adapters) + + def reset_vera_parameters(self, adapter_name, d_initial: float = 0.1): + if adapter_name in self.vera_lambda_d.keys(): + with torch.no_grad(): + nn.init.zeros_(self.vera_lambda_d[adapter_name]).fill_(d_initial) + nn.init.zeros_(self.vera_lambda_b[adapter_name]) + + +class Linear(nn.Linear, VeraLayer): + # Vera implemented in a dense layer + def __init__( + self, + base_layer, + vera_A: BufferDict, + vera_B: BufferDict, + adapter_name: str, + r: int = 0, + vera_dropout: float = 0.0, + fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) + is_target_conv_1d_layer: bool = False, + init_weights: bool = True, + d_initial: float = 0.1, + **kwargs, + ) -> None: + # this gets the init from nn.Linear's super perspective, i.e. nn.Module.__init__, which should always be called + super(nn.Linear, self).__init__() + VeraLayer.__init__(self, base_layer, **kwargs) + self.fan_in_fan_out = fan_in_fan_out + + self._active_adapter = adapter_name + self.update_layer(adapter_name, vera_A, vera_B, r, vera_dropout, init_weights, d_initial=d_initial) + self.is_target_conv_1d_layer = is_target_conv_1d_layer + + def merge(self, safe_merge: bool = False, adapter_names: Optional[List[str]] = None) -> None: + """ + Merge the active adapter weights into the base weights + + Args: + safe_merge (`bool`, *optional*): + If True, the merge operation will be performed in a copy of the original weights and check for NaNs + before merging the weights. This is useful if you want to check if the merge operation will produce + NaNs. Defaults to `False`. + adapter_names (`List[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + """ + adapter_names = check_adapters_to_merge(self, adapter_names) + if not adapter_names: + # no adapter to merge + return + + for active_adapter in adapter_names: + if active_adapter in self.vera_lambda_d.keys(): + base_layer = self.get_base_layer() + if safe_merge: + # Note that safe_merge will be slower than the normal merge + # because of the copy operation. + orig_weights = base_layer.weight.data.clone() + + orig_weights += self.get_delta_weight(active_adapter) + + if not torch.isfinite(orig_weights).all(): + raise ValueError( + f"NaNs detected in the merged weights. The adapter {active_adapter} seems to be broken" + ) + + base_layer.weight.data = orig_weights + else: + base_layer.weight.data += self.get_delta_weight(active_adapter) + self.merged_adapters.append(active_adapter) + + def unmerge(self) -> None: + if not self.merged: + warnings.warn("Already unmerged. Nothing to do.") + return + + while len(self.merged_adapters) > 0: + active_adapter = self.merged_adapters.pop() + if active_adapter in self.vera_lambda_d.keys(): + self.get_base_layer().weight.data -= self.get_delta_weight(active_adapter) + + def get_delta_weight(self, adapter) -> torch.Tensor: + """ + Compute the delta weight for the given adapter. + + Args: + adapter (str): + The name of the adapter for which the delta weight should be computed. + """ + vera_A = self.vera_A[adapter] + vera_B = self.vera_B[adapter] + + device = vera_B.device + dtype = vera_B.dtype + + # In case users wants to merge the adapter weights that are in + # (b)float16 while being on CPU, we need to cast the weights to float32, perform the merge and then cast back to + # (b)float16 because some CPUs have slow bf16/fp16 matmuls. + cast_to_fp32 = device.type == "cpu" and (dtype == torch.float16 or dtype == torch.bfloat16) + + lambda_d = self.vera_lambda_d[adapter] + lambda_b = self.vera_lambda_b[adapter] + + if cast_to_fp32: + vera_A = vera_A.float() + vera_B = vera_B.float() + lambda_d = lambda_d.float() + lambda_b = lambda_b.float() + + sliced_A = vera_A[:, : self.in_features].to(lambda_d.device) + sliced_B = vera_B[: self.out_features, :].to(lambda_d.device) + lambda_b = lambda_b.unsqueeze(-1) + lambda_d = lambda_d.unsqueeze(-1) + output_tensor = transpose((lambda_b * sliced_B) @ (lambda_d * sliced_A), self.fan_in_fan_out) + + if cast_to_fp32: + output_tensor = output_tensor.to(dtype=dtype) + + # cast back the weights + # TODO: why? + self.vera_lambda_d[adapter].data = lambda_d.to(dtype) + self.vera_lambda_b[adapter].data = lambda_b.to(dtype) + + return output_tensor + + def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + previous_dtype = x.dtype + + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + result = self.base_layer(x, *args, **kwargs) + for active_adapter in self.active_adapters: + if active_adapter not in self.vera_lambda_d.keys(): + continue + + lambda_d = self.vera_lambda_d[active_adapter] + lambda_b = self.vera_lambda_b[active_adapter] + + vera_A = self.vera_A[active_adapter] + vera_B = self.vera_B[active_adapter] + + # As adapted layers may have different shapes and VeRA contains a single shared pair of A and B matrices, + # we initialize these matrices with the largest required size for each dimension. + # During the forward pass, required submatrices are sliced out from the shared vera_A and vera_B. + sliced_A = vera_A[:, : self.in_features].to(x.device) + sliced_B = vera_B[: self.out_features, :].to(x.device) + + dropout = self.vera_dropout[active_adapter] + x = x.to(lambda_d.dtype) + result = result + lambda_b * F.linear(lambda_d * F.linear(dropout(x), sliced_A), sliced_B) + + result = result.to(previous_dtype) + return result + + def __repr__(self) -> str: + rep = super().__repr__() + return "vera." + rep diff --git a/lib/python3.12/site-packages/peft/tuners/vera/model.py b/lib/python3.12/site-packages/peft/tuners/vera/model.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf1dc57d61f19a90014f4cb85ee738cb44fa84b --- /dev/null +++ b/lib/python3.12/site-packages/peft/tuners/vera/model.py @@ -0,0 +1,507 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 math +import warnings +from dataclasses import asdict +from enum import Enum +from typing import Optional, Union + +import torch +import torch.nn as nn +from torch.nn.init import _calculate_correct_fan +from tqdm import tqdm +from transformers.pytorch_utils import Conv1D + +from peft.import_utils import is_bnb_4bit_available, is_bnb_available +from peft.tuners.tuners_utils import BaseTuner, BaseTunerLayer, check_target_module_exists +from peft.utils import ( + TRANSFORMERS_MODELS_TO_VERA_TARGET_MODULES_MAPPING, + ModulesToSaveWrapper, + _get_submodules, +) + +from .._buffer_dict import BufferDict +from ..tuners_utils import _maybe_include_all_linear_layers +from .config import VeraConfig +from .layer import Linear, VeraLayer + + +def _kaiming_init( + tensor_or_shape: Union[torch.Tensor, tuple[int, ...]], + generator: torch.Generator, +) -> torch.Tensor: + """ + Kaiming Uniform Initialisation adapted to accept a `torch.Generator` object for PRNG. + + Args: + tensor_or_shape (`Union[torch.Tensor, tuple[int, ...]]`): + Tensor to initialise, or shape of new tensor to create and then initialise. + generator: (`torch.Generator`): + Generator object that manages the state of the PRNG algorithm in use. + + Returns: + `torch.Tensor`: The initialised tensor. + """ + if isinstance(tensor_or_shape, tuple): + tensor = torch.empty(tensor_or_shape) + else: + tensor = tensor_or_shape + fan = _calculate_correct_fan(tensor, "fan_in") + gain = math.sqrt(2) + std = gain / math.sqrt(fan) + bound = math.sqrt(3.0) * std + + with torch.no_grad(): + return tensor.uniform_(-bound, bound, generator=generator) + + +class VeraModel(BaseTuner): + """ + Creates Vector-based Random Matrix Adaptation (Vera) model from a pretrained transformers model. + + Args: + model ([`~transformers.PreTrainedModel`]): The model to be adapted. + config ([`VeraConfig`]): The configuration of the Vera model. + adapter_name (`str`): The name of the adapter, defaults to `"default"`. + low_cpu_mem_usage (`bool`, `optional`, defaults to `False`): + Create empty adapter weights on meta device. Useful to speed up the loading process. + + Returns: + `torch.nn.Module`: The Vera model. + + Example: + + ```py + >>> from transformers import AutoModelForCausalLM + >>> from peft import VeraConfig, get_peft_model + + >>> base_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m") + >>> config = VeraConfig(r=128) + >>> model = get_peft_model(base_model, config) + ``` + + **Attributes**: + - **model** ([`~transformers.PreTrainedModel`]) -- The model to be adapted. + - **peft_config** ([`VeraConfig`]): The configuration of the Vera model. + """ + + prefix: str = "vera_lambda_" + + def __init__(self, model, config, adapter_name, low_cpu_mem_usage: bool = False) -> None: + super().__init__(model, config, adapter_name, low_cpu_mem_usage=low_cpu_mem_usage) + + def _find_dim(self, config) -> tuple[int, int]: + """ + Finds the largest input and output dimensions across linear layers that have been wrapped with VeRA. + + This will be used for determining the size of the shared vera_A and vera_B matrices. + """ + model_config = self.get_model_config(self.model) + + peft_config = self._prepare_adapter_config(config, model_config) + peft_config = _maybe_include_all_linear_layers(peft_config, self.model) + + largest_shape = None + for key, module in self.model.named_modules(): + if not self._check_target_module_exists(peft_config, key): + continue + + if isinstance(module, nn.Linear): + module_shape = module.out_features, module.in_features + elif isinstance(module, Conv1D): + module_shape = module.weight.ds_shape if hasattr(module.weight, "ds_shape") else module.weight.shape + module_shape = module_shape[::-1] + else: + continue + + if largest_shape is None: + largest_shape = module_shape + continue + + if module_shape != largest_shape: + largest_shape = tuple(max(a, b) for a, b in zip(largest_shape, module_shape)) + + if largest_shape is None: + msg = "No layers types compatible with VeRA were found. Please check `peft_config.target_modules`." + raise ValueError(msg) + + return largest_shape + + def _init_vera_A_vera_B(self, config: VeraConfig, adapter_name: str) -> None: + linear_out_dim, linear_in_dim = self._find_dim(config) + + # use of persistent to exclude vera_A and vera_B from the state dict if we choose not to save them. + self.vera_A = BufferDict({}, persistent=config.save_projection) + self.vera_B = BufferDict({}, persistent=config.save_projection) + + # deterministic init of vera_A and vera_B if we know the key + generator = torch.Generator(device="cpu").manual_seed(config.projection_prng_key) + vera_A = _kaiming_init((config.r, linear_in_dim), generator=generator) + vera_B = _kaiming_init((linear_out_dim, config.r), generator=generator) + + self.vera_A[adapter_name] = vera_A + self.vera_B[adapter_name] = vera_B + + def _pre_injection_hook(self, model: nn.Module, config: VeraConfig, adapter_name: str) -> None: + self._init_vera_A_vera_B(config, adapter_name) + + def _check_new_adapter_config(self, config: VeraConfig) -> None: + """ + A helper method to check the config when a new adapter is being added. + + Raise a ValueError if there is something wrong with the config or if it conflicts with existing adapters. + + """ + # the below todo is copied from LoRA + # TODO: there should be a check if any of the existing adapters actually has bias != "none", or else the check + # does not fully correspond to the error message. + if (len(self.peft_config) > 1) and (config.bias != "none"): + raise ValueError( + f"{self.__class__.__name__} supports only 1 adapter with bias. When using multiple adapters, " + "set bias to 'none' for all adapters." + ) + + for existing_config in self.peft_config.values(): + if existing_config is config: + # skip the current config + continue + + if existing_config.projection_prng_key != config.projection_prng_key: + raise ValueError( + f"Vera PRNG initialisation key must be the same for all adapters. Got {config.projection_prng_key=} but " + f"previous config had {existing_config.projection_prng_key}." + ) + + save_project_unique_values = sorted({config.save_projection for config in self.peft_config.values()}) + if len(save_project_unique_values) > 1: + raise ValueError( + "VeRA projection weights must be saved for all adapters or none, but got multiple different values: " + f"{save_project_unique_values}" + ) + + @staticmethod + def _check_target_module_exists(vera_config, key): + return check_target_module_exists(vera_config, key) + + def _create_and_replace( + self, + vera_config, + adapter_name, + target, + target_name, + parent, + current_key, + **optional_kwargs, + ): + if current_key is None: + raise ValueError("Current Key shouldn't be `None`") + + r = vera_config.r + bias = hasattr(target, "bias") and target.bias is not None + kwargs = { + "r": r, + "vera_dropout": vera_config.vera_dropout, + "fan_in_fan_out": vera_config.fan_in_fan_out, + "init_weights": vera_config.init_weights, + "loaded_in_8bit": getattr(self.model, "is_loaded_in_8bit", False), + "loaded_in_4bit": getattr(self.model, "is_loaded_in_4bit", False), + } + kwargs["bias"] = bias + + if isinstance(target, Linear): + target.update_layer( + adapter_name, + self.vera_A, + self.vera_B, + r, + vera_config.vera_dropout, + vera_config.init_weights, + d_initial=vera_config.d_initial, + ) + else: + new_module = self._create_new_module(vera_config, self.vera_A, self.vera_B, adapter_name, target, **kwargs) + if adapter_name not in self.active_adapter: + # adding an additional adapter: it is not automatically trainable + new_module.requires_grad_(False) + self._replace_module(parent, target_name, new_module, target) + + @staticmethod + def _replace_module(parent, child_name, new_module, child): + setattr(parent, child_name, new_module) + # It's not necessary to set requires_grad here, as that is handled by + # _mark_only_adapters_as_trainable + + # child layer wraps the original module, unpack it + if hasattr(child, "base_layer"): + child = child.base_layer + + if not hasattr(new_module, "base_layer"): + new_module.weight = child.weight + if hasattr(child, "bias"): + new_module.bias = child.bias + + if getattr(child, "state", None) is not None: + if hasattr(new_module, "base_layer"): + new_module.base_layer.state = child.state + else: + new_module.state = child.state + new_module.to(child.weight.device) + + meta = torch.device("meta") + # dispatch to correct device + for name, module in new_module.named_modules(): + if "vera_" in name: + if not any(p.device == meta for p in module.parameters()): + module.to(child.weight.device) + + def _mark_only_adapters_as_trainable(self, model: nn.Module) -> None: + for n, p in model.named_parameters(): + if self.prefix not in n: + p.requires_grad = False + + for active_adapter in self.active_adapters: + bias = self.peft_config[active_adapter].bias + if bias == "none": + continue + + if bias == "all": + for n, p in model.named_parameters(): + if "bias" in n: + p.requires_grad = True + elif bias == "vera_only": + for m in model.modules(): + if isinstance(m, VeraLayer) and hasattr(m, "bias") and m.bias is not None: + m.bias.requires_grad = True + else: + raise NotImplementedError(f"Requested bias: {bias}, is not implemented.") + + @staticmethod + def _create_new_module(vera_config, vera_A, vera_B, adapter_name, target, **kwargs): + # avoid eager bnb import + if is_bnb_available(): + import bitsandbytes as bnb + + from .bnb import Linear8bitLt + + if is_bnb_4bit_available(): + from .bnb import Linear4bit + + bias = kwargs.pop("bias", False) + loaded_in_8bit = kwargs.get("loaded_in_8bit", False) + loaded_in_4bit = kwargs.get("loaded_in_4bit", False) + + if isinstance(target, BaseTunerLayer): + target_base_layer = target.get_base_layer() + else: + target_base_layer = target + + if loaded_in_8bit and isinstance(target_base_layer, bnb.nn.Linear8bitLt): + eightbit_kwargs = kwargs.copy() + eightbit_kwargs.update( + { + "has_fp16_weights": target_base_layer.state.has_fp16_weights, + "threshold": target_base_layer.state.threshold, + "index": target_base_layer.index, + } + ) + return Linear8bitLt(target, adapter_name, vera_A, vera_B, **eightbit_kwargs) + elif loaded_in_4bit and isinstance(target_base_layer, bnb.nn.Linear4bit): + fourbit_kwargs = kwargs.copy() + fourbit_kwargs.update( + { + "compute_dtype": target_base_layer.compute_dtype, + "compress_statistics": target_base_layer.weight.compress_statistics, + "quant_type": target_base_layer.weight.quant_type, + } + ) + return Linear4bit(target, adapter_name, vera_A, vera_B, **fourbit_kwargs) + elif isinstance(target_base_layer, torch.nn.Linear): + if kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to True but the target module is `torch.nn.Linear`. " + "Setting fan_in_fan_out to False." + ) + kwargs["fan_in_fan_out"] = vera_config.fan_in_fan_out = False + elif isinstance(target_base_layer, Conv1D): + kwargs["is_target_conv_1d_layer"] = True + if not kwargs["fan_in_fan_out"]: + warnings.warn( + "fan_in_fan_out is set to False but the target module is `Conv1D`. Setting fan_in_fan_out to True." + ) + kwargs["fan_in_fan_out"] = vera_config.fan_in_fan_out = True + else: + raise ValueError( + f"Target module {target} is not supported. Currently, only the following modules are supported: " + "`torch.nn.Linear`, `transformers.pytorch_utils.Conv1D`." + ) + new_module = Linear( + target, + vera_A, + vera_B, + adapter_name, + bias=bias, + d_initial=vera_config.d_initial, + **kwargs, + ) + + return new_module + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + if name == "model": # see #1892: prevent infinite recursion if class is not initialized + raise + return getattr(self.model, name) + + def get_peft_config_as_dict(self, inference: bool = False): + config_dict = {} + for key, value in self.peft_config.items(): + config = {k: v.value if isinstance(v, Enum) else v for k, v in asdict(value).items()} + if inference: + config["inference_mode"] = True + config_dict[key] = config + return config + + def _set_adapter_layers(self, enabled=True): + for module in self.model.modules(): + if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)): + module.enable_adapters(enabled) + + def enable_adapter_layers(self): + self._set_adapter_layers(enabled=True) + + def disable_adapter_layers(self): + for active_adapter in self.active_adapters: + val = self.peft_config[active_adapter].bias + if val != "none": + msg = ( + f"Careful, disabling adapter layers with bias configured to be '{val}' does not produce the same " + "output as the the base model would without adaption." + ) + warnings.warn(msg) + self._set_adapter_layers(enabled=False) + + def set_adapter(self, adapter_name): + for module in self.model.modules(): + if isinstance(module, VeraLayer): + if module.merged: + warnings.warn("Adapter cannot be set when the model is merged. Unmerging the model first.") + module.unmerge() + module.set_adapter(adapter_name) + self.active_adapter = adapter_name + + @staticmethod + def _prepare_adapter_config(peft_config, model_config): + if peft_config.target_modules is None: + if model_config["model_type"] not in TRANSFORMERS_MODELS_TO_VERA_TARGET_MODULES_MAPPING: + raise ValueError("Please specify `target_modules` in `peft_config`") + peft_config.target_modules = set( + TRANSFORMERS_MODELS_TO_VERA_TARGET_MODULES_MAPPING[model_config["model_type"]] + ) + return peft_config + + def _unload_and_optionally_merge( + self, + merge=True, + progressbar: bool = False, + safe_merge: bool = False, + adapter_names: Optional[list[str]] = None, + ): + # we cannot use self.prefix as we want to include non-trainable vera parameters + key_list = [key for key, _ in self.model.named_modules() if "vera" not in key] + desc = "Unloading " + ("and merging " if merge else "") + "model" + for key in tqdm(key_list, disable=not progressbar, desc=desc): + try: + parent, target, target_name = _get_submodules(self.model, key) + except AttributeError: + continue + + if hasattr(target, "base_layer"): + if merge: + target.merge(safe_merge=safe_merge, adapter_names=adapter_names) + + self._replace_module(parent, target_name, target.get_base_layer(), target) + elif isinstance(target, ModulesToSaveWrapper): + # save any additional trainable modules part of `modules_to_save` + setattr(parent, target_name, target.modules_to_save[target.active_adapter]) + + return self.model + + def delete_adapter(self, adapter_name: str): + """ + Deletes an existing adapter. + + Args: + adapter_name (str): Name of the adapter to be deleted. + """ + if adapter_name not in list(self.peft_config.keys()): + raise ValueError(f"Adapter {adapter_name} does not exist") + del self.peft_config[adapter_name] + + # we cannot use self.prefix as we want to include non-trainable vera parameters + key_list = [key for key, _ in self.model.named_modules() if "vera" not in key] + new_adapter = None + for key in key_list: + _, target, _ = _get_submodules(self.model, key) + if isinstance(target, VeraLayer): + target.delete_adapter(adapter_name) + if new_adapter is None: + new_adapter = target.active_adapter[:] + + self.active_adapter = new_adapter or [] + + def merge_and_unload( + self, progressbar: bool = False, safe_merge: bool = False, adapter_names: Optional[list[str]] = None + ): + r""" + This method merges the Vera layers into the base model. This is needed if someone wants to use the base model + as a standalone model. + + Args: + progressbar (`bool`): + whether to show a progressbar indicating the unload and merge process + safe_merge (`bool`): + whether to activate the safe merging check to check if there is any potential Nan in the adapter + weights + adapter_names (`list[str]`, *optional*): + The list of adapter names that should be merged. If None, all active adapters will be merged. Defaults + to `None`. + + Example: + + ```py + >>> from transformers import AutoModelForCausalLM + >>> from peft import PeftModel + + >>> base_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-40b") + >>> peft_model_id = "smangrul/falcon-40B-int4-peft-lora-sfttrainer-sample" + >>> model = PeftModel.from_pretrained(base_model, peft_model_id) + >>> merged_model = model.merge_and_unload() + ``` + """ + return self._unload_and_optionally_merge( + progressbar=progressbar, safe_merge=safe_merge, adapter_names=adapter_names + ) + + def unload(self): + """ + Gets back the base model by removing all the Vera modules without merging. This gives back the original base + model. + """ + return self._unload_and_optionally_merge(merge=False) diff --git a/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/INSTALLER b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/METADATA b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..46c806e797e8b8147a51873e637b0a7d3a670e4c --- /dev/null +++ b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/METADATA @@ -0,0 +1,245 @@ +Metadata-Version: 2.4 +Name: tensordict +Version: 0.11.0 +Summary: TensorDict is a pytorch dedicated tensor container. +Author-email: Vincent Moens +License: BSD +Project-URL: homepage, https://github.com/pytorch/tensordict +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Classifier: Development Status :: 4 - Beta +Requires-Python: >=3.10 +Description-Content-Type: text/markdown +License-File: LICENSE +Requires-Dist: torch +Requires-Dist: numpy +Requires-Dist: cloudpickle +Requires-Dist: packaging +Requires-Dist: importlib_metadata +Requires-Dist: orjson; python_version < "3.13" +Requires-Dist: pyvers<0.3.0,>=0.2.0 +Provides-Extra: tests +Requires-Dist: pytest; extra == "tests" +Requires-Dist: pyyaml; extra == "tests" +Requires-Dist: pytest-instafail; extra == "tests" +Requires-Dist: pytest-rerunfailures; extra == "tests" +Requires-Dist: pytest-benchmark; extra == "tests" +Provides-Extra: h5 +Requires-Dist: h5py>=3.8; extra == "h5" +Provides-Extra: dev +Requires-Dist: pybind11; extra == "dev" +Requires-Dist: ninja; extra == "dev" +Provides-Extra: typecheck +Requires-Dist: mypy>=1.0.0; extra == "typecheck" +Provides-Extra: onnx +Requires-Dist: onnx; extra == "onnx" +Requires-Dist: onnxscript; extra == "onnx" +Requires-Dist: onnxruntime; extra == "onnx" +Dynamic: license-file + + + +[![Docs - GitHub.io](https://img.shields.io/static/v1?logo=github&style=flat&color=pink&label=docs&message=tensordict)][#docs-package] +[![Discord Shield](https://dcbadge.vercel.app/api/server/tz3TgTAe3D)](https://discord.gg/tz3TgTAe3D) +[![Benchmarks](https://img.shields.io/badge/Benchmarks-blue.svg)][#docs-package-benchmark] +[![Python version](https://img.shields.io/pypi/pyversions/tensordict.svg)](https://www.python.org/downloads/) +[![GitHub license](https://img.shields.io/badge/license-MIT-blue.svg)][#github-license] +pypi version +pypi nightly version +[![Downloads](https://static.pepy.tech/personalized-badge/tensordict?period=total&units=international_system&left_color=blue&right_color=orange&left_text=Downloads)][#pepy-package] +[![Downloads](https://static.pepy.tech/personalized-badge/tensordict-nightly?period=total&units=international_system&left_color=blue&right_color=orange&left_text=Downloads%20(nightly))][#pepy-package-nightly] +[![codecov](https://codecov.io/gh/pytorch/tensordict/branch/main/graph/badge.svg?token=9QTUG6NAGQ)][#codecov-package] +[![circleci](https://circleci.com/gh/pytorch/tensordict.svg?style=shield)][#circleci-package] +[![Conda - Platform](https://img.shields.io/conda/pn/conda-forge/tensordict?logo=anaconda&style=flat)][#conda-forge-package] +[![Conda (channel only)](https://img.shields.io/conda/vn/conda-forge/tensordict?logo=anaconda&style=flat&color=orange)][#conda-forge-package] + +[#docs-package]: https://pytorch.github.io/tensordict/ +[#docs-package-benchmark]: https://pytorch.github.io/tensordict/dev/bench/ +[#github-license]: https://github.com/pytorch/tensordict/blob/main/LICENSE +[#pepy-package]: https://pepy.tech/project/tensordict +[#pepy-package-nightly]: https://pepy.tech/project/tensordict-nightly +[#codecov-package]: https://codecov.io/gh/pytorch/tensordict +[#circleci-package]: https://circleci.com/gh/pytorch/tensordict +[#conda-forge-package]: https://anaconda.org/conda-forge/tensordict + + + +# 📖 TensorDict + +TensorDict is a dictionary-like class that inherits properties from tensors, making it easy to work with collections of +tensors in PyTorch. It provides a simple and intuitive way to manipulate and process tensors, allowing you to focus on +building and training your models. + +[**Key Features**](#key-features) | +[**Examples**](#examples) | +[**Installation**](#installation) | +[**Citation**](#citation) | +[**License**](#license) + +## Key Features +TensorDict makes your code-bases more _readable_, _compact_, _modular_ and _fast_. +It abstracts away tailored operations, making your code less error-prone as it takes care of +dispatching the operation on the leaves for you. + +The key features are: + +- 🧮 **Composability**: `TensorDict` generalizes `torch.Tensor` operations to collection of tensors. +- ⚡️ **Speed**: asynchronous transfer to device, fast node-to-node communication through `consolidate`, compatible with `torch.compile`. +- ✂️ **Shape operations**: Perform tensor-like operations on TensorDict instances, such as indexing, slicing or + concatenation. +- 🌐 **Distributed / multiprocessed capabilities**: Easily distribute TensorDict instances across multiple workers, + devices and machines. +- 💾 **Serialization** and memory-mapping +- λ **Functional programming** and compatibility with `torch.vmap` +- 📦 **Nesting**: Nest TensorDict instances to create hierarchical structures. +- ⏰ **Lazy preallocation**: Preallocate memory for TensorDict instances without initializing the tensors. +- 📝 **Specialized dataclass** for torch.Tensor ([`@tensorclass`](#tensorclass)) + +![tensordict.png](docs%2Ftensordict.png) + +## Examples + +This section presents a couple of stand-out applications of the library. +Check our [**Getting Started**](GETTING_STARTED.md) guide for an overview of TensorDict's features! + +### Fast copy on device +`TensorDict` optimizes transfers from/to device to make them safe and fast. +By default, data transfers will be made asynchronously and synchronizations will be called whenever needed. +```python +# Fast and safe asynchronous copy to 'cuda' +td_cuda = TensorDict(**dict_of_tensor, device="cuda") +# Fast and safe asynchronous copy to 'cpu' +td_cpu = td_cuda.to("cpu") +# Force synchronous copy +td_cpu = td_cuda.to("cpu", non_blocking=False) +``` + +### Coding an optimizer +For instance, using `TensorDict` you can code the Adam optimizer as you would for a single `torch.Tensor` and apply +that to a `TensorDict` input as well. On `cuda`, these operations will rely on fused kernels, making it very fast to +execute: +```python +class Adam: + def __init__(self, weights: TensorDict, alpha: float=1e-3, + beta1: float=0.9, beta2: float=0.999, + eps: float = 1e-6): + # Lock for efficiency + weights = weights.lock_() + self.weights = weights + self.t = 0 + + self._mu = weights.data.clone() + self._sigma = weights.data.mul(0.0) + self.beta1 = beta1 + self.beta2 = beta2 + self.alpha = alpha + self.eps = eps + + def step(self): + self._mu.mul_(self.beta1).add_(self.weights.grad, 1 - self.beta1) + self._sigma.mul_(self.beta2).add_(self.weights.grad.pow(2), 1 - self.beta2) + self.t += 1 + mu = self._mu.div_(1-self.beta1**self.t) + sigma = self._sigma.div_(1 - self.beta2 ** self.t) + self.weights.data.add_(mu.div_(sigma.sqrt_().add_(self.eps)).mul_(-self.alpha)) +``` + +### Training a model + +Using tensordict primitives, most supervised training loops can be rewritten in a generic way: +```python +for i, data in enumerate(dataset): + # the model reads and writes tensordicts + data = model(data) + loss = loss_module(data) + loss.backward() + optimizer.step() + optimizer.zero_grad() +``` + +With this level of abstraction, one can recycle a training loop for highly heterogeneous task. +Each individual step of the training loop (data collection and transform, model prediction, loss computation etc.) +can be tailored to the use case at hand without impacting the others. +For instance, the above example can be easily used across classification and segmentation tasks, among many others. + +## Installation + +**With Pip**: + +To install the latest stable version of tensordict, simply run + +```bash +pip install tensordict +``` + +This will work with Python 3.7 and upward as well as PyTorch 1.12 and upward. + +To enjoy the latest features, one can use + +```bash +pip install tensordict-nightly +``` + +**With uv + PyTorch nightlies**: + +If you're using a **PyTorch nightly** (e.g. installed from the PyTorch nightly wheel index), then for **editable** +installs you should install tensordict with **`--no-deps`**. + +This avoids uv re-resolving `torch` from its configured indexes (by default: PyPI), which can otherwise replace an +existing nightly with the latest stable `torch`. + +This is an **uv-specific** behavior; plain `pip install -e .` typically won’t replace an already-installed PyTorch +nightly. + +To keep a nightly `torch` with an editable install: + +- install without resolving deps (**recommended**): + +```bash +uv pip install -e . --no-deps +``` + +- or (less recommended) explicitly point uv at the PyTorch nightly wheel index (CPU shown; use the appropriate backend directory like `cu126/`): + +```bash +uv pip install -e . --prerelease=allow -f "https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html" +``` + +**With Conda**: + +Install `tensordict` from `conda-forge` channel. + +```sh +conda install -c conda-forge tensordict +``` + +## Citation + +If you're using TensorDict, please refer to this BibTeX entry to cite this work: +``` +@misc{bou2023torchrl, + title={TorchRL: A data-driven decision-making library for PyTorch}, + author={Albert Bou and Matteo Bettini and Sebastian Dittert and Vikash Kumar and Shagun Sodhani and Xiaomeng Yang and Gianni De Fabritiis and Vincent Moens}, + year={2023}, + eprint={2306.00577}, + archivePrefix={arXiv}, + primaryClass={cs.LG} +} +``` + +## Disclaimer + +TensorDict is at the *beta*-stage, meaning that there may be bc-breaking changes introduced, but +they should come with a warranty. +Hopefully these should not happen too often, as the current roadmap mostly +involves adding new features and building compatibility with the broader +PyTorch ecosystem. + +## License + +TensorDict is licensed under the MIT License. See [LICENSE](LICENSE) for details. diff --git a/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/RECORD b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..a9043ea18cc1e69cdce67c16f8562a8c73dca7e4 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/RECORD @@ -0,0 +1,89 @@ +tensordict-0.11.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +tensordict-0.11.0.dist-info/METADATA,sha256=bdgz6Hi4buW9Sa5PAr_a0FE1xDeKIXA15AKekuMoOyw,10267 +tensordict-0.11.0.dist-info/RECORD,, +tensordict-0.11.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +tensordict-0.11.0.dist-info/WHEEL,sha256=APLcm4xnO3gxtPYIxN0qfo8NUFHpgCLQlh8uXFGE2eA,114 +tensordict-0.11.0.dist-info/licenses/LICENSE,sha256=xdjS4_xk-IwnLuIFCvTYTl9Y8aXRejqpmke3dGam_nI,1098 +tensordict-0.11.0.dist-info/top_level.txt,sha256=4EMgKpsnmq04uFLPJXf8tMQb1POT8QqR1pR74y5wycc,11 +tensordict/_C.so,sha256=QD8VoQ4zDSPFSQeGgKmzgYp0hu7M_JTs8KE75bxvMjU,202224 +tensordict/_C/__init__.pyi,sha256=13Uw6DBr5yALVmGcEQM41EBsNJsr4Cfqsy_pPajywlw,310 +tensordict/__init__.py,sha256=eyE7syu5vPr1V5Efdrc3JHZmq9FWcuKwAXlGuDRV9Wo,3963 +tensordict/__init__.pyi,sha256=BoPXGoWsyGXaLenaWLDf3MW4lIhNoCKKEBKWz4QM9a0,3264 +tensordict/__pycache__/__init__.cpython-312.pyc,, +tensordict/__pycache__/_contextlib.cpython-312.pyc,, +tensordict/__pycache__/_datasets.cpython-312.pyc,, +tensordict/__pycache__/_lazy.cpython-312.pyc,, +tensordict/__pycache__/_nestedkey.cpython-312.pyc,, +tensordict/__pycache__/_pytree.cpython-312.pyc,, +tensordict/__pycache__/_reductions.cpython-312.pyc,, +tensordict/__pycache__/_td.cpython-312.pyc,, +tensordict/__pycache__/_tensorcollection.cpython-312.pyc,, +tensordict/__pycache__/_torch_func.cpython-312.pyc,, +tensordict/__pycache__/_unbatched.cpython-312.pyc,, +tensordict/__pycache__/_version.cpython-312.pyc,, +tensordict/__pycache__/base.cpython-312.pyc,, +tensordict/__pycache__/functional.cpython-312.pyc,, +tensordict/__pycache__/memmap.cpython-312.pyc,, +tensordict/__pycache__/persistent.cpython-312.pyc,, +tensordict/__pycache__/tensorclass.cpython-312.pyc,, +tensordict/__pycache__/tensordict.cpython-312.pyc,, +tensordict/__pycache__/utils.cpython-312.pyc,, +tensordict/_contextlib.py,sha256=4Gu6Y-ncQj1IbGhc9eU5dztXUvVB5WGdZuy1XwZO17U,14281 +tensordict/_datasets.py,sha256=xEE4gSKA_ro6dbANDboAIV1wBSP3OLRWzeLuqpCa4y4,7932 +tensordict/_lazy.py,sha256=kjGWDAwKfGyRb8qVlkz1OGYoJLlJuIQMnZa8mNWvmSA,184867 +tensordict/_nestedkey.py,sha256=DG2owzUNtMaMOUCg7vNn0AKguBQpa4OFsnJ-V5kTqz4,776 +tensordict/_nestedkey.pyi,sha256=M5rCTUHeHEThMW8gYc1EXEcS-WedrL9BOG3wX49uM2c,352 +tensordict/_pytree.py,sha256=-3zVvgrQsWDPqUfcenMWGXlaI_oCrt1uOKW9K6sAnE8,7718 +tensordict/_reductions.py,sha256=o3IJL4gmzxSayPI2fwGL3K270i4ldblFE3i0S0MZhxo,6371 +tensordict/_td.py,sha256=vq0EtYiUeg8FV0sDNWffLzj-vLcqKuB_yfFaSdcdk6w,199870 +tensordict/_tensorcollection.py,sha256=qx-JY6K6bEjESUNIW47m86KPIJVyqSn3GWfmYKR6L_8,879 +tensordict/_tensorcollection.pyi,sha256=Dm_UkKjKX_gSQARUxcwIbyplQ6PSE9eCRP-FV_JY2zI,47003 +tensordict/_tensordict/__init__.py,sha256=bF5EFt_es-0Db0DMPZZR3SZzl_gk_JowW9F2weZ0W4s,558 +tensordict/_tensordict/__pycache__/__init__.cpython-312.pyc,, +tensordict/_torch_func.py,sha256=OE4m6ZAzpr-MZEjgD5X_4bCALMZElJNPlQhGRoufbyI,34822 +tensordict/_unbatched.py,sha256=_vO2Aae6Xo4J7x8WodylOKonGxTJvbUB8B3pWIj5rP0,8068 +tensordict/_version.py,sha256=9eKRDJ72C44i2IPiti-C7phzF429SwV2Nogzt0etpr0,706 +tensordict/base.py,sha256=HRVPevfUzF-BwqASTj-z-RAFPaLhUxjFBkXAWTH1ZzY,627689 +tensordict/functional.py,sha256=0zXEF_qoj4zqihD2balJC0vLHuZrjEq-wMCLevr-50M,21309 +tensordict/memmap.py,sha256=1PqvKca7tA95vaNWVlxt_QdFuyqL5vtb_yxPPTVASpw,41535 +tensordict/nn/__init__.py,sha256=CzMouxwlD113X8A03d27GHFc6d8nRxBMOIBCDKZ1FIY,2451 +tensordict/nn/__pycache__/__init__.cpython-312.pyc,, +tensordict/nn/__pycache__/common.cpython-312.pyc,, +tensordict/nn/__pycache__/cudagraphs.cpython-312.pyc,, +tensordict/nn/__pycache__/ensemble.cpython-312.pyc,, +tensordict/nn/__pycache__/functional_modules.cpython-312.pyc,, +tensordict/nn/__pycache__/params.cpython-312.pyc,, +tensordict/nn/__pycache__/probabilistic.cpython-312.pyc,, +tensordict/nn/__pycache__/sequence.cpython-312.pyc,, +tensordict/nn/__pycache__/tensorclass_module.cpython-312.pyc,, +tensordict/nn/__pycache__/utils.cpython-312.pyc,, +tensordict/nn/common.py,sha256=4ka5l33R0gdPu7_k-OxTubugWZXlcoGLPYrzgmhtbSc,61225 +tensordict/nn/cudagraphs.py,sha256=lvDN3drLmYuUR6wUGZb4cKXloVkcoH3NI4BXNGJ3Sh0,23075 +tensordict/nn/distributions/__init__.py,sha256=P8cETAKzBw9QwtpNlQtJXx_eAHwHxmAP9vH3Xc6Ahno,1317 +tensordict/nn/distributions/__pycache__/__init__.cpython-312.pyc,, +tensordict/nn/distributions/__pycache__/composite.cpython-312.pyc,, +tensordict/nn/distributions/__pycache__/continuous.cpython-312.pyc,, +tensordict/nn/distributions/__pycache__/discrete.cpython-312.pyc,, +tensordict/nn/distributions/__pycache__/truncated_normal.cpython-312.pyc,, +tensordict/nn/distributions/__pycache__/utils.cpython-312.pyc,, +tensordict/nn/distributions/composite.py,sha256=KYBBk0kCJ6ASlvDg9OWE2cr_jPU3Py0mnqphg8mhoNE,24395 +tensordict/nn/distributions/continuous.py,sha256=9dP-AimQjaW_d9TLyPId5C2s57YuQ1xS4pk8od0mb_A,9883 +tensordict/nn/distributions/discrete.py,sha256=-DVZe2K7JCWgejUFT7XZb53yS0DIg9d2ZfRq3RmJOzg,3002 +tensordict/nn/distributions/truncated_normal.py,sha256=Q_pc3h7ZeM3LtmvK5fO0H6acJTjO-jFZ4f7JbcIje8k,6505 +tensordict/nn/distributions/utils.py,sha256=fX6NUeNnWKK6kDdOp8NTM43Lls3vMCF5h-S2teVZw3E,1226 +tensordict/nn/ensemble.py,sha256=ftSNir6yyLYU23_aO2iHckjWYW1jf9-mxMkZ0GWdjE4,5870 +tensordict/nn/functional_modules.py,sha256=cUdAoYYYocu_XS4ut1ruqAofABUZiGMNnM4m4k2V3M8,12696 +tensordict/nn/params.py,sha256=f1w_Rga8KuMjpvycLXOouPu8BCnA-QBjTAsZ2u5JC84,43154 +tensordict/nn/probabilistic.py,sha256=smMXJ-T1CZeUTuyGTrOVHu0UqCCL9J0Hi0E-Zm0B0sk,65494 +tensordict/nn/sequence.py,sha256=UuLDCm6WAMyo36ant9HUM5v2sfmLHQo6fGDAWSm1QNU,29694 +tensordict/nn/tensorclass_module.py,sha256=HU2Or9NQAwgTpF_5MfruZHWhJeILRsbBNPHcWhE7kkk,8717 +tensordict/nn/utils.py,sha256=aSHm7VokNPxll37_uNdX1hpwUvu-mKcVH-UAkqJCff4,20600 +tensordict/persistent.py,sha256=n95ALnjYsvgYNz_rAwVWE91J5qGuG9gIDJ0e4iAuJnA,54555 +tensordict/prototype/__init__.py,sha256=qgMM907GHcCaf_8kG4hdhkymRx4X6zZnduthh6bxT3A,267 +tensordict/prototype/__pycache__/__init__.cpython-312.pyc,, +tensordict/prototype/__pycache__/fx.cpython-312.pyc,, +tensordict/prototype/fx.py,sha256=XCGyfEQT1bHijPYpB7A4rDY_B8WQoaBsYfAdadaAPAo,7749 +tensordict/tensorclass.py,sha256=cFd7LLSc4bPC6b7r8IUnelew4yY4TPYFHPEFLcXAaDw,167132 +tensordict/tensorclass.pyi,sha256=den5CGFsa34nMd1oxj6ghyCM7PxQwzEGeRryrdtu7VY,48342 +tensordict/tensordict.py,sha256=itgrIY3BaQbrfx09IGVND4tr7kWEU85AFBZWUyLVEXY,936 +tensordict/utils.py,sha256=GzdO_mOZNONJKT008JhCyquFwQntRJtsR7yRHPH7MpY,103474 diff --git a/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/REQUESTED b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/WHEEL b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..507c97f39c8650b6401b43cc1b2ef7148643d3e4 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: setuptools (80.10.2) +Root-Is-Purelib: false +Tag: cp312-cp312-manylinux_2_28_x86_64 + diff --git a/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/licenses/LICENSE b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..426ba7a3fc3dacaa7bed9159c49b5a989654a14a --- /dev/null +++ b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/licenses/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) Meta Platforms, Inc. and affiliates. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE. diff --git a/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/top_level.txt b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..dd2b7ca72785306f6cfeaaa44798525ee5c0a506 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict-0.11.0.dist-info/top_level.txt @@ -0,0 +1 @@ +tensordict diff --git a/lib/python3.12/site-packages/transformers/__init__.py b/lib/python3.12/site-packages/transformers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a5f027b9bb7b04e3f119ef98d736565be0e855e3 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/__init__.py @@ -0,0 +1,9849 @@ +# Copyright 2020 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. + +# When adding a new object to this init, remember to add it twice: once inside the `_import_structure` dictionary and +# once inside the `if TYPE_CHECKING` branch. The `TYPE_CHECKING` should have import statements as usual, but they are +# only there for type checking. The `_import_structure` is a dictionary submodule to list of object names, and is used +# to defer the actual importing for when the objects are requested. This way `import transformers` provides the names +# in the namespace without actually importing anything (and especially none of the backends). + +__version__ = "4.51.1" + +from typing import TYPE_CHECKING + +# Check the dependencies satisfy the minimal versions required. +from . import dependency_versions_check +from .utils import ( + OptionalDependencyNotAvailable, + _LazyModule, + is_bitsandbytes_available, + is_essentia_available, + is_flax_available, + is_g2p_en_available, + is_keras_nlp_available, + is_librosa_available, + is_pretty_midi_available, + is_scipy_available, + is_sentencepiece_available, + is_speech_available, + is_tensorflow_text_available, + is_tf_available, + is_timm_available, + is_tokenizers_available, + is_torch_available, + is_torchaudio_available, + is_torchvision_available, + is_vision_available, + logging, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Base objects, independent of any specific backend +_import_structure = { + "agents": [ + "Agent", + "CodeAgent", + "HfApiEngine", + "ManagedAgent", + "PipelineTool", + "ReactAgent", + "ReactCodeAgent", + "ReactJsonAgent", + "Tool", + "Toolbox", + "ToolCollection", + "TransformersEngine", + "launch_gradio_demo", + "load_tool", + "stream_to_gradio", + "tool", + ], + "audio_utils": [], + "commands": [], + "configuration_utils": ["PretrainedConfig"], + "convert_graph_to_onnx": [], + "convert_slow_tokenizers_checkpoints_to_fast": [], + "convert_tf_hub_seq_to_seq_bert_to_pytorch": [], + "data": [ + "DataProcessor", + "InputExample", + "InputFeatures", + "SingleSentenceClassificationProcessor", + "SquadExample", + "SquadFeatures", + "SquadV1Processor", + "SquadV2Processor", + "glue_compute_metrics", + "glue_convert_examples_to_features", + "glue_output_modes", + "glue_processors", + "glue_tasks_num_labels", + "squad_convert_examples_to_features", + "xnli_compute_metrics", + "xnli_output_modes", + "xnli_processors", + "xnli_tasks_num_labels", + ], + "data.data_collator": [ + "DataCollator", + "DataCollatorForLanguageModeling", + "DataCollatorForMultipleChoice", + "DataCollatorForPermutationLanguageModeling", + "DataCollatorForSeq2Seq", + "DataCollatorForSOP", + "DataCollatorForTokenClassification", + "DataCollatorForWholeWordMask", + "DataCollatorWithFlattening", + "DataCollatorWithPadding", + "DefaultDataCollator", + "default_data_collator", + ], + "data.metrics": [], + "data.processors": [], + "debug_utils": [], + "dependency_versions_check": [], + "dependency_versions_table": [], + "dynamic_module_utils": [], + "feature_extraction_sequence_utils": ["SequenceFeatureExtractor"], + "feature_extraction_utils": ["BatchFeature", "FeatureExtractionMixin"], + "file_utils": [], + "generation": [ + "AsyncTextIteratorStreamer", + "CompileConfig", + "GenerationConfig", + "TextIteratorStreamer", + "TextStreamer", + "WatermarkingConfig", + ], + "hf_argparser": ["HfArgumentParser"], + "hyperparameter_search": [], + "image_transforms": [], + "integrations": [ + "is_clearml_available", + "is_comet_available", + "is_dvclive_available", + "is_neptune_available", + "is_optuna_available", + "is_ray_available", + "is_ray_tune_available", + "is_sigopt_available", + "is_swanlab_available", + "is_tensorboard_available", + "is_wandb_available", + ], + "loss": [], + "modelcard": ["ModelCard"], + # Losses + "modeling_tf_pytorch_utils": [ + "convert_tf_weight_name_to_pt_weight_name", + "load_pytorch_checkpoint_in_tf2_model", + "load_pytorch_model_in_tf2_model", + "load_pytorch_weights_in_tf2_model", + "load_tf2_checkpoint_in_pytorch_model", + "load_tf2_model_in_pytorch_model", + "load_tf2_weights_in_pytorch_model", + ], + # Models + "models": [], + "models.albert": ["AlbertConfig"], + "models.align": [ + "AlignConfig", + "AlignProcessor", + "AlignTextConfig", + "AlignVisionConfig", + ], + "models.altclip": [ + "AltCLIPConfig", + "AltCLIPProcessor", + "AltCLIPTextConfig", + "AltCLIPVisionConfig", + ], + "models.aria": [ + "AriaConfig", + "AriaProcessor", + "AriaTextConfig", + ], + "models.audio_spectrogram_transformer": [ + "ASTConfig", + "ASTFeatureExtractor", + ], + "models.auto": [ + "CONFIG_MAPPING", + "FEATURE_EXTRACTOR_MAPPING", + "IMAGE_PROCESSOR_MAPPING", + "MODEL_NAMES_MAPPING", + "PROCESSOR_MAPPING", + "TOKENIZER_MAPPING", + "AutoConfig", + "AutoFeatureExtractor", + "AutoImageProcessor", + "AutoProcessor", + "AutoTokenizer", + ], + "models.autoformer": ["AutoformerConfig"], + "models.aya_vision": ["AyaVisionConfig", "AyaVisionProcessor"], + "models.bamba": ["BambaConfig"], + "models.bark": [ + "BarkCoarseConfig", + "BarkConfig", + "BarkFineConfig", + "BarkProcessor", + "BarkSemanticConfig", + ], + "models.bart": ["BartConfig", "BartTokenizer"], + "models.barthez": [], + "models.bartpho": [], + "models.beit": ["BeitConfig"], + "models.bert": [ + "BasicTokenizer", + "BertConfig", + "BertTokenizer", + "WordpieceTokenizer", + ], + "models.bert_generation": ["BertGenerationConfig"], + "models.bert_japanese": [ + "BertJapaneseTokenizer", + "CharacterTokenizer", + "MecabTokenizer", + ], + "models.bertweet": ["BertweetTokenizer"], + "models.big_bird": ["BigBirdConfig"], + "models.bigbird_pegasus": ["BigBirdPegasusConfig"], + "models.biogpt": [ + "BioGptConfig", + "BioGptTokenizer", + ], + "models.bit": ["BitConfig"], + "models.blenderbot": [ + "BlenderbotConfig", + "BlenderbotTokenizer", + ], + "models.blenderbot_small": [ + "BlenderbotSmallConfig", + "BlenderbotSmallTokenizer", + ], + "models.blip": [ + "BlipConfig", + "BlipProcessor", + "BlipTextConfig", + "BlipVisionConfig", + ], + "models.blip_2": [ + "Blip2Config", + "Blip2Processor", + "Blip2QFormerConfig", + "Blip2VisionConfig", + ], + "models.bloom": ["BloomConfig"], + "models.bridgetower": [ + "BridgeTowerConfig", + "BridgeTowerProcessor", + "BridgeTowerTextConfig", + "BridgeTowerVisionConfig", + ], + "models.bros": [ + "BrosConfig", + "BrosProcessor", + ], + "models.byt5": ["ByT5Tokenizer"], + "models.camembert": ["CamembertConfig"], + "models.canine": [ + "CanineConfig", + "CanineTokenizer", + ], + "models.chameleon": [ + "ChameleonConfig", + "ChameleonProcessor", + "ChameleonVQVAEConfig", + ], + "models.chinese_clip": [ + "ChineseCLIPConfig", + "ChineseCLIPProcessor", + "ChineseCLIPTextConfig", + "ChineseCLIPVisionConfig", + ], + "models.clap": [ + "ClapAudioConfig", + "ClapConfig", + "ClapProcessor", + "ClapTextConfig", + ], + "models.clip": [ + "CLIPConfig", + "CLIPProcessor", + "CLIPTextConfig", + "CLIPTokenizer", + "CLIPVisionConfig", + ], + "models.clipseg": [ + "CLIPSegConfig", + "CLIPSegProcessor", + "CLIPSegTextConfig", + "CLIPSegVisionConfig", + ], + "models.clvp": [ + "ClvpConfig", + "ClvpDecoderConfig", + "ClvpEncoderConfig", + "ClvpFeatureExtractor", + "ClvpProcessor", + "ClvpTokenizer", + ], + "models.code_llama": [], + "models.codegen": [ + "CodeGenConfig", + "CodeGenTokenizer", + ], + "models.cohere": ["CohereConfig"], + "models.cohere2": ["Cohere2Config"], + "models.colpali": [ + "ColPaliConfig", + "ColPaliProcessor", + ], + "models.conditional_detr": ["ConditionalDetrConfig"], + "models.convbert": [ + "ConvBertConfig", + "ConvBertTokenizer", + ], + "models.convnext": ["ConvNextConfig"], + "models.convnextv2": ["ConvNextV2Config"], + "models.cpm": [], + "models.cpmant": [ + "CpmAntConfig", + "CpmAntTokenizer", + ], + "models.ctrl": [ + "CTRLConfig", + "CTRLTokenizer", + ], + "models.cvt": ["CvtConfig"], + "models.dab_detr": ["DabDetrConfig"], + "models.dac": ["DacConfig", "DacFeatureExtractor"], + "models.data2vec": [ + "Data2VecAudioConfig", + "Data2VecTextConfig", + "Data2VecVisionConfig", + ], + "models.dbrx": ["DbrxConfig"], + "models.deberta": [ + "DebertaConfig", + "DebertaTokenizer", + ], + "models.deberta_v2": ["DebertaV2Config"], + "models.decision_transformer": ["DecisionTransformerConfig"], + "models.deepseek_v3": ["DeepseekV3Config"], + "models.deformable_detr": ["DeformableDetrConfig"], + "models.deit": ["DeiTConfig"], + "models.deprecated": [], + "models.deprecated.bort": [], + "models.deprecated.deta": ["DetaConfig"], + "models.deprecated.efficientformer": ["EfficientFormerConfig"], + "models.deprecated.ernie_m": ["ErnieMConfig"], + "models.deprecated.gptsan_japanese": [ + "GPTSanJapaneseConfig", + "GPTSanJapaneseTokenizer", + ], + "models.deprecated.graphormer": ["GraphormerConfig"], + "models.deprecated.jukebox": [ + "JukeboxConfig", + "JukeboxPriorConfig", + "JukeboxTokenizer", + "JukeboxVQVAEConfig", + ], + "models.deprecated.mctct": [ + "MCTCTConfig", + "MCTCTFeatureExtractor", + "MCTCTProcessor", + ], + "models.deprecated.mega": ["MegaConfig"], + "models.deprecated.mmbt": ["MMBTConfig"], + "models.deprecated.nat": ["NatConfig"], + "models.deprecated.nezha": ["NezhaConfig"], + "models.deprecated.open_llama": ["OpenLlamaConfig"], + "models.deprecated.qdqbert": ["QDQBertConfig"], + "models.deprecated.realm": [ + "RealmConfig", + "RealmTokenizer", + ], + "models.deprecated.retribert": [ + "RetriBertConfig", + "RetriBertTokenizer", + ], + "models.deprecated.speech_to_text_2": [ + "Speech2Text2Config", + "Speech2Text2Processor", + "Speech2Text2Tokenizer", + ], + "models.deprecated.tapex": ["TapexTokenizer"], + "models.deprecated.trajectory_transformer": ["TrajectoryTransformerConfig"], + "models.deprecated.transfo_xl": [ + "TransfoXLConfig", + "TransfoXLCorpus", + "TransfoXLTokenizer", + ], + "models.deprecated.tvlt": [ + "TvltConfig", + "TvltFeatureExtractor", + "TvltProcessor", + ], + "models.deprecated.van": ["VanConfig"], + "models.deprecated.vit_hybrid": ["ViTHybridConfig"], + "models.deprecated.xlm_prophetnet": ["XLMProphetNetConfig"], + "models.depth_anything": ["DepthAnythingConfig"], + "models.depth_pro": ["DepthProConfig"], + "models.detr": ["DetrConfig"], + "models.dialogpt": [], + "models.diffllama": ["DiffLlamaConfig"], + "models.dinat": ["DinatConfig"], + "models.dinov2": ["Dinov2Config"], + "models.dinov2_with_registers": ["Dinov2WithRegistersConfig"], + "models.distilbert": [ + "DistilBertConfig", + "DistilBertTokenizer", + ], + "models.dit": [], + "models.donut": [ + "DonutProcessor", + "DonutSwinConfig", + ], + "models.dpr": [ + "DPRConfig", + "DPRContextEncoderTokenizer", + "DPRQuestionEncoderTokenizer", + "DPRReaderOutput", + "DPRReaderTokenizer", + ], + "models.dpt": ["DPTConfig"], + "models.efficientnet": ["EfficientNetConfig"], + "models.electra": [ + "ElectraConfig", + "ElectraTokenizer", + ], + "models.emu3": [ + "Emu3Config", + "Emu3Processor", + "Emu3TextConfig", + "Emu3VQVAEConfig", + ], + "models.encodec": [ + "EncodecConfig", + "EncodecFeatureExtractor", + ], + "models.encoder_decoder": ["EncoderDecoderConfig"], + "models.ernie": ["ErnieConfig"], + "models.esm": ["EsmConfig", "EsmTokenizer"], + "models.falcon": ["FalconConfig"], + "models.falcon_mamba": ["FalconMambaConfig"], + "models.fastspeech2_conformer": [ + "FastSpeech2ConformerConfig", + "FastSpeech2ConformerHifiGanConfig", + "FastSpeech2ConformerTokenizer", + "FastSpeech2ConformerWithHifiGanConfig", + ], + "models.flaubert": ["FlaubertConfig", "FlaubertTokenizer"], + "models.flava": [ + "FlavaConfig", + "FlavaImageCodebookConfig", + "FlavaImageConfig", + "FlavaMultimodalConfig", + "FlavaTextConfig", + ], + "models.fnet": ["FNetConfig"], + "models.focalnet": ["FocalNetConfig"], + "models.fsmt": [ + "FSMTConfig", + "FSMTTokenizer", + ], + "models.funnel": [ + "FunnelConfig", + "FunnelTokenizer", + ], + "models.fuyu": ["FuyuConfig"], + "models.gemma": ["GemmaConfig"], + "models.gemma2": ["Gemma2Config"], + "models.gemma3": ["Gemma3Config", "Gemma3Processor", "Gemma3TextConfig"], + "models.git": [ + "GitConfig", + "GitProcessor", + "GitVisionConfig", + ], + "models.glm": ["GlmConfig"], + "models.glpn": ["GLPNConfig"], + "models.got_ocr2": [ + "GotOcr2Config", + "GotOcr2Processor", + "GotOcr2VisionConfig", + ], + "models.gpt2": [ + "GPT2Config", + "GPT2Tokenizer", + ], + "models.gpt_bigcode": ["GPTBigCodeConfig"], + "models.gpt_neo": ["GPTNeoConfig"], + "models.gpt_neox": ["GPTNeoXConfig"], + "models.gpt_neox_japanese": ["GPTNeoXJapaneseConfig"], + "models.gpt_sw3": [], + "models.gptj": ["GPTJConfig"], + "models.granite": ["GraniteConfig"], + "models.granitemoe": ["GraniteMoeConfig"], + "models.granitemoeshared": ["GraniteMoeSharedConfig"], + "models.grounding_dino": [ + "GroundingDinoConfig", + "GroundingDinoProcessor", + ], + "models.groupvit": [ + "GroupViTConfig", + "GroupViTTextConfig", + "GroupViTVisionConfig", + ], + "models.helium": ["HeliumConfig"], + "models.herbert": ["HerbertTokenizer"], + "models.hiera": ["HieraConfig"], + "models.hubert": ["HubertConfig"], + "models.ibert": ["IBertConfig"], + "models.idefics": ["IdeficsConfig"], + "models.idefics2": ["Idefics2Config"], + "models.idefics3": ["Idefics3Config"], + "models.ijepa": ["IJepaConfig"], + "models.imagegpt": ["ImageGPTConfig"], + "models.informer": ["InformerConfig"], + "models.instructblip": [ + "InstructBlipConfig", + "InstructBlipProcessor", + "InstructBlipQFormerConfig", + "InstructBlipVisionConfig", + ], + "models.instructblipvideo": [ + "InstructBlipVideoConfig", + "InstructBlipVideoProcessor", + "InstructBlipVideoQFormerConfig", + "InstructBlipVideoVisionConfig", + ], + "models.jamba": ["JambaConfig"], + "models.jetmoe": ["JetMoeConfig"], + "models.kosmos2": [ + "Kosmos2Config", + "Kosmos2Processor", + ], + "models.layoutlm": [ + "LayoutLMConfig", + "LayoutLMTokenizer", + ], + "models.layoutlmv2": [ + "LayoutLMv2Config", + "LayoutLMv2FeatureExtractor", + "LayoutLMv2ImageProcessor", + "LayoutLMv2Processor", + "LayoutLMv2Tokenizer", + ], + "models.layoutlmv3": [ + "LayoutLMv3Config", + "LayoutLMv3FeatureExtractor", + "LayoutLMv3ImageProcessor", + "LayoutLMv3Processor", + "LayoutLMv3Tokenizer", + ], + "models.layoutxlm": ["LayoutXLMProcessor"], + "models.led": ["LEDConfig", "LEDTokenizer"], + "models.levit": ["LevitConfig"], + "models.lilt": ["LiltConfig"], + "models.llama": ["LlamaConfig"], + "models.llama4": [ + "Llama4Config", + "Llama4Processor", + "Llama4TextConfig", + "Llama4VisionConfig", + ], + "models.llava": [ + "LlavaConfig", + "LlavaProcessor", + ], + "models.llava_next": [ + "LlavaNextConfig", + "LlavaNextProcessor", + ], + "models.llava_next_video": [ + "LlavaNextVideoConfig", + "LlavaNextVideoProcessor", + ], + "models.llava_onevision": ["LlavaOnevisionConfig", "LlavaOnevisionProcessor"], + "models.longformer": [ + "LongformerConfig", + "LongformerTokenizer", + ], + "models.longt5": ["LongT5Config"], + "models.luke": [ + "LukeConfig", + "LukeTokenizer", + ], + "models.lxmert": [ + "LxmertConfig", + "LxmertTokenizer", + ], + "models.m2m_100": ["M2M100Config"], + "models.mamba": ["MambaConfig"], + "models.mamba2": ["Mamba2Config"], + "models.marian": ["MarianConfig"], + "models.markuplm": [ + "MarkupLMConfig", + "MarkupLMFeatureExtractor", + "MarkupLMProcessor", + "MarkupLMTokenizer", + ], + "models.mask2former": ["Mask2FormerConfig"], + "models.maskformer": [ + "MaskFormerConfig", + "MaskFormerSwinConfig", + ], + "models.mbart": ["MBartConfig"], + "models.mbart50": [], + "models.megatron_bert": ["MegatronBertConfig"], + "models.megatron_gpt2": [], + "models.mgp_str": [ + "MgpstrConfig", + "MgpstrProcessor", + "MgpstrTokenizer", + ], + "models.mimi": ["MimiConfig"], + "models.mistral": ["MistralConfig"], + "models.mistral3": ["Mistral3Config"], + "models.mixtral": ["MixtralConfig"], + "models.mllama": [ + "MllamaConfig", + "MllamaProcessor", + ], + "models.mluke": [], + "models.mobilebert": [ + "MobileBertConfig", + "MobileBertTokenizer", + ], + "models.mobilenet_v1": ["MobileNetV1Config"], + "models.mobilenet_v2": ["MobileNetV2Config"], + "models.mobilevit": ["MobileViTConfig"], + "models.mobilevitv2": ["MobileViTV2Config"], + "models.modernbert": ["ModernBertConfig"], + "models.moonshine": ["MoonshineConfig"], + "models.moshi": [ + "MoshiConfig", + "MoshiDepthConfig", + ], + "models.mpnet": [ + "MPNetConfig", + "MPNetTokenizer", + ], + "models.mpt": ["MptConfig"], + "models.mra": ["MraConfig"], + "models.mt5": ["MT5Config"], + "models.musicgen": [ + "MusicgenConfig", + "MusicgenDecoderConfig", + ], + "models.musicgen_melody": [ + "MusicgenMelodyConfig", + "MusicgenMelodyDecoderConfig", + ], + "models.mvp": ["MvpConfig", "MvpTokenizer"], + "models.myt5": ["MyT5Tokenizer"], + "models.nemotron": ["NemotronConfig"], + "models.nllb": [], + "models.nllb_moe": ["NllbMoeConfig"], + "models.nougat": ["NougatProcessor"], + "models.nystromformer": ["NystromformerConfig"], + "models.olmo": ["OlmoConfig"], + "models.olmo2": ["Olmo2Config"], + "models.olmoe": ["OlmoeConfig"], + "models.omdet_turbo": [ + "OmDetTurboConfig", + "OmDetTurboProcessor", + ], + "models.oneformer": [ + "OneFormerConfig", + "OneFormerProcessor", + ], + "models.openai": [ + "OpenAIGPTConfig", + "OpenAIGPTTokenizer", + ], + "models.opt": ["OPTConfig"], + "models.owlv2": [ + "Owlv2Config", + "Owlv2Processor", + "Owlv2TextConfig", + "Owlv2VisionConfig", + ], + "models.owlvit": [ + "OwlViTConfig", + "OwlViTProcessor", + "OwlViTTextConfig", + "OwlViTVisionConfig", + ], + "models.paligemma": ["PaliGemmaConfig"], + "models.patchtsmixer": ["PatchTSMixerConfig"], + "models.patchtst": ["PatchTSTConfig"], + "models.pegasus": [ + "PegasusConfig", + "PegasusTokenizer", + ], + "models.pegasus_x": ["PegasusXConfig"], + "models.perceiver": [ + "PerceiverConfig", + "PerceiverTokenizer", + ], + "models.persimmon": ["PersimmonConfig"], + "models.phi": ["PhiConfig"], + "models.phi3": ["Phi3Config"], + "models.phi4_multimodal": [ + "Phi4MultimodalAudioConfig", + "Phi4MultimodalConfig", + "Phi4MultimodalFeatureExtractor", + "Phi4MultimodalProcessor", + "Phi4MultimodalVisionConfig", + ], + "models.phimoe": ["PhimoeConfig"], + "models.phobert": ["PhobertTokenizer"], + "models.pix2struct": [ + "Pix2StructConfig", + "Pix2StructProcessor", + "Pix2StructTextConfig", + "Pix2StructVisionConfig", + ], + "models.pixtral": ["PixtralProcessor", "PixtralVisionConfig"], + "models.plbart": ["PLBartConfig"], + "models.poolformer": ["PoolFormerConfig"], + "models.pop2piano": ["Pop2PianoConfig"], + "models.prompt_depth_anything": ["PromptDepthAnythingConfig"], + "models.prophetnet": [ + "ProphetNetConfig", + "ProphetNetTokenizer", + ], + "models.pvt": ["PvtConfig"], + "models.pvt_v2": ["PvtV2Config"], + "models.qwen2": [ + "Qwen2Config", + "Qwen2Tokenizer", + ], + "models.qwen2_5_vl": [ + "Qwen2_5_VLConfig", + "Qwen2_5_VLProcessor", + ], + "models.qwen2_audio": [ + "Qwen2AudioConfig", + "Qwen2AudioEncoderConfig", + "Qwen2AudioProcessor", + ], + "models.qwen2_moe": ["Qwen2MoeConfig"], + "models.qwen2_vl": [ + "Qwen2VLConfig", + "Qwen2VLProcessor", + ], + "models.qwen3": ["Qwen3Config"], + "models.qwen3_moe": ["Qwen3MoeConfig"], + "models.rag": ["RagConfig", "RagRetriever", "RagTokenizer"], + "models.recurrent_gemma": ["RecurrentGemmaConfig"], + "models.reformer": ["ReformerConfig"], + "models.regnet": ["RegNetConfig"], + "models.rembert": ["RemBertConfig"], + "models.resnet": ["ResNetConfig"], + "models.roberta": [ + "RobertaConfig", + "RobertaTokenizer", + ], + "models.roberta_prelayernorm": ["RobertaPreLayerNormConfig"], + "models.roc_bert": [ + "RoCBertConfig", + "RoCBertTokenizer", + ], + "models.roformer": [ + "RoFormerConfig", + "RoFormerTokenizer", + ], + "models.rt_detr": ["RTDetrConfig", "RTDetrResNetConfig"], + "models.rt_detr_v2": ["RTDetrV2Config"], + "models.rwkv": ["RwkvConfig"], + "models.sam": [ + "SamConfig", + "SamMaskDecoderConfig", + "SamProcessor", + "SamPromptEncoderConfig", + "SamVisionConfig", + ], + "models.seamless_m4t": [ + "SeamlessM4TConfig", + "SeamlessM4TFeatureExtractor", + "SeamlessM4TProcessor", + ], + "models.seamless_m4t_v2": ["SeamlessM4Tv2Config"], + "models.segformer": ["SegformerConfig"], + "models.seggpt": ["SegGptConfig"], + "models.sew": ["SEWConfig"], + "models.sew_d": ["SEWDConfig"], + "models.shieldgemma2": [ + "ShieldGemma2Config", + "ShieldGemma2Processor", + ], + "models.siglip": [ + "SiglipConfig", + "SiglipProcessor", + "SiglipTextConfig", + "SiglipVisionConfig", + ], + "models.siglip2": [ + "Siglip2Config", + "Siglip2Processor", + "Siglip2TextConfig", + "Siglip2VisionConfig", + ], + "models.smolvlm": ["SmolVLMConfig"], + "models.speech_encoder_decoder": ["SpeechEncoderDecoderConfig"], + "models.speech_to_text": [ + "Speech2TextConfig", + "Speech2TextFeatureExtractor", + "Speech2TextProcessor", + ], + "models.speecht5": [ + "SpeechT5Config", + "SpeechT5FeatureExtractor", + "SpeechT5HifiGanConfig", + "SpeechT5Processor", + ], + "models.splinter": [ + "SplinterConfig", + "SplinterTokenizer", + ], + "models.squeezebert": [ + "SqueezeBertConfig", + "SqueezeBertTokenizer", + ], + "models.stablelm": ["StableLmConfig"], + "models.starcoder2": ["Starcoder2Config"], + "models.superglue": ["SuperGlueConfig"], + "models.superpoint": ["SuperPointConfig"], + "models.swiftformer": ["SwiftFormerConfig"], + "models.swin": ["SwinConfig"], + "models.swin2sr": ["Swin2SRConfig"], + "models.swinv2": ["Swinv2Config"], + "models.switch_transformers": ["SwitchTransformersConfig"], + "models.t5": ["T5Config"], + "models.table_transformer": ["TableTransformerConfig"], + "models.tapas": [ + "TapasConfig", + "TapasTokenizer", + ], + "models.textnet": ["TextNetConfig"], + "models.time_series_transformer": ["TimeSeriesTransformerConfig"], + "models.timesformer": ["TimesformerConfig"], + "models.timm_backbone": ["TimmBackboneConfig"], + "models.timm_wrapper": ["TimmWrapperConfig"], + "models.trocr": [ + "TrOCRConfig", + "TrOCRProcessor", + ], + "models.tvp": [ + "TvpConfig", + "TvpProcessor", + ], + "models.udop": [ + "UdopConfig", + "UdopProcessor", + ], + "models.umt5": ["UMT5Config"], + "models.unispeech": ["UniSpeechConfig"], + "models.unispeech_sat": ["UniSpeechSatConfig"], + "models.univnet": [ + "UnivNetConfig", + "UnivNetFeatureExtractor", + ], + "models.upernet": ["UperNetConfig"], + "models.video_llava": ["VideoLlavaConfig"], + "models.videomae": ["VideoMAEConfig"], + "models.vilt": [ + "ViltConfig", + "ViltFeatureExtractor", + "ViltImageProcessor", + "ViltProcessor", + ], + "models.vipllava": ["VipLlavaConfig"], + "models.vision_encoder_decoder": ["VisionEncoderDecoderConfig"], + "models.vision_text_dual_encoder": [ + "VisionTextDualEncoderConfig", + "VisionTextDualEncoderProcessor", + ], + "models.visual_bert": ["VisualBertConfig"], + "models.vit": ["ViTConfig"], + "models.vit_mae": ["ViTMAEConfig"], + "models.vit_msn": ["ViTMSNConfig"], + "models.vitdet": ["VitDetConfig"], + "models.vitmatte": ["VitMatteConfig"], + "models.vitpose": ["VitPoseConfig"], + "models.vitpose_backbone": ["VitPoseBackboneConfig"], + "models.vits": [ + "VitsConfig", + "VitsTokenizer", + ], + "models.vivit": ["VivitConfig"], + "models.wav2vec2": [ + "Wav2Vec2Config", + "Wav2Vec2CTCTokenizer", + "Wav2Vec2FeatureExtractor", + "Wav2Vec2Processor", + "Wav2Vec2Tokenizer", + ], + "models.wav2vec2_bert": [ + "Wav2Vec2BertConfig", + "Wav2Vec2BertProcessor", + ], + "models.wav2vec2_conformer": ["Wav2Vec2ConformerConfig"], + "models.wav2vec2_phoneme": ["Wav2Vec2PhonemeCTCTokenizer"], + "models.wav2vec2_with_lm": ["Wav2Vec2ProcessorWithLM"], + "models.wavlm": ["WavLMConfig"], + "models.whisper": [ + "WhisperConfig", + "WhisperFeatureExtractor", + "WhisperProcessor", + "WhisperTokenizer", + ], + "models.x_clip": [ + "XCLIPConfig", + "XCLIPProcessor", + "XCLIPTextConfig", + "XCLIPVisionConfig", + ], + "models.xglm": ["XGLMConfig"], + "models.xlm": ["XLMConfig", "XLMTokenizer"], + "models.xlm_roberta": ["XLMRobertaConfig"], + "models.xlm_roberta_xl": ["XLMRobertaXLConfig"], + "models.xlnet": ["XLNetConfig"], + "models.xmod": ["XmodConfig"], + "models.yolos": ["YolosConfig"], + "models.yoso": ["YosoConfig"], + "models.zamba": ["ZambaConfig"], + "models.zamba2": ["Zamba2Config"], + "models.zoedepth": ["ZoeDepthConfig"], + "onnx": [], + "pipelines": [ + "AudioClassificationPipeline", + "AutomaticSpeechRecognitionPipeline", + "CsvPipelineDataFormat", + "DepthEstimationPipeline", + "DocumentQuestionAnsweringPipeline", + "FeatureExtractionPipeline", + "FillMaskPipeline", + "ImageClassificationPipeline", + "ImageFeatureExtractionPipeline", + "ImageSegmentationPipeline", + "ImageTextToTextPipeline", + "ImageToImagePipeline", + "ImageToTextPipeline", + "JsonPipelineDataFormat", + "MaskGenerationPipeline", + "NerPipeline", + "ObjectDetectionPipeline", + "PipedPipelineDataFormat", + "Pipeline", + "PipelineDataFormat", + "QuestionAnsweringPipeline", + "SummarizationPipeline", + "TableQuestionAnsweringPipeline", + "Text2TextGenerationPipeline", + "TextClassificationPipeline", + "TextGenerationPipeline", + "TextToAudioPipeline", + "TokenClassificationPipeline", + "TranslationPipeline", + "VideoClassificationPipeline", + "VisualQuestionAnsweringPipeline", + "ZeroShotAudioClassificationPipeline", + "ZeroShotClassificationPipeline", + "ZeroShotImageClassificationPipeline", + "ZeroShotObjectDetectionPipeline", + "pipeline", + ], + "processing_utils": ["ProcessorMixin"], + "quantizers": [], + "testing_utils": [], + "tokenization_utils": ["PreTrainedTokenizer"], + "tokenization_utils_base": [ + "AddedToken", + "BatchEncoding", + "CharSpan", + "PreTrainedTokenizerBase", + "SpecialTokensMixin", + "TokenSpan", + ], + "trainer_callback": [ + "DefaultFlowCallback", + "EarlyStoppingCallback", + "PrinterCallback", + "ProgressCallback", + "TrainerCallback", + "TrainerControl", + "TrainerState", + ], + "trainer_utils": [ + "EvalPrediction", + "IntervalStrategy", + "SchedulerType", + "enable_full_determinism", + "set_seed", + ], + "training_args": ["TrainingArguments"], + "training_args_seq2seq": ["Seq2SeqTrainingArguments"], + "training_args_tf": ["TFTrainingArguments"], + "utils": [ + "CONFIG_NAME", + "MODEL_CARD_NAME", + "PYTORCH_PRETRAINED_BERT_CACHE", + "PYTORCH_TRANSFORMERS_CACHE", + "SPIECE_UNDERLINE", + "TF2_WEIGHTS_NAME", + "TF_WEIGHTS_NAME", + "TRANSFORMERS_CACHE", + "WEIGHTS_NAME", + "TensorType", + "add_end_docstrings", + "add_start_docstrings", + "is_apex_available", + "is_av_available", + "is_bitsandbytes_available", + "is_datasets_available", + "is_faiss_available", + "is_flax_available", + "is_keras_nlp_available", + "is_phonemizer_available", + "is_psutil_available", + "is_py3nvml_available", + "is_pyctcdecode_available", + "is_sacremoses_available", + "is_safetensors_available", + "is_scipy_available", + "is_sentencepiece_available", + "is_sklearn_available", + "is_speech_available", + "is_tensorflow_text_available", + "is_tf_available", + "is_timm_available", + "is_tokenizers_available", + "is_torch_available", + "is_torch_hpu_available", + "is_torch_mlu_available", + "is_torch_musa_available", + "is_torch_neuroncore_available", + "is_torch_npu_available", + "is_torchvision_available", + "is_torch_xla_available", + "is_torch_xpu_available", + "is_vision_available", + "logging", + ], + "utils.quantization_config": [ + "AqlmConfig", + "AwqConfig", + "BitNetConfig", + "BitsAndBytesConfig", + "CompressedTensorsConfig", + "EetqConfig", + "FbgemmFp8Config", + "FineGrainedFP8Config", + "GPTQConfig", + "HiggsConfig", + "HqqConfig", + "QuantoConfig", + "QuarkConfig", + "SpQRConfig", + "TorchAoConfig", + "VptqConfig", + ], +} + +# sentencepiece-backed objects +try: + if not is_sentencepiece_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_sentencepiece_objects + + _import_structure["utils.dummy_sentencepiece_objects"] = [ + name for name in dir(dummy_sentencepiece_objects) if not name.startswith("_") + ] +else: + _import_structure["models.albert"].append("AlbertTokenizer") + _import_structure["models.barthez"].append("BarthezTokenizer") + _import_structure["models.bartpho"].append("BartphoTokenizer") + _import_structure["models.bert_generation"].append("BertGenerationTokenizer") + _import_structure["models.big_bird"].append("BigBirdTokenizer") + _import_structure["models.camembert"].append("CamembertTokenizer") + _import_structure["models.code_llama"].append("CodeLlamaTokenizer") + _import_structure["models.cpm"].append("CpmTokenizer") + _import_structure["models.deberta_v2"].append("DebertaV2Tokenizer") + _import_structure["models.deprecated.ernie_m"].append("ErnieMTokenizer") + _import_structure["models.deprecated.xlm_prophetnet"].append("XLMProphetNetTokenizer") + _import_structure["models.fnet"].append("FNetTokenizer") + _import_structure["models.gemma"].append("GemmaTokenizer") + _import_structure["models.gpt_sw3"].append("GPTSw3Tokenizer") + _import_structure["models.layoutxlm"].append("LayoutXLMTokenizer") + _import_structure["models.llama"].append("LlamaTokenizer") + _import_structure["models.m2m_100"].append("M2M100Tokenizer") + _import_structure["models.marian"].append("MarianTokenizer") + _import_structure["models.mbart"].append("MBartTokenizer") + _import_structure["models.mbart50"].append("MBart50Tokenizer") + _import_structure["models.mluke"].append("MLukeTokenizer") + _import_structure["models.mt5"].append("MT5Tokenizer") + _import_structure["models.nllb"].append("NllbTokenizer") + _import_structure["models.pegasus"].append("PegasusTokenizer") + _import_structure["models.plbart"].append("PLBartTokenizer") + _import_structure["models.reformer"].append("ReformerTokenizer") + _import_structure["models.rembert"].append("RemBertTokenizer") + _import_structure["models.seamless_m4t"].append("SeamlessM4TTokenizer") + _import_structure["models.siglip"].append("SiglipTokenizer") + _import_structure["models.speech_to_text"].append("Speech2TextTokenizer") + _import_structure["models.speecht5"].append("SpeechT5Tokenizer") + _import_structure["models.t5"].append("T5Tokenizer") + _import_structure["models.udop"].append("UdopTokenizer") + _import_structure["models.xglm"].append("XGLMTokenizer") + _import_structure["models.xlm_roberta"].append("XLMRobertaTokenizer") + _import_structure["models.xlnet"].append("XLNetTokenizer") + +# tokenizers-backed objects +try: + if not is_tokenizers_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_tokenizers_objects + + _import_structure["utils.dummy_tokenizers_objects"] = [ + name for name in dir(dummy_tokenizers_objects) if not name.startswith("_") + ] +else: + # Fast tokenizers structure + _import_structure["models.albert"].append("AlbertTokenizerFast") + _import_structure["models.bart"].append("BartTokenizerFast") + _import_structure["models.barthez"].append("BarthezTokenizerFast") + _import_structure["models.bert"].append("BertTokenizerFast") + _import_structure["models.big_bird"].append("BigBirdTokenizerFast") + _import_structure["models.blenderbot"].append("BlenderbotTokenizerFast") + _import_structure["models.blenderbot_small"].append("BlenderbotSmallTokenizerFast") + _import_structure["models.bloom"].append("BloomTokenizerFast") + _import_structure["models.camembert"].append("CamembertTokenizerFast") + _import_structure["models.clip"].append("CLIPTokenizerFast") + _import_structure["models.code_llama"].append("CodeLlamaTokenizerFast") + _import_structure["models.codegen"].append("CodeGenTokenizerFast") + _import_structure["models.cohere"].append("CohereTokenizerFast") + _import_structure["models.convbert"].append("ConvBertTokenizerFast") + _import_structure["models.cpm"].append("CpmTokenizerFast") + _import_structure["models.deberta"].append("DebertaTokenizerFast") + _import_structure["models.deberta_v2"].append("DebertaV2TokenizerFast") + _import_structure["models.deprecated.realm"].append("RealmTokenizerFast") + _import_structure["models.deprecated.retribert"].append("RetriBertTokenizerFast") + _import_structure["models.distilbert"].append("DistilBertTokenizerFast") + _import_structure["models.dpr"].extend( + [ + "DPRContextEncoderTokenizerFast", + "DPRQuestionEncoderTokenizerFast", + "DPRReaderTokenizerFast", + ] + ) + _import_structure["models.electra"].append("ElectraTokenizerFast") + _import_structure["models.fnet"].append("FNetTokenizerFast") + _import_structure["models.funnel"].append("FunnelTokenizerFast") + _import_structure["models.gemma"].append("GemmaTokenizerFast") + _import_structure["models.gpt2"].append("GPT2TokenizerFast") + _import_structure["models.gpt_neox"].append("GPTNeoXTokenizerFast") + _import_structure["models.gpt_neox_japanese"].append("GPTNeoXJapaneseTokenizer") + _import_structure["models.herbert"].append("HerbertTokenizerFast") + _import_structure["models.layoutlm"].append("LayoutLMTokenizerFast") + _import_structure["models.layoutlmv2"].append("LayoutLMv2TokenizerFast") + _import_structure["models.layoutlmv3"].append("LayoutLMv3TokenizerFast") + _import_structure["models.layoutxlm"].append("LayoutXLMTokenizerFast") + _import_structure["models.led"].append("LEDTokenizerFast") + _import_structure["models.llama"].append("LlamaTokenizerFast") + _import_structure["models.longformer"].append("LongformerTokenizerFast") + _import_structure["models.lxmert"].append("LxmertTokenizerFast") + _import_structure["models.markuplm"].append("MarkupLMTokenizerFast") + _import_structure["models.mbart"].append("MBartTokenizerFast") + _import_structure["models.mbart50"].append("MBart50TokenizerFast") + _import_structure["models.mobilebert"].append("MobileBertTokenizerFast") + _import_structure["models.mpnet"].append("MPNetTokenizerFast") + _import_structure["models.mt5"].append("MT5TokenizerFast") + _import_structure["models.mvp"].append("MvpTokenizerFast") + _import_structure["models.nllb"].append("NllbTokenizerFast") + _import_structure["models.nougat"].append("NougatTokenizerFast") + _import_structure["models.openai"].append("OpenAIGPTTokenizerFast") + _import_structure["models.pegasus"].append("PegasusTokenizerFast") + _import_structure["models.qwen2"].append("Qwen2TokenizerFast") + _import_structure["models.reformer"].append("ReformerTokenizerFast") + _import_structure["models.rembert"].append("RemBertTokenizerFast") + _import_structure["models.roberta"].append("RobertaTokenizerFast") + _import_structure["models.roformer"].append("RoFormerTokenizerFast") + _import_structure["models.seamless_m4t"].append("SeamlessM4TTokenizerFast") + _import_structure["models.splinter"].append("SplinterTokenizerFast") + _import_structure["models.squeezebert"].append("SqueezeBertTokenizerFast") + _import_structure["models.t5"].append("T5TokenizerFast") + _import_structure["models.udop"].append("UdopTokenizerFast") + _import_structure["models.whisper"].append("WhisperTokenizerFast") + _import_structure["models.xglm"].append("XGLMTokenizerFast") + _import_structure["models.xlm_roberta"].append("XLMRobertaTokenizerFast") + _import_structure["models.xlnet"].append("XLNetTokenizerFast") + _import_structure["tokenization_utils_fast"] = ["PreTrainedTokenizerFast"] + + +try: + if not (is_sentencepiece_available() and is_tokenizers_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_sentencepiece_and_tokenizers_objects + + _import_structure["utils.dummy_sentencepiece_and_tokenizers_objects"] = [ + name for name in dir(dummy_sentencepiece_and_tokenizers_objects) if not name.startswith("_") + ] +else: + _import_structure["convert_slow_tokenizer"] = [ + "SLOW_TO_FAST_CONVERTERS", + "convert_slow_tokenizer", + ] + +# Tensorflow-text-specific objects +try: + if not is_tensorflow_text_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_tensorflow_text_objects + + _import_structure["utils.dummy_tensorflow_text_objects"] = [ + name for name in dir(dummy_tensorflow_text_objects) if not name.startswith("_") + ] +else: + _import_structure["models.bert"].append("TFBertTokenizer") + +# keras-nlp-specific objects +try: + if not is_keras_nlp_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_keras_nlp_objects + + _import_structure["utils.dummy_keras_nlp_objects"] = [ + name for name in dir(dummy_keras_nlp_objects) if not name.startswith("_") + ] +else: + _import_structure["models.gpt2"].append("TFGPT2Tokenizer") + +# Vision-specific objects +try: + if not is_vision_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_vision_objects + + _import_structure["utils.dummy_vision_objects"] = [ + name for name in dir(dummy_vision_objects) if not name.startswith("_") + ] +else: + _import_structure["image_processing_base"] = ["ImageProcessingMixin"] + _import_structure["image_processing_utils"] = ["BaseImageProcessor"] + _import_structure["image_utils"] = ["ImageFeatureExtractionMixin"] + _import_structure["models.aria"].extend(["AriaImageProcessor"]) + _import_structure["models.beit"].extend(["BeitFeatureExtractor", "BeitImageProcessor"]) + _import_structure["models.bit"].extend(["BitImageProcessor"]) + _import_structure["models.blip"].extend(["BlipImageProcessor"]) + _import_structure["models.bridgetower"].append("BridgeTowerImageProcessor") + _import_structure["models.chameleon"].append("ChameleonImageProcessor") + _import_structure["models.chinese_clip"].extend(["ChineseCLIPFeatureExtractor", "ChineseCLIPImageProcessor"]) + _import_structure["models.clip"].extend(["CLIPFeatureExtractor", "CLIPImageProcessor"]) + _import_structure["models.conditional_detr"].extend( + ["ConditionalDetrFeatureExtractor", "ConditionalDetrImageProcessor"] + ) + _import_structure["models.convnext"].extend(["ConvNextFeatureExtractor", "ConvNextImageProcessor"]) + _import_structure["models.deformable_detr"].extend( + ["DeformableDetrFeatureExtractor", "DeformableDetrImageProcessor"] + ) + _import_structure["models.deit"].extend(["DeiTFeatureExtractor", "DeiTImageProcessor"]) + _import_structure["models.deprecated.deta"].append("DetaImageProcessor") + _import_structure["models.deprecated.efficientformer"].append("EfficientFormerImageProcessor") + _import_structure["models.deprecated.tvlt"].append("TvltImageProcessor") + _import_structure["models.deprecated.vit_hybrid"].extend(["ViTHybridImageProcessor"]) + _import_structure["models.depth_pro"].extend(["DepthProImageProcessor", "DepthProImageProcessorFast"]) + _import_structure["models.detr"].extend(["DetrFeatureExtractor", "DetrImageProcessor"]) + _import_structure["models.donut"].extend(["DonutFeatureExtractor", "DonutImageProcessor"]) + _import_structure["models.dpt"].extend(["DPTFeatureExtractor", "DPTImageProcessor"]) + _import_structure["models.efficientnet"].append("EfficientNetImageProcessor") + _import_structure["models.emu3"].append("Emu3ImageProcessor") + _import_structure["models.flava"].extend(["FlavaFeatureExtractor", "FlavaImageProcessor", "FlavaProcessor"]) + _import_structure["models.fuyu"].extend(["FuyuImageProcessor", "FuyuProcessor"]) + _import_structure["models.gemma3"].append("Gemma3ImageProcessor") + _import_structure["models.glpn"].extend(["GLPNFeatureExtractor", "GLPNImageProcessor"]) + _import_structure["models.got_ocr2"].extend(["GotOcr2ImageProcessor"]) + _import_structure["models.grounding_dino"].extend(["GroundingDinoImageProcessor"]) + _import_structure["models.idefics"].extend(["IdeficsImageProcessor"]) + _import_structure["models.idefics2"].extend(["Idefics2ImageProcessor"]) + _import_structure["models.idefics3"].extend(["Idefics3ImageProcessor"]) + _import_structure["models.imagegpt"].extend(["ImageGPTFeatureExtractor", "ImageGPTImageProcessor"]) + _import_structure["models.instructblipvideo"].extend(["InstructBlipVideoImageProcessor"]) + _import_structure["models.layoutlmv2"].extend(["LayoutLMv2FeatureExtractor", "LayoutLMv2ImageProcessor"]) + _import_structure["models.layoutlmv3"].extend(["LayoutLMv3FeatureExtractor", "LayoutLMv3ImageProcessor"]) + _import_structure["models.levit"].extend(["LevitFeatureExtractor", "LevitImageProcessor"]) + _import_structure["models.llava"].append("LlavaImageProcessor") + _import_structure["models.llava_next"].append("LlavaNextImageProcessor") + _import_structure["models.llava_next_video"].append("LlavaNextVideoImageProcessor") + _import_structure["models.llava_onevision"].extend( + ["LlavaOnevisionImageProcessor", "LlavaOnevisionVideoProcessor"] + ) + _import_structure["models.mask2former"].append("Mask2FormerImageProcessor") + _import_structure["models.maskformer"].extend(["MaskFormerFeatureExtractor", "MaskFormerImageProcessor"]) + _import_structure["models.mllama"].extend(["MllamaImageProcessor"]) + _import_structure["models.mobilenet_v1"].extend(["MobileNetV1FeatureExtractor", "MobileNetV1ImageProcessor"]) + _import_structure["models.mobilenet_v2"].extend(["MobileNetV2FeatureExtractor", "MobileNetV2ImageProcessor"]) + _import_structure["models.mobilevit"].extend(["MobileViTFeatureExtractor", "MobileViTImageProcessor"]) + _import_structure["models.nougat"].append("NougatImageProcessor") + _import_structure["models.oneformer"].extend(["OneFormerImageProcessor"]) + _import_structure["models.owlv2"].append("Owlv2ImageProcessor") + _import_structure["models.owlvit"].extend(["OwlViTFeatureExtractor", "OwlViTImageProcessor"]) + _import_structure["models.perceiver"].extend(["PerceiverFeatureExtractor", "PerceiverImageProcessor"]) + _import_structure["models.pix2struct"].extend(["Pix2StructImageProcessor"]) + _import_structure["models.pixtral"].append("PixtralImageProcessor") + _import_structure["models.poolformer"].extend(["PoolFormerFeatureExtractor", "PoolFormerImageProcessor"]) + _import_structure["models.prompt_depth_anything"].extend(["PromptDepthAnythingImageProcessor"]) + _import_structure["models.pvt"].extend(["PvtImageProcessor"]) + _import_structure["models.qwen2_vl"].extend(["Qwen2VLImageProcessor"]) + _import_structure["models.rt_detr"].extend(["RTDetrImageProcessor"]) + _import_structure["models.sam"].extend(["SamImageProcessor"]) + _import_structure["models.segformer"].extend(["SegformerFeatureExtractor", "SegformerImageProcessor"]) + _import_structure["models.seggpt"].extend(["SegGptImageProcessor"]) + _import_structure["models.siglip"].append("SiglipImageProcessor") + _import_structure["models.siglip2"].append("Siglip2ImageProcessor") + _import_structure["models.smolvlm"].extend(["SmolVLMImageProcessor"]) + _import_structure["models.superglue"].extend(["SuperGlueImageProcessor"]) + _import_structure["models.superpoint"].extend(["SuperPointImageProcessor"]) + _import_structure["models.swin2sr"].append("Swin2SRImageProcessor") + _import_structure["models.textnet"].extend(["TextNetImageProcessor"]) + _import_structure["models.tvp"].append("TvpImageProcessor") + _import_structure["models.video_llava"].append("VideoLlavaImageProcessor") + _import_structure["models.videomae"].extend(["VideoMAEFeatureExtractor", "VideoMAEImageProcessor"]) + _import_structure["models.vilt"].extend(["ViltFeatureExtractor", "ViltImageProcessor", "ViltProcessor"]) + _import_structure["models.vit"].extend(["ViTFeatureExtractor", "ViTImageProcessor"]) + _import_structure["models.vitmatte"].append("VitMatteImageProcessor") + _import_structure["models.vitpose"].append("VitPoseImageProcessor") + _import_structure["models.vivit"].append("VivitImageProcessor") + _import_structure["models.yolos"].extend(["YolosFeatureExtractor", "YolosImageProcessor"]) + _import_structure["models.zoedepth"].append("ZoeDepthImageProcessor") + +try: + if not is_torchvision_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_torchvision_objects + + _import_structure["utils.dummy_torchvision_objects"] = [ + name for name in dir(dummy_torchvision_objects) if not name.startswith("_") + ] +else: + _import_structure["image_processing_utils_fast"] = ["BaseImageProcessorFast"] + _import_structure["models.blip"].append("BlipImageProcessorFast") + _import_structure["models.clip"].append("CLIPImageProcessorFast") + _import_structure["models.convnext"].append("ConvNextImageProcessorFast") + _import_structure["models.deformable_detr"].append("DeformableDetrImageProcessorFast") + _import_structure["models.deit"].append("DeiTImageProcessorFast") + _import_structure["models.depth_pro"].append("DepthProImageProcessorFast") + _import_structure["models.detr"].append("DetrImageProcessorFast") + _import_structure["models.gemma3"].append("Gemma3ImageProcessorFast") + _import_structure["models.got_ocr2"].append("GotOcr2ImageProcessorFast") + _import_structure["models.llama4"].append("Llama4ImageProcessorFast") + _import_structure["models.llava"].append("LlavaImageProcessorFast") + _import_structure["models.llava_next"].append("LlavaNextImageProcessorFast") + _import_structure["models.llava_onevision"].append("LlavaOnevisionImageProcessorFast") + _import_structure["models.phi4_multimodal"].append("Phi4MultimodalImageProcessorFast") + _import_structure["models.pixtral"].append("PixtralImageProcessorFast") + _import_structure["models.qwen2_vl"].append("Qwen2VLImageProcessorFast") + _import_structure["models.rt_detr"].append("RTDetrImageProcessorFast") + _import_structure["models.siglip"].append("SiglipImageProcessorFast") + _import_structure["models.siglip2"].append("Siglip2ImageProcessorFast") + _import_structure["models.vit"].append("ViTImageProcessorFast") + +try: + if not (is_torchvision_available() and is_timm_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_timm_and_torchvision_objects + + _import_structure["utils.dummy_timm_and_torchvision_objects"] = [ + name for name in dir(dummy_timm_and_torchvision_objects) if not name.startswith("_") + ] +else: + _import_structure["models.timm_wrapper"].extend(["TimmWrapperImageProcessor"]) + +# PyTorch-backed objects +try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_pt_objects + + _import_structure["utils.dummy_pt_objects"] = [name for name in dir(dummy_pt_objects) if not name.startswith("_")] +else: + _import_structure["model_debugging_utils"] = [ + "model_addition_debugger", + "model_addition_debugger_context", + ] + _import_structure["activations"] = [] + _import_structure["cache_utils"] = [ + "Cache", + "CacheConfig", + "DynamicCache", + "EncoderDecoderCache", + "HQQQuantizedCache", + "HybridCache", + "MambaCache", + "OffloadedCache", + "OffloadedStaticCache", + "QuantizedCache", + "QuantizedCacheConfig", + "QuantoQuantizedCache", + "SinkCache", + "SlidingWindowCache", + "StaticCache", + ] + _import_structure["data.datasets"] = [ + "GlueDataset", + "GlueDataTrainingArguments", + "LineByLineTextDataset", + "LineByLineWithRefDataset", + "LineByLineWithSOPTextDataset", + "SquadDataset", + "SquadDataTrainingArguments", + "TextDataset", + "TextDatasetForNextSentencePrediction", + ] + _import_structure["generation"].extend( + [ + "AlternatingCodebooksLogitsProcessor", + "BayesianDetectorConfig", + "BayesianDetectorModel", + "BeamScorer", + "BeamSearchScorer", + "ClassifierFreeGuidanceLogitsProcessor", + "ConstrainedBeamSearchScorer", + "Constraint", + "ConstraintListState", + "DisjunctiveConstraint", + "EncoderNoRepeatNGramLogitsProcessor", + "EncoderRepetitionPenaltyLogitsProcessor", + "EosTokenCriteria", + "EpsilonLogitsWarper", + "EtaLogitsWarper", + "ExponentialDecayLengthPenalty", + "ForcedBOSTokenLogitsProcessor", + "ForcedEOSTokenLogitsProcessor", + "GenerationMixin", + "HammingDiversityLogitsProcessor", + "InfNanRemoveLogitsProcessor", + "LogitNormalization", + "LogitsProcessor", + "LogitsProcessorList", + "MaxLengthCriteria", + "MaxTimeCriteria", + "MinLengthLogitsProcessor", + "MinNewTokensLengthLogitsProcessor", + "MinPLogitsWarper", + "NoBadWordsLogitsProcessor", + "NoRepeatNGramLogitsProcessor", + "PhrasalConstraint", + "PrefixConstrainedLogitsProcessor", + "RepetitionPenaltyLogitsProcessor", + "SequenceBiasLogitsProcessor", + "StoppingCriteria", + "StoppingCriteriaList", + "StopStringCriteria", + "SuppressTokensAtBeginLogitsProcessor", + "SuppressTokensLogitsProcessor", + "SynthIDTextWatermarkDetector", + "SynthIDTextWatermarkingConfig", + "SynthIDTextWatermarkLogitsProcessor", + "TemperatureLogitsWarper", + "TopKLogitsWarper", + "TopPLogitsWarper", + "TypicalLogitsWarper", + "UnbatchedClassifierFreeGuidanceLogitsProcessor", + "WatermarkDetector", + "WatermarkLogitsProcessor", + "WhisperTimeStampLogitsProcessor", + ] + ) + + # PyTorch domain libraries integration + _import_structure["integrations.executorch"] = [ + "TorchExportableModuleWithStaticCache", + "convert_and_export_with_cache", + ] + + _import_structure["modeling_flash_attention_utils"] = [] + _import_structure["modeling_outputs"] = [] + _import_structure["modeling_rope_utils"] = ["ROPE_INIT_FUNCTIONS", "dynamic_rope_update"] + _import_structure["modeling_utils"] = ["PreTrainedModel", "AttentionInterface"] + + # PyTorch models structure + + _import_structure["models.albert"].extend( + [ + "AlbertForMaskedLM", + "AlbertForMultipleChoice", + "AlbertForPreTraining", + "AlbertForQuestionAnswering", + "AlbertForSequenceClassification", + "AlbertForTokenClassification", + "AlbertModel", + "AlbertPreTrainedModel", + "load_tf_weights_in_albert", + ] + ) + + _import_structure["models.align"].extend( + [ + "AlignModel", + "AlignPreTrainedModel", + "AlignTextModel", + "AlignVisionModel", + ] + ) + _import_structure["models.altclip"].extend( + [ + "AltCLIPModel", + "AltCLIPPreTrainedModel", + "AltCLIPTextModel", + "AltCLIPVisionModel", + ] + ) + _import_structure["models.aria"].extend( + [ + "AriaForConditionalGeneration", + "AriaPreTrainedModel", + "AriaTextForCausalLM", + "AriaTextModel", + "AriaTextPreTrainedModel", + ] + ) + _import_structure["models.audio_spectrogram_transformer"].extend( + [ + "ASTForAudioClassification", + "ASTModel", + "ASTPreTrainedModel", + ] + ) + _import_structure["models.auto"].extend( + [ + "MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING", + "MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING", + "MODEL_FOR_AUDIO_XVECTOR_MAPPING", + "MODEL_FOR_BACKBONE_MAPPING", + "MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING", + "MODEL_FOR_CAUSAL_LM_MAPPING", + "MODEL_FOR_CTC_MAPPING", + "MODEL_FOR_DEPTH_ESTIMATION_MAPPING", + "MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING", + "MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING", + "MODEL_FOR_IMAGE_MAPPING", + "MODEL_FOR_IMAGE_SEGMENTATION_MAPPING", + "MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING", + "MODEL_FOR_IMAGE_TO_IMAGE_MAPPING", + "MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING", + "MODEL_FOR_KEYPOINT_DETECTION_MAPPING", + "MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING", + "MODEL_FOR_MASKED_LM_MAPPING", + "MODEL_FOR_MASK_GENERATION_MAPPING", + "MODEL_FOR_MULTIPLE_CHOICE_MAPPING", + "MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING", + "MODEL_FOR_OBJECT_DETECTION_MAPPING", + "MODEL_FOR_PRETRAINING_MAPPING", + "MODEL_FOR_QUESTION_ANSWERING_MAPPING", + "MODEL_FOR_RETRIEVAL_MAPPING", + "MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING", + "MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING", + "MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING", + "MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING", + "MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING", + "MODEL_FOR_TEXT_ENCODING_MAPPING", + "MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING", + "MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING", + "MODEL_FOR_TIME_SERIES_CLASSIFICATION_MAPPING", + "MODEL_FOR_TIME_SERIES_REGRESSION_MAPPING", + "MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING", + "MODEL_FOR_UNIVERSAL_SEGMENTATION_MAPPING", + "MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING", + "MODEL_FOR_VISION_2_SEQ_MAPPING", + "MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING", + "MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING", + "MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING", + "MODEL_MAPPING", + "MODEL_WITH_LM_HEAD_MAPPING", + "AutoBackbone", + "AutoModel", + "AutoModelForAudioClassification", + "AutoModelForAudioFrameClassification", + "AutoModelForAudioXVector", + "AutoModelForCausalLM", + "AutoModelForCTC", + "AutoModelForDepthEstimation", + "AutoModelForDocumentQuestionAnswering", + "AutoModelForImageClassification", + "AutoModelForImageSegmentation", + "AutoModelForImageTextToText", + "AutoModelForImageToImage", + "AutoModelForInstanceSegmentation", + "AutoModelForKeypointDetection", + "AutoModelForMaskedImageModeling", + "AutoModelForMaskedLM", + "AutoModelForMaskGeneration", + "AutoModelForMultipleChoice", + "AutoModelForNextSentencePrediction", + "AutoModelForObjectDetection", + "AutoModelForPreTraining", + "AutoModelForQuestionAnswering", + "AutoModelForSemanticSegmentation", + "AutoModelForSeq2SeqLM", + "AutoModelForSequenceClassification", + "AutoModelForSpeechSeq2Seq", + "AutoModelForTableQuestionAnswering", + "AutoModelForTextEncoding", + "AutoModelForTextToSpectrogram", + "AutoModelForTextToWaveform", + "AutoModelForTokenClassification", + "AutoModelForUniversalSegmentation", + "AutoModelForVideoClassification", + "AutoModelForVision2Seq", + "AutoModelForVisualQuestionAnswering", + "AutoModelForZeroShotImageClassification", + "AutoModelForZeroShotObjectDetection", + "AutoModelWithLMHead", + ] + ) + _import_structure["models.autoformer"].extend( + [ + "AutoformerForPrediction", + "AutoformerModel", + "AutoformerPreTrainedModel", + ] + ) + _import_structure["models.aya_vision"].extend(["AyaVisionForConditionalGeneration", "AyaVisionPreTrainedModel"]) + _import_structure["models.bamba"].extend( + [ + "BambaForCausalLM", + "BambaModel", + "BambaPreTrainedModel", + ] + ) + _import_structure["models.bark"].extend( + [ + "BarkCausalModel", + "BarkCoarseModel", + "BarkFineModel", + "BarkModel", + "BarkPreTrainedModel", + "BarkSemanticModel", + ] + ) + _import_structure["models.bart"].extend( + [ + "BartForCausalLM", + "BartForConditionalGeneration", + "BartForQuestionAnswering", + "BartForSequenceClassification", + "BartModel", + "BartPretrainedModel", + "BartPreTrainedModel", + "PretrainedBartModel", + ] + ) + _import_structure["models.beit"].extend( + [ + "BeitBackbone", + "BeitForImageClassification", + "BeitForMaskedImageModeling", + "BeitForSemanticSegmentation", + "BeitModel", + "BeitPreTrainedModel", + ] + ) + _import_structure["models.bert"].extend( + [ + "BertForMaskedLM", + "BertForMultipleChoice", + "BertForNextSentencePrediction", + "BertForPreTraining", + "BertForQuestionAnswering", + "BertForSequenceClassification", + "BertForTokenClassification", + "BertLMHeadModel", + "BertModel", + "BertPreTrainedModel", + "load_tf_weights_in_bert", + ] + ) + _import_structure["models.bert_generation"].extend( + [ + "BertGenerationDecoder", + "BertGenerationEncoder", + "BertGenerationPreTrainedModel", + "load_tf_weights_in_bert_generation", + ] + ) + _import_structure["models.big_bird"].extend( + [ + "BigBirdForCausalLM", + "BigBirdForMaskedLM", + "BigBirdForMultipleChoice", + "BigBirdForPreTraining", + "BigBirdForQuestionAnswering", + "BigBirdForSequenceClassification", + "BigBirdForTokenClassification", + "BigBirdModel", + "BigBirdPreTrainedModel", + "load_tf_weights_in_big_bird", + ] + ) + _import_structure["models.bigbird_pegasus"].extend( + [ + "BigBirdPegasusForCausalLM", + "BigBirdPegasusForConditionalGeneration", + "BigBirdPegasusForQuestionAnswering", + "BigBirdPegasusForSequenceClassification", + "BigBirdPegasusModel", + "BigBirdPegasusPreTrainedModel", + ] + ) + _import_structure["models.biogpt"].extend( + [ + "BioGptForCausalLM", + "BioGptForSequenceClassification", + "BioGptForTokenClassification", + "BioGptModel", + "BioGptPreTrainedModel", + ] + ) + _import_structure["models.bit"].extend( + [ + "BitBackbone", + "BitForImageClassification", + "BitModel", + "BitPreTrainedModel", + ] + ) + _import_structure["models.blenderbot"].extend( + [ + "BlenderbotForCausalLM", + "BlenderbotForConditionalGeneration", + "BlenderbotModel", + "BlenderbotPreTrainedModel", + ] + ) + _import_structure["models.blenderbot_small"].extend( + [ + "BlenderbotSmallForCausalLM", + "BlenderbotSmallForConditionalGeneration", + "BlenderbotSmallModel", + "BlenderbotSmallPreTrainedModel", + ] + ) + _import_structure["models.blip"].extend( + [ + "BlipForConditionalGeneration", + "BlipForImageTextRetrieval", + "BlipForQuestionAnswering", + "BlipModel", + "BlipPreTrainedModel", + "BlipTextModel", + "BlipVisionModel", + ] + ) + _import_structure["models.blip_2"].extend( + [ + "Blip2ForConditionalGeneration", + "Blip2ForImageTextRetrieval", + "Blip2Model", + "Blip2PreTrainedModel", + "Blip2QFormerModel", + "Blip2TextModelWithProjection", + "Blip2VisionModel", + "Blip2VisionModelWithProjection", + ] + ) + _import_structure["models.bloom"].extend( + [ + "BloomForCausalLM", + "BloomForQuestionAnswering", + "BloomForSequenceClassification", + "BloomForTokenClassification", + "BloomModel", + "BloomPreTrainedModel", + ] + ) + _import_structure["models.bridgetower"].extend( + [ + "BridgeTowerForContrastiveLearning", + "BridgeTowerForImageAndTextRetrieval", + "BridgeTowerForMaskedLM", + "BridgeTowerModel", + "BridgeTowerPreTrainedModel", + ] + ) + _import_structure["models.bros"].extend( + [ + "BrosForTokenClassification", + "BrosModel", + "BrosPreTrainedModel", + "BrosProcessor", + "BrosSpadeEEForTokenClassification", + "BrosSpadeELForTokenClassification", + ] + ) + _import_structure["models.camembert"].extend( + [ + "CamembertForCausalLM", + "CamembertForMaskedLM", + "CamembertForMultipleChoice", + "CamembertForQuestionAnswering", + "CamembertForSequenceClassification", + "CamembertForTokenClassification", + "CamembertModel", + "CamembertPreTrainedModel", + ] + ) + _import_structure["models.canine"].extend( + [ + "CanineForMultipleChoice", + "CanineForQuestionAnswering", + "CanineForSequenceClassification", + "CanineForTokenClassification", + "CanineModel", + "CaninePreTrainedModel", + "load_tf_weights_in_canine", + ] + ) + _import_structure["models.chameleon"].extend( + [ + "ChameleonForConditionalGeneration", + "ChameleonModel", + "ChameleonPreTrainedModel", + "ChameleonProcessor", + "ChameleonVQVAE", + ] + ) + _import_structure["models.chinese_clip"].extend( + [ + "ChineseCLIPModel", + "ChineseCLIPPreTrainedModel", + "ChineseCLIPTextModel", + "ChineseCLIPVisionModel", + ] + ) + _import_structure["models.clap"].extend( + [ + "ClapAudioModel", + "ClapAudioModelWithProjection", + "ClapFeatureExtractor", + "ClapModel", + "ClapPreTrainedModel", + "ClapTextModel", + "ClapTextModelWithProjection", + ] + ) + _import_structure["models.clip"].extend( + [ + "CLIPForImageClassification", + "CLIPModel", + "CLIPPreTrainedModel", + "CLIPTextModel", + "CLIPTextModelWithProjection", + "CLIPVisionModel", + "CLIPVisionModelWithProjection", + ] + ) + _import_structure["models.clipseg"].extend( + [ + "CLIPSegForImageSegmentation", + "CLIPSegModel", + "CLIPSegPreTrainedModel", + "CLIPSegTextModel", + "CLIPSegVisionModel", + ] + ) + _import_structure["models.clvp"].extend( + [ + "ClvpDecoder", + "ClvpEncoder", + "ClvpForCausalLM", + "ClvpModel", + "ClvpModelForConditionalGeneration", + "ClvpPreTrainedModel", + ] + ) + _import_structure["models.codegen"].extend( + [ + "CodeGenForCausalLM", + "CodeGenModel", + "CodeGenPreTrainedModel", + ] + ) + _import_structure["models.cohere"].extend(["CohereForCausalLM", "CohereModel", "CoherePreTrainedModel"]) + _import_structure["models.cohere2"].extend(["Cohere2ForCausalLM", "Cohere2Model", "Cohere2PreTrainedModel"]) + _import_structure["models.colpali"].extend( + [ + "ColPaliForRetrieval", + "ColPaliPreTrainedModel", + ] + ) + _import_structure["models.conditional_detr"].extend( + [ + "ConditionalDetrForObjectDetection", + "ConditionalDetrForSegmentation", + "ConditionalDetrModel", + "ConditionalDetrPreTrainedModel", + ] + ) + _import_structure["models.convbert"].extend( + [ + "ConvBertForMaskedLM", + "ConvBertForMultipleChoice", + "ConvBertForQuestionAnswering", + "ConvBertForSequenceClassification", + "ConvBertForTokenClassification", + "ConvBertModel", + "ConvBertPreTrainedModel", + "load_tf_weights_in_convbert", + ] + ) + _import_structure["models.convnext"].extend( + [ + "ConvNextBackbone", + "ConvNextForImageClassification", + "ConvNextModel", + "ConvNextPreTrainedModel", + ] + ) + _import_structure["models.convnextv2"].extend( + [ + "ConvNextV2Backbone", + "ConvNextV2ForImageClassification", + "ConvNextV2Model", + "ConvNextV2PreTrainedModel", + ] + ) + _import_structure["models.cpmant"].extend( + [ + "CpmAntForCausalLM", + "CpmAntModel", + "CpmAntPreTrainedModel", + ] + ) + _import_structure["models.ctrl"].extend( + [ + "CTRLForSequenceClassification", + "CTRLLMHeadModel", + "CTRLModel", + "CTRLPreTrainedModel", + ] + ) + _import_structure["models.cvt"].extend( + [ + "CvtForImageClassification", + "CvtModel", + "CvtPreTrainedModel", + ] + ) + _import_structure["models.dab_detr"].extend( + [ + "DabDetrForObjectDetection", + "DabDetrModel", + "DabDetrPreTrainedModel", + ] + ) + _import_structure["models.dac"].extend( + [ + "DacModel", + "DacPreTrainedModel", + ] + ) + _import_structure["models.data2vec"].extend( + [ + "Data2VecAudioForAudioFrameClassification", + "Data2VecAudioForCTC", + "Data2VecAudioForSequenceClassification", + "Data2VecAudioForXVector", + "Data2VecAudioModel", + "Data2VecAudioPreTrainedModel", + "Data2VecTextForCausalLM", + "Data2VecTextForMaskedLM", + "Data2VecTextForMultipleChoice", + "Data2VecTextForQuestionAnswering", + "Data2VecTextForSequenceClassification", + "Data2VecTextForTokenClassification", + "Data2VecTextModel", + "Data2VecTextPreTrainedModel", + "Data2VecVisionForImageClassification", + "Data2VecVisionForSemanticSegmentation", + "Data2VecVisionModel", + "Data2VecVisionPreTrainedModel", + ] + ) + _import_structure["models.dbrx"].extend( + [ + "DbrxForCausalLM", + "DbrxModel", + "DbrxPreTrainedModel", + ] + ) + _import_structure["models.deberta"].extend( + [ + "DebertaForMaskedLM", + "DebertaForQuestionAnswering", + "DebertaForSequenceClassification", + "DebertaForTokenClassification", + "DebertaModel", + "DebertaPreTrainedModel", + ] + ) + _import_structure["models.deberta_v2"].extend( + [ + "DebertaV2ForMaskedLM", + "DebertaV2ForMultipleChoice", + "DebertaV2ForQuestionAnswering", + "DebertaV2ForSequenceClassification", + "DebertaV2ForTokenClassification", + "DebertaV2Model", + "DebertaV2PreTrainedModel", + ] + ) + _import_structure["models.decision_transformer"].extend( + [ + "DecisionTransformerGPT2Model", + "DecisionTransformerGPT2PreTrainedModel", + "DecisionTransformerModel", + "DecisionTransformerPreTrainedModel", + ] + ) + _import_structure["models.deepseek_v3"].extend( + [ + "DeepseekV3ForCausalLM", + "DeepseekV3Model", + "DeepseekV3PreTrainedModel", + ] + ) + _import_structure["models.deformable_detr"].extend( + [ + "DeformableDetrForObjectDetection", + "DeformableDetrModel", + "DeformableDetrPreTrainedModel", + ] + ) + _import_structure["models.deit"].extend( + [ + "DeiTForImageClassification", + "DeiTForImageClassificationWithTeacher", + "DeiTForMaskedImageModeling", + "DeiTModel", + "DeiTPreTrainedModel", + ] + ) + _import_structure["models.deprecated.deta"].extend( + [ + "DetaForObjectDetection", + "DetaModel", + "DetaPreTrainedModel", + ] + ) + _import_structure["models.deprecated.efficientformer"].extend( + [ + "EfficientFormerForImageClassification", + "EfficientFormerForImageClassificationWithTeacher", + "EfficientFormerModel", + "EfficientFormerPreTrainedModel", + ] + ) + _import_structure["models.deprecated.ernie_m"].extend( + [ + "ErnieMForInformationExtraction", + "ErnieMForMultipleChoice", + "ErnieMForQuestionAnswering", + "ErnieMForSequenceClassification", + "ErnieMForTokenClassification", + "ErnieMModel", + "ErnieMPreTrainedModel", + ] + ) + _import_structure["models.deprecated.gptsan_japanese"].extend( + [ + "GPTSanJapaneseForConditionalGeneration", + "GPTSanJapaneseModel", + "GPTSanJapanesePreTrainedModel", + ] + ) + _import_structure["models.deprecated.graphormer"].extend( + [ + "GraphormerForGraphClassification", + "GraphormerModel", + "GraphormerPreTrainedModel", + ] + ) + _import_structure["models.deprecated.jukebox"].extend( + [ + "JukeboxModel", + "JukeboxPreTrainedModel", + "JukeboxPrior", + "JukeboxVQVAE", + ] + ) + _import_structure["models.deprecated.mctct"].extend( + [ + "MCTCTForCTC", + "MCTCTModel", + "MCTCTPreTrainedModel", + ] + ) + _import_structure["models.deprecated.mega"].extend( + [ + "MegaForCausalLM", + "MegaForMaskedLM", + "MegaForMultipleChoice", + "MegaForQuestionAnswering", + "MegaForSequenceClassification", + "MegaForTokenClassification", + "MegaModel", + "MegaPreTrainedModel", + ] + ) + _import_structure["models.deprecated.mmbt"].extend(["MMBTForClassification", "MMBTModel", "ModalEmbeddings"]) + _import_structure["models.deprecated.nat"].extend( + [ + "NatBackbone", + "NatForImageClassification", + "NatModel", + "NatPreTrainedModel", + ] + ) + _import_structure["models.deprecated.nezha"].extend( + [ + "NezhaForMaskedLM", + "NezhaForMultipleChoice", + "NezhaForNextSentencePrediction", + "NezhaForPreTraining", + "NezhaForQuestionAnswering", + "NezhaForSequenceClassification", + "NezhaForTokenClassification", + "NezhaModel", + "NezhaPreTrainedModel", + ] + ) + _import_structure["models.deprecated.open_llama"].extend( + [ + "OpenLlamaForCausalLM", + "OpenLlamaForSequenceClassification", + "OpenLlamaModel", + "OpenLlamaPreTrainedModel", + ] + ) + _import_structure["models.deprecated.qdqbert"].extend( + [ + "QDQBertForMaskedLM", + "QDQBertForMultipleChoice", + "QDQBertForNextSentencePrediction", + "QDQBertForQuestionAnswering", + "QDQBertForSequenceClassification", + "QDQBertForTokenClassification", + "QDQBertLMHeadModel", + "QDQBertModel", + "QDQBertPreTrainedModel", + "load_tf_weights_in_qdqbert", + ] + ) + _import_structure["models.deprecated.realm"].extend( + [ + "RealmEmbedder", + "RealmForOpenQA", + "RealmKnowledgeAugEncoder", + "RealmPreTrainedModel", + "RealmReader", + "RealmRetriever", + "RealmScorer", + "load_tf_weights_in_realm", + ] + ) + _import_structure["models.deprecated.retribert"].extend( + [ + "RetriBertModel", + "RetriBertPreTrainedModel", + ] + ) + _import_structure["models.deprecated.speech_to_text_2"].extend( + ["Speech2Text2ForCausalLM", "Speech2Text2PreTrainedModel"] + ) + _import_structure["models.deprecated.trajectory_transformer"].extend( + [ + "TrajectoryTransformerModel", + "TrajectoryTransformerPreTrainedModel", + ] + ) + _import_structure["models.deprecated.transfo_xl"].extend( + [ + "AdaptiveEmbedding", + "TransfoXLForSequenceClassification", + "TransfoXLLMHeadModel", + "TransfoXLModel", + "TransfoXLPreTrainedModel", + "load_tf_weights_in_transfo_xl", + ] + ) + _import_structure["models.deprecated.tvlt"].extend( + [ + "TvltForAudioVisualClassification", + "TvltForPreTraining", + "TvltModel", + "TvltPreTrainedModel", + ] + ) + _import_structure["models.deprecated.van"].extend( + [ + "VanForImageClassification", + "VanModel", + "VanPreTrainedModel", + ] + ) + _import_structure["models.deprecated.vit_hybrid"].extend( + [ + "ViTHybridForImageClassification", + "ViTHybridModel", + "ViTHybridPreTrainedModel", + ] + ) + _import_structure["models.deprecated.xlm_prophetnet"].extend( + [ + "XLMProphetNetDecoder", + "XLMProphetNetEncoder", + "XLMProphetNetForCausalLM", + "XLMProphetNetForConditionalGeneration", + "XLMProphetNetModel", + "XLMProphetNetPreTrainedModel", + ] + ) + _import_structure["models.depth_anything"].extend( + [ + "DepthAnythingForDepthEstimation", + "DepthAnythingPreTrainedModel", + ] + ) + _import_structure["models.depth_pro"].extend( + [ + "DepthProForDepthEstimation", + "DepthProModel", + "DepthProPreTrainedModel", + ] + ) + _import_structure["models.detr"].extend( + [ + "DetrForObjectDetection", + "DetrForSegmentation", + "DetrModel", + "DetrPreTrainedModel", + ] + ) + _import_structure["models.diffllama"].extend( + [ + "DiffLlamaForCausalLM", + "DiffLlamaForQuestionAnswering", + "DiffLlamaForSequenceClassification", + "DiffLlamaForTokenClassification", + "DiffLlamaModel", + "DiffLlamaPreTrainedModel", + ] + ) + _import_structure["models.dinat"].extend( + [ + "DinatBackbone", + "DinatForImageClassification", + "DinatModel", + "DinatPreTrainedModel", + ] + ) + _import_structure["models.dinov2"].extend( + [ + "Dinov2Backbone", + "Dinov2ForImageClassification", + "Dinov2Model", + "Dinov2PreTrainedModel", + ] + ) + _import_structure["models.dinov2_with_registers"].extend( + [ + "Dinov2WithRegistersBackbone", + "Dinov2WithRegistersForImageClassification", + "Dinov2WithRegistersModel", + "Dinov2WithRegistersPreTrainedModel", + ] + ) + _import_structure["models.distilbert"].extend( + [ + "DistilBertForMaskedLM", + "DistilBertForMultipleChoice", + "DistilBertForQuestionAnswering", + "DistilBertForSequenceClassification", + "DistilBertForTokenClassification", + "DistilBertModel", + "DistilBertPreTrainedModel", + ] + ) + _import_structure["models.donut"].extend( + [ + "DonutSwinModel", + "DonutSwinPreTrainedModel", + ] + ) + _import_structure["models.dpr"].extend( + [ + "DPRContextEncoder", + "DPRPretrainedContextEncoder", + "DPRPreTrainedModel", + "DPRPretrainedQuestionEncoder", + "DPRPretrainedReader", + "DPRQuestionEncoder", + "DPRReader", + ] + ) + _import_structure["models.dpt"].extend( + [ + "DPTForDepthEstimation", + "DPTForSemanticSegmentation", + "DPTModel", + "DPTPreTrainedModel", + ] + ) + _import_structure["models.efficientnet"].extend( + [ + "EfficientNetForImageClassification", + "EfficientNetModel", + "EfficientNetPreTrainedModel", + ] + ) + _import_structure["models.electra"].extend( + [ + "ElectraForCausalLM", + "ElectraForMaskedLM", + "ElectraForMultipleChoice", + "ElectraForPreTraining", + "ElectraForQuestionAnswering", + "ElectraForSequenceClassification", + "ElectraForTokenClassification", + "ElectraModel", + "ElectraPreTrainedModel", + "load_tf_weights_in_electra", + ] + ) + _import_structure["models.emu3"].extend( + [ + "Emu3ForCausalLM", + "Emu3ForConditionalGeneration", + "Emu3PreTrainedModel", + "Emu3TextModel", + "Emu3VQVAE", + ] + ) + _import_structure["models.encodec"].extend( + [ + "EncodecModel", + "EncodecPreTrainedModel", + ] + ) + _import_structure["models.encoder_decoder"].append("EncoderDecoderModel") + _import_structure["models.ernie"].extend( + [ + "ErnieForCausalLM", + "ErnieForMaskedLM", + "ErnieForMultipleChoice", + "ErnieForNextSentencePrediction", + "ErnieForPreTraining", + "ErnieForQuestionAnswering", + "ErnieForSequenceClassification", + "ErnieForTokenClassification", + "ErnieModel", + "ErniePreTrainedModel", + ] + ) + _import_structure["models.esm"].extend( + [ + "EsmFoldPreTrainedModel", + "EsmForMaskedLM", + "EsmForProteinFolding", + "EsmForSequenceClassification", + "EsmForTokenClassification", + "EsmModel", + "EsmPreTrainedModel", + ] + ) + _import_structure["models.falcon"].extend( + [ + "FalconForCausalLM", + "FalconForQuestionAnswering", + "FalconForSequenceClassification", + "FalconForTokenClassification", + "FalconModel", + "FalconPreTrainedModel", + ] + ) + _import_structure["models.falcon_mamba"].extend( + [ + "FalconMambaForCausalLM", + "FalconMambaModel", + "FalconMambaPreTrainedModel", + ] + ) + _import_structure["models.fastspeech2_conformer"].extend( + [ + "FastSpeech2ConformerHifiGan", + "FastSpeech2ConformerModel", + "FastSpeech2ConformerPreTrainedModel", + "FastSpeech2ConformerWithHifiGan", + ] + ) + _import_structure["models.flaubert"].extend( + [ + "FlaubertForMultipleChoice", + "FlaubertForQuestionAnswering", + "FlaubertForQuestionAnsweringSimple", + "FlaubertForSequenceClassification", + "FlaubertForTokenClassification", + "FlaubertModel", + "FlaubertPreTrainedModel", + "FlaubertWithLMHeadModel", + ] + ) + _import_structure["models.flava"].extend( + [ + "FlavaForPreTraining", + "FlavaImageCodebook", + "FlavaImageModel", + "FlavaModel", + "FlavaMultimodalModel", + "FlavaPreTrainedModel", + "FlavaTextModel", + ] + ) + _import_structure["models.fnet"].extend( + [ + "FNetForMaskedLM", + "FNetForMultipleChoice", + "FNetForNextSentencePrediction", + "FNetForPreTraining", + "FNetForQuestionAnswering", + "FNetForSequenceClassification", + "FNetForTokenClassification", + "FNetModel", + "FNetPreTrainedModel", + ] + ) + _import_structure["models.focalnet"].extend( + [ + "FocalNetBackbone", + "FocalNetForImageClassification", + "FocalNetForMaskedImageModeling", + "FocalNetModel", + "FocalNetPreTrainedModel", + ] + ) + _import_structure["models.fsmt"].extend(["FSMTForConditionalGeneration", "FSMTModel", "PretrainedFSMTModel"]) + _import_structure["models.funnel"].extend( + [ + "FunnelBaseModel", + "FunnelForMaskedLM", + "FunnelForMultipleChoice", + "FunnelForPreTraining", + "FunnelForQuestionAnswering", + "FunnelForSequenceClassification", + "FunnelForTokenClassification", + "FunnelModel", + "FunnelPreTrainedModel", + "load_tf_weights_in_funnel", + ] + ) + _import_structure["models.fuyu"].extend(["FuyuForCausalLM", "FuyuPreTrainedModel"]) + _import_structure["models.gemma"].extend( + [ + "GemmaForCausalLM", + "GemmaForSequenceClassification", + "GemmaForTokenClassification", + "GemmaModel", + "GemmaPreTrainedModel", + ] + ) + _import_structure["models.gemma2"].extend( + [ + "Gemma2ForCausalLM", + "Gemma2ForSequenceClassification", + "Gemma2ForTokenClassification", + "Gemma2Model", + "Gemma2PreTrainedModel", + ] + ) + _import_structure["models.gemma3"].extend( + [ + "Gemma3ForCausalLM", + "Gemma3ForConditionalGeneration", + "Gemma3PreTrainedModel", + "Gemma3TextModel", + ] + ) + _import_structure["models.git"].extend( + [ + "GitForCausalLM", + "GitModel", + "GitPreTrainedModel", + "GitVisionModel", + ] + ) + _import_structure["models.glm"].extend( + [ + "GlmForCausalLM", + "GlmForSequenceClassification", + "GlmForTokenClassification", + "GlmModel", + "GlmPreTrainedModel", + ] + ) + _import_structure["models.llama4"].extend( + [ + "Llama4ForCausalLM", + "Llama4ForConditionalGeneration", + "Llama4TextModel", + "Llama4VisionModel", + "Llama4PreTrainedModel", + ] + ) + _import_structure["models.glpn"].extend( + [ + "GLPNForDepthEstimation", + "GLPNModel", + "GLPNPreTrainedModel", + ] + ) + _import_structure["models.got_ocr2"].extend( + [ + "GotOcr2ForConditionalGeneration", + "GotOcr2PreTrainedModel", + ] + ) + _import_structure["models.gpt2"].extend( + [ + "GPT2DoubleHeadsModel", + "GPT2ForQuestionAnswering", + "GPT2ForSequenceClassification", + "GPT2ForTokenClassification", + "GPT2LMHeadModel", + "GPT2Model", + "GPT2PreTrainedModel", + "load_tf_weights_in_gpt2", + ] + ) + _import_structure["models.gpt_bigcode"].extend( + [ + "GPTBigCodeForCausalLM", + "GPTBigCodeForSequenceClassification", + "GPTBigCodeForTokenClassification", + "GPTBigCodeModel", + "GPTBigCodePreTrainedModel", + ] + ) + _import_structure["models.gpt_neo"].extend( + [ + "GPTNeoForCausalLM", + "GPTNeoForQuestionAnswering", + "GPTNeoForSequenceClassification", + "GPTNeoForTokenClassification", + "GPTNeoModel", + "GPTNeoPreTrainedModel", + "load_tf_weights_in_gpt_neo", + ] + ) + _import_structure["models.gpt_neox"].extend( + [ + "GPTNeoXForCausalLM", + "GPTNeoXForQuestionAnswering", + "GPTNeoXForSequenceClassification", + "GPTNeoXForTokenClassification", + "GPTNeoXModel", + "GPTNeoXPreTrainedModel", + ] + ) + _import_structure["models.gpt_neox_japanese"].extend( + [ + "GPTNeoXJapaneseForCausalLM", + "GPTNeoXJapaneseModel", + "GPTNeoXJapanesePreTrainedModel", + ] + ) + _import_structure["models.gptj"].extend( + [ + "GPTJForCausalLM", + "GPTJForQuestionAnswering", + "GPTJForSequenceClassification", + "GPTJModel", + "GPTJPreTrainedModel", + ] + ) + _import_structure["models.granite"].extend( + [ + "GraniteForCausalLM", + "GraniteModel", + "GranitePreTrainedModel", + ] + ) + _import_structure["models.granitemoe"].extend( + [ + "GraniteMoeForCausalLM", + "GraniteMoeModel", + "GraniteMoePreTrainedModel", + ] + ) + + _import_structure["models.granitemoeshared"].extend( + [ + "GraniteMoeSharedForCausalLM", + "GraniteMoeSharedModel", + "GraniteMoeSharedPreTrainedModel", + ] + ) + _import_structure["models.grounding_dino"].extend( + [ + "GroundingDinoForObjectDetection", + "GroundingDinoModel", + "GroundingDinoPreTrainedModel", + ] + ) + _import_structure["models.groupvit"].extend( + [ + "GroupViTModel", + "GroupViTPreTrainedModel", + "GroupViTTextModel", + "GroupViTVisionModel", + ] + ) + _import_structure["models.helium"].extend( + [ + "HeliumForCausalLM", + "HeliumForSequenceClassification", + "HeliumForTokenClassification", + "HeliumModel", + "HeliumPreTrainedModel", + ] + ) + _import_structure["models.hiera"].extend( + [ + "HieraBackbone", + "HieraForImageClassification", + "HieraForPreTraining", + "HieraModel", + "HieraPreTrainedModel", + ] + ) + _import_structure["models.hubert"].extend( + [ + "HubertForCTC", + "HubertForSequenceClassification", + "HubertModel", + "HubertPreTrainedModel", + ] + ) + _import_structure["models.ibert"].extend( + [ + "IBertForMaskedLM", + "IBertForMultipleChoice", + "IBertForQuestionAnswering", + "IBertForSequenceClassification", + "IBertForTokenClassification", + "IBertModel", + "IBertPreTrainedModel", + ] + ) + _import_structure["models.idefics"].extend( + [ + "IdeficsForVisionText2Text", + "IdeficsModel", + "IdeficsPreTrainedModel", + "IdeficsProcessor", + ] + ) + _import_structure["models.idefics2"].extend( + [ + "Idefics2ForConditionalGeneration", + "Idefics2Model", + "Idefics2PreTrainedModel", + "Idefics2Processor", + ] + ) + _import_structure["models.idefics3"].extend( + [ + "Idefics3ForConditionalGeneration", + "Idefics3Model", + "Idefics3PreTrainedModel", + "Idefics3Processor", + "Idefics3VisionConfig", + "Idefics3VisionTransformer", + ] + ) + _import_structure["models.ijepa"].extend( + [ + "IJepaForImageClassification", + "IJepaModel", + "IJepaPreTrainedModel", + ] + ) + _import_structure["models.imagegpt"].extend( + [ + "ImageGPTForCausalImageModeling", + "ImageGPTForImageClassification", + "ImageGPTModel", + "ImageGPTPreTrainedModel", + "load_tf_weights_in_imagegpt", + ] + ) + _import_structure["models.informer"].extend( + [ + "InformerForPrediction", + "InformerModel", + "InformerPreTrainedModel", + ] + ) + _import_structure["models.instructblip"].extend( + [ + "InstructBlipForConditionalGeneration", + "InstructBlipPreTrainedModel", + "InstructBlipQFormerModel", + "InstructBlipVisionModel", + ] + ) + _import_structure["models.instructblipvideo"].extend( + [ + "InstructBlipVideoForConditionalGeneration", + "InstructBlipVideoPreTrainedModel", + "InstructBlipVideoQFormerModel", + "InstructBlipVideoVisionModel", + ] + ) + _import_structure["models.jamba"].extend( + [ + "JambaForCausalLM", + "JambaForSequenceClassification", + "JambaModel", + "JambaPreTrainedModel", + ] + ) + _import_structure["models.jetmoe"].extend( + [ + "JetMoeForCausalLM", + "JetMoeForSequenceClassification", + "JetMoeModel", + "JetMoePreTrainedModel", + ] + ) + _import_structure["models.kosmos2"].extend( + [ + "Kosmos2ForConditionalGeneration", + "Kosmos2Model", + "Kosmos2PreTrainedModel", + ] + ) + _import_structure["models.layoutlm"].extend( + [ + "LayoutLMForMaskedLM", + "LayoutLMForQuestionAnswering", + "LayoutLMForSequenceClassification", + "LayoutLMForTokenClassification", + "LayoutLMModel", + "LayoutLMPreTrainedModel", + ] + ) + _import_structure["models.layoutlmv2"].extend( + [ + "LayoutLMv2ForQuestionAnswering", + "LayoutLMv2ForSequenceClassification", + "LayoutLMv2ForTokenClassification", + "LayoutLMv2Model", + "LayoutLMv2PreTrainedModel", + ] + ) + _import_structure["models.layoutlmv3"].extend( + [ + "LayoutLMv3ForQuestionAnswering", + "LayoutLMv3ForSequenceClassification", + "LayoutLMv3ForTokenClassification", + "LayoutLMv3Model", + "LayoutLMv3PreTrainedModel", + ] + ) + _import_structure["models.led"].extend( + [ + "LEDForConditionalGeneration", + "LEDForQuestionAnswering", + "LEDForSequenceClassification", + "LEDModel", + "LEDPreTrainedModel", + ] + ) + _import_structure["models.levit"].extend( + [ + "LevitForImageClassification", + "LevitForImageClassificationWithTeacher", + "LevitModel", + "LevitPreTrainedModel", + ] + ) + _import_structure["models.lilt"].extend( + [ + "LiltForQuestionAnswering", + "LiltForSequenceClassification", + "LiltForTokenClassification", + "LiltModel", + "LiltPreTrainedModel", + ] + ) + _import_structure["models.llama"].extend( + [ + "LlamaForCausalLM", + "LlamaForQuestionAnswering", + "LlamaForSequenceClassification", + "LlamaForTokenClassification", + "LlamaModel", + "LlamaPreTrainedModel", + ] + ) + _import_structure["models.llava"].extend( + [ + "LlavaForConditionalGeneration", + "LlavaPreTrainedModel", + ] + ) + _import_structure["models.llava_next"].extend( + [ + "LlavaNextForConditionalGeneration", + "LlavaNextPreTrainedModel", + ] + ) + _import_structure["models.phi4_multimodal"].extend( + [ + "Phi4MultimodalForCausalLM", + "Phi4MultimodalPreTrainedModel", + "Phi4MultimodalAudioModel", + "Phi4MultimodalAudioPreTrainedModel", + "Phi4MultimodalModel", + "Phi4MultimodalVisionModel", + "Phi4MultimodalVisionPreTrainedModel", + ] + ) + _import_structure["models.llava_next_video"].extend( + [ + "LlavaNextVideoForConditionalGeneration", + "LlavaNextVideoPreTrainedModel", + ] + ) + _import_structure["models.llava_onevision"].extend( + [ + "LlavaOnevisionForConditionalGeneration", + "LlavaOnevisionPreTrainedModel", + ] + ) + _import_structure["models.longformer"].extend( + [ + "LongformerForMaskedLM", + "LongformerForMultipleChoice", + "LongformerForQuestionAnswering", + "LongformerForSequenceClassification", + "LongformerForTokenClassification", + "LongformerModel", + "LongformerPreTrainedModel", + ] + ) + _import_structure["models.longt5"].extend( + [ + "LongT5EncoderModel", + "LongT5ForConditionalGeneration", + "LongT5Model", + "LongT5PreTrainedModel", + ] + ) + _import_structure["models.luke"].extend( + [ + "LukeForEntityClassification", + "LukeForEntityPairClassification", + "LukeForEntitySpanClassification", + "LukeForMaskedLM", + "LukeForMultipleChoice", + "LukeForQuestionAnswering", + "LukeForSequenceClassification", + "LukeForTokenClassification", + "LukeModel", + "LukePreTrainedModel", + ] + ) + _import_structure["models.lxmert"].extend( + [ + "LxmertEncoder", + "LxmertForPreTraining", + "LxmertForQuestionAnswering", + "LxmertModel", + "LxmertPreTrainedModel", + "LxmertVisualFeatureEncoder", + ] + ) + _import_structure["models.m2m_100"].extend( + [ + "M2M100ForConditionalGeneration", + "M2M100Model", + "M2M100PreTrainedModel", + ] + ) + _import_structure["models.mamba"].extend( + [ + "MambaForCausalLM", + "MambaModel", + "MambaPreTrainedModel", + ] + ) + _import_structure["models.mamba2"].extend( + [ + "Mamba2ForCausalLM", + "Mamba2Model", + "Mamba2PreTrainedModel", + ] + ) + _import_structure["models.marian"].extend( + ["MarianForCausalLM", "MarianModel", "MarianMTModel", "MarianPreTrainedModel"] + ) + _import_structure["models.markuplm"].extend( + [ + "MarkupLMForQuestionAnswering", + "MarkupLMForSequenceClassification", + "MarkupLMForTokenClassification", + "MarkupLMModel", + "MarkupLMPreTrainedModel", + ] + ) + _import_structure["models.mask2former"].extend( + [ + "Mask2FormerForUniversalSegmentation", + "Mask2FormerModel", + "Mask2FormerPreTrainedModel", + ] + ) + _import_structure["models.maskformer"].extend( + [ + "MaskFormerForInstanceSegmentation", + "MaskFormerModel", + "MaskFormerPreTrainedModel", + "MaskFormerSwinBackbone", + ] + ) + _import_structure["models.mbart"].extend( + [ + "MBartForCausalLM", + "MBartForConditionalGeneration", + "MBartForQuestionAnswering", + "MBartForSequenceClassification", + "MBartModel", + "MBartPreTrainedModel", + ] + ) + _import_structure["models.megatron_bert"].extend( + [ + "MegatronBertForCausalLM", + "MegatronBertForMaskedLM", + "MegatronBertForMultipleChoice", + "MegatronBertForNextSentencePrediction", + "MegatronBertForPreTraining", + "MegatronBertForQuestionAnswering", + "MegatronBertForSequenceClassification", + "MegatronBertForTokenClassification", + "MegatronBertModel", + "MegatronBertPreTrainedModel", + ] + ) + _import_structure["models.mgp_str"].extend( + [ + "MgpstrForSceneTextRecognition", + "MgpstrModel", + "MgpstrPreTrainedModel", + ] + ) + _import_structure["models.mimi"].extend( + [ + "MimiModel", + "MimiPreTrainedModel", + ] + ) + _import_structure["models.mistral"].extend( + [ + "MistralForCausalLM", + "MistralForQuestionAnswering", + "MistralForSequenceClassification", + "MistralForTokenClassification", + "MistralModel", + "MistralPreTrainedModel", + ] + ) + _import_structure["models.mistral3"].extend( + [ + "Mistral3ForConditionalGeneration", + "Mistral3PreTrainedModel", + ] + ) + _import_structure["models.mixtral"].extend( + [ + "MixtralForCausalLM", + "MixtralForQuestionAnswering", + "MixtralForSequenceClassification", + "MixtralForTokenClassification", + "MixtralModel", + "MixtralPreTrainedModel", + ] + ) + _import_structure["models.mllama"].extend( + [ + "MllamaForCausalLM", + "MllamaForConditionalGeneration", + "MllamaPreTrainedModel", + "MllamaProcessor", + "MllamaTextModel", + "MllamaVisionModel", + ] + ) + _import_structure["models.mobilebert"].extend( + [ + "MobileBertForMaskedLM", + "MobileBertForMultipleChoice", + "MobileBertForNextSentencePrediction", + "MobileBertForPreTraining", + "MobileBertForQuestionAnswering", + "MobileBertForSequenceClassification", + "MobileBertForTokenClassification", + "MobileBertModel", + "MobileBertPreTrainedModel", + "load_tf_weights_in_mobilebert", + ] + ) + _import_structure["models.mobilenet_v1"].extend( + [ + "MobileNetV1ForImageClassification", + "MobileNetV1Model", + "MobileNetV1PreTrainedModel", + "load_tf_weights_in_mobilenet_v1", + ] + ) + _import_structure["models.mobilenet_v2"].extend( + [ + "MobileNetV2ForImageClassification", + "MobileNetV2ForSemanticSegmentation", + "MobileNetV2Model", + "MobileNetV2PreTrainedModel", + "load_tf_weights_in_mobilenet_v2", + ] + ) + _import_structure["models.mobilevit"].extend( + [ + "MobileViTForImageClassification", + "MobileViTForSemanticSegmentation", + "MobileViTModel", + "MobileViTPreTrainedModel", + ] + ) + _import_structure["models.mobilevitv2"].extend( + [ + "MobileViTV2ForImageClassification", + "MobileViTV2ForSemanticSegmentation", + "MobileViTV2Model", + "MobileViTV2PreTrainedModel", + ] + ) + _import_structure["models.modernbert"].extend( + [ + "ModernBertForMaskedLM", + "ModernBertForQuestionAnswering", + "ModernBertForSequenceClassification", + "ModernBertForTokenClassification", + "ModernBertModel", + "ModernBertPreTrainedModel", + ] + ) + _import_structure["models.moonshine"].extend( + [ + "MoonshineForConditionalGeneration", + "MoonshineModel", + "MoonshinePreTrainedModel", + ] + ) + _import_structure["models.moshi"].extend( + [ + "MoshiForCausalLM", + "MoshiForConditionalGeneration", + "MoshiModel", + "MoshiPreTrainedModel", + ] + ) + _import_structure["models.mpnet"].extend( + [ + "MPNetForMaskedLM", + "MPNetForMultipleChoice", + "MPNetForQuestionAnswering", + "MPNetForSequenceClassification", + "MPNetForTokenClassification", + "MPNetModel", + "MPNetPreTrainedModel", + ] + ) + _import_structure["models.mpt"].extend( + [ + "MptForCausalLM", + "MptForQuestionAnswering", + "MptForSequenceClassification", + "MptForTokenClassification", + "MptModel", + "MptPreTrainedModel", + ] + ) + _import_structure["models.mra"].extend( + [ + "MraForMaskedLM", + "MraForMultipleChoice", + "MraForQuestionAnswering", + "MraForSequenceClassification", + "MraForTokenClassification", + "MraModel", + "MraPreTrainedModel", + ] + ) + _import_structure["models.mt5"].extend( + [ + "MT5EncoderModel", + "MT5ForConditionalGeneration", + "MT5ForQuestionAnswering", + "MT5ForSequenceClassification", + "MT5ForTokenClassification", + "MT5Model", + "MT5PreTrainedModel", + ] + ) + _import_structure["models.musicgen"].extend( + [ + "MusicgenForCausalLM", + "MusicgenForConditionalGeneration", + "MusicgenModel", + "MusicgenPreTrainedModel", + "MusicgenProcessor", + ] + ) + _import_structure["models.musicgen_melody"].extend( + [ + "MusicgenMelodyForCausalLM", + "MusicgenMelodyForConditionalGeneration", + "MusicgenMelodyModel", + "MusicgenMelodyPreTrainedModel", + ] + ) + _import_structure["models.mvp"].extend( + [ + "MvpForCausalLM", + "MvpForConditionalGeneration", + "MvpForQuestionAnswering", + "MvpForSequenceClassification", + "MvpModel", + "MvpPreTrainedModel", + ] + ) + _import_structure["models.nemotron"].extend( + [ + "NemotronForCausalLM", + "NemotronForQuestionAnswering", + "NemotronForSequenceClassification", + "NemotronForTokenClassification", + "NemotronModel", + "NemotronPreTrainedModel", + ] + ) + _import_structure["models.nllb_moe"].extend( + [ + "NllbMoeForConditionalGeneration", + "NllbMoeModel", + "NllbMoePreTrainedModel", + "NllbMoeSparseMLP", + "NllbMoeTop2Router", + ] + ) + _import_structure["models.nystromformer"].extend( + [ + "NystromformerForMaskedLM", + "NystromformerForMultipleChoice", + "NystromformerForQuestionAnswering", + "NystromformerForSequenceClassification", + "NystromformerForTokenClassification", + "NystromformerModel", + "NystromformerPreTrainedModel", + ] + ) + _import_structure["models.olmo"].extend( + [ + "OlmoForCausalLM", + "OlmoModel", + "OlmoPreTrainedModel", + ] + ) + _import_structure["models.olmo2"].extend( + [ + "Olmo2ForCausalLM", + "Olmo2Model", + "Olmo2PreTrainedModel", + ] + ) + _import_structure["models.olmoe"].extend( + [ + "OlmoeForCausalLM", + "OlmoeModel", + "OlmoePreTrainedModel", + ] + ) + _import_structure["models.omdet_turbo"].extend( + [ + "OmDetTurboForObjectDetection", + "OmDetTurboPreTrainedModel", + ] + ) + _import_structure["models.oneformer"].extend( + [ + "OneFormerForUniversalSegmentation", + "OneFormerModel", + "OneFormerPreTrainedModel", + ] + ) + _import_structure["models.openai"].extend( + [ + "OpenAIGPTDoubleHeadsModel", + "OpenAIGPTForSequenceClassification", + "OpenAIGPTLMHeadModel", + "OpenAIGPTModel", + "OpenAIGPTPreTrainedModel", + "load_tf_weights_in_openai_gpt", + ] + ) + _import_structure["models.opt"].extend( + [ + "OPTForCausalLM", + "OPTForQuestionAnswering", + "OPTForSequenceClassification", + "OPTModel", + "OPTPreTrainedModel", + ] + ) + _import_structure["models.owlv2"].extend( + [ + "Owlv2ForObjectDetection", + "Owlv2Model", + "Owlv2PreTrainedModel", + "Owlv2TextModel", + "Owlv2VisionModel", + ] + ) + _import_structure["models.owlvit"].extend( + [ + "OwlViTForObjectDetection", + "OwlViTModel", + "OwlViTPreTrainedModel", + "OwlViTTextModel", + "OwlViTVisionModel", + ] + ) + _import_structure["models.paligemma"].extend( + [ + "PaliGemmaForConditionalGeneration", + "PaliGemmaPreTrainedModel", + "PaliGemmaProcessor", + ] + ) + _import_structure["models.patchtsmixer"].extend( + [ + "PatchTSMixerForPrediction", + "PatchTSMixerForPretraining", + "PatchTSMixerForRegression", + "PatchTSMixerForTimeSeriesClassification", + "PatchTSMixerModel", + "PatchTSMixerPreTrainedModel", + ] + ) + _import_structure["models.patchtst"].extend( + [ + "PatchTSTForClassification", + "PatchTSTForPrediction", + "PatchTSTForPretraining", + "PatchTSTForRegression", + "PatchTSTModel", + "PatchTSTPreTrainedModel", + ] + ) + _import_structure["models.pegasus"].extend( + [ + "PegasusForCausalLM", + "PegasusForConditionalGeneration", + "PegasusModel", + "PegasusPreTrainedModel", + ] + ) + _import_structure["models.pegasus_x"].extend( + [ + "PegasusXForConditionalGeneration", + "PegasusXModel", + "PegasusXPreTrainedModel", + ] + ) + _import_structure["models.perceiver"].extend( + [ + "PerceiverForImageClassificationConvProcessing", + "PerceiverForImageClassificationFourier", + "PerceiverForImageClassificationLearned", + "PerceiverForMaskedLM", + "PerceiverForMultimodalAutoencoding", + "PerceiverForOpticalFlow", + "PerceiverForSequenceClassification", + "PerceiverModel", + "PerceiverPreTrainedModel", + ] + ) + _import_structure["models.persimmon"].extend( + [ + "PersimmonForCausalLM", + "PersimmonForSequenceClassification", + "PersimmonForTokenClassification", + "PersimmonModel", + "PersimmonPreTrainedModel", + ] + ) + _import_structure["models.phi"].extend( + [ + "PhiForCausalLM", + "PhiForSequenceClassification", + "PhiForTokenClassification", + "PhiModel", + "PhiPreTrainedModel", + ] + ) + _import_structure["models.phi3"].extend( + [ + "Phi3ForCausalLM", + "Phi3ForSequenceClassification", + "Phi3ForTokenClassification", + "Phi3Model", + "Phi3PreTrainedModel", + ] + ) + _import_structure["models.phimoe"].extend( + [ + "PhimoeForCausalLM", + "PhimoeForSequenceClassification", + "PhimoeModel", + "PhimoePreTrainedModel", + ] + ) + _import_structure["models.pix2struct"].extend( + [ + "Pix2StructForConditionalGeneration", + "Pix2StructPreTrainedModel", + "Pix2StructTextModel", + "Pix2StructVisionModel", + ] + ) + _import_structure["models.pixtral"].extend(["PixtralPreTrainedModel", "PixtralVisionModel"]) + _import_structure["models.plbart"].extend( + [ + "PLBartForCausalLM", + "PLBartForConditionalGeneration", + "PLBartForSequenceClassification", + "PLBartModel", + "PLBartPreTrainedModel", + ] + ) + _import_structure["models.poolformer"].extend( + [ + "PoolFormerForImageClassification", + "PoolFormerModel", + "PoolFormerPreTrainedModel", + ] + ) + _import_structure["models.pop2piano"].extend( + [ + "Pop2PianoForConditionalGeneration", + "Pop2PianoPreTrainedModel", + ] + ) + _import_structure["models.prompt_depth_anything"].extend( + [ + "PromptDepthAnythingForDepthEstimation", + "PromptDepthAnythingPreTrainedModel", + ] + ) + _import_structure["models.prophetnet"].extend( + [ + "ProphetNetDecoder", + "ProphetNetEncoder", + "ProphetNetForCausalLM", + "ProphetNetForConditionalGeneration", + "ProphetNetModel", + "ProphetNetPreTrainedModel", + ] + ) + _import_structure["models.pvt"].extend( + [ + "PvtForImageClassification", + "PvtModel", + "PvtPreTrainedModel", + ] + ) + _import_structure["models.pvt_v2"].extend( + [ + "PvtV2Backbone", + "PvtV2ForImageClassification", + "PvtV2Model", + "PvtV2PreTrainedModel", + ] + ) + _import_structure["models.qwen2"].extend( + [ + "Qwen2ForCausalLM", + "Qwen2ForQuestionAnswering", + "Qwen2ForSequenceClassification", + "Qwen2ForTokenClassification", + "Qwen2Model", + "Qwen2PreTrainedModel", + ] + ) + _import_structure["models.qwen2_5_vl"].extend( + [ + "Qwen2_5_VLForConditionalGeneration", + "Qwen2_5_VLModel", + "Qwen2_5_VLPreTrainedModel", + ] + ) + _import_structure["models.qwen2_audio"].extend( + [ + "Qwen2AudioEncoder", + "Qwen2AudioForConditionalGeneration", + "Qwen2AudioPreTrainedModel", + ] + ) + _import_structure["models.qwen2_moe"].extend( + [ + "Qwen2MoeForCausalLM", + "Qwen2MoeForQuestionAnswering", + "Qwen2MoeForSequenceClassification", + "Qwen2MoeForTokenClassification", + "Qwen2MoeModel", + "Qwen2MoePreTrainedModel", + ] + ) + _import_structure["models.qwen2_vl"].extend( + [ + "Qwen2VLForConditionalGeneration", + "Qwen2VLModel", + "Qwen2VLPreTrainedModel", + ] + ) + _import_structure["models.qwen3"].extend( + [ + "Qwen3ForCausalLM", + "Qwen3ForQuestionAnswering", + "Qwen3ForSequenceClassification", + "Qwen3ForTokenClassification", + "Qwen3Model", + "Qwen3PreTrainedModel", + ] + ) + _import_structure["models.qwen3_moe"].extend( + [ + "Qwen3MoeForCausalLM", + "Qwen3MoeForQuestionAnswering", + "Qwen3MoeForSequenceClassification", + "Qwen3MoeForTokenClassification", + "Qwen3MoeModel", + "Qwen3MoePreTrainedModel", + ] + ) + _import_structure["models.rag"].extend( + [ + "RagModel", + "RagPreTrainedModel", + "RagSequenceForGeneration", + "RagTokenForGeneration", + ] + ) + _import_structure["models.recurrent_gemma"].extend( + [ + "RecurrentGemmaForCausalLM", + "RecurrentGemmaModel", + "RecurrentGemmaPreTrainedModel", + ] + ) + _import_structure["models.reformer"].extend( + [ + "ReformerForMaskedLM", + "ReformerForQuestionAnswering", + "ReformerForSequenceClassification", + "ReformerModel", + "ReformerModelWithLMHead", + "ReformerPreTrainedModel", + ] + ) + _import_structure["models.regnet"].extend( + [ + "RegNetForImageClassification", + "RegNetModel", + "RegNetPreTrainedModel", + ] + ) + _import_structure["models.rembert"].extend( + [ + "RemBertForCausalLM", + "RemBertForMaskedLM", + "RemBertForMultipleChoice", + "RemBertForQuestionAnswering", + "RemBertForSequenceClassification", + "RemBertForTokenClassification", + "RemBertModel", + "RemBertPreTrainedModel", + "load_tf_weights_in_rembert", + ] + ) + _import_structure["models.resnet"].extend( + [ + "ResNetBackbone", + "ResNetForImageClassification", + "ResNetModel", + "ResNetPreTrainedModel", + ] + ) + _import_structure["models.roberta"].extend( + [ + "RobertaForCausalLM", + "RobertaForMaskedLM", + "RobertaForMultipleChoice", + "RobertaForQuestionAnswering", + "RobertaForSequenceClassification", + "RobertaForTokenClassification", + "RobertaModel", + "RobertaPreTrainedModel", + ] + ) + _import_structure["models.roberta_prelayernorm"].extend( + [ + "RobertaPreLayerNormForCausalLM", + "RobertaPreLayerNormForMaskedLM", + "RobertaPreLayerNormForMultipleChoice", + "RobertaPreLayerNormForQuestionAnswering", + "RobertaPreLayerNormForSequenceClassification", + "RobertaPreLayerNormForTokenClassification", + "RobertaPreLayerNormModel", + "RobertaPreLayerNormPreTrainedModel", + ] + ) + _import_structure["models.roc_bert"].extend( + [ + "RoCBertForCausalLM", + "RoCBertForMaskedLM", + "RoCBertForMultipleChoice", + "RoCBertForPreTraining", + "RoCBertForQuestionAnswering", + "RoCBertForSequenceClassification", + "RoCBertForTokenClassification", + "RoCBertModel", + "RoCBertPreTrainedModel", + "load_tf_weights_in_roc_bert", + ] + ) + _import_structure["models.roformer"].extend( + [ + "RoFormerForCausalLM", + "RoFormerForMaskedLM", + "RoFormerForMultipleChoice", + "RoFormerForQuestionAnswering", + "RoFormerForSequenceClassification", + "RoFormerForTokenClassification", + "RoFormerModel", + "RoFormerPreTrainedModel", + "load_tf_weights_in_roformer", + ] + ) + _import_structure["models.rt_detr"].extend( + [ + "RTDetrForObjectDetection", + "RTDetrModel", + "RTDetrPreTrainedModel", + "RTDetrResNetBackbone", + "RTDetrResNetPreTrainedModel", + ] + ) + _import_structure["models.rt_detr_v2"].extend( + ["RTDetrV2ForObjectDetection", "RTDetrV2Model", "RTDetrV2PreTrainedModel"] + ) + _import_structure["models.rwkv"].extend( + [ + "RwkvForCausalLM", + "RwkvModel", + "RwkvPreTrainedModel", + ] + ) + _import_structure["models.sam"].extend( + [ + "SamModel", + "SamPreTrainedModel", + "SamVisionModel", + ] + ) + _import_structure["models.seamless_m4t"].extend( + [ + "SeamlessM4TCodeHifiGan", + "SeamlessM4TForSpeechToSpeech", + "SeamlessM4TForSpeechToText", + "SeamlessM4TForTextToSpeech", + "SeamlessM4TForTextToText", + "SeamlessM4THifiGan", + "SeamlessM4TModel", + "SeamlessM4TPreTrainedModel", + "SeamlessM4TTextToUnitForConditionalGeneration", + "SeamlessM4TTextToUnitModel", + ] + ) + _import_structure["models.seamless_m4t_v2"].extend( + [ + "SeamlessM4Tv2ForSpeechToSpeech", + "SeamlessM4Tv2ForSpeechToText", + "SeamlessM4Tv2ForTextToSpeech", + "SeamlessM4Tv2ForTextToText", + "SeamlessM4Tv2Model", + "SeamlessM4Tv2PreTrainedModel", + ] + ) + _import_structure["models.segformer"].extend( + [ + "SegformerDecodeHead", + "SegformerForImageClassification", + "SegformerForSemanticSegmentation", + "SegformerModel", + "SegformerPreTrainedModel", + ] + ) + _import_structure["models.seggpt"].extend( + [ + "SegGptForImageSegmentation", + "SegGptModel", + "SegGptPreTrainedModel", + ] + ) + _import_structure["models.sew"].extend( + [ + "SEWForCTC", + "SEWForSequenceClassification", + "SEWModel", + "SEWPreTrainedModel", + ] + ) + _import_structure["models.sew_d"].extend( + [ + "SEWDForCTC", + "SEWDForSequenceClassification", + "SEWDModel", + "SEWDPreTrainedModel", + ] + ) + _import_structure["models.shieldgemma2"].append("ShieldGemma2ForImageClassification") + _import_structure["models.siglip"].extend( + [ + "SiglipForImageClassification", + "SiglipModel", + "SiglipPreTrainedModel", + "SiglipTextModel", + "SiglipVisionModel", + ] + ) + _import_structure["models.siglip2"].extend( + [ + "Siglip2ForImageClassification", + "Siglip2Model", + "Siglip2PreTrainedModel", + "Siglip2TextModel", + "Siglip2VisionModel", + ] + ) + _import_structure["models.smolvlm"].extend( + [ + "SmolVLMForConditionalGeneration", + "SmolVLMModel", + "SmolVLMPreTrainedModel", + "SmolVLMProcessor", + "SmolVLMVisionConfig", + "SmolVLMVisionTransformer", + ] + ) + _import_structure["models.speech_encoder_decoder"].extend(["SpeechEncoderDecoderModel"]) + _import_structure["models.speech_to_text"].extend( + [ + "Speech2TextForConditionalGeneration", + "Speech2TextModel", + "Speech2TextPreTrainedModel", + ] + ) + _import_structure["models.speecht5"].extend( + [ + "SpeechT5ForSpeechToSpeech", + "SpeechT5ForSpeechToText", + "SpeechT5ForTextToSpeech", + "SpeechT5HifiGan", + "SpeechT5Model", + "SpeechT5PreTrainedModel", + ] + ) + _import_structure["models.splinter"].extend( + [ + "SplinterForPreTraining", + "SplinterForQuestionAnswering", + "SplinterModel", + "SplinterPreTrainedModel", + ] + ) + _import_structure["models.squeezebert"].extend( + [ + "SqueezeBertForMaskedLM", + "SqueezeBertForMultipleChoice", + "SqueezeBertForQuestionAnswering", + "SqueezeBertForSequenceClassification", + "SqueezeBertForTokenClassification", + "SqueezeBertModel", + "SqueezeBertPreTrainedModel", + ] + ) + _import_structure["models.stablelm"].extend( + [ + "StableLmForCausalLM", + "StableLmForSequenceClassification", + "StableLmForTokenClassification", + "StableLmModel", + "StableLmPreTrainedModel", + ] + ) + _import_structure["models.starcoder2"].extend( + [ + "Starcoder2ForCausalLM", + "Starcoder2ForSequenceClassification", + "Starcoder2ForTokenClassification", + "Starcoder2Model", + "Starcoder2PreTrainedModel", + ] + ) + _import_structure["models.superglue"].extend( + [ + "SuperGlueForKeypointMatching", + "SuperGluePreTrainedModel", + ] + ) + _import_structure["models.superpoint"].extend( + [ + "SuperPointForKeypointDetection", + "SuperPointPreTrainedModel", + ] + ) + _import_structure["models.swiftformer"].extend( + [ + "SwiftFormerForImageClassification", + "SwiftFormerModel", + "SwiftFormerPreTrainedModel", + ] + ) + _import_structure["models.swin"].extend( + [ + "SwinBackbone", + "SwinForImageClassification", + "SwinForMaskedImageModeling", + "SwinModel", + "SwinPreTrainedModel", + ] + ) + _import_structure["models.swin2sr"].extend( + [ + "Swin2SRForImageSuperResolution", + "Swin2SRModel", + "Swin2SRPreTrainedModel", + ] + ) + _import_structure["models.swinv2"].extend( + [ + "Swinv2Backbone", + "Swinv2ForImageClassification", + "Swinv2ForMaskedImageModeling", + "Swinv2Model", + "Swinv2PreTrainedModel", + ] + ) + _import_structure["models.switch_transformers"].extend( + [ + "SwitchTransformersEncoderModel", + "SwitchTransformersForConditionalGeneration", + "SwitchTransformersModel", + "SwitchTransformersPreTrainedModel", + "SwitchTransformersSparseMLP", + "SwitchTransformersTop1Router", + ] + ) + _import_structure["models.t5"].extend( + [ + "T5EncoderModel", + "T5ForConditionalGeneration", + "T5ForQuestionAnswering", + "T5ForSequenceClassification", + "T5ForTokenClassification", + "T5Model", + "T5PreTrainedModel", + "load_tf_weights_in_t5", + ] + ) + _import_structure["models.table_transformer"].extend( + [ + "TableTransformerForObjectDetection", + "TableTransformerModel", + "TableTransformerPreTrainedModel", + ] + ) + _import_structure["models.tapas"].extend( + [ + "TapasForMaskedLM", + "TapasForQuestionAnswering", + "TapasForSequenceClassification", + "TapasModel", + "TapasPreTrainedModel", + "load_tf_weights_in_tapas", + ] + ) + _import_structure["models.textnet"].extend( + [ + "TextNetBackbone", + "TextNetForImageClassification", + "TextNetModel", + "TextNetPreTrainedModel", + ] + ) + _import_structure["models.time_series_transformer"].extend( + [ + "TimeSeriesTransformerForPrediction", + "TimeSeriesTransformerModel", + "TimeSeriesTransformerPreTrainedModel", + ] + ) + _import_structure["models.timesformer"].extend( + [ + "TimesformerForVideoClassification", + "TimesformerModel", + "TimesformerPreTrainedModel", + ] + ) + _import_structure["models.timm_backbone"].extend(["TimmBackbone"]) + _import_structure["models.timm_wrapper"].extend( + ["TimmWrapperForImageClassification", "TimmWrapperModel", "TimmWrapperPreTrainedModel"] + ) + _import_structure["models.trocr"].extend( + [ + "TrOCRForCausalLM", + "TrOCRPreTrainedModel", + ] + ) + _import_structure["models.tvp"].extend( + [ + "TvpForVideoGrounding", + "TvpModel", + "TvpPreTrainedModel", + ] + ) + _import_structure["models.udop"].extend( + [ + "UdopEncoderModel", + "UdopForConditionalGeneration", + "UdopModel", + "UdopPreTrainedModel", + ], + ) + _import_structure["models.umt5"].extend( + [ + "UMT5EncoderModel", + "UMT5ForConditionalGeneration", + "UMT5ForQuestionAnswering", + "UMT5ForSequenceClassification", + "UMT5ForTokenClassification", + "UMT5Model", + "UMT5PreTrainedModel", + ] + ) + _import_structure["models.unispeech"].extend( + [ + "UniSpeechForCTC", + "UniSpeechForPreTraining", + "UniSpeechForSequenceClassification", + "UniSpeechModel", + "UniSpeechPreTrainedModel", + ] + ) + _import_structure["models.unispeech_sat"].extend( + [ + "UniSpeechSatForAudioFrameClassification", + "UniSpeechSatForCTC", + "UniSpeechSatForPreTraining", + "UniSpeechSatForSequenceClassification", + "UniSpeechSatForXVector", + "UniSpeechSatModel", + "UniSpeechSatPreTrainedModel", + ] + ) + _import_structure["models.univnet"].extend( + [ + "UnivNetModel", + ] + ) + _import_structure["models.upernet"].extend( + [ + "UperNetForSemanticSegmentation", + "UperNetPreTrainedModel", + ] + ) + _import_structure["models.video_llava"].extend( + [ + "VideoLlavaForConditionalGeneration", + "VideoLlavaPreTrainedModel", + "VideoLlavaProcessor", + ] + ) + _import_structure["models.videomae"].extend( + [ + "VideoMAEForPreTraining", + "VideoMAEForVideoClassification", + "VideoMAEModel", + "VideoMAEPreTrainedModel", + ] + ) + _import_structure["models.vilt"].extend( + [ + "ViltForImageAndTextRetrieval", + "ViltForImagesAndTextClassification", + "ViltForMaskedLM", + "ViltForQuestionAnswering", + "ViltForTokenClassification", + "ViltModel", + "ViltPreTrainedModel", + ] + ) + _import_structure["models.vipllava"].extend( + [ + "VipLlavaForConditionalGeneration", + "VipLlavaPreTrainedModel", + ] + ) + _import_structure["models.vision_encoder_decoder"].extend(["VisionEncoderDecoderModel"]) + _import_structure["models.vision_text_dual_encoder"].extend(["VisionTextDualEncoderModel"]) + _import_structure["models.visual_bert"].extend( + [ + "VisualBertForMultipleChoice", + "VisualBertForPreTraining", + "VisualBertForQuestionAnswering", + "VisualBertForRegionToPhraseAlignment", + "VisualBertForVisualReasoning", + "VisualBertModel", + "VisualBertPreTrainedModel", + ] + ) + _import_structure["models.vit"].extend( + [ + "ViTForImageClassification", + "ViTForMaskedImageModeling", + "ViTModel", + "ViTPreTrainedModel", + ] + ) + _import_structure["models.vit_mae"].extend( + [ + "ViTMAEForPreTraining", + "ViTMAEModel", + "ViTMAEPreTrainedModel", + ] + ) + _import_structure["models.vit_msn"].extend( + [ + "ViTMSNForImageClassification", + "ViTMSNModel", + "ViTMSNPreTrainedModel", + ] + ) + _import_structure["models.vitdet"].extend( + [ + "VitDetBackbone", + "VitDetModel", + "VitDetPreTrainedModel", + ] + ) + _import_structure["models.vitmatte"].extend( + [ + "VitMatteForImageMatting", + "VitMattePreTrainedModel", + ] + ) + _import_structure["models.vitpose"].extend( + [ + "VitPoseForPoseEstimation", + "VitPosePreTrainedModel", + ] + ) + _import_structure["models.vitpose_backbone"].extend( + [ + "VitPoseBackbone", + "VitPoseBackbonePreTrainedModel", + ] + ) + _import_structure["models.vits"].extend( + [ + "VitsModel", + "VitsPreTrainedModel", + ] + ) + _import_structure["models.vivit"].extend( + [ + "VivitForVideoClassification", + "VivitModel", + "VivitPreTrainedModel", + ] + ) + _import_structure["models.wav2vec2"].extend( + [ + "Wav2Vec2ForAudioFrameClassification", + "Wav2Vec2ForCTC", + "Wav2Vec2ForMaskedLM", + "Wav2Vec2ForPreTraining", + "Wav2Vec2ForSequenceClassification", + "Wav2Vec2ForXVector", + "Wav2Vec2Model", + "Wav2Vec2PreTrainedModel", + ] + ) + _import_structure["models.wav2vec2_bert"].extend( + [ + "Wav2Vec2BertForAudioFrameClassification", + "Wav2Vec2BertForCTC", + "Wav2Vec2BertForSequenceClassification", + "Wav2Vec2BertForXVector", + "Wav2Vec2BertModel", + "Wav2Vec2BertPreTrainedModel", + ] + ) + _import_structure["models.wav2vec2_conformer"].extend( + [ + "Wav2Vec2ConformerForAudioFrameClassification", + "Wav2Vec2ConformerForCTC", + "Wav2Vec2ConformerForPreTraining", + "Wav2Vec2ConformerForSequenceClassification", + "Wav2Vec2ConformerForXVector", + "Wav2Vec2ConformerModel", + "Wav2Vec2ConformerPreTrainedModel", + ] + ) + _import_structure["models.wavlm"].extend( + [ + "WavLMForAudioFrameClassification", + "WavLMForCTC", + "WavLMForSequenceClassification", + "WavLMForXVector", + "WavLMModel", + "WavLMPreTrainedModel", + ] + ) + _import_structure["models.whisper"].extend( + [ + "WhisperForAudioClassification", + "WhisperForCausalLM", + "WhisperForConditionalGeneration", + "WhisperModel", + "WhisperPreTrainedModel", + ] + ) + _import_structure["models.x_clip"].extend( + [ + "XCLIPModel", + "XCLIPPreTrainedModel", + "XCLIPTextModel", + "XCLIPVisionModel", + ] + ) + _import_structure["models.xglm"].extend( + [ + "XGLMForCausalLM", + "XGLMModel", + "XGLMPreTrainedModel", + ] + ) + _import_structure["models.xlm"].extend( + [ + "XLMForMultipleChoice", + "XLMForQuestionAnswering", + "XLMForQuestionAnsweringSimple", + "XLMForSequenceClassification", + "XLMForTokenClassification", + "XLMModel", + "XLMPreTrainedModel", + "XLMWithLMHeadModel", + ] + ) + _import_structure["models.xlm_roberta"].extend( + [ + "XLMRobertaForCausalLM", + "XLMRobertaForMaskedLM", + "XLMRobertaForMultipleChoice", + "XLMRobertaForQuestionAnswering", + "XLMRobertaForSequenceClassification", + "XLMRobertaForTokenClassification", + "XLMRobertaModel", + "XLMRobertaPreTrainedModel", + ] + ) + _import_structure["models.xlm_roberta_xl"].extend( + [ + "XLMRobertaXLForCausalLM", + "XLMRobertaXLForMaskedLM", + "XLMRobertaXLForMultipleChoice", + "XLMRobertaXLForQuestionAnswering", + "XLMRobertaXLForSequenceClassification", + "XLMRobertaXLForTokenClassification", + "XLMRobertaXLModel", + "XLMRobertaXLPreTrainedModel", + ] + ) + _import_structure["models.xlnet"].extend( + [ + "XLNetForMultipleChoice", + "XLNetForQuestionAnswering", + "XLNetForQuestionAnsweringSimple", + "XLNetForSequenceClassification", + "XLNetForTokenClassification", + "XLNetLMHeadModel", + "XLNetModel", + "XLNetPreTrainedModel", + "load_tf_weights_in_xlnet", + ] + ) + _import_structure["models.xmod"].extend( + [ + "XmodForCausalLM", + "XmodForMaskedLM", + "XmodForMultipleChoice", + "XmodForQuestionAnswering", + "XmodForSequenceClassification", + "XmodForTokenClassification", + "XmodModel", + "XmodPreTrainedModel", + ] + ) + _import_structure["models.yolos"].extend( + [ + "YolosForObjectDetection", + "YolosModel", + "YolosPreTrainedModel", + ] + ) + _import_structure["models.yoso"].extend( + [ + "YosoForMaskedLM", + "YosoForMultipleChoice", + "YosoForQuestionAnswering", + "YosoForSequenceClassification", + "YosoForTokenClassification", + "YosoModel", + "YosoPreTrainedModel", + ] + ) + _import_structure["models.zamba"].extend( + [ + "ZambaForCausalLM", + "ZambaForSequenceClassification", + "ZambaModel", + "ZambaPreTrainedModel", + ] + ) + _import_structure["models.zamba2"].extend( + [ + "Zamba2ForCausalLM", + "Zamba2ForSequenceClassification", + "Zamba2Model", + "Zamba2PreTrainedModel", + ] + ) + _import_structure["models.zoedepth"].extend( + [ + "ZoeDepthForDepthEstimation", + "ZoeDepthPreTrainedModel", + ] + ) + _import_structure["optimization"] = [ + "Adafactor", + "get_constant_schedule", + "get_constant_schedule_with_warmup", + "get_cosine_schedule_with_warmup", + "get_cosine_with_hard_restarts_schedule_with_warmup", + "get_inverse_sqrt_schedule", + "get_linear_schedule_with_warmup", + "get_polynomial_decay_schedule_with_warmup", + "get_scheduler", + "get_wsd_schedule", + ] + _import_structure["pytorch_utils"] = [ + "Conv1D", + "apply_chunking_to_forward", + "prune_layer", + ] + _import_structure["sagemaker"] = [] + _import_structure["time_series_utils"] = [] + _import_structure["trainer"] = ["Trainer"] + _import_structure["trainer_pt_utils"] = ["torch_distributed_zero_first"] + _import_structure["trainer_seq2seq"] = ["Seq2SeqTrainer"] + +# TensorFlow-backed objects +try: + if not is_tf_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_tf_objects + + _import_structure["utils.dummy_tf_objects"] = [name for name in dir(dummy_tf_objects) if not name.startswith("_")] +else: + _import_structure["activations_tf"] = [] + _import_structure["generation"].extend( + [ + "TFForcedBOSTokenLogitsProcessor", + "TFForcedEOSTokenLogitsProcessor", + "TFForceTokensLogitsProcessor", + "TFGenerationMixin", + "TFLogitsProcessor", + "TFLogitsProcessorList", + "TFLogitsWarper", + "TFMinLengthLogitsProcessor", + "TFNoBadWordsLogitsProcessor", + "TFNoRepeatNGramLogitsProcessor", + "TFRepetitionPenaltyLogitsProcessor", + "TFSuppressTokensAtBeginLogitsProcessor", + "TFSuppressTokensLogitsProcessor", + "TFTemperatureLogitsWarper", + "TFTopKLogitsWarper", + "TFTopPLogitsWarper", + ] + ) + _import_structure["keras_callbacks"] = ["KerasMetricCallback", "PushToHubCallback"] + _import_structure["modeling_tf_outputs"] = [] + _import_structure["modeling_tf_utils"] = [ + "TFPreTrainedModel", + "TFSequenceSummary", + "TFSharedEmbeddings", + "shape_list", + ] + # TensorFlow models structure + _import_structure["models.albert"].extend( + [ + "TFAlbertForMaskedLM", + "TFAlbertForMultipleChoice", + "TFAlbertForPreTraining", + "TFAlbertForQuestionAnswering", + "TFAlbertForSequenceClassification", + "TFAlbertForTokenClassification", + "TFAlbertMainLayer", + "TFAlbertModel", + "TFAlbertPreTrainedModel", + ] + ) + _import_structure["models.auto"].extend( + [ + "TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING", + "TF_MODEL_FOR_CAUSAL_LM_MAPPING", + "TF_MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING", + "TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING", + "TF_MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING", + "TF_MODEL_FOR_MASKED_LM_MAPPING", + "TF_MODEL_FOR_MASK_GENERATION_MAPPING", + "TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING", + "TF_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING", + "TF_MODEL_FOR_PRETRAINING_MAPPING", + "TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING", + "TF_MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING", + "TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING", + "TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING", + "TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING", + "TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING", + "TF_MODEL_FOR_TEXT_ENCODING_MAPPING", + "TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING", + "TF_MODEL_FOR_VISION_2_SEQ_MAPPING", + "TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING", + "TF_MODEL_MAPPING", + "TF_MODEL_WITH_LM_HEAD_MAPPING", + "TFAutoModel", + "TFAutoModelForAudioClassification", + "TFAutoModelForCausalLM", + "TFAutoModelForDocumentQuestionAnswering", + "TFAutoModelForImageClassification", + "TFAutoModelForMaskedImageModeling", + "TFAutoModelForMaskedLM", + "TFAutoModelForMaskGeneration", + "TFAutoModelForMultipleChoice", + "TFAutoModelForNextSentencePrediction", + "TFAutoModelForPreTraining", + "TFAutoModelForQuestionAnswering", + "TFAutoModelForSemanticSegmentation", + "TFAutoModelForSeq2SeqLM", + "TFAutoModelForSequenceClassification", + "TFAutoModelForSpeechSeq2Seq", + "TFAutoModelForTableQuestionAnswering", + "TFAutoModelForTextEncoding", + "TFAutoModelForTokenClassification", + "TFAutoModelForVision2Seq", + "TFAutoModelForZeroShotImageClassification", + "TFAutoModelWithLMHead", + ] + ) + _import_structure["models.bart"].extend( + [ + "TFBartForConditionalGeneration", + "TFBartForSequenceClassification", + "TFBartModel", + "TFBartPretrainedModel", + ] + ) + _import_structure["models.bert"].extend( + [ + "TFBertForMaskedLM", + "TFBertForMultipleChoice", + "TFBertForNextSentencePrediction", + "TFBertForPreTraining", + "TFBertForQuestionAnswering", + "TFBertForSequenceClassification", + "TFBertForTokenClassification", + "TFBertLMHeadModel", + "TFBertMainLayer", + "TFBertModel", + "TFBertPreTrainedModel", + ] + ) + _import_structure["models.blenderbot"].extend( + [ + "TFBlenderbotForConditionalGeneration", + "TFBlenderbotModel", + "TFBlenderbotPreTrainedModel", + ] + ) + _import_structure["models.blenderbot_small"].extend( + [ + "TFBlenderbotSmallForConditionalGeneration", + "TFBlenderbotSmallModel", + "TFBlenderbotSmallPreTrainedModel", + ] + ) + _import_structure["models.blip"].extend( + [ + "TFBlipForConditionalGeneration", + "TFBlipForImageTextRetrieval", + "TFBlipForQuestionAnswering", + "TFBlipModel", + "TFBlipPreTrainedModel", + "TFBlipTextModel", + "TFBlipVisionModel", + ] + ) + _import_structure["models.camembert"].extend( + [ + "TFCamembertForCausalLM", + "TFCamembertForMaskedLM", + "TFCamembertForMultipleChoice", + "TFCamembertForQuestionAnswering", + "TFCamembertForSequenceClassification", + "TFCamembertForTokenClassification", + "TFCamembertModel", + "TFCamembertPreTrainedModel", + ] + ) + _import_structure["models.clip"].extend( + [ + "TFCLIPModel", + "TFCLIPPreTrainedModel", + "TFCLIPTextModel", + "TFCLIPVisionModel", + ] + ) + _import_structure["models.convbert"].extend( + [ + "TFConvBertForMaskedLM", + "TFConvBertForMultipleChoice", + "TFConvBertForQuestionAnswering", + "TFConvBertForSequenceClassification", + "TFConvBertForTokenClassification", + "TFConvBertModel", + "TFConvBertPreTrainedModel", + ] + ) + _import_structure["models.convnext"].extend( + [ + "TFConvNextForImageClassification", + "TFConvNextModel", + "TFConvNextPreTrainedModel", + ] + ) + _import_structure["models.convnextv2"].extend( + [ + "TFConvNextV2ForImageClassification", + "TFConvNextV2Model", + "TFConvNextV2PreTrainedModel", + ] + ) + _import_structure["models.ctrl"].extend( + [ + "TFCTRLForSequenceClassification", + "TFCTRLLMHeadModel", + "TFCTRLModel", + "TFCTRLPreTrainedModel", + ] + ) + _import_structure["models.cvt"].extend( + [ + "TFCvtForImageClassification", + "TFCvtModel", + "TFCvtPreTrainedModel", + ] + ) + _import_structure["models.data2vec"].extend( + [ + "TFData2VecVisionForImageClassification", + "TFData2VecVisionForSemanticSegmentation", + "TFData2VecVisionModel", + "TFData2VecVisionPreTrainedModel", + ] + ) + _import_structure["models.deberta"].extend( + [ + "TFDebertaForMaskedLM", + "TFDebertaForQuestionAnswering", + "TFDebertaForSequenceClassification", + "TFDebertaForTokenClassification", + "TFDebertaModel", + "TFDebertaPreTrainedModel", + ] + ) + _import_structure["models.deberta_v2"].extend( + [ + "TFDebertaV2ForMaskedLM", + "TFDebertaV2ForMultipleChoice", + "TFDebertaV2ForQuestionAnswering", + "TFDebertaV2ForSequenceClassification", + "TFDebertaV2ForTokenClassification", + "TFDebertaV2Model", + "TFDebertaV2PreTrainedModel", + ] + ) + _import_structure["models.deit"].extend( + [ + "TFDeiTForImageClassification", + "TFDeiTForImageClassificationWithTeacher", + "TFDeiTForMaskedImageModeling", + "TFDeiTModel", + "TFDeiTPreTrainedModel", + ] + ) + _import_structure["models.deprecated.efficientformer"].extend( + [ + "TFEfficientFormerForImageClassification", + "TFEfficientFormerForImageClassificationWithTeacher", + "TFEfficientFormerModel", + "TFEfficientFormerPreTrainedModel", + ] + ) + _import_structure["models.deprecated.transfo_xl"].extend( + [ + "TFAdaptiveEmbedding", + "TFTransfoXLForSequenceClassification", + "TFTransfoXLLMHeadModel", + "TFTransfoXLMainLayer", + "TFTransfoXLModel", + "TFTransfoXLPreTrainedModel", + ] + ) + _import_structure["models.distilbert"].extend( + [ + "TFDistilBertForMaskedLM", + "TFDistilBertForMultipleChoice", + "TFDistilBertForQuestionAnswering", + "TFDistilBertForSequenceClassification", + "TFDistilBertForTokenClassification", + "TFDistilBertMainLayer", + "TFDistilBertModel", + "TFDistilBertPreTrainedModel", + ] + ) + _import_structure["models.dpr"].extend( + [ + "TFDPRContextEncoder", + "TFDPRPretrainedContextEncoder", + "TFDPRPretrainedQuestionEncoder", + "TFDPRPretrainedReader", + "TFDPRQuestionEncoder", + "TFDPRReader", + ] + ) + _import_structure["models.electra"].extend( + [ + "TFElectraForMaskedLM", + "TFElectraForMultipleChoice", + "TFElectraForPreTraining", + "TFElectraForQuestionAnswering", + "TFElectraForSequenceClassification", + "TFElectraForTokenClassification", + "TFElectraModel", + "TFElectraPreTrainedModel", + ] + ) + _import_structure["models.encoder_decoder"].append("TFEncoderDecoderModel") + _import_structure["models.esm"].extend( + [ + "TFEsmForMaskedLM", + "TFEsmForSequenceClassification", + "TFEsmForTokenClassification", + "TFEsmModel", + "TFEsmPreTrainedModel", + ] + ) + _import_structure["models.flaubert"].extend( + [ + "TFFlaubertForMultipleChoice", + "TFFlaubertForQuestionAnsweringSimple", + "TFFlaubertForSequenceClassification", + "TFFlaubertForTokenClassification", + "TFFlaubertModel", + "TFFlaubertPreTrainedModel", + "TFFlaubertWithLMHeadModel", + ] + ) + _import_structure["models.funnel"].extend( + [ + "TFFunnelBaseModel", + "TFFunnelForMaskedLM", + "TFFunnelForMultipleChoice", + "TFFunnelForPreTraining", + "TFFunnelForQuestionAnswering", + "TFFunnelForSequenceClassification", + "TFFunnelForTokenClassification", + "TFFunnelModel", + "TFFunnelPreTrainedModel", + ] + ) + _import_structure["models.gpt2"].extend( + [ + "TFGPT2DoubleHeadsModel", + "TFGPT2ForSequenceClassification", + "TFGPT2LMHeadModel", + "TFGPT2MainLayer", + "TFGPT2Model", + "TFGPT2PreTrainedModel", + ] + ) + _import_structure["models.gptj"].extend( + [ + "TFGPTJForCausalLM", + "TFGPTJForQuestionAnswering", + "TFGPTJForSequenceClassification", + "TFGPTJModel", + "TFGPTJPreTrainedModel", + ] + ) + _import_structure["models.groupvit"].extend( + [ + "TFGroupViTModel", + "TFGroupViTPreTrainedModel", + "TFGroupViTTextModel", + "TFGroupViTVisionModel", + ] + ) + _import_structure["models.hubert"].extend( + [ + "TFHubertForCTC", + "TFHubertModel", + "TFHubertPreTrainedModel", + ] + ) + + _import_structure["models.idefics"].extend( + [ + "TFIdeficsForVisionText2Text", + "TFIdeficsModel", + "TFIdeficsPreTrainedModel", + ] + ) + + _import_structure["models.layoutlm"].extend( + [ + "TFLayoutLMForMaskedLM", + "TFLayoutLMForQuestionAnswering", + "TFLayoutLMForSequenceClassification", + "TFLayoutLMForTokenClassification", + "TFLayoutLMMainLayer", + "TFLayoutLMModel", + "TFLayoutLMPreTrainedModel", + ] + ) + _import_structure["models.layoutlmv3"].extend( + [ + "TFLayoutLMv3ForQuestionAnswering", + "TFLayoutLMv3ForSequenceClassification", + "TFLayoutLMv3ForTokenClassification", + "TFLayoutLMv3Model", + "TFLayoutLMv3PreTrainedModel", + ] + ) + _import_structure["models.led"].extend(["TFLEDForConditionalGeneration", "TFLEDModel", "TFLEDPreTrainedModel"]) + _import_structure["models.longformer"].extend( + [ + "TFLongformerForMaskedLM", + "TFLongformerForMultipleChoice", + "TFLongformerForQuestionAnswering", + "TFLongformerForSequenceClassification", + "TFLongformerForTokenClassification", + "TFLongformerModel", + "TFLongformerPreTrainedModel", + ] + ) + _import_structure["models.lxmert"].extend( + [ + "TFLxmertForPreTraining", + "TFLxmertMainLayer", + "TFLxmertModel", + "TFLxmertPreTrainedModel", + "TFLxmertVisualFeatureEncoder", + ] + ) + _import_structure["models.marian"].extend(["TFMarianModel", "TFMarianMTModel", "TFMarianPreTrainedModel"]) + _import_structure["models.mbart"].extend( + ["TFMBartForConditionalGeneration", "TFMBartModel", "TFMBartPreTrainedModel"] + ) + _import_structure["models.mistral"].extend( + ["TFMistralForCausalLM", "TFMistralForSequenceClassification", "TFMistralModel", "TFMistralPreTrainedModel"] + ) + _import_structure["models.mobilebert"].extend( + [ + "TFMobileBertForMaskedLM", + "TFMobileBertForMultipleChoice", + "TFMobileBertForNextSentencePrediction", + "TFMobileBertForPreTraining", + "TFMobileBertForQuestionAnswering", + "TFMobileBertForSequenceClassification", + "TFMobileBertForTokenClassification", + "TFMobileBertMainLayer", + "TFMobileBertModel", + "TFMobileBertPreTrainedModel", + ] + ) + _import_structure["models.mobilevit"].extend( + [ + "TFMobileViTForImageClassification", + "TFMobileViTForSemanticSegmentation", + "TFMobileViTModel", + "TFMobileViTPreTrainedModel", + ] + ) + _import_structure["models.mpnet"].extend( + [ + "TFMPNetForMaskedLM", + "TFMPNetForMultipleChoice", + "TFMPNetForQuestionAnswering", + "TFMPNetForSequenceClassification", + "TFMPNetForTokenClassification", + "TFMPNetMainLayer", + "TFMPNetModel", + "TFMPNetPreTrainedModel", + ] + ) + _import_structure["models.mt5"].extend(["TFMT5EncoderModel", "TFMT5ForConditionalGeneration", "TFMT5Model"]) + _import_structure["models.openai"].extend( + [ + "TFOpenAIGPTDoubleHeadsModel", + "TFOpenAIGPTForSequenceClassification", + "TFOpenAIGPTLMHeadModel", + "TFOpenAIGPTMainLayer", + "TFOpenAIGPTModel", + "TFOpenAIGPTPreTrainedModel", + ] + ) + _import_structure["models.opt"].extend( + [ + "TFOPTForCausalLM", + "TFOPTModel", + "TFOPTPreTrainedModel", + ] + ) + _import_structure["models.pegasus"].extend( + [ + "TFPegasusForConditionalGeneration", + "TFPegasusModel", + "TFPegasusPreTrainedModel", + ] + ) + _import_structure["models.rag"].extend( + [ + "TFRagModel", + "TFRagPreTrainedModel", + "TFRagSequenceForGeneration", + "TFRagTokenForGeneration", + ] + ) + _import_structure["models.regnet"].extend( + [ + "TFRegNetForImageClassification", + "TFRegNetModel", + "TFRegNetPreTrainedModel", + ] + ) + _import_structure["models.rembert"].extend( + [ + "TFRemBertForCausalLM", + "TFRemBertForMaskedLM", + "TFRemBertForMultipleChoice", + "TFRemBertForQuestionAnswering", + "TFRemBertForSequenceClassification", + "TFRemBertForTokenClassification", + "TFRemBertModel", + "TFRemBertPreTrainedModel", + ] + ) + _import_structure["models.resnet"].extend( + [ + "TFResNetForImageClassification", + "TFResNetModel", + "TFResNetPreTrainedModel", + ] + ) + _import_structure["models.roberta"].extend( + [ + "TFRobertaForCausalLM", + "TFRobertaForMaskedLM", + "TFRobertaForMultipleChoice", + "TFRobertaForQuestionAnswering", + "TFRobertaForSequenceClassification", + "TFRobertaForTokenClassification", + "TFRobertaMainLayer", + "TFRobertaModel", + "TFRobertaPreTrainedModel", + ] + ) + _import_structure["models.roberta_prelayernorm"].extend( + [ + "TFRobertaPreLayerNormForCausalLM", + "TFRobertaPreLayerNormForMaskedLM", + "TFRobertaPreLayerNormForMultipleChoice", + "TFRobertaPreLayerNormForQuestionAnswering", + "TFRobertaPreLayerNormForSequenceClassification", + "TFRobertaPreLayerNormForTokenClassification", + "TFRobertaPreLayerNormMainLayer", + "TFRobertaPreLayerNormModel", + "TFRobertaPreLayerNormPreTrainedModel", + ] + ) + _import_structure["models.roformer"].extend( + [ + "TFRoFormerForCausalLM", + "TFRoFormerForMaskedLM", + "TFRoFormerForMultipleChoice", + "TFRoFormerForQuestionAnswering", + "TFRoFormerForSequenceClassification", + "TFRoFormerForTokenClassification", + "TFRoFormerModel", + "TFRoFormerPreTrainedModel", + ] + ) + _import_structure["models.sam"].extend( + [ + "TFSamModel", + "TFSamPreTrainedModel", + "TFSamVisionModel", + ] + ) + _import_structure["models.segformer"].extend( + [ + "TFSegformerDecodeHead", + "TFSegformerForImageClassification", + "TFSegformerForSemanticSegmentation", + "TFSegformerModel", + "TFSegformerPreTrainedModel", + ] + ) + _import_structure["models.speech_to_text"].extend( + [ + "TFSpeech2TextForConditionalGeneration", + "TFSpeech2TextModel", + "TFSpeech2TextPreTrainedModel", + ] + ) + _import_structure["models.swiftformer"].extend( + [ + "TFSwiftFormerForImageClassification", + "TFSwiftFormerModel", + "TFSwiftFormerPreTrainedModel", + ] + ) + _import_structure["models.swin"].extend( + [ + "TFSwinForImageClassification", + "TFSwinForMaskedImageModeling", + "TFSwinModel", + "TFSwinPreTrainedModel", + ] + ) + _import_structure["models.t5"].extend( + [ + "TFT5EncoderModel", + "TFT5ForConditionalGeneration", + "TFT5Model", + "TFT5PreTrainedModel", + ] + ) + _import_structure["models.tapas"].extend( + [ + "TFTapasForMaskedLM", + "TFTapasForQuestionAnswering", + "TFTapasForSequenceClassification", + "TFTapasModel", + "TFTapasPreTrainedModel", + ] + ) + _import_structure["models.vision_encoder_decoder"].extend(["TFVisionEncoderDecoderModel"]) + _import_structure["models.vision_text_dual_encoder"].extend(["TFVisionTextDualEncoderModel"]) + _import_structure["models.vit"].extend( + [ + "TFViTForImageClassification", + "TFViTModel", + "TFViTPreTrainedModel", + ] + ) + _import_structure["models.vit_mae"].extend( + [ + "TFViTMAEForPreTraining", + "TFViTMAEModel", + "TFViTMAEPreTrainedModel", + ] + ) + _import_structure["models.wav2vec2"].extend( + [ + "TFWav2Vec2ForCTC", + "TFWav2Vec2ForSequenceClassification", + "TFWav2Vec2Model", + "TFWav2Vec2PreTrainedModel", + ] + ) + _import_structure["models.whisper"].extend( + [ + "TFWhisperForConditionalGeneration", + "TFWhisperModel", + "TFWhisperPreTrainedModel", + ] + ) + _import_structure["models.xglm"].extend( + [ + "TFXGLMForCausalLM", + "TFXGLMModel", + "TFXGLMPreTrainedModel", + ] + ) + _import_structure["models.xlm"].extend( + [ + "TFXLMForMultipleChoice", + "TFXLMForQuestionAnsweringSimple", + "TFXLMForSequenceClassification", + "TFXLMForTokenClassification", + "TFXLMMainLayer", + "TFXLMModel", + "TFXLMPreTrainedModel", + "TFXLMWithLMHeadModel", + ] + ) + _import_structure["models.xlm_roberta"].extend( + [ + "TFXLMRobertaForCausalLM", + "TFXLMRobertaForMaskedLM", + "TFXLMRobertaForMultipleChoice", + "TFXLMRobertaForQuestionAnswering", + "TFXLMRobertaForSequenceClassification", + "TFXLMRobertaForTokenClassification", + "TFXLMRobertaModel", + "TFXLMRobertaPreTrainedModel", + ] + ) + _import_structure["models.xlnet"].extend( + [ + "TFXLNetForMultipleChoice", + "TFXLNetForQuestionAnsweringSimple", + "TFXLNetForSequenceClassification", + "TFXLNetForTokenClassification", + "TFXLNetLMHeadModel", + "TFXLNetMainLayer", + "TFXLNetModel", + "TFXLNetPreTrainedModel", + ] + ) + _import_structure["optimization_tf"] = [ + "AdamWeightDecay", + "GradientAccumulator", + "WarmUp", + "create_optimizer", + ] + _import_structure["tf_utils"] = [] + + +try: + if not ( + is_librosa_available() + and is_essentia_available() + and is_scipy_available() + and is_torch_available() + and is_pretty_midi_available() + ): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import ( + dummy_essentia_and_librosa_and_pretty_midi_and_scipy_and_torch_objects, + ) + + _import_structure["utils.dummy_essentia_and_librosa_and_pretty_midi_and_scipy_and_torch_objects"] = [ + name + for name in dir(dummy_essentia_and_librosa_and_pretty_midi_and_scipy_and_torch_objects) + if not name.startswith("_") + ] +else: + _import_structure["models.pop2piano"].append("Pop2PianoFeatureExtractor") + _import_structure["models.pop2piano"].append("Pop2PianoTokenizer") + _import_structure["models.pop2piano"].append("Pop2PianoProcessor") + +try: + if not is_torchaudio_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import ( + dummy_torchaudio_objects, + ) + + _import_structure["utils.dummy_torchaudio_objects"] = [ + name for name in dir(dummy_torchaudio_objects) if not name.startswith("_") + ] +else: + _import_structure["models.musicgen_melody"].append("MusicgenMelodyFeatureExtractor") + _import_structure["models.musicgen_melody"].append("MusicgenMelodyProcessor") + + +# FLAX-backed objects +try: + if not is_flax_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_flax_objects + + _import_structure["utils.dummy_flax_objects"] = [ + name for name in dir(dummy_flax_objects) if not name.startswith("_") + ] +else: + _import_structure["generation"].extend( + [ + "FlaxForcedBOSTokenLogitsProcessor", + "FlaxForcedEOSTokenLogitsProcessor", + "FlaxForceTokensLogitsProcessor", + "FlaxGenerationMixin", + "FlaxLogitsProcessor", + "FlaxLogitsProcessorList", + "FlaxLogitsWarper", + "FlaxMinLengthLogitsProcessor", + "FlaxTemperatureLogitsWarper", + "FlaxSuppressTokensAtBeginLogitsProcessor", + "FlaxSuppressTokensLogitsProcessor", + "FlaxTopKLogitsWarper", + "FlaxTopPLogitsWarper", + "FlaxWhisperTimeStampLogitsProcessor", + ] + ) + _import_structure["modeling_flax_outputs"] = [] + _import_structure["modeling_flax_utils"] = ["FlaxPreTrainedModel"] + _import_structure["models.albert"].extend( + [ + "FlaxAlbertForMaskedLM", + "FlaxAlbertForMultipleChoice", + "FlaxAlbertForPreTraining", + "FlaxAlbertForQuestionAnswering", + "FlaxAlbertForSequenceClassification", + "FlaxAlbertForTokenClassification", + "FlaxAlbertModel", + "FlaxAlbertPreTrainedModel", + ] + ) + _import_structure["models.auto"].extend( + [ + "FLAX_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING", + "FLAX_MODEL_FOR_CAUSAL_LM_MAPPING", + "FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING", + "FLAX_MODEL_FOR_MASKED_LM_MAPPING", + "FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING", + "FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING", + "FLAX_MODEL_FOR_PRETRAINING_MAPPING", + "FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING", + "FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING", + "FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING", + "FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING", + "FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING", + "FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING", + "FLAX_MODEL_MAPPING", + "FlaxAutoModel", + "FlaxAutoModelForCausalLM", + "FlaxAutoModelForImageClassification", + "FlaxAutoModelForMaskedLM", + "FlaxAutoModelForMultipleChoice", + "FlaxAutoModelForNextSentencePrediction", + "FlaxAutoModelForPreTraining", + "FlaxAutoModelForQuestionAnswering", + "FlaxAutoModelForSeq2SeqLM", + "FlaxAutoModelForSequenceClassification", + "FlaxAutoModelForSpeechSeq2Seq", + "FlaxAutoModelForTokenClassification", + "FlaxAutoModelForVision2Seq", + ] + ) + + # Flax models structure + + _import_structure["models.bart"].extend( + [ + "FlaxBartDecoderPreTrainedModel", + "FlaxBartForCausalLM", + "FlaxBartForConditionalGeneration", + "FlaxBartForQuestionAnswering", + "FlaxBartForSequenceClassification", + "FlaxBartModel", + "FlaxBartPreTrainedModel", + ] + ) + _import_structure["models.beit"].extend( + [ + "FlaxBeitForImageClassification", + "FlaxBeitForMaskedImageModeling", + "FlaxBeitModel", + "FlaxBeitPreTrainedModel", + ] + ) + + _import_structure["models.bert"].extend( + [ + "FlaxBertForCausalLM", + "FlaxBertForMaskedLM", + "FlaxBertForMultipleChoice", + "FlaxBertForNextSentencePrediction", + "FlaxBertForPreTraining", + "FlaxBertForQuestionAnswering", + "FlaxBertForSequenceClassification", + "FlaxBertForTokenClassification", + "FlaxBertModel", + "FlaxBertPreTrainedModel", + ] + ) + _import_structure["models.big_bird"].extend( + [ + "FlaxBigBirdForCausalLM", + "FlaxBigBirdForMaskedLM", + "FlaxBigBirdForMultipleChoice", + "FlaxBigBirdForPreTraining", + "FlaxBigBirdForQuestionAnswering", + "FlaxBigBirdForSequenceClassification", + "FlaxBigBirdForTokenClassification", + "FlaxBigBirdModel", + "FlaxBigBirdPreTrainedModel", + ] + ) + _import_structure["models.blenderbot"].extend( + [ + "FlaxBlenderbotForConditionalGeneration", + "FlaxBlenderbotModel", + "FlaxBlenderbotPreTrainedModel", + ] + ) + _import_structure["models.blenderbot_small"].extend( + [ + "FlaxBlenderbotSmallForConditionalGeneration", + "FlaxBlenderbotSmallModel", + "FlaxBlenderbotSmallPreTrainedModel", + ] + ) + _import_structure["models.bloom"].extend( + [ + "FlaxBloomForCausalLM", + "FlaxBloomModel", + "FlaxBloomPreTrainedModel", + ] + ) + _import_structure["models.clip"].extend( + [ + "FlaxCLIPModel", + "FlaxCLIPPreTrainedModel", + "FlaxCLIPTextModel", + "FlaxCLIPTextPreTrainedModel", + "FlaxCLIPTextModelWithProjection", + "FlaxCLIPVisionModel", + "FlaxCLIPVisionPreTrainedModel", + ] + ) + _import_structure["models.dinov2"].extend( + [ + "FlaxDinov2Model", + "FlaxDinov2ForImageClassification", + "FlaxDinov2PreTrainedModel", + ] + ) + _import_structure["models.distilbert"].extend( + [ + "FlaxDistilBertForMaskedLM", + "FlaxDistilBertForMultipleChoice", + "FlaxDistilBertForQuestionAnswering", + "FlaxDistilBertForSequenceClassification", + "FlaxDistilBertForTokenClassification", + "FlaxDistilBertModel", + "FlaxDistilBertPreTrainedModel", + ] + ) + _import_structure["models.electra"].extend( + [ + "FlaxElectraForCausalLM", + "FlaxElectraForMaskedLM", + "FlaxElectraForMultipleChoice", + "FlaxElectraForPreTraining", + "FlaxElectraForQuestionAnswering", + "FlaxElectraForSequenceClassification", + "FlaxElectraForTokenClassification", + "FlaxElectraModel", + "FlaxElectraPreTrainedModel", + ] + ) + _import_structure["models.encoder_decoder"].append("FlaxEncoderDecoderModel") + _import_structure["models.gpt2"].extend(["FlaxGPT2LMHeadModel", "FlaxGPT2Model", "FlaxGPT2PreTrainedModel"]) + _import_structure["models.gpt_neo"].extend( + ["FlaxGPTNeoForCausalLM", "FlaxGPTNeoModel", "FlaxGPTNeoPreTrainedModel"] + ) + _import_structure["models.gptj"].extend(["FlaxGPTJForCausalLM", "FlaxGPTJModel", "FlaxGPTJPreTrainedModel"]) + _import_structure["models.llama"].extend(["FlaxLlamaForCausalLM", "FlaxLlamaModel", "FlaxLlamaPreTrainedModel"]) + _import_structure["models.gemma"].extend(["FlaxGemmaForCausalLM", "FlaxGemmaModel", "FlaxGemmaPreTrainedModel"]) + _import_structure["models.longt5"].extend( + [ + "FlaxLongT5ForConditionalGeneration", + "FlaxLongT5Model", + "FlaxLongT5PreTrainedModel", + ] + ) + _import_structure["models.marian"].extend( + [ + "FlaxMarianModel", + "FlaxMarianMTModel", + "FlaxMarianPreTrainedModel", + ] + ) + _import_structure["models.mbart"].extend( + [ + "FlaxMBartForConditionalGeneration", + "FlaxMBartForQuestionAnswering", + "FlaxMBartForSequenceClassification", + "FlaxMBartModel", + "FlaxMBartPreTrainedModel", + ] + ) + _import_structure["models.mistral"].extend( + [ + "FlaxMistralForCausalLM", + "FlaxMistralModel", + "FlaxMistralPreTrainedModel", + ] + ) + _import_structure["models.mt5"].extend(["FlaxMT5EncoderModel", "FlaxMT5ForConditionalGeneration", "FlaxMT5Model"]) + _import_structure["models.opt"].extend( + [ + "FlaxOPTForCausalLM", + "FlaxOPTModel", + "FlaxOPTPreTrainedModel", + ] + ) + _import_structure["models.pegasus"].extend( + [ + "FlaxPegasusForConditionalGeneration", + "FlaxPegasusModel", + "FlaxPegasusPreTrainedModel", + ] + ) + _import_structure["models.regnet"].extend( + [ + "FlaxRegNetForImageClassification", + "FlaxRegNetModel", + "FlaxRegNetPreTrainedModel", + ] + ) + _import_structure["models.resnet"].extend( + [ + "FlaxResNetForImageClassification", + "FlaxResNetModel", + "FlaxResNetPreTrainedModel", + ] + ) + _import_structure["models.roberta"].extend( + [ + "FlaxRobertaForCausalLM", + "FlaxRobertaForMaskedLM", + "FlaxRobertaForMultipleChoice", + "FlaxRobertaForQuestionAnswering", + "FlaxRobertaForSequenceClassification", + "FlaxRobertaForTokenClassification", + "FlaxRobertaModel", + "FlaxRobertaPreTrainedModel", + ] + ) + _import_structure["models.roberta_prelayernorm"].extend( + [ + "FlaxRobertaPreLayerNormForCausalLM", + "FlaxRobertaPreLayerNormForMaskedLM", + "FlaxRobertaPreLayerNormForMultipleChoice", + "FlaxRobertaPreLayerNormForQuestionAnswering", + "FlaxRobertaPreLayerNormForSequenceClassification", + "FlaxRobertaPreLayerNormForTokenClassification", + "FlaxRobertaPreLayerNormModel", + "FlaxRobertaPreLayerNormPreTrainedModel", + ] + ) + _import_structure["models.roformer"].extend( + [ + "FlaxRoFormerForMaskedLM", + "FlaxRoFormerForMultipleChoice", + "FlaxRoFormerForQuestionAnswering", + "FlaxRoFormerForSequenceClassification", + "FlaxRoFormerForTokenClassification", + "FlaxRoFormerModel", + "FlaxRoFormerPreTrainedModel", + ] + ) + _import_structure["models.speech_encoder_decoder"].append("FlaxSpeechEncoderDecoderModel") + _import_structure["models.t5"].extend( + [ + "FlaxT5EncoderModel", + "FlaxT5ForConditionalGeneration", + "FlaxT5Model", + "FlaxT5PreTrainedModel", + ] + ) + _import_structure["models.vision_encoder_decoder"].append("FlaxVisionEncoderDecoderModel") + _import_structure["models.vision_text_dual_encoder"].extend(["FlaxVisionTextDualEncoderModel"]) + _import_structure["models.vit"].extend(["FlaxViTForImageClassification", "FlaxViTModel", "FlaxViTPreTrainedModel"]) + _import_structure["models.wav2vec2"].extend( + [ + "FlaxWav2Vec2ForCTC", + "FlaxWav2Vec2ForPreTraining", + "FlaxWav2Vec2Model", + "FlaxWav2Vec2PreTrainedModel", + ] + ) + _import_structure["models.whisper"].extend( + [ + "FlaxWhisperForConditionalGeneration", + "FlaxWhisperModel", + "FlaxWhisperPreTrainedModel", + "FlaxWhisperForAudioClassification", + ] + ) + _import_structure["models.xglm"].extend( + [ + "FlaxXGLMForCausalLM", + "FlaxXGLMModel", + "FlaxXGLMPreTrainedModel", + ] + ) + _import_structure["models.xlm_roberta"].extend( + [ + "FlaxXLMRobertaForMaskedLM", + "FlaxXLMRobertaForMultipleChoice", + "FlaxXLMRobertaForQuestionAnswering", + "FlaxXLMRobertaForSequenceClassification", + "FlaxXLMRobertaForTokenClassification", + "FlaxXLMRobertaModel", + "FlaxXLMRobertaForCausalLM", + "FlaxXLMRobertaPreTrainedModel", + ] + ) + + +# Direct imports for type-checking +if TYPE_CHECKING: + # Configuration + # Agents + from .agents import ( + Agent, + CodeAgent, + HfApiEngine, + ManagedAgent, + PipelineTool, + ReactAgent, + ReactCodeAgent, + ReactJsonAgent, + Tool, + Toolbox, + ToolCollection, + TransformersEngine, + launch_gradio_demo, + load_tool, + stream_to_gradio, + tool, + ) + from .configuration_utils import PretrainedConfig + + # Data + from .data import ( + DataProcessor, + InputExample, + InputFeatures, + SingleSentenceClassificationProcessor, + SquadExample, + SquadFeatures, + SquadV1Processor, + SquadV2Processor, + glue_compute_metrics, + glue_convert_examples_to_features, + glue_output_modes, + glue_processors, + glue_tasks_num_labels, + squad_convert_examples_to_features, + xnli_compute_metrics, + xnli_output_modes, + xnli_processors, + xnli_tasks_num_labels, + ) + from .data.data_collator import ( + DataCollator, + DataCollatorForLanguageModeling, + DataCollatorForMultipleChoice, + DataCollatorForPermutationLanguageModeling, + DataCollatorForSeq2Seq, + DataCollatorForSOP, + DataCollatorForTokenClassification, + DataCollatorForWholeWordMask, + DataCollatorWithFlattening, + DataCollatorWithPadding, + DefaultDataCollator, + default_data_collator, + ) + from .feature_extraction_sequence_utils import SequenceFeatureExtractor + + # Feature Extractor + from .feature_extraction_utils import BatchFeature, FeatureExtractionMixin + + # Generation + from .generation import ( + AsyncTextIteratorStreamer, + CompileConfig, + GenerationConfig, + TextIteratorStreamer, + TextStreamer, + WatermarkingConfig, + ) + from .hf_argparser import HfArgumentParser + + # Integrations + from .integrations import ( + is_clearml_available, + is_comet_available, + is_dvclive_available, + is_neptune_available, + is_optuna_available, + is_ray_available, + is_ray_tune_available, + is_sigopt_available, + is_swanlab_available, + is_tensorboard_available, + is_wandb_available, + ) + + # Model Cards + from .modelcard import ModelCard + + # TF 2.0 <=> PyTorch conversion utilities + from .modeling_tf_pytorch_utils import ( + convert_tf_weight_name_to_pt_weight_name, + load_pytorch_checkpoint_in_tf2_model, + load_pytorch_model_in_tf2_model, + load_pytorch_weights_in_tf2_model, + load_tf2_checkpoint_in_pytorch_model, + load_tf2_model_in_pytorch_model, + load_tf2_weights_in_pytorch_model, + ) + from .models.albert import AlbertConfig + from .models.align import ( + AlignConfig, + AlignProcessor, + AlignTextConfig, + AlignVisionConfig, + ) + from .models.altclip import ( + AltCLIPConfig, + AltCLIPProcessor, + AltCLIPTextConfig, + AltCLIPVisionConfig, + ) + from .models.aria import ( + AriaConfig, + AriaProcessor, + AriaTextConfig, + ) + from .models.audio_spectrogram_transformer import ( + ASTConfig, + ASTFeatureExtractor, + ) + from .models.auto import ( + CONFIG_MAPPING, + FEATURE_EXTRACTOR_MAPPING, + IMAGE_PROCESSOR_MAPPING, + MODEL_NAMES_MAPPING, + PROCESSOR_MAPPING, + TOKENIZER_MAPPING, + AutoConfig, + AutoFeatureExtractor, + AutoImageProcessor, + AutoProcessor, + AutoTokenizer, + ) + from .models.autoformer import ( + AutoformerConfig, + ) + from .models.aya_vision import ( + AyaVisionConfig, + AyaVisionProcessor, + ) + from .models.bamba import BambaConfig + from .models.bark import ( + BarkCoarseConfig, + BarkConfig, + BarkFineConfig, + BarkProcessor, + BarkSemanticConfig, + ) + from .models.bart import BartConfig, BartTokenizer + from .models.beit import BeitConfig + from .models.bert import ( + BasicTokenizer, + BertConfig, + BertTokenizer, + WordpieceTokenizer, + ) + from .models.bert_generation import BertGenerationConfig + from .models.bert_japanese import ( + BertJapaneseTokenizer, + CharacterTokenizer, + MecabTokenizer, + ) + from .models.bertweet import BertweetTokenizer + from .models.big_bird import BigBirdConfig + from .models.bigbird_pegasus import ( + BigBirdPegasusConfig, + ) + from .models.biogpt import ( + BioGptConfig, + BioGptTokenizer, + ) + from .models.bit import BitConfig + from .models.blenderbot import ( + BlenderbotConfig, + BlenderbotTokenizer, + ) + from .models.blenderbot_small import ( + BlenderbotSmallConfig, + BlenderbotSmallTokenizer, + ) + from .models.blip import ( + BlipConfig, + BlipProcessor, + BlipTextConfig, + BlipVisionConfig, + ) + from .models.blip_2 import ( + Blip2Config, + Blip2Processor, + Blip2QFormerConfig, + Blip2VisionConfig, + ) + from .models.bloom import BloomConfig + from .models.bridgetower import ( + BridgeTowerConfig, + BridgeTowerProcessor, + BridgeTowerTextConfig, + BridgeTowerVisionConfig, + ) + from .models.bros import ( + BrosConfig, + BrosProcessor, + ) + from .models.byt5 import ByT5Tokenizer + from .models.camembert import ( + CamembertConfig, + ) + from .models.canine import ( + CanineConfig, + CanineTokenizer, + ) + from .models.chameleon import ( + ChameleonConfig, + ChameleonProcessor, + ChameleonVQVAEConfig, + ) + from .models.chinese_clip import ( + ChineseCLIPConfig, + ChineseCLIPProcessor, + ChineseCLIPTextConfig, + ChineseCLIPVisionConfig, + ) + from .models.clap import ( + ClapAudioConfig, + ClapConfig, + ClapProcessor, + ClapTextConfig, + ) + from .models.clip import ( + CLIPConfig, + CLIPProcessor, + CLIPTextConfig, + CLIPTokenizer, + CLIPVisionConfig, + ) + from .models.clipseg import ( + CLIPSegConfig, + CLIPSegProcessor, + CLIPSegTextConfig, + CLIPSegVisionConfig, + ) + from .models.clvp import ( + ClvpConfig, + ClvpDecoderConfig, + ClvpEncoderConfig, + ClvpFeatureExtractor, + ClvpProcessor, + ClvpTokenizer, + ) + from .models.codegen import ( + CodeGenConfig, + CodeGenTokenizer, + ) + from .models.cohere import CohereConfig + from .models.cohere2 import Cohere2Config + from .models.colpali import ( + ColPaliConfig, + ColPaliProcessor, + ) + from .models.conditional_detr import ( + ConditionalDetrConfig, + ) + from .models.convbert import ( + ConvBertConfig, + ConvBertTokenizer, + ) + from .models.convnext import ConvNextConfig + from .models.convnextv2 import ( + ConvNextV2Config, + ) + from .models.cpmant import ( + CpmAntConfig, + CpmAntTokenizer, + ) + from .models.ctrl import ( + CTRLConfig, + CTRLTokenizer, + ) + from .models.cvt import CvtConfig + from .models.dab_detr import ( + DabDetrConfig, + ) + from .models.dac import ( + DacConfig, + DacFeatureExtractor, + ) + from .models.data2vec import ( + Data2VecAudioConfig, + Data2VecTextConfig, + Data2VecVisionConfig, + ) + from .models.dbrx import DbrxConfig + from .models.deberta import ( + DebertaConfig, + DebertaTokenizer, + ) + from .models.deberta_v2 import ( + DebertaV2Config, + ) + from .models.decision_transformer import ( + DecisionTransformerConfig, + ) + from .models.deepseek_v3 import ( + DeepseekV3Config, + ) + from .models.deformable_detr import ( + DeformableDetrConfig, + ) + from .models.deit import DeiTConfig + from .models.deprecated.deta import DetaConfig + from .models.deprecated.efficientformer import ( + EfficientFormerConfig, + ) + from .models.deprecated.ernie_m import ErnieMConfig + from .models.deprecated.gptsan_japanese import ( + GPTSanJapaneseConfig, + GPTSanJapaneseTokenizer, + ) + from .models.deprecated.graphormer import GraphormerConfig + from .models.deprecated.jukebox import ( + JukeboxConfig, + JukeboxPriorConfig, + JukeboxTokenizer, + JukeboxVQVAEConfig, + ) + from .models.deprecated.mctct import ( + MCTCTConfig, + MCTCTFeatureExtractor, + MCTCTProcessor, + ) + from .models.deprecated.mega import MegaConfig + from .models.deprecated.mmbt import MMBTConfig + from .models.deprecated.nat import NatConfig + from .models.deprecated.nezha import NezhaConfig + from .models.deprecated.open_llama import ( + OpenLlamaConfig, + ) + from .models.deprecated.qdqbert import QDQBertConfig + from .models.deprecated.realm import ( + RealmConfig, + RealmTokenizer, + ) + from .models.deprecated.retribert import ( + RetriBertConfig, + RetriBertTokenizer, + ) + from .models.deprecated.speech_to_text_2 import ( + Speech2Text2Config, + Speech2Text2Processor, + Speech2Text2Tokenizer, + ) + from .models.deprecated.tapex import TapexTokenizer + from .models.deprecated.trajectory_transformer import ( + TrajectoryTransformerConfig, + ) + from .models.deprecated.transfo_xl import ( + TransfoXLConfig, + TransfoXLCorpus, + TransfoXLTokenizer, + ) + from .models.deprecated.tvlt import ( + TvltConfig, + TvltFeatureExtractor, + TvltProcessor, + ) + from .models.deprecated.van import VanConfig + from .models.deprecated.vit_hybrid import ( + ViTHybridConfig, + ) + from .models.deprecated.xlm_prophetnet import ( + XLMProphetNetConfig, + ) + from .models.depth_anything import DepthAnythingConfig + from .models.depth_pro import DepthProConfig + from .models.detr import DetrConfig + from .models.diffllama import DiffLlamaConfig + from .models.dinat import DinatConfig + from .models.dinov2 import Dinov2Config + from .models.dinov2_with_registers import Dinov2WithRegistersConfig + from .models.distilbert import ( + DistilBertConfig, + DistilBertTokenizer, + ) + from .models.donut import ( + DonutProcessor, + DonutSwinConfig, + ) + from .models.dpr import ( + DPRConfig, + DPRContextEncoderTokenizer, + DPRQuestionEncoderTokenizer, + DPRReaderOutput, + DPRReaderTokenizer, + ) + from .models.dpt import DPTConfig + from .models.efficientnet import ( + EfficientNetConfig, + ) + from .models.electra import ( + ElectraConfig, + ElectraTokenizer, + ) + from .models.emu3 import ( + Emu3Config, + Emu3Processor, + Emu3TextConfig, + Emu3VQVAEConfig, + ) + from .models.encodec import ( + EncodecConfig, + EncodecFeatureExtractor, + ) + from .models.encoder_decoder import EncoderDecoderConfig + from .models.ernie import ErnieConfig + from .models.esm import EsmConfig, EsmTokenizer + from .models.falcon import FalconConfig + from .models.falcon_mamba import FalconMambaConfig + from .models.fastspeech2_conformer import ( + FastSpeech2ConformerConfig, + FastSpeech2ConformerHifiGanConfig, + FastSpeech2ConformerTokenizer, + FastSpeech2ConformerWithHifiGanConfig, + ) + from .models.flaubert import FlaubertConfig, FlaubertTokenizer + from .models.flava import ( + FlavaConfig, + FlavaImageCodebookConfig, + FlavaImageConfig, + FlavaMultimodalConfig, + FlavaTextConfig, + ) + from .models.fnet import FNetConfig + from .models.focalnet import FocalNetConfig + from .models.fsmt import ( + FSMTConfig, + FSMTTokenizer, + ) + from .models.funnel import ( + FunnelConfig, + FunnelTokenizer, + ) + from .models.fuyu import FuyuConfig + from .models.gemma import GemmaConfig + from .models.gemma2 import Gemma2Config + from .models.gemma3 import Gemma3Config, Gemma3Processor, Gemma3TextConfig + from .models.git import ( + GitConfig, + GitProcessor, + GitVisionConfig, + ) + from .models.glm import GlmConfig + from .models.glpn import GLPNConfig + from .models.got_ocr2 import GotOcr2Config, GotOcr2Processor, GotOcr2VisionConfig + from .models.gpt2 import ( + GPT2Config, + GPT2Tokenizer, + ) + from .models.gpt_bigcode import ( + GPTBigCodeConfig, + ) + from .models.gpt_neo import GPTNeoConfig + from .models.gpt_neox import GPTNeoXConfig + from .models.gpt_neox_japanese import ( + GPTNeoXJapaneseConfig, + ) + from .models.gptj import GPTJConfig + from .models.granite import GraniteConfig + from .models.granitemoe import GraniteMoeConfig + from .models.granitemoeshared import GraniteMoeSharedConfig + from .models.grounding_dino import ( + GroundingDinoConfig, + GroundingDinoProcessor, + ) + from .models.groupvit import ( + GroupViTConfig, + GroupViTTextConfig, + GroupViTVisionConfig, + ) + from .models.helium import HeliumConfig + from .models.herbert import HerbertTokenizer + from .models.hiera import HieraConfig + from .models.hubert import HubertConfig + from .models.ibert import IBertConfig + from .models.idefics import ( + IdeficsConfig, + ) + from .models.idefics2 import Idefics2Config + from .models.idefics3 import Idefics3Config + from .models.ijepa import IJepaConfig + from .models.imagegpt import ImageGPTConfig + from .models.informer import InformerConfig + from .models.instructblip import ( + InstructBlipConfig, + InstructBlipProcessor, + InstructBlipQFormerConfig, + InstructBlipVisionConfig, + ) + from .models.instructblipvideo import ( + InstructBlipVideoConfig, + InstructBlipVideoProcessor, + InstructBlipVideoQFormerConfig, + InstructBlipVideoVisionConfig, + ) + from .models.jamba import JambaConfig + from .models.jetmoe import JetMoeConfig + from .models.kosmos2 import ( + Kosmos2Config, + Kosmos2Processor, + ) + from .models.layoutlm import ( + LayoutLMConfig, + LayoutLMTokenizer, + ) + from .models.layoutlmv2 import ( + LayoutLMv2Config, + LayoutLMv2FeatureExtractor, + LayoutLMv2ImageProcessor, + LayoutLMv2Processor, + LayoutLMv2Tokenizer, + ) + from .models.layoutlmv3 import ( + LayoutLMv3Config, + LayoutLMv3FeatureExtractor, + LayoutLMv3ImageProcessor, + LayoutLMv3Processor, + LayoutLMv3Tokenizer, + ) + from .models.layoutxlm import LayoutXLMProcessor + from .models.led import LEDConfig, LEDTokenizer + from .models.levit import LevitConfig + from .models.lilt import LiltConfig + from .models.llama import LlamaConfig + from .models.llama4 import ( + Llama4Config, + Llama4Processor, + Llama4TextConfig, + Llama4VisionConfig, + ) + from .models.llava import ( + LlavaConfig, + LlavaProcessor, + ) + from .models.llava_next import ( + LlavaNextConfig, + LlavaNextProcessor, + ) + from .models.llava_next_video import ( + LlavaNextVideoConfig, + LlavaNextVideoProcessor, + ) + from .models.llava_onevision import ( + LlavaOnevisionConfig, + LlavaOnevisionProcessor, + ) + from .models.longformer import ( + LongformerConfig, + LongformerTokenizer, + ) + from .models.longt5 import LongT5Config + from .models.luke import ( + LukeConfig, + LukeTokenizer, + ) + from .models.lxmert import ( + LxmertConfig, + LxmertTokenizer, + ) + from .models.m2m_100 import M2M100Config + from .models.mamba import MambaConfig + from .models.mamba2 import Mamba2Config + from .models.marian import MarianConfig + from .models.markuplm import ( + MarkupLMConfig, + MarkupLMFeatureExtractor, + MarkupLMProcessor, + MarkupLMTokenizer, + ) + from .models.mask2former import ( + Mask2FormerConfig, + ) + from .models.maskformer import ( + MaskFormerConfig, + MaskFormerSwinConfig, + ) + from .models.mbart import MBartConfig + from .models.megatron_bert import ( + MegatronBertConfig, + ) + from .models.mgp_str import ( + MgpstrConfig, + MgpstrProcessor, + MgpstrTokenizer, + ) + from .models.mimi import ( + MimiConfig, + ) + from .models.mistral import MistralConfig + from .models.mistral3 import Mistral3Config + from .models.mixtral import MixtralConfig + from .models.mllama import ( + MllamaConfig, + MllamaProcessor, + ) + from .models.mobilebert import ( + MobileBertConfig, + MobileBertTokenizer, + ) + from .models.mobilenet_v1 import ( + MobileNetV1Config, + ) + from .models.mobilenet_v2 import ( + MobileNetV2Config, + ) + from .models.mobilevit import ( + MobileViTConfig, + ) + from .models.mobilevitv2 import ( + MobileViTV2Config, + ) + from .models.modernbert import ModernBertConfig + from .models.moonshine import MoonshineConfig + from .models.moshi import ( + MoshiConfig, + MoshiDepthConfig, + ) + from .models.mpnet import ( + MPNetConfig, + MPNetTokenizer, + ) + from .models.mpt import MptConfig + from .models.mra import MraConfig + from .models.mt5 import MT5Config + from .models.musicgen import ( + MusicgenConfig, + MusicgenDecoderConfig, + ) + from .models.musicgen_melody import ( + MusicgenMelodyConfig, + MusicgenMelodyDecoderConfig, + ) + from .models.mvp import MvpConfig, MvpTokenizer + from .models.myt5 import MyT5Tokenizer + from .models.nemotron import NemotronConfig + from .models.nllb_moe import NllbMoeConfig + from .models.nougat import NougatProcessor + from .models.nystromformer import ( + NystromformerConfig, + ) + from .models.olmo import OlmoConfig + from .models.olmo2 import Olmo2Config + from .models.olmoe import OlmoeConfig + from .models.omdet_turbo import ( + OmDetTurboConfig, + OmDetTurboProcessor, + ) + from .models.oneformer import ( + OneFormerConfig, + OneFormerProcessor, + ) + from .models.openai import ( + OpenAIGPTConfig, + OpenAIGPTTokenizer, + ) + from .models.opt import OPTConfig + from .models.owlv2 import ( + Owlv2Config, + Owlv2Processor, + Owlv2TextConfig, + Owlv2VisionConfig, + ) + from .models.owlvit import ( + OwlViTConfig, + OwlViTProcessor, + OwlViTTextConfig, + OwlViTVisionConfig, + ) + from .models.paligemma import ( + PaliGemmaConfig, + ) + from .models.patchtsmixer import ( + PatchTSMixerConfig, + ) + from .models.patchtst import PatchTSTConfig + from .models.pegasus import ( + PegasusConfig, + PegasusTokenizer, + ) + from .models.pegasus_x import ( + PegasusXConfig, + ) + from .models.perceiver import ( + PerceiverConfig, + PerceiverTokenizer, + ) + from .models.persimmon import ( + PersimmonConfig, + ) + from .models.phi import PhiConfig + from .models.phi3 import Phi3Config + from .models.phi4_multimodal import ( + Phi4MultimodalAudioConfig, + Phi4MultimodalConfig, + Phi4MultimodalFeatureExtractor, + Phi4MultimodalProcessor, + Phi4MultimodalVisionConfig, + ) + from .models.phimoe import PhimoeConfig + from .models.phobert import PhobertTokenizer + from .models.pix2struct import ( + Pix2StructConfig, + Pix2StructProcessor, + Pix2StructTextConfig, + Pix2StructVisionConfig, + ) + from .models.pixtral import ( + PixtralProcessor, + PixtralVisionConfig, + ) + from .models.plbart import PLBartConfig + from .models.poolformer import ( + PoolFormerConfig, + ) + from .models.pop2piano import ( + Pop2PianoConfig, + ) + from .models.prompt_depth_anything import PromptDepthAnythingConfig + from .models.prophetnet import ( + ProphetNetConfig, + ProphetNetTokenizer, + ) + from .models.pvt import PvtConfig + from .models.pvt_v2 import PvtV2Config + from .models.qwen2 import Qwen2Config, Qwen2Tokenizer + from .models.qwen2_5_vl import ( + Qwen2_5_VLConfig, + Qwen2_5_VLProcessor, + ) + from .models.qwen2_audio import ( + Qwen2AudioConfig, + Qwen2AudioEncoderConfig, + Qwen2AudioProcessor, + ) + from .models.qwen2_moe import Qwen2MoeConfig + from .models.qwen2_vl import ( + Qwen2VLConfig, + Qwen2VLProcessor, + ) + from .models.qwen3 import Qwen3Config + from .models.qwen3_moe import Qwen3MoeConfig + from .models.rag import RagConfig, RagRetriever, RagTokenizer + from .models.recurrent_gemma import RecurrentGemmaConfig + from .models.reformer import ReformerConfig + from .models.regnet import RegNetConfig + from .models.rembert import RemBertConfig + from .models.resnet import ResNetConfig + from .models.roberta import ( + RobertaConfig, + RobertaTokenizer, + ) + from .models.roberta_prelayernorm import ( + RobertaPreLayerNormConfig, + ) + from .models.roc_bert import ( + RoCBertConfig, + RoCBertTokenizer, + ) + from .models.roformer import ( + RoFormerConfig, + RoFormerTokenizer, + ) + from .models.rt_detr import ( + RTDetrConfig, + RTDetrResNetConfig, + ) + from .models.rt_detr_v2 import RTDetrV2Config + from .models.rwkv import RwkvConfig + from .models.sam import ( + SamConfig, + SamMaskDecoderConfig, + SamProcessor, + SamPromptEncoderConfig, + SamVisionConfig, + ) + from .models.seamless_m4t import ( + SeamlessM4TConfig, + SeamlessM4TFeatureExtractor, + SeamlessM4TProcessor, + ) + from .models.seamless_m4t_v2 import ( + SeamlessM4Tv2Config, + ) + from .models.segformer import SegformerConfig + from .models.seggpt import SegGptConfig + from .models.sew import SEWConfig + from .models.sew_d import SEWDConfig + from .models.shieldgemma2 import ( + ShieldGemma2Config, + ShieldGemma2Processor, + ) + from .models.siglip import ( + SiglipConfig, + SiglipProcessor, + SiglipTextConfig, + SiglipVisionConfig, + ) + from .models.siglip2 import ( + Siglip2Config, + Siglip2Processor, + Siglip2TextConfig, + Siglip2VisionConfig, + ) + from .models.smolvlm import SmolVLMConfig + from .models.speech_encoder_decoder import SpeechEncoderDecoderConfig + from .models.speech_to_text import ( + Speech2TextConfig, + Speech2TextFeatureExtractor, + Speech2TextProcessor, + ) + from .models.speecht5 import ( + SpeechT5Config, + SpeechT5FeatureExtractor, + SpeechT5HifiGanConfig, + SpeechT5Processor, + ) + from .models.splinter import ( + SplinterConfig, + SplinterTokenizer, + ) + from .models.squeezebert import ( + SqueezeBertConfig, + SqueezeBertTokenizer, + ) + from .models.stablelm import StableLmConfig + from .models.starcoder2 import Starcoder2Config + from .models.superglue import SuperGlueConfig + from .models.superpoint import SuperPointConfig + from .models.swiftformer import ( + SwiftFormerConfig, + ) + from .models.swin import SwinConfig + from .models.swin2sr import Swin2SRConfig + from .models.swinv2 import Swinv2Config + from .models.switch_transformers import ( + SwitchTransformersConfig, + ) + from .models.t5 import T5Config + from .models.table_transformer import ( + TableTransformerConfig, + ) + from .models.tapas import ( + TapasConfig, + TapasTokenizer, + ) + from .models.textnet import TextNetConfig + from .models.time_series_transformer import ( + TimeSeriesTransformerConfig, + ) + from .models.timesformer import ( + TimesformerConfig, + ) + from .models.timm_backbone import TimmBackboneConfig + from .models.timm_wrapper import TimmWrapperConfig + from .models.trocr import ( + TrOCRConfig, + TrOCRProcessor, + ) + from .models.tvp import ( + TvpConfig, + TvpProcessor, + ) + from .models.udop import UdopConfig, UdopProcessor + from .models.umt5 import UMT5Config + from .models.unispeech import ( + UniSpeechConfig, + ) + from .models.unispeech_sat import ( + UniSpeechSatConfig, + ) + from .models.univnet import ( + UnivNetConfig, + UnivNetFeatureExtractor, + ) + from .models.upernet import UperNetConfig + from .models.video_llava import VideoLlavaConfig + from .models.videomae import VideoMAEConfig + from .models.vilt import ( + ViltConfig, + ViltFeatureExtractor, + ViltImageProcessor, + ViltProcessor, + ) + from .models.vipllava import ( + VipLlavaConfig, + ) + from .models.vision_encoder_decoder import VisionEncoderDecoderConfig + from .models.vision_text_dual_encoder import ( + VisionTextDualEncoderConfig, + VisionTextDualEncoderProcessor, + ) + from .models.visual_bert import ( + VisualBertConfig, + ) + from .models.vit import ViTConfig + from .models.vit_mae import ViTMAEConfig + from .models.vit_msn import ViTMSNConfig + from .models.vitdet import VitDetConfig + from .models.vitmatte import VitMatteConfig + from .models.vitpose import VitPoseConfig + from .models.vitpose_backbone import VitPoseBackboneConfig + from .models.vits import ( + VitsConfig, + VitsTokenizer, + ) + from .models.vivit import VivitConfig + from .models.wav2vec2 import ( + Wav2Vec2Config, + Wav2Vec2CTCTokenizer, + Wav2Vec2FeatureExtractor, + Wav2Vec2Processor, + Wav2Vec2Tokenizer, + ) + from .models.wav2vec2_bert import ( + Wav2Vec2BertConfig, + Wav2Vec2BertProcessor, + ) + from .models.wav2vec2_conformer import ( + Wav2Vec2ConformerConfig, + ) + from .models.wav2vec2_phoneme import Wav2Vec2PhonemeCTCTokenizer + from .models.wav2vec2_with_lm import Wav2Vec2ProcessorWithLM + from .models.wavlm import WavLMConfig + from .models.whisper import ( + WhisperConfig, + WhisperFeatureExtractor, + WhisperProcessor, + WhisperTokenizer, + ) + from .models.x_clip import ( + XCLIPConfig, + XCLIPProcessor, + XCLIPTextConfig, + XCLIPVisionConfig, + ) + from .models.xglm import XGLMConfig + from .models.xlm import XLMConfig, XLMTokenizer + from .models.xlm_roberta import ( + XLMRobertaConfig, + ) + from .models.xlm_roberta_xl import ( + XLMRobertaXLConfig, + ) + from .models.xlnet import XLNetConfig + from .models.xmod import XmodConfig + from .models.yolos import YolosConfig + from .models.yoso import YosoConfig + from .models.zamba import ZambaConfig + from .models.zamba2 import Zamba2Config + from .models.zoedepth import ZoeDepthConfig + + # Pipelines + from .pipelines import ( + AudioClassificationPipeline, + AutomaticSpeechRecognitionPipeline, + CsvPipelineDataFormat, + DepthEstimationPipeline, + DocumentQuestionAnsweringPipeline, + FeatureExtractionPipeline, + FillMaskPipeline, + ImageClassificationPipeline, + ImageFeatureExtractionPipeline, + ImageSegmentationPipeline, + ImageTextToTextPipeline, + ImageToImagePipeline, + ImageToTextPipeline, + JsonPipelineDataFormat, + MaskGenerationPipeline, + NerPipeline, + ObjectDetectionPipeline, + PipedPipelineDataFormat, + Pipeline, + PipelineDataFormat, + QuestionAnsweringPipeline, + SummarizationPipeline, + TableQuestionAnsweringPipeline, + Text2TextGenerationPipeline, + TextClassificationPipeline, + TextGenerationPipeline, + TextToAudioPipeline, + TokenClassificationPipeline, + TranslationPipeline, + VideoClassificationPipeline, + VisualQuestionAnsweringPipeline, + ZeroShotAudioClassificationPipeline, + ZeroShotClassificationPipeline, + ZeroShotImageClassificationPipeline, + ZeroShotObjectDetectionPipeline, + pipeline, + ) + from .processing_utils import ProcessorMixin + + # Tokenization + from .tokenization_utils import PreTrainedTokenizer + from .tokenization_utils_base import ( + AddedToken, + BatchEncoding, + CharSpan, + PreTrainedTokenizerBase, + SpecialTokensMixin, + TokenSpan, + ) + + # Trainer + from .trainer_callback import ( + DefaultFlowCallback, + EarlyStoppingCallback, + PrinterCallback, + ProgressCallback, + TrainerCallback, + TrainerControl, + TrainerState, + ) + from .trainer_utils import ( + EvalPrediction, + IntervalStrategy, + SchedulerType, + enable_full_determinism, + set_seed, + ) + from .training_args import TrainingArguments + from .training_args_seq2seq import Seq2SeqTrainingArguments + from .training_args_tf import TFTrainingArguments + + # Files and general utilities + from .utils import ( + CONFIG_NAME, + MODEL_CARD_NAME, + PYTORCH_PRETRAINED_BERT_CACHE, + PYTORCH_TRANSFORMERS_CACHE, + SPIECE_UNDERLINE, + TF2_WEIGHTS_NAME, + TF_WEIGHTS_NAME, + TRANSFORMERS_CACHE, + WEIGHTS_NAME, + TensorType, + add_end_docstrings, + add_start_docstrings, + is_apex_available, + is_av_available, + is_bitsandbytes_available, + is_datasets_available, + is_faiss_available, + is_flax_available, + is_keras_nlp_available, + is_phonemizer_available, + is_psutil_available, + is_py3nvml_available, + is_pyctcdecode_available, + is_sacremoses_available, + is_safetensors_available, + is_scipy_available, + is_sentencepiece_available, + is_sklearn_available, + is_speech_available, + is_tensorflow_text_available, + is_tf_available, + is_timm_available, + is_tokenizers_available, + is_torch_available, + is_torch_hpu_available, + is_torch_mlu_available, + is_torch_musa_available, + is_torch_neuroncore_available, + is_torch_npu_available, + is_torch_xla_available, + is_torch_xpu_available, + is_torchvision_available, + is_vision_available, + logging, + ) + + # bitsandbytes config + from .utils.quantization_config import ( + AqlmConfig, + AwqConfig, + BitNetConfig, + BitsAndBytesConfig, + CompressedTensorsConfig, + EetqConfig, + FbgemmFp8Config, + FineGrainedFP8Config, + GPTQConfig, + HiggsConfig, + HqqConfig, + QuantoConfig, + QuarkConfig, + SpQRConfig, + TorchAoConfig, + VptqConfig, + ) + + try: + if not is_sentencepiece_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_sentencepiece_objects import * + else: + from .models.albert import AlbertTokenizer + from .models.barthez import BarthezTokenizer + from .models.bartpho import BartphoTokenizer + from .models.bert_generation import BertGenerationTokenizer + from .models.big_bird import BigBirdTokenizer + from .models.camembert import CamembertTokenizer + from .models.code_llama import CodeLlamaTokenizer + from .models.cpm import CpmTokenizer + from .models.deberta_v2 import DebertaV2Tokenizer + from .models.deprecated.ernie_m import ErnieMTokenizer + from .models.deprecated.xlm_prophetnet import XLMProphetNetTokenizer + from .models.fnet import FNetTokenizer + from .models.gemma import GemmaTokenizer + from .models.gpt_sw3 import GPTSw3Tokenizer + from .models.layoutxlm import LayoutXLMTokenizer + from .models.llama import LlamaTokenizer + from .models.m2m_100 import M2M100Tokenizer + from .models.marian import MarianTokenizer + from .models.mbart import MBartTokenizer + from .models.mbart50 import MBart50Tokenizer + from .models.mluke import MLukeTokenizer + from .models.mt5 import MT5Tokenizer + from .models.nllb import NllbTokenizer + from .models.pegasus import PegasusTokenizer + from .models.plbart import PLBartTokenizer + from .models.reformer import ReformerTokenizer + from .models.rembert import RemBertTokenizer + from .models.seamless_m4t import SeamlessM4TTokenizer + from .models.siglip import SiglipTokenizer + from .models.speech_to_text import Speech2TextTokenizer + from .models.speecht5 import SpeechT5Tokenizer + from .models.t5 import T5Tokenizer + from .models.udop import UdopTokenizer + from .models.xglm import XGLMTokenizer + from .models.xlm_roberta import XLMRobertaTokenizer + from .models.xlnet import XLNetTokenizer + + try: + if not is_tokenizers_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_tokenizers_objects import * + else: + # Fast tokenizers imports + from .models.albert import AlbertTokenizerFast + from .models.bart import BartTokenizerFast + from .models.barthez import BarthezTokenizerFast + from .models.bert import BertTokenizerFast + from .models.big_bird import BigBirdTokenizerFast + from .models.blenderbot import BlenderbotTokenizerFast + from .models.blenderbot_small import BlenderbotSmallTokenizerFast + from .models.bloom import BloomTokenizerFast + from .models.camembert import CamembertTokenizerFast + from .models.clip import CLIPTokenizerFast + from .models.code_llama import CodeLlamaTokenizerFast + from .models.codegen import CodeGenTokenizerFast + from .models.cohere import CohereTokenizerFast + from .models.convbert import ConvBertTokenizerFast + from .models.cpm import CpmTokenizerFast + from .models.deberta import DebertaTokenizerFast + from .models.deberta_v2 import DebertaV2TokenizerFast + from .models.deprecated.realm import RealmTokenizerFast + from .models.deprecated.retribert import RetriBertTokenizerFast + from .models.distilbert import DistilBertTokenizerFast + from .models.dpr import ( + DPRContextEncoderTokenizerFast, + DPRQuestionEncoderTokenizerFast, + DPRReaderTokenizerFast, + ) + from .models.electra import ElectraTokenizerFast + from .models.fnet import FNetTokenizerFast + from .models.funnel import FunnelTokenizerFast + from .models.gemma import GemmaTokenizerFast + from .models.gpt2 import GPT2TokenizerFast + from .models.gpt_neox import GPTNeoXTokenizerFast + from .models.gpt_neox_japanese import GPTNeoXJapaneseTokenizer + from .models.herbert import HerbertTokenizerFast + from .models.layoutlm import LayoutLMTokenizerFast + from .models.layoutlmv2 import LayoutLMv2TokenizerFast + from .models.layoutlmv3 import LayoutLMv3TokenizerFast + from .models.layoutxlm import LayoutXLMTokenizerFast + from .models.led import LEDTokenizerFast + from .models.llama import LlamaTokenizerFast + from .models.longformer import LongformerTokenizerFast + from .models.lxmert import LxmertTokenizerFast + from .models.markuplm import MarkupLMTokenizerFast + from .models.mbart import MBartTokenizerFast + from .models.mbart50 import MBart50TokenizerFast + from .models.mobilebert import MobileBertTokenizerFast + from .models.mpnet import MPNetTokenizerFast + from .models.mt5 import MT5TokenizerFast + from .models.mvp import MvpTokenizerFast + from .models.nllb import NllbTokenizerFast + from .models.nougat import NougatTokenizerFast + from .models.openai import OpenAIGPTTokenizerFast + from .models.pegasus import PegasusTokenizerFast + from .models.qwen2 import Qwen2TokenizerFast + from .models.reformer import ReformerTokenizerFast + from .models.rembert import RemBertTokenizerFast + from .models.roberta import RobertaTokenizerFast + from .models.roformer import RoFormerTokenizerFast + from .models.seamless_m4t import SeamlessM4TTokenizerFast + from .models.splinter import SplinterTokenizerFast + from .models.squeezebert import SqueezeBertTokenizerFast + from .models.t5 import T5TokenizerFast + from .models.udop import UdopTokenizerFast + from .models.whisper import WhisperTokenizerFast + from .models.xglm import XGLMTokenizerFast + from .models.xlm_roberta import XLMRobertaTokenizerFast + from .models.xlnet import XLNetTokenizerFast + from .tokenization_utils_fast import PreTrainedTokenizerFast + + try: + if not (is_sentencepiece_available() and is_tokenizers_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummies_sentencepiece_and_tokenizers_objects import * + else: + from .convert_slow_tokenizer import ( + SLOW_TO_FAST_CONVERTERS, + convert_slow_tokenizer, + ) + + try: + if not is_tensorflow_text_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_tensorflow_text_objects import * + else: + from .models.bert import TFBertTokenizer + + try: + if not is_keras_nlp_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_keras_nlp_objects import * + else: + from .models.gpt2 import TFGPT2Tokenizer + + try: + if not is_vision_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_vision_objects import * + else: + from .image_processing_base import ImageProcessingMixin + from .image_processing_utils import BaseImageProcessor + from .image_utils import ImageFeatureExtractionMixin + from .models.aria import AriaImageProcessor + from .models.beit import BeitFeatureExtractor, BeitImageProcessor + from .models.bit import BitImageProcessor + from .models.blip import BlipImageProcessor + from .models.bridgetower import BridgeTowerImageProcessor + from .models.chameleon import ChameleonImageProcessor + from .models.chinese_clip import ( + ChineseCLIPFeatureExtractor, + ChineseCLIPImageProcessor, + ) + from .models.clip import CLIPFeatureExtractor, CLIPImageProcessor + from .models.conditional_detr import ( + ConditionalDetrFeatureExtractor, + ConditionalDetrImageProcessor, + ) + from .models.convnext import ConvNextFeatureExtractor, ConvNextImageProcessor + from .models.deformable_detr import DeformableDetrFeatureExtractor, DeformableDetrImageProcessor + from .models.deit import DeiTFeatureExtractor, DeiTImageProcessor + from .models.deprecated.deta import DetaImageProcessor + from .models.deprecated.efficientformer import EfficientFormerImageProcessor + from .models.deprecated.tvlt import TvltImageProcessor + from .models.deprecated.vit_hybrid import ViTHybridImageProcessor + from .models.depth_pro import DepthProImageProcessor, DepthProImageProcessorFast + from .models.detr import DetrFeatureExtractor, DetrImageProcessor + from .models.donut import DonutFeatureExtractor, DonutImageProcessor + from .models.dpt import DPTFeatureExtractor, DPTImageProcessor + from .models.efficientnet import EfficientNetImageProcessor + from .models.emu3 import Emu3ImageProcessor + from .models.flava import ( + FlavaFeatureExtractor, + FlavaImageProcessor, + FlavaProcessor, + ) + from .models.fuyu import FuyuImageProcessor, FuyuProcessor + from .models.gemma3 import Gemma3ImageProcessor + from .models.glpn import GLPNFeatureExtractor, GLPNImageProcessor + from .models.got_ocr2 import GotOcr2ImageProcessor + from .models.grounding_dino import GroundingDinoImageProcessor + from .models.idefics import IdeficsImageProcessor + from .models.idefics2 import Idefics2ImageProcessor + from .models.idefics3 import Idefics3ImageProcessor + from .models.imagegpt import ImageGPTFeatureExtractor, ImageGPTImageProcessor + from .models.instructblipvideo import InstructBlipVideoImageProcessor + from .models.layoutlmv2 import ( + LayoutLMv2FeatureExtractor, + LayoutLMv2ImageProcessor, + ) + from .models.layoutlmv3 import ( + LayoutLMv3FeatureExtractor, + LayoutLMv3ImageProcessor, + ) + from .models.levit import LevitFeatureExtractor, LevitImageProcessor + from .models.llava import LlavaImageProcessor + from .models.llava_next import LlavaNextImageProcessor + from .models.llava_next_video import LlavaNextVideoImageProcessor + from .models.llava_onevision import LlavaOnevisionImageProcessor, LlavaOnevisionVideoProcessor + from .models.mask2former import Mask2FormerImageProcessor + from .models.maskformer import ( + MaskFormerFeatureExtractor, + MaskFormerImageProcessor, + ) + from .models.mllama import MllamaImageProcessor + from .models.mobilenet_v1 import ( + MobileNetV1FeatureExtractor, + MobileNetV1ImageProcessor, + ) + from .models.mobilenet_v2 import ( + MobileNetV2FeatureExtractor, + MobileNetV2ImageProcessor, + ) + from .models.mobilevit import MobileViTFeatureExtractor, MobileViTImageProcessor + from .models.nougat import NougatImageProcessor + from .models.oneformer import OneFormerImageProcessor + from .models.owlv2 import Owlv2ImageProcessor + from .models.owlvit import OwlViTFeatureExtractor, OwlViTImageProcessor + from .models.perceiver import PerceiverFeatureExtractor, PerceiverImageProcessor + from .models.pix2struct import Pix2StructImageProcessor + from .models.pixtral import PixtralImageProcessor + from .models.poolformer import ( + PoolFormerFeatureExtractor, + PoolFormerImageProcessor, + ) + from .models.prompt_depth_anything import PromptDepthAnythingImageProcessor + from .models.pvt import PvtImageProcessor + from .models.qwen2_vl import Qwen2VLImageProcessor + from .models.rt_detr import RTDetrImageProcessor + from .models.sam import SamImageProcessor + from .models.segformer import SegformerFeatureExtractor, SegformerImageProcessor + from .models.seggpt import SegGptImageProcessor + from .models.siglip import SiglipImageProcessor + from .models.siglip2 import Siglip2ImageProcessor + from .models.smolvlm import SmolVLMImageProcessor + from .models.superglue import SuperGlueImageProcessor + from .models.superpoint import SuperPointImageProcessor + from .models.swin2sr import Swin2SRImageProcessor + from .models.textnet import TextNetImageProcessor + from .models.tvp import TvpImageProcessor + from .models.video_llava import VideoLlavaImageProcessor + from .models.videomae import VideoMAEFeatureExtractor, VideoMAEImageProcessor + from .models.vilt import ViltFeatureExtractor, ViltImageProcessor, ViltProcessor + from .models.vit import ViTFeatureExtractor, ViTImageProcessor + from .models.vitmatte import VitMatteImageProcessor + from .models.vitpose import VitPoseImageProcessor + from .models.vivit import VivitImageProcessor + from .models.yolos import YolosFeatureExtractor, YolosImageProcessor + from .models.zoedepth import ZoeDepthImageProcessor + + try: + if not is_torchvision_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_torchvision_objects import * + else: + from .image_processing_utils_fast import BaseImageProcessorFast + from .models.blip import BlipImageProcessorFast + from .models.clip import CLIPImageProcessorFast + from .models.convnext import ConvNextImageProcessorFast + from .models.deformable_detr import DeformableDetrImageProcessorFast + from .models.deit import DeiTImageProcessorFast + from .models.depth_pro import DepthProImageProcessorFast + from .models.detr import DetrImageProcessorFast + from .models.gemma3 import Gemma3ImageProcessorFast + from .models.got_ocr2 import GotOcr2ImageProcessorFast + from .models.llama4 import Llama4ImageProcessorFast + from .models.llava import LlavaImageProcessorFast + from .models.llava_next import LlavaNextImageProcessorFast + from .models.llava_onevision import LlavaOnevisionImageProcessorFast + from .models.phi4_multimodal import Phi4MultimodalImageProcessorFast + from .models.pixtral import PixtralImageProcessorFast + from .models.qwen2_vl import Qwen2VLImageProcessorFast + from .models.rt_detr import RTDetrImageProcessorFast + from .models.siglip import SiglipImageProcessorFast + from .models.siglip2 import Siglip2ImageProcessorFast + from .models.vit import ViTImageProcessorFast + + try: + if not (is_torchvision_available() and is_timm_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_timm_and_torchvision_objects import * + else: + from .models.timm_wrapper import TimmWrapperImageProcessor + + # Modeling + try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_pt_objects import * + else: + # Debugging + from .cache_utils import ( + Cache, + CacheConfig, + DynamicCache, + EncoderDecoderCache, + HQQQuantizedCache, + HybridCache, + MambaCache, + OffloadedCache, + OffloadedStaticCache, + QuantizedCache, + QuantizedCacheConfig, + QuantoQuantizedCache, + SinkCache, + SlidingWindowCache, + StaticCache, + ) + from .data.datasets import ( + GlueDataset, + GlueDataTrainingArguments, + LineByLineTextDataset, + LineByLineWithRefDataset, + LineByLineWithSOPTextDataset, + SquadDataset, + SquadDataTrainingArguments, + TextDataset, + TextDatasetForNextSentencePrediction, + ) + from .generation import ( + AlternatingCodebooksLogitsProcessor, + BayesianDetectorConfig, + BayesianDetectorModel, + BeamScorer, + BeamSearchScorer, + ClassifierFreeGuidanceLogitsProcessor, + ConstrainedBeamSearchScorer, + Constraint, + ConstraintListState, + DisjunctiveConstraint, + EncoderNoRepeatNGramLogitsProcessor, + EncoderRepetitionPenaltyLogitsProcessor, + EosTokenCriteria, + EpsilonLogitsWarper, + EtaLogitsWarper, + ExponentialDecayLengthPenalty, + ForcedBOSTokenLogitsProcessor, + ForcedEOSTokenLogitsProcessor, + GenerationMixin, + HammingDiversityLogitsProcessor, + InfNanRemoveLogitsProcessor, + LogitNormalization, + LogitsProcessor, + LogitsProcessorList, + MaxLengthCriteria, + MaxTimeCriteria, + MinLengthLogitsProcessor, + MinNewTokensLengthLogitsProcessor, + MinPLogitsWarper, + NoBadWordsLogitsProcessor, + NoRepeatNGramLogitsProcessor, + PhrasalConstraint, + PrefixConstrainedLogitsProcessor, + RepetitionPenaltyLogitsProcessor, + SequenceBiasLogitsProcessor, + StoppingCriteria, + StoppingCriteriaList, + StopStringCriteria, + SuppressTokensAtBeginLogitsProcessor, + SuppressTokensLogitsProcessor, + SynthIDTextWatermarkDetector, + SynthIDTextWatermarkingConfig, + SynthIDTextWatermarkLogitsProcessor, + TemperatureLogitsWarper, + TopKLogitsWarper, + TopPLogitsWarper, + TypicalLogitsWarper, + UnbatchedClassifierFreeGuidanceLogitsProcessor, + WatermarkDetector, + WatermarkLogitsProcessor, + WhisperTimeStampLogitsProcessor, + ) + from .integrations.executorch import ( + TorchExportableModuleWithStaticCache, + convert_and_export_with_cache, + ) + from .model_debugging_utils import ( + model_addition_debugger, + model_addition_debugger_context, + ) + from .modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update + from .modeling_utils import AttentionInterface, PreTrainedModel + from .models.albert import ( + AlbertForMaskedLM, + AlbertForMultipleChoice, + AlbertForPreTraining, + AlbertForQuestionAnswering, + AlbertForSequenceClassification, + AlbertForTokenClassification, + AlbertModel, + AlbertPreTrainedModel, + load_tf_weights_in_albert, + ) + from .models.align import ( + AlignModel, + AlignPreTrainedModel, + AlignTextModel, + AlignVisionModel, + ) + from .models.altclip import ( + AltCLIPModel, + AltCLIPPreTrainedModel, + AltCLIPTextModel, + AltCLIPVisionModel, + ) + from .models.aria import ( + AriaForConditionalGeneration, + AriaPreTrainedModel, + AriaTextForCausalLM, + AriaTextModel, + AriaTextPreTrainedModel, + ) + from .models.audio_spectrogram_transformer import ( + ASTForAudioClassification, + ASTModel, + ASTPreTrainedModel, + ) + from .models.auto import ( + MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING, + MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING, + MODEL_FOR_AUDIO_XVECTOR_MAPPING, + MODEL_FOR_BACKBONE_MAPPING, + MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING, + MODEL_FOR_CAUSAL_LM_MAPPING, + MODEL_FOR_CTC_MAPPING, + MODEL_FOR_DEPTH_ESTIMATION_MAPPING, + MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING, + MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING, + MODEL_FOR_IMAGE_MAPPING, + MODEL_FOR_IMAGE_SEGMENTATION_MAPPING, + MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING, + MODEL_FOR_IMAGE_TO_IMAGE_MAPPING, + MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING, + MODEL_FOR_KEYPOINT_DETECTION_MAPPING, + MODEL_FOR_MASK_GENERATION_MAPPING, + MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING, + MODEL_FOR_MASKED_LM_MAPPING, + MODEL_FOR_MULTIPLE_CHOICE_MAPPING, + MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING, + MODEL_FOR_OBJECT_DETECTION_MAPPING, + MODEL_FOR_PRETRAINING_MAPPING, + MODEL_FOR_QUESTION_ANSWERING_MAPPING, + MODEL_FOR_RETRIEVAL_MAPPING, + MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING, + MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, + MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING, + MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, + MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING, + MODEL_FOR_TEXT_ENCODING_MAPPING, + MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING, + MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING, + MODEL_FOR_TIME_SERIES_CLASSIFICATION_MAPPING, + MODEL_FOR_TIME_SERIES_REGRESSION_MAPPING, + MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, + MODEL_FOR_UNIVERSAL_SEGMENTATION_MAPPING, + MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, + MODEL_FOR_VISION_2_SEQ_MAPPING, + MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING, + MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING, + MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING, + MODEL_MAPPING, + MODEL_WITH_LM_HEAD_MAPPING, + AutoBackbone, + AutoModel, + AutoModelForAudioClassification, + AutoModelForAudioFrameClassification, + AutoModelForAudioXVector, + AutoModelForCausalLM, + AutoModelForCTC, + AutoModelForDepthEstimation, + AutoModelForDocumentQuestionAnswering, + AutoModelForImageClassification, + AutoModelForImageSegmentation, + AutoModelForImageTextToText, + AutoModelForImageToImage, + AutoModelForInstanceSegmentation, + AutoModelForKeypointDetection, + AutoModelForMaskedImageModeling, + AutoModelForMaskedLM, + AutoModelForMaskGeneration, + AutoModelForMultipleChoice, + AutoModelForNextSentencePrediction, + AutoModelForObjectDetection, + AutoModelForPreTraining, + AutoModelForQuestionAnswering, + AutoModelForSemanticSegmentation, + AutoModelForSeq2SeqLM, + AutoModelForSequenceClassification, + AutoModelForSpeechSeq2Seq, + AutoModelForTableQuestionAnswering, + AutoModelForTextEncoding, + AutoModelForTextToSpectrogram, + AutoModelForTextToWaveform, + AutoModelForTokenClassification, + AutoModelForUniversalSegmentation, + AutoModelForVideoClassification, + AutoModelForVision2Seq, + AutoModelForVisualQuestionAnswering, + AutoModelForZeroShotImageClassification, + AutoModelForZeroShotObjectDetection, + AutoModelWithLMHead, + ) + from .models.autoformer import ( + AutoformerForPrediction, + AutoformerModel, + AutoformerPreTrainedModel, + ) + from .models.aya_vision import AyaVisionForConditionalGeneration, AyaVisionPreTrainedModel + from .models.bamba import BambaForCausalLM, BambaModel, BambaPreTrainedModel + from .models.bark import ( + BarkCausalModel, + BarkCoarseModel, + BarkFineModel, + BarkModel, + BarkPreTrainedModel, + BarkSemanticModel, + ) + from .models.bart import ( + BartForCausalLM, + BartForConditionalGeneration, + BartForQuestionAnswering, + BartForSequenceClassification, + BartModel, + BartPreTrainedModel, + BartPretrainedModel, + PretrainedBartModel, + ) + from .models.beit import ( + BeitBackbone, + BeitForImageClassification, + BeitForMaskedImageModeling, + BeitForSemanticSegmentation, + BeitModel, + BeitPreTrainedModel, + ) + from .models.bert import ( + BertForMaskedLM, + BertForMultipleChoice, + BertForNextSentencePrediction, + BertForPreTraining, + BertForQuestionAnswering, + BertForSequenceClassification, + BertForTokenClassification, + BertLMHeadModel, + BertModel, + BertPreTrainedModel, + load_tf_weights_in_bert, + ) + from .models.bert_generation import ( + BertGenerationDecoder, + BertGenerationEncoder, + BertGenerationPreTrainedModel, + load_tf_weights_in_bert_generation, + ) + from .models.big_bird import ( + BigBirdForCausalLM, + BigBirdForMaskedLM, + BigBirdForMultipleChoice, + BigBirdForPreTraining, + BigBirdForQuestionAnswering, + BigBirdForSequenceClassification, + BigBirdForTokenClassification, + BigBirdModel, + BigBirdPreTrainedModel, + load_tf_weights_in_big_bird, + ) + from .models.bigbird_pegasus import ( + BigBirdPegasusForCausalLM, + BigBirdPegasusForConditionalGeneration, + BigBirdPegasusForQuestionAnswering, + BigBirdPegasusForSequenceClassification, + BigBirdPegasusModel, + BigBirdPegasusPreTrainedModel, + ) + from .models.biogpt import ( + BioGptForCausalLM, + BioGptForSequenceClassification, + BioGptForTokenClassification, + BioGptModel, + BioGptPreTrainedModel, + ) + from .models.bit import ( + BitBackbone, + BitForImageClassification, + BitModel, + BitPreTrainedModel, + ) + from .models.blenderbot import ( + BlenderbotForCausalLM, + BlenderbotForConditionalGeneration, + BlenderbotModel, + BlenderbotPreTrainedModel, + ) + from .models.blenderbot_small import ( + BlenderbotSmallForCausalLM, + BlenderbotSmallForConditionalGeneration, + BlenderbotSmallModel, + BlenderbotSmallPreTrainedModel, + ) + from .models.blip import ( + BlipForConditionalGeneration, + BlipForImageTextRetrieval, + BlipForQuestionAnswering, + BlipModel, + BlipPreTrainedModel, + BlipTextModel, + BlipVisionModel, + ) + from .models.blip_2 import ( + Blip2ForConditionalGeneration, + Blip2ForImageTextRetrieval, + Blip2Model, + Blip2PreTrainedModel, + Blip2QFormerModel, + Blip2TextModelWithProjection, + Blip2VisionModel, + Blip2VisionModelWithProjection, + ) + from .models.bloom import ( + BloomForCausalLM, + BloomForQuestionAnswering, + BloomForSequenceClassification, + BloomForTokenClassification, + BloomModel, + BloomPreTrainedModel, + ) + from .models.bridgetower import ( + BridgeTowerForContrastiveLearning, + BridgeTowerForImageAndTextRetrieval, + BridgeTowerForMaskedLM, + BridgeTowerModel, + BridgeTowerPreTrainedModel, + ) + from .models.bros import ( + BrosForTokenClassification, + BrosModel, + BrosPreTrainedModel, + BrosProcessor, + BrosSpadeEEForTokenClassification, + BrosSpadeELForTokenClassification, + ) + from .models.camembert import ( + CamembertForCausalLM, + CamembertForMaskedLM, + CamembertForMultipleChoice, + CamembertForQuestionAnswering, + CamembertForSequenceClassification, + CamembertForTokenClassification, + CamembertModel, + CamembertPreTrainedModel, + ) + from .models.canine import ( + CanineForMultipleChoice, + CanineForQuestionAnswering, + CanineForSequenceClassification, + CanineForTokenClassification, + CanineModel, + CaninePreTrainedModel, + load_tf_weights_in_canine, + ) + from .models.chameleon import ( + ChameleonForConditionalGeneration, + ChameleonModel, + ChameleonPreTrainedModel, + ChameleonProcessor, + ChameleonVQVAE, + ) + from .models.chinese_clip import ( + ChineseCLIPModel, + ChineseCLIPPreTrainedModel, + ChineseCLIPTextModel, + ChineseCLIPVisionModel, + ) + from .models.clap import ( + ClapAudioModel, + ClapAudioModelWithProjection, + ClapFeatureExtractor, + ClapModel, + ClapPreTrainedModel, + ClapTextModel, + ClapTextModelWithProjection, + ) + from .models.clip import ( + CLIPForImageClassification, + CLIPModel, + CLIPPreTrainedModel, + CLIPTextModel, + CLIPTextModelWithProjection, + CLIPVisionModel, + CLIPVisionModelWithProjection, + ) + from .models.clipseg import ( + CLIPSegForImageSegmentation, + CLIPSegModel, + CLIPSegPreTrainedModel, + CLIPSegTextModel, + CLIPSegVisionModel, + ) + from .models.clvp import ( + ClvpDecoder, + ClvpEncoder, + ClvpForCausalLM, + ClvpModel, + ClvpModelForConditionalGeneration, + ClvpPreTrainedModel, + ) + from .models.codegen import ( + CodeGenForCausalLM, + CodeGenModel, + CodeGenPreTrainedModel, + ) + from .models.cohere import ( + CohereForCausalLM, + CohereModel, + CoherePreTrainedModel, + ) + from .models.cohere2 import ( + Cohere2ForCausalLM, + Cohere2Model, + Cohere2PreTrainedModel, + ) + from .models.colpali import ( + ColPaliForRetrieval, + ColPaliPreTrainedModel, + ) + from .models.conditional_detr import ( + ConditionalDetrForObjectDetection, + ConditionalDetrForSegmentation, + ConditionalDetrModel, + ConditionalDetrPreTrainedModel, + ) + from .models.convbert import ( + ConvBertForMaskedLM, + ConvBertForMultipleChoice, + ConvBertForQuestionAnswering, + ConvBertForSequenceClassification, + ConvBertForTokenClassification, + ConvBertModel, + ConvBertPreTrainedModel, + load_tf_weights_in_convbert, + ) + from .models.convnext import ( + ConvNextBackbone, + ConvNextForImageClassification, + ConvNextModel, + ConvNextPreTrainedModel, + ) + from .models.convnextv2 import ( + ConvNextV2Backbone, + ConvNextV2ForImageClassification, + ConvNextV2Model, + ConvNextV2PreTrainedModel, + ) + from .models.cpmant import ( + CpmAntForCausalLM, + CpmAntModel, + CpmAntPreTrainedModel, + ) + from .models.ctrl import ( + CTRLForSequenceClassification, + CTRLLMHeadModel, + CTRLModel, + CTRLPreTrainedModel, + ) + from .models.cvt import ( + CvtForImageClassification, + CvtModel, + CvtPreTrainedModel, + ) + from .models.dab_detr import ( + DabDetrForObjectDetection, + DabDetrModel, + DabDetrPreTrainedModel, + ) + from .models.dac import ( + DacModel, + DacPreTrainedModel, + ) + from .models.data2vec import ( + Data2VecAudioForAudioFrameClassification, + Data2VecAudioForCTC, + Data2VecAudioForSequenceClassification, + Data2VecAudioForXVector, + Data2VecAudioModel, + Data2VecAudioPreTrainedModel, + Data2VecTextForCausalLM, + Data2VecTextForMaskedLM, + Data2VecTextForMultipleChoice, + Data2VecTextForQuestionAnswering, + Data2VecTextForSequenceClassification, + Data2VecTextForTokenClassification, + Data2VecTextModel, + Data2VecTextPreTrainedModel, + Data2VecVisionForImageClassification, + Data2VecVisionForSemanticSegmentation, + Data2VecVisionModel, + Data2VecVisionPreTrainedModel, + ) + + # PyTorch model imports + from .models.dbrx import ( + DbrxForCausalLM, + DbrxModel, + DbrxPreTrainedModel, + ) + from .models.deberta import ( + DebertaForMaskedLM, + DebertaForQuestionAnswering, + DebertaForSequenceClassification, + DebertaForTokenClassification, + DebertaModel, + DebertaPreTrainedModel, + ) + from .models.deberta_v2 import ( + DebertaV2ForMaskedLM, + DebertaV2ForMultipleChoice, + DebertaV2ForQuestionAnswering, + DebertaV2ForSequenceClassification, + DebertaV2ForTokenClassification, + DebertaV2Model, + DebertaV2PreTrainedModel, + ) + from .models.decision_transformer import ( + DecisionTransformerGPT2Model, + DecisionTransformerGPT2PreTrainedModel, + DecisionTransformerModel, + DecisionTransformerPreTrainedModel, + ) + from .models.deepseek_v3 import ( + DeepseekV3ForCausalLM, + DeepseekV3Model, + DeepseekV3PreTrainedModel, + ) + from .models.deformable_detr import ( + DeformableDetrForObjectDetection, + DeformableDetrModel, + DeformableDetrPreTrainedModel, + ) + from .models.deit import ( + DeiTForImageClassification, + DeiTForImageClassificationWithTeacher, + DeiTForMaskedImageModeling, + DeiTModel, + DeiTPreTrainedModel, + ) + from .models.deprecated.deta import ( + DetaForObjectDetection, + DetaModel, + DetaPreTrainedModel, + ) + from .models.deprecated.efficientformer import ( + EfficientFormerForImageClassification, + EfficientFormerForImageClassificationWithTeacher, + EfficientFormerModel, + EfficientFormerPreTrainedModel, + ) + from .models.deprecated.ernie_m import ( + ErnieMForInformationExtraction, + ErnieMForMultipleChoice, + ErnieMForQuestionAnswering, + ErnieMForSequenceClassification, + ErnieMForTokenClassification, + ErnieMModel, + ErnieMPreTrainedModel, + ) + from .models.deprecated.gptsan_japanese import ( + GPTSanJapaneseForConditionalGeneration, + GPTSanJapaneseModel, + GPTSanJapanesePreTrainedModel, + ) + from .models.deprecated.graphormer import ( + GraphormerForGraphClassification, + GraphormerModel, + GraphormerPreTrainedModel, + ) + from .models.deprecated.jukebox import ( + JukeboxModel, + JukeboxPreTrainedModel, + JukeboxPrior, + JukeboxVQVAE, + ) + from .models.deprecated.mctct import ( + MCTCTForCTC, + MCTCTModel, + MCTCTPreTrainedModel, + ) + from .models.deprecated.mega import ( + MegaForCausalLM, + MegaForMaskedLM, + MegaForMultipleChoice, + MegaForQuestionAnswering, + MegaForSequenceClassification, + MegaForTokenClassification, + MegaModel, + MegaPreTrainedModel, + ) + from .models.deprecated.mmbt import ( + MMBTForClassification, + MMBTModel, + ModalEmbeddings, + ) + from .models.deprecated.nat import ( + NatBackbone, + NatForImageClassification, + NatModel, + NatPreTrainedModel, + ) + from .models.deprecated.nezha import ( + NezhaForMaskedLM, + NezhaForMultipleChoice, + NezhaForNextSentencePrediction, + NezhaForPreTraining, + NezhaForQuestionAnswering, + NezhaForSequenceClassification, + NezhaForTokenClassification, + NezhaModel, + NezhaPreTrainedModel, + ) + from .models.deprecated.open_llama import ( + OpenLlamaForCausalLM, + OpenLlamaForSequenceClassification, + OpenLlamaModel, + OpenLlamaPreTrainedModel, + ) + from .models.deprecated.qdqbert import ( + QDQBertForMaskedLM, + QDQBertForMultipleChoice, + QDQBertForNextSentencePrediction, + QDQBertForQuestionAnswering, + QDQBertForSequenceClassification, + QDQBertForTokenClassification, + QDQBertLMHeadModel, + QDQBertModel, + QDQBertPreTrainedModel, + load_tf_weights_in_qdqbert, + ) + from .models.deprecated.realm import ( + RealmEmbedder, + RealmForOpenQA, + RealmKnowledgeAugEncoder, + RealmPreTrainedModel, + RealmReader, + RealmRetriever, + RealmScorer, + load_tf_weights_in_realm, + ) + from .models.deprecated.retribert import ( + RetriBertModel, + RetriBertPreTrainedModel, + ) + from .models.deprecated.speech_to_text_2 import ( + Speech2Text2ForCausalLM, + Speech2Text2PreTrainedModel, + ) + from .models.deprecated.trajectory_transformer import ( + TrajectoryTransformerModel, + TrajectoryTransformerPreTrainedModel, + ) + from .models.deprecated.transfo_xl import ( + AdaptiveEmbedding, + TransfoXLForSequenceClassification, + TransfoXLLMHeadModel, + TransfoXLModel, + TransfoXLPreTrainedModel, + load_tf_weights_in_transfo_xl, + ) + from .models.deprecated.tvlt import ( + TvltForAudioVisualClassification, + TvltForPreTraining, + TvltModel, + TvltPreTrainedModel, + ) + from .models.deprecated.van import ( + VanForImageClassification, + VanModel, + VanPreTrainedModel, + ) + from .models.deprecated.vit_hybrid import ( + ViTHybridForImageClassification, + ViTHybridModel, + ViTHybridPreTrainedModel, + ) + from .models.deprecated.xlm_prophetnet import ( + XLMProphetNetDecoder, + XLMProphetNetEncoder, + XLMProphetNetForCausalLM, + XLMProphetNetForConditionalGeneration, + XLMProphetNetModel, + XLMProphetNetPreTrainedModel, + ) + from .models.depth_anything import ( + DepthAnythingForDepthEstimation, + DepthAnythingPreTrainedModel, + ) + from .models.depth_pro import ( + DepthProForDepthEstimation, + DepthProModel, + DepthProPreTrainedModel, + ) + from .models.detr import ( + DetrForObjectDetection, + DetrForSegmentation, + DetrModel, + DetrPreTrainedModel, + ) + from .models.diffllama import ( + DiffLlamaForCausalLM, + DiffLlamaForQuestionAnswering, + DiffLlamaForSequenceClassification, + DiffLlamaForTokenClassification, + DiffLlamaModel, + DiffLlamaPreTrainedModel, + ) + from .models.dinat import ( + DinatBackbone, + DinatForImageClassification, + DinatModel, + DinatPreTrainedModel, + ) + from .models.dinov2 import ( + Dinov2Backbone, + Dinov2ForImageClassification, + Dinov2Model, + Dinov2PreTrainedModel, + ) + from .models.dinov2_with_registers import ( + Dinov2WithRegistersBackbone, + Dinov2WithRegistersForImageClassification, + Dinov2WithRegistersModel, + Dinov2WithRegistersPreTrainedModel, + ) + from .models.distilbert import ( + DistilBertForMaskedLM, + DistilBertForMultipleChoice, + DistilBertForQuestionAnswering, + DistilBertForSequenceClassification, + DistilBertForTokenClassification, + DistilBertModel, + DistilBertPreTrainedModel, + ) + from .models.donut import ( + DonutSwinModel, + DonutSwinPreTrainedModel, + ) + from .models.dpr import ( + DPRContextEncoder, + DPRPretrainedContextEncoder, + DPRPreTrainedModel, + DPRPretrainedQuestionEncoder, + DPRPretrainedReader, + DPRQuestionEncoder, + DPRReader, + ) + from .models.dpt import ( + DPTForDepthEstimation, + DPTForSemanticSegmentation, + DPTModel, + DPTPreTrainedModel, + ) + from .models.efficientnet import ( + EfficientNetForImageClassification, + EfficientNetModel, + EfficientNetPreTrainedModel, + ) + from .models.electra import ( + ElectraForCausalLM, + ElectraForMaskedLM, + ElectraForMultipleChoice, + ElectraForPreTraining, + ElectraForQuestionAnswering, + ElectraForSequenceClassification, + ElectraForTokenClassification, + ElectraModel, + ElectraPreTrainedModel, + load_tf_weights_in_electra, + ) + from .models.emu3 import ( + Emu3ForCausalLM, + Emu3ForConditionalGeneration, + Emu3PreTrainedModel, + Emu3TextModel, + Emu3VQVAE, + ) + from .models.encodec import ( + EncodecModel, + EncodecPreTrainedModel, + ) + from .models.encoder_decoder import EncoderDecoderModel + from .models.ernie import ( + ErnieForCausalLM, + ErnieForMaskedLM, + ErnieForMultipleChoice, + ErnieForNextSentencePrediction, + ErnieForPreTraining, + ErnieForQuestionAnswering, + ErnieForSequenceClassification, + ErnieForTokenClassification, + ErnieModel, + ErniePreTrainedModel, + ) + from .models.esm import ( + EsmFoldPreTrainedModel, + EsmForMaskedLM, + EsmForProteinFolding, + EsmForSequenceClassification, + EsmForTokenClassification, + EsmModel, + EsmPreTrainedModel, + ) + from .models.falcon import ( + FalconForCausalLM, + FalconForQuestionAnswering, + FalconForSequenceClassification, + FalconForTokenClassification, + FalconModel, + FalconPreTrainedModel, + ) + from .models.falcon_mamba import ( + FalconMambaForCausalLM, + FalconMambaModel, + FalconMambaPreTrainedModel, + ) + from .models.fastspeech2_conformer import ( + FastSpeech2ConformerHifiGan, + FastSpeech2ConformerModel, + FastSpeech2ConformerPreTrainedModel, + FastSpeech2ConformerWithHifiGan, + ) + from .models.flaubert import ( + FlaubertForMultipleChoice, + FlaubertForQuestionAnswering, + FlaubertForQuestionAnsweringSimple, + FlaubertForSequenceClassification, + FlaubertForTokenClassification, + FlaubertModel, + FlaubertPreTrainedModel, + FlaubertWithLMHeadModel, + ) + from .models.flava import ( + FlavaForPreTraining, + FlavaImageCodebook, + FlavaImageModel, + FlavaModel, + FlavaMultimodalModel, + FlavaPreTrainedModel, + FlavaTextModel, + ) + from .models.fnet import ( + FNetForMaskedLM, + FNetForMultipleChoice, + FNetForNextSentencePrediction, + FNetForPreTraining, + FNetForQuestionAnswering, + FNetForSequenceClassification, + FNetForTokenClassification, + FNetModel, + FNetPreTrainedModel, + ) + from .models.focalnet import ( + FocalNetBackbone, + FocalNetForImageClassification, + FocalNetForMaskedImageModeling, + FocalNetModel, + FocalNetPreTrainedModel, + ) + from .models.fsmt import ( + FSMTForConditionalGeneration, + FSMTModel, + PretrainedFSMTModel, + ) + from .models.funnel import ( + FunnelBaseModel, + FunnelForMaskedLM, + FunnelForMultipleChoice, + FunnelForPreTraining, + FunnelForQuestionAnswering, + FunnelForSequenceClassification, + FunnelForTokenClassification, + FunnelModel, + FunnelPreTrainedModel, + load_tf_weights_in_funnel, + ) + from .models.fuyu import ( + FuyuForCausalLM, + FuyuPreTrainedModel, + ) + from .models.gemma import ( + GemmaForCausalLM, + GemmaForSequenceClassification, + GemmaForTokenClassification, + GemmaModel, + GemmaPreTrainedModel, + ) + from .models.gemma2 import ( + Gemma2ForCausalLM, + Gemma2ForSequenceClassification, + Gemma2ForTokenClassification, + Gemma2Model, + Gemma2PreTrainedModel, + ) + from .models.gemma3 import ( + Gemma3ForCausalLM, + Gemma3ForConditionalGeneration, + Gemma3PreTrainedModel, + Gemma3TextModel, + ) + from .models.git import ( + GitForCausalLM, + GitModel, + GitPreTrainedModel, + GitVisionModel, + ) + from .models.glm import ( + GlmForCausalLM, + GlmForSequenceClassification, + GlmForTokenClassification, + GlmModel, + GlmPreTrainedModel, + ) + from .models.glpn import ( + GLPNForDepthEstimation, + GLPNModel, + GLPNPreTrainedModel, + ) + from .models.got_ocr2 import ( + GotOcr2ForConditionalGeneration, + GotOcr2PreTrainedModel, + ) + from .models.gpt2 import ( + GPT2DoubleHeadsModel, + GPT2ForQuestionAnswering, + GPT2ForSequenceClassification, + GPT2ForTokenClassification, + GPT2LMHeadModel, + GPT2Model, + GPT2PreTrainedModel, + load_tf_weights_in_gpt2, + ) + from .models.gpt_bigcode import ( + GPTBigCodeForCausalLM, + GPTBigCodeForSequenceClassification, + GPTBigCodeForTokenClassification, + GPTBigCodeModel, + GPTBigCodePreTrainedModel, + ) + from .models.gpt_neo import ( + GPTNeoForCausalLM, + GPTNeoForQuestionAnswering, + GPTNeoForSequenceClassification, + GPTNeoForTokenClassification, + GPTNeoModel, + GPTNeoPreTrainedModel, + load_tf_weights_in_gpt_neo, + ) + from .models.gpt_neox import ( + GPTNeoXForCausalLM, + GPTNeoXForQuestionAnswering, + GPTNeoXForSequenceClassification, + GPTNeoXForTokenClassification, + GPTNeoXModel, + GPTNeoXPreTrainedModel, + ) + from .models.gpt_neox_japanese import ( + GPTNeoXJapaneseForCausalLM, + GPTNeoXJapaneseModel, + GPTNeoXJapanesePreTrainedModel, + ) + from .models.gptj import ( + GPTJForCausalLM, + GPTJForQuestionAnswering, + GPTJForSequenceClassification, + GPTJModel, + GPTJPreTrainedModel, + ) + from .models.granite import ( + GraniteForCausalLM, + GraniteModel, + GranitePreTrainedModel, + ) + from .models.granitemoe import ( + GraniteMoeForCausalLM, + GraniteMoeModel, + GraniteMoePreTrainedModel, + ) + from .models.granitemoeshared import ( + GraniteMoeSharedForCausalLM, + GraniteMoeSharedModel, + GraniteMoeSharedPreTrainedModel, + ) + from .models.grounding_dino import ( + GroundingDinoForObjectDetection, + GroundingDinoModel, + GroundingDinoPreTrainedModel, + ) + from .models.groupvit import ( + GroupViTModel, + GroupViTPreTrainedModel, + GroupViTTextModel, + GroupViTVisionModel, + ) + from .models.helium import ( + HeliumForCausalLM, + HeliumForSequenceClassification, + HeliumForTokenClassification, + HeliumModel, + HeliumPreTrainedModel, + ) + from .models.hiera import ( + HieraBackbone, + HieraForImageClassification, + HieraForPreTraining, + HieraModel, + HieraPreTrainedModel, + ) + from .models.hubert import ( + HubertForCTC, + HubertForSequenceClassification, + HubertModel, + HubertPreTrainedModel, + ) + from .models.ibert import ( + IBertForMaskedLM, + IBertForMultipleChoice, + IBertForQuestionAnswering, + IBertForSequenceClassification, + IBertForTokenClassification, + IBertModel, + IBertPreTrainedModel, + ) + from .models.idefics import ( + IdeficsForVisionText2Text, + IdeficsModel, + IdeficsPreTrainedModel, + IdeficsProcessor, + ) + from .models.idefics2 import ( + Idefics2ForConditionalGeneration, + Idefics2Model, + Idefics2PreTrainedModel, + Idefics2Processor, + ) + from .models.idefics3 import ( + Idefics3ForConditionalGeneration, + Idefics3Model, + Idefics3PreTrainedModel, + Idefics3Processor, + Idefics3VisionConfig, + Idefics3VisionTransformer, + ) + from .models.ijepa import ( + IJepaForImageClassification, + IJepaModel, + IJepaPreTrainedModel, + ) + from .models.imagegpt import ( + ImageGPTForCausalImageModeling, + ImageGPTForImageClassification, + ImageGPTModel, + ImageGPTPreTrainedModel, + load_tf_weights_in_imagegpt, + ) + from .models.informer import ( + InformerForPrediction, + InformerModel, + InformerPreTrainedModel, + ) + from .models.instructblip import ( + InstructBlipForConditionalGeneration, + InstructBlipPreTrainedModel, + InstructBlipQFormerModel, + InstructBlipVisionModel, + ) + from .models.instructblipvideo import ( + InstructBlipVideoForConditionalGeneration, + InstructBlipVideoPreTrainedModel, + InstructBlipVideoQFormerModel, + InstructBlipVideoVisionModel, + ) + from .models.jamba import ( + JambaForCausalLM, + JambaForSequenceClassification, + JambaModel, + JambaPreTrainedModel, + ) + from .models.jetmoe import ( + JetMoeForCausalLM, + JetMoeForSequenceClassification, + JetMoeModel, + JetMoePreTrainedModel, + ) + from .models.kosmos2 import ( + Kosmos2ForConditionalGeneration, + Kosmos2Model, + Kosmos2PreTrainedModel, + ) + from .models.layoutlm import ( + LayoutLMForMaskedLM, + LayoutLMForQuestionAnswering, + LayoutLMForSequenceClassification, + LayoutLMForTokenClassification, + LayoutLMModel, + LayoutLMPreTrainedModel, + ) + from .models.layoutlmv2 import ( + LayoutLMv2ForQuestionAnswering, + LayoutLMv2ForSequenceClassification, + LayoutLMv2ForTokenClassification, + LayoutLMv2Model, + LayoutLMv2PreTrainedModel, + ) + from .models.layoutlmv3 import ( + LayoutLMv3ForQuestionAnswering, + LayoutLMv3ForSequenceClassification, + LayoutLMv3ForTokenClassification, + LayoutLMv3Model, + LayoutLMv3PreTrainedModel, + ) + from .models.led import ( + LEDForConditionalGeneration, + LEDForQuestionAnswering, + LEDForSequenceClassification, + LEDModel, + LEDPreTrainedModel, + ) + from .models.levit import ( + LevitForImageClassification, + LevitForImageClassificationWithTeacher, + LevitModel, + LevitPreTrainedModel, + ) + from .models.lilt import ( + LiltForQuestionAnswering, + LiltForSequenceClassification, + LiltForTokenClassification, + LiltModel, + LiltPreTrainedModel, + ) + from .models.llama import ( + LlamaForCausalLM, + LlamaForQuestionAnswering, + LlamaForSequenceClassification, + LlamaForTokenClassification, + LlamaModel, + LlamaPreTrainedModel, + ) + from .models.llama4 import ( + Llama4ForCausalLM, + Llama4ForConditionalGeneration, + Llama4PreTrainedModel, + Llama4TextModel, + Llama4VisionModel, + ) + from .models.llava import ( + LlavaForConditionalGeneration, + LlavaPreTrainedModel, + ) + from .models.llava_next import ( + LlavaNextForConditionalGeneration, + LlavaNextPreTrainedModel, + ) + from .models.llava_next_video import ( + LlavaNextVideoForConditionalGeneration, + LlavaNextVideoPreTrainedModel, + ) + from .models.llava_onevision import ( + LlavaOnevisionForConditionalGeneration, + LlavaOnevisionPreTrainedModel, + ) + from .models.longformer import ( + LongformerForMaskedLM, + LongformerForMultipleChoice, + LongformerForQuestionAnswering, + LongformerForSequenceClassification, + LongformerForTokenClassification, + LongformerModel, + LongformerPreTrainedModel, + ) + from .models.longt5 import ( + LongT5EncoderModel, + LongT5ForConditionalGeneration, + LongT5Model, + LongT5PreTrainedModel, + ) + from .models.luke import ( + LukeForEntityClassification, + LukeForEntityPairClassification, + LukeForEntitySpanClassification, + LukeForMaskedLM, + LukeForMultipleChoice, + LukeForQuestionAnswering, + LukeForSequenceClassification, + LukeForTokenClassification, + LukeModel, + LukePreTrainedModel, + ) + from .models.lxmert import ( + LxmertEncoder, + LxmertForPreTraining, + LxmertForQuestionAnswering, + LxmertModel, + LxmertPreTrainedModel, + LxmertVisualFeatureEncoder, + ) + from .models.m2m_100 import ( + M2M100ForConditionalGeneration, + M2M100Model, + M2M100PreTrainedModel, + ) + from .models.mamba import ( + MambaForCausalLM, + MambaModel, + MambaPreTrainedModel, + ) + from .models.mamba2 import ( + Mamba2ForCausalLM, + Mamba2Model, + Mamba2PreTrainedModel, + ) + from .models.marian import MarianForCausalLM, MarianModel, MarianMTModel, MarianPreTrainedModel + from .models.markuplm import ( + MarkupLMForQuestionAnswering, + MarkupLMForSequenceClassification, + MarkupLMForTokenClassification, + MarkupLMModel, + MarkupLMPreTrainedModel, + ) + from .models.mask2former import ( + Mask2FormerForUniversalSegmentation, + Mask2FormerModel, + Mask2FormerPreTrainedModel, + ) + from .models.maskformer import ( + MaskFormerForInstanceSegmentation, + MaskFormerModel, + MaskFormerPreTrainedModel, + MaskFormerSwinBackbone, + ) + from .models.mbart import ( + MBartForCausalLM, + MBartForConditionalGeneration, + MBartForQuestionAnswering, + MBartForSequenceClassification, + MBartModel, + MBartPreTrainedModel, + ) + from .models.megatron_bert import ( + MegatronBertForCausalLM, + MegatronBertForMaskedLM, + MegatronBertForMultipleChoice, + MegatronBertForNextSentencePrediction, + MegatronBertForPreTraining, + MegatronBertForQuestionAnswering, + MegatronBertForSequenceClassification, + MegatronBertForTokenClassification, + MegatronBertModel, + MegatronBertPreTrainedModel, + ) + from .models.mgp_str import ( + MgpstrForSceneTextRecognition, + MgpstrModel, + MgpstrPreTrainedModel, + ) + from .models.mimi import ( + MimiModel, + MimiPreTrainedModel, + ) + from .models.mistral import ( + MistralForCausalLM, + MistralForQuestionAnswering, + MistralForSequenceClassification, + MistralForTokenClassification, + MistralModel, + MistralPreTrainedModel, + ) + from .models.mistral3 import ( + Mistral3ForConditionalGeneration, + Mistral3PreTrainedModel, + ) + from .models.mixtral import ( + MixtralForCausalLM, + MixtralForQuestionAnswering, + MixtralForSequenceClassification, + MixtralForTokenClassification, + MixtralModel, + MixtralPreTrainedModel, + ) + from .models.mllama import ( + MllamaForCausalLM, + MllamaForConditionalGeneration, + MllamaPreTrainedModel, + MllamaProcessor, + MllamaTextModel, + MllamaVisionModel, + ) + from .models.mobilebert import ( + MobileBertForMaskedLM, + MobileBertForMultipleChoice, + MobileBertForNextSentencePrediction, + MobileBertForPreTraining, + MobileBertForQuestionAnswering, + MobileBertForSequenceClassification, + MobileBertForTokenClassification, + MobileBertModel, + MobileBertPreTrainedModel, + load_tf_weights_in_mobilebert, + ) + from .models.mobilenet_v1 import ( + MobileNetV1ForImageClassification, + MobileNetV1Model, + MobileNetV1PreTrainedModel, + load_tf_weights_in_mobilenet_v1, + ) + from .models.mobilenet_v2 import ( + MobileNetV2ForImageClassification, + MobileNetV2ForSemanticSegmentation, + MobileNetV2Model, + MobileNetV2PreTrainedModel, + load_tf_weights_in_mobilenet_v2, + ) + from .models.mobilevit import ( + MobileViTForImageClassification, + MobileViTForSemanticSegmentation, + MobileViTModel, + MobileViTPreTrainedModel, + ) + from .models.mobilevitv2 import ( + MobileViTV2ForImageClassification, + MobileViTV2ForSemanticSegmentation, + MobileViTV2Model, + MobileViTV2PreTrainedModel, + ) + from .models.modernbert import ( + ModernBertForMaskedLM, + ModernBertForQuestionAnswering, + ModernBertForSequenceClassification, + ModernBertForTokenClassification, + ModernBertModel, + ModernBertPreTrainedModel, + ) + from .models.moonshine import ( + MoonshineForConditionalGeneration, + MoonshineModel, + MoonshinePreTrainedModel, + ) + from .models.moshi import ( + MoshiForCausalLM, + MoshiForConditionalGeneration, + MoshiModel, + MoshiPreTrainedModel, + ) + from .models.mpnet import ( + MPNetForMaskedLM, + MPNetForMultipleChoice, + MPNetForQuestionAnswering, + MPNetForSequenceClassification, + MPNetForTokenClassification, + MPNetModel, + MPNetPreTrainedModel, + ) + from .models.mpt import ( + MptForCausalLM, + MptForQuestionAnswering, + MptForSequenceClassification, + MptForTokenClassification, + MptModel, + MptPreTrainedModel, + ) + from .models.mra import ( + MraForMaskedLM, + MraForMultipleChoice, + MraForQuestionAnswering, + MraForSequenceClassification, + MraForTokenClassification, + MraModel, + MraPreTrainedModel, + ) + from .models.mt5 import ( + MT5EncoderModel, + MT5ForConditionalGeneration, + MT5ForQuestionAnswering, + MT5ForSequenceClassification, + MT5ForTokenClassification, + MT5Model, + MT5PreTrainedModel, + ) + from .models.musicgen import ( + MusicgenForCausalLM, + MusicgenForConditionalGeneration, + MusicgenModel, + MusicgenPreTrainedModel, + MusicgenProcessor, + ) + from .models.musicgen_melody import ( + MusicgenMelodyForCausalLM, + MusicgenMelodyForConditionalGeneration, + MusicgenMelodyModel, + MusicgenMelodyPreTrainedModel, + ) + from .models.mvp import ( + MvpForCausalLM, + MvpForConditionalGeneration, + MvpForQuestionAnswering, + MvpForSequenceClassification, + MvpModel, + MvpPreTrainedModel, + ) + from .models.nemotron import ( + NemotronForCausalLM, + NemotronForQuestionAnswering, + NemotronForSequenceClassification, + NemotronForTokenClassification, + NemotronModel, + NemotronPreTrainedModel, + ) + from .models.nllb_moe import ( + NllbMoeForConditionalGeneration, + NllbMoeModel, + NllbMoePreTrainedModel, + NllbMoeSparseMLP, + NllbMoeTop2Router, + ) + from .models.nystromformer import ( + NystromformerForMaskedLM, + NystromformerForMultipleChoice, + NystromformerForQuestionAnswering, + NystromformerForSequenceClassification, + NystromformerForTokenClassification, + NystromformerModel, + NystromformerPreTrainedModel, + ) + from .models.olmo import ( + OlmoForCausalLM, + OlmoModel, + OlmoPreTrainedModel, + ) + from .models.olmo2 import ( + Olmo2ForCausalLM, + Olmo2Model, + Olmo2PreTrainedModel, + ) + from .models.olmoe import ( + OlmoeForCausalLM, + OlmoeModel, + OlmoePreTrainedModel, + ) + from .models.omdet_turbo import ( + OmDetTurboForObjectDetection, + OmDetTurboPreTrainedModel, + ) + from .models.oneformer import ( + OneFormerForUniversalSegmentation, + OneFormerModel, + OneFormerPreTrainedModel, + ) + from .models.openai import ( + OpenAIGPTDoubleHeadsModel, + OpenAIGPTForSequenceClassification, + OpenAIGPTLMHeadModel, + OpenAIGPTModel, + OpenAIGPTPreTrainedModel, + load_tf_weights_in_openai_gpt, + ) + from .models.opt import ( + OPTForCausalLM, + OPTForQuestionAnswering, + OPTForSequenceClassification, + OPTModel, + OPTPreTrainedModel, + ) + from .models.owlv2 import ( + Owlv2ForObjectDetection, + Owlv2Model, + Owlv2PreTrainedModel, + Owlv2TextModel, + Owlv2VisionModel, + ) + from .models.owlvit import ( + OwlViTForObjectDetection, + OwlViTModel, + OwlViTPreTrainedModel, + OwlViTTextModel, + OwlViTVisionModel, + ) + from .models.paligemma import ( + PaliGemmaForConditionalGeneration, + PaliGemmaPreTrainedModel, + PaliGemmaProcessor, + ) + from .models.patchtsmixer import ( + PatchTSMixerForPrediction, + PatchTSMixerForPretraining, + PatchTSMixerForRegression, + PatchTSMixerForTimeSeriesClassification, + PatchTSMixerModel, + PatchTSMixerPreTrainedModel, + ) + from .models.patchtst import ( + PatchTSTForClassification, + PatchTSTForPrediction, + PatchTSTForPretraining, + PatchTSTForRegression, + PatchTSTModel, + PatchTSTPreTrainedModel, + ) + from .models.pegasus import ( + PegasusForCausalLM, + PegasusForConditionalGeneration, + PegasusModel, + PegasusPreTrainedModel, + ) + from .models.pegasus_x import ( + PegasusXForConditionalGeneration, + PegasusXModel, + PegasusXPreTrainedModel, + ) + from .models.perceiver import ( + PerceiverForImageClassificationConvProcessing, + PerceiverForImageClassificationFourier, + PerceiverForImageClassificationLearned, + PerceiverForMaskedLM, + PerceiverForMultimodalAutoencoding, + PerceiverForOpticalFlow, + PerceiverForSequenceClassification, + PerceiverModel, + PerceiverPreTrainedModel, + ) + from .models.persimmon import ( + PersimmonForCausalLM, + PersimmonForSequenceClassification, + PersimmonForTokenClassification, + PersimmonModel, + PersimmonPreTrainedModel, + ) + from .models.phi import ( + PhiForCausalLM, + PhiForSequenceClassification, + PhiForTokenClassification, + PhiModel, + PhiPreTrainedModel, + ) + from .models.phi3 import ( + Phi3ForCausalLM, + Phi3ForSequenceClassification, + Phi3ForTokenClassification, + Phi3Model, + Phi3PreTrainedModel, + ) + from .models.phi4_multimodal import ( + Phi4MultimodalAudioModel, + Phi4MultimodalAudioPreTrainedModel, + Phi4MultimodalForCausalLM, + Phi4MultimodalModel, + Phi4MultimodalPreTrainedModel, + Phi4MultimodalVisionModel, + Phi4MultimodalVisionPreTrainedModel, + ) + from .models.phimoe import ( + PhimoeForCausalLM, + PhimoeForSequenceClassification, + PhimoeModel, + PhimoePreTrainedModel, + ) + from .models.pix2struct import ( + Pix2StructForConditionalGeneration, + Pix2StructPreTrainedModel, + Pix2StructTextModel, + Pix2StructVisionModel, + ) + from .models.pixtral import ( + PixtralPreTrainedModel, + PixtralVisionModel, + ) + from .models.plbart import ( + PLBartForCausalLM, + PLBartForConditionalGeneration, + PLBartForSequenceClassification, + PLBartModel, + PLBartPreTrainedModel, + ) + from .models.poolformer import ( + PoolFormerForImageClassification, + PoolFormerModel, + PoolFormerPreTrainedModel, + ) + from .models.pop2piano import ( + Pop2PianoForConditionalGeneration, + Pop2PianoPreTrainedModel, + ) + from .models.prompt_depth_anything import ( + PromptDepthAnythingForDepthEstimation, + PromptDepthAnythingPreTrainedModel, + ) + from .models.prophetnet import ( + ProphetNetDecoder, + ProphetNetEncoder, + ProphetNetForCausalLM, + ProphetNetForConditionalGeneration, + ProphetNetModel, + ProphetNetPreTrainedModel, + ) + from .models.pvt import ( + PvtForImageClassification, + PvtModel, + PvtPreTrainedModel, + ) + from .models.pvt_v2 import ( + PvtV2Backbone, + PvtV2ForImageClassification, + PvtV2Model, + PvtV2PreTrainedModel, + ) + from .models.qwen2 import ( + Qwen2ForCausalLM, + Qwen2ForQuestionAnswering, + Qwen2ForSequenceClassification, + Qwen2ForTokenClassification, + Qwen2Model, + Qwen2PreTrainedModel, + ) + from .models.qwen2_5_vl import ( + Qwen2_5_VLForConditionalGeneration, + Qwen2_5_VLModel, + Qwen2_5_VLPreTrainedModel, + ) + from .models.qwen2_audio import ( + Qwen2AudioEncoder, + Qwen2AudioForConditionalGeneration, + Qwen2AudioPreTrainedModel, + ) + from .models.qwen2_moe import ( + Qwen2MoeForCausalLM, + Qwen2MoeForQuestionAnswering, + Qwen2MoeForSequenceClassification, + Qwen2MoeForTokenClassification, + Qwen2MoeModel, + Qwen2MoePreTrainedModel, + ) + from .models.qwen2_vl import ( + Qwen2VLForConditionalGeneration, + Qwen2VLModel, + Qwen2VLPreTrainedModel, + ) + from .models.qwen3 import ( + Qwen3ForCausalLM, + Qwen3ForQuestionAnswering, + Qwen3ForSequenceClassification, + Qwen3ForTokenClassification, + Qwen3Model, + Qwen3PreTrainedModel, + ) + from .models.qwen3_moe import ( + Qwen3MoeForCausalLM, + Qwen3MoeForQuestionAnswering, + Qwen3MoeForSequenceClassification, + Qwen3MoeForTokenClassification, + Qwen3MoeModel, + Qwen3MoePreTrainedModel, + ) + from .models.rag import ( + RagModel, + RagPreTrainedModel, + RagSequenceForGeneration, + RagTokenForGeneration, + ) + from .models.recurrent_gemma import ( + RecurrentGemmaForCausalLM, + RecurrentGemmaModel, + RecurrentGemmaPreTrainedModel, + ) + from .models.reformer import ( + ReformerForMaskedLM, + ReformerForQuestionAnswering, + ReformerForSequenceClassification, + ReformerModel, + ReformerModelWithLMHead, + ReformerPreTrainedModel, + ) + from .models.regnet import ( + RegNetForImageClassification, + RegNetModel, + RegNetPreTrainedModel, + ) + from .models.rembert import ( + RemBertForCausalLM, + RemBertForMaskedLM, + RemBertForMultipleChoice, + RemBertForQuestionAnswering, + RemBertForSequenceClassification, + RemBertForTokenClassification, + RemBertModel, + RemBertPreTrainedModel, + load_tf_weights_in_rembert, + ) + from .models.resnet import ( + ResNetBackbone, + ResNetForImageClassification, + ResNetModel, + ResNetPreTrainedModel, + ) + from .models.roberta import ( + RobertaForCausalLM, + RobertaForMaskedLM, + RobertaForMultipleChoice, + RobertaForQuestionAnswering, + RobertaForSequenceClassification, + RobertaForTokenClassification, + RobertaModel, + RobertaPreTrainedModel, + ) + from .models.roberta_prelayernorm import ( + RobertaPreLayerNormForCausalLM, + RobertaPreLayerNormForMaskedLM, + RobertaPreLayerNormForMultipleChoice, + RobertaPreLayerNormForQuestionAnswering, + RobertaPreLayerNormForSequenceClassification, + RobertaPreLayerNormForTokenClassification, + RobertaPreLayerNormModel, + RobertaPreLayerNormPreTrainedModel, + ) + from .models.roc_bert import ( + RoCBertForCausalLM, + RoCBertForMaskedLM, + RoCBertForMultipleChoice, + RoCBertForPreTraining, + RoCBertForQuestionAnswering, + RoCBertForSequenceClassification, + RoCBertForTokenClassification, + RoCBertModel, + RoCBertPreTrainedModel, + load_tf_weights_in_roc_bert, + ) + from .models.roformer import ( + RoFormerForCausalLM, + RoFormerForMaskedLM, + RoFormerForMultipleChoice, + RoFormerForQuestionAnswering, + RoFormerForSequenceClassification, + RoFormerForTokenClassification, + RoFormerModel, + RoFormerPreTrainedModel, + load_tf_weights_in_roformer, + ) + from .models.rt_detr import ( + RTDetrForObjectDetection, + RTDetrModel, + RTDetrPreTrainedModel, + RTDetrResNetBackbone, + RTDetrResNetPreTrainedModel, + ) + from .models.rt_detr_v2 import RTDetrV2ForObjectDetection, RTDetrV2Model, RTDetrV2PreTrainedModel + from .models.rwkv import ( + RwkvForCausalLM, + RwkvModel, + RwkvPreTrainedModel, + ) + from .models.sam import ( + SamModel, + SamPreTrainedModel, + SamVisionModel, + ) + from .models.seamless_m4t import ( + SeamlessM4TCodeHifiGan, + SeamlessM4TForSpeechToSpeech, + SeamlessM4TForSpeechToText, + SeamlessM4TForTextToSpeech, + SeamlessM4TForTextToText, + SeamlessM4THifiGan, + SeamlessM4TModel, + SeamlessM4TPreTrainedModel, + SeamlessM4TTextToUnitForConditionalGeneration, + SeamlessM4TTextToUnitModel, + ) + from .models.seamless_m4t_v2 import ( + SeamlessM4Tv2ForSpeechToSpeech, + SeamlessM4Tv2ForSpeechToText, + SeamlessM4Tv2ForTextToSpeech, + SeamlessM4Tv2ForTextToText, + SeamlessM4Tv2Model, + SeamlessM4Tv2PreTrainedModel, + ) + from .models.segformer import ( + SegformerDecodeHead, + SegformerForImageClassification, + SegformerForSemanticSegmentation, + SegformerModel, + SegformerPreTrainedModel, + ) + from .models.seggpt import ( + SegGptForImageSegmentation, + SegGptModel, + SegGptPreTrainedModel, + ) + from .models.sew import ( + SEWForCTC, + SEWForSequenceClassification, + SEWModel, + SEWPreTrainedModel, + ) + from .models.sew_d import ( + SEWDForCTC, + SEWDForSequenceClassification, + SEWDModel, + SEWDPreTrainedModel, + ) + from .models.shieldgemma2 import ( + ShieldGemma2ForImageClassification, + ) + from .models.siglip import ( + SiglipForImageClassification, + SiglipModel, + SiglipPreTrainedModel, + SiglipTextModel, + SiglipVisionModel, + ) + from .models.siglip2 import ( + Siglip2ForImageClassification, + Siglip2Model, + Siglip2PreTrainedModel, + Siglip2TextModel, + Siglip2VisionModel, + ) + from .models.smolvlm import ( + SmolVLMForConditionalGeneration, + SmolVLMModel, + SmolVLMPreTrainedModel, + SmolVLMProcessor, + SmolVLMVisionConfig, + SmolVLMVisionTransformer, + ) + from .models.speech_encoder_decoder import SpeechEncoderDecoderModel + from .models.speech_to_text import ( + Speech2TextForConditionalGeneration, + Speech2TextModel, + Speech2TextPreTrainedModel, + ) + from .models.speecht5 import ( + SpeechT5ForSpeechToSpeech, + SpeechT5ForSpeechToText, + SpeechT5ForTextToSpeech, + SpeechT5HifiGan, + SpeechT5Model, + SpeechT5PreTrainedModel, + ) + from .models.splinter import ( + SplinterForPreTraining, + SplinterForQuestionAnswering, + SplinterModel, + SplinterPreTrainedModel, + ) + from .models.squeezebert import ( + SqueezeBertForMaskedLM, + SqueezeBertForMultipleChoice, + SqueezeBertForQuestionAnswering, + SqueezeBertForSequenceClassification, + SqueezeBertForTokenClassification, + SqueezeBertModel, + SqueezeBertPreTrainedModel, + ) + from .models.stablelm import ( + StableLmForCausalLM, + StableLmForSequenceClassification, + StableLmForTokenClassification, + StableLmModel, + StableLmPreTrainedModel, + ) + from .models.starcoder2 import ( + Starcoder2ForCausalLM, + Starcoder2ForSequenceClassification, + Starcoder2ForTokenClassification, + Starcoder2Model, + Starcoder2PreTrainedModel, + ) + from .models.superglue import ( + SuperGlueForKeypointMatching, + SuperGluePreTrainedModel, + ) + from .models.superpoint import ( + SuperPointForKeypointDetection, + SuperPointPreTrainedModel, + ) + from .models.swiftformer import ( + SwiftFormerForImageClassification, + SwiftFormerModel, + SwiftFormerPreTrainedModel, + ) + from .models.swin import ( + SwinBackbone, + SwinForImageClassification, + SwinForMaskedImageModeling, + SwinModel, + SwinPreTrainedModel, + ) + from .models.swin2sr import ( + Swin2SRForImageSuperResolution, + Swin2SRModel, + Swin2SRPreTrainedModel, + ) + from .models.swinv2 import ( + Swinv2Backbone, + Swinv2ForImageClassification, + Swinv2ForMaskedImageModeling, + Swinv2Model, + Swinv2PreTrainedModel, + ) + from .models.switch_transformers import ( + SwitchTransformersEncoderModel, + SwitchTransformersForConditionalGeneration, + SwitchTransformersModel, + SwitchTransformersPreTrainedModel, + SwitchTransformersSparseMLP, + SwitchTransformersTop1Router, + ) + from .models.t5 import ( + T5EncoderModel, + T5ForConditionalGeneration, + T5ForQuestionAnswering, + T5ForSequenceClassification, + T5ForTokenClassification, + T5Model, + T5PreTrainedModel, + load_tf_weights_in_t5, + ) + from .models.table_transformer import ( + TableTransformerForObjectDetection, + TableTransformerModel, + TableTransformerPreTrainedModel, + ) + from .models.tapas import ( + TapasForMaskedLM, + TapasForQuestionAnswering, + TapasForSequenceClassification, + TapasModel, + TapasPreTrainedModel, + load_tf_weights_in_tapas, + ) + from .models.textnet import ( + TextNetBackbone, + TextNetForImageClassification, + TextNetModel, + TextNetPreTrainedModel, + ) + from .models.time_series_transformer import ( + TimeSeriesTransformerForPrediction, + TimeSeriesTransformerModel, + TimeSeriesTransformerPreTrainedModel, + ) + from .models.timesformer import ( + TimesformerForVideoClassification, + TimesformerModel, + TimesformerPreTrainedModel, + ) + from .models.timm_backbone import TimmBackbone + from .models.timm_wrapper import ( + TimmWrapperForImageClassification, + TimmWrapperModel, + TimmWrapperPreTrainedModel, + ) + from .models.trocr import ( + TrOCRForCausalLM, + TrOCRPreTrainedModel, + ) + from .models.tvp import ( + TvpForVideoGrounding, + TvpModel, + TvpPreTrainedModel, + ) + from .models.udop import ( + UdopEncoderModel, + UdopForConditionalGeneration, + UdopModel, + UdopPreTrainedModel, + ) + from .models.umt5 import ( + UMT5EncoderModel, + UMT5ForConditionalGeneration, + UMT5ForQuestionAnswering, + UMT5ForSequenceClassification, + UMT5ForTokenClassification, + UMT5Model, + UMT5PreTrainedModel, + ) + from .models.unispeech import ( + UniSpeechForCTC, + UniSpeechForPreTraining, + UniSpeechForSequenceClassification, + UniSpeechModel, + UniSpeechPreTrainedModel, + ) + from .models.unispeech_sat import ( + UniSpeechSatForAudioFrameClassification, + UniSpeechSatForCTC, + UniSpeechSatForPreTraining, + UniSpeechSatForSequenceClassification, + UniSpeechSatForXVector, + UniSpeechSatModel, + UniSpeechSatPreTrainedModel, + ) + from .models.univnet import UnivNetModel + from .models.upernet import ( + UperNetForSemanticSegmentation, + UperNetPreTrainedModel, + ) + from .models.video_llava import ( + VideoLlavaForConditionalGeneration, + VideoLlavaPreTrainedModel, + VideoLlavaProcessor, + ) + from .models.videomae import ( + VideoMAEForPreTraining, + VideoMAEForVideoClassification, + VideoMAEModel, + VideoMAEPreTrainedModel, + ) + from .models.vilt import ( + ViltForImageAndTextRetrieval, + ViltForImagesAndTextClassification, + ViltForMaskedLM, + ViltForQuestionAnswering, + ViltForTokenClassification, + ViltModel, + ViltPreTrainedModel, + ) + from .models.vipllava import ( + VipLlavaForConditionalGeneration, + VipLlavaPreTrainedModel, + ) + from .models.vision_encoder_decoder import VisionEncoderDecoderModel + from .models.vision_text_dual_encoder import VisionTextDualEncoderModel + from .models.visual_bert import ( + VisualBertForMultipleChoice, + VisualBertForPreTraining, + VisualBertForQuestionAnswering, + VisualBertForRegionToPhraseAlignment, + VisualBertForVisualReasoning, + VisualBertModel, + VisualBertPreTrainedModel, + ) + from .models.vit import ( + ViTForImageClassification, + ViTForMaskedImageModeling, + ViTModel, + ViTPreTrainedModel, + ) + from .models.vit_mae import ( + ViTMAEForPreTraining, + ViTMAEModel, + ViTMAEPreTrainedModel, + ) + from .models.vit_msn import ( + ViTMSNForImageClassification, + ViTMSNModel, + ViTMSNPreTrainedModel, + ) + from .models.vitdet import ( + VitDetBackbone, + VitDetModel, + VitDetPreTrainedModel, + ) + from .models.vitmatte import ( + VitMatteForImageMatting, + VitMattePreTrainedModel, + ) + from .models.vitpose import ( + VitPoseForPoseEstimation, + VitPosePreTrainedModel, + ) + from .models.vitpose_backbone import VitPoseBackbone, VitPoseBackbonePreTrainedModel + from .models.vits import ( + VitsModel, + VitsPreTrainedModel, + ) + from .models.vivit import ( + VivitForVideoClassification, + VivitModel, + VivitPreTrainedModel, + ) + from .models.wav2vec2 import ( + Wav2Vec2ForAudioFrameClassification, + Wav2Vec2ForCTC, + Wav2Vec2ForMaskedLM, + Wav2Vec2ForPreTraining, + Wav2Vec2ForSequenceClassification, + Wav2Vec2ForXVector, + Wav2Vec2Model, + Wav2Vec2PreTrainedModel, + ) + from .models.wav2vec2_bert import ( + Wav2Vec2BertForAudioFrameClassification, + Wav2Vec2BertForCTC, + Wav2Vec2BertForSequenceClassification, + Wav2Vec2BertForXVector, + Wav2Vec2BertModel, + Wav2Vec2BertPreTrainedModel, + ) + from .models.wav2vec2_conformer import ( + Wav2Vec2ConformerForAudioFrameClassification, + Wav2Vec2ConformerForCTC, + Wav2Vec2ConformerForPreTraining, + Wav2Vec2ConformerForSequenceClassification, + Wav2Vec2ConformerForXVector, + Wav2Vec2ConformerModel, + Wav2Vec2ConformerPreTrainedModel, + ) + from .models.wavlm import ( + WavLMForAudioFrameClassification, + WavLMForCTC, + WavLMForSequenceClassification, + WavLMForXVector, + WavLMModel, + WavLMPreTrainedModel, + ) + from .models.whisper import ( + WhisperForAudioClassification, + WhisperForCausalLM, + WhisperForConditionalGeneration, + WhisperModel, + WhisperPreTrainedModel, + ) + from .models.x_clip import ( + XCLIPModel, + XCLIPPreTrainedModel, + XCLIPTextModel, + XCLIPVisionModel, + ) + from .models.xglm import ( + XGLMForCausalLM, + XGLMModel, + XGLMPreTrainedModel, + ) + from .models.xlm import ( + XLMForMultipleChoice, + XLMForQuestionAnswering, + XLMForQuestionAnsweringSimple, + XLMForSequenceClassification, + XLMForTokenClassification, + XLMModel, + XLMPreTrainedModel, + XLMWithLMHeadModel, + ) + from .models.xlm_roberta import ( + XLMRobertaForCausalLM, + XLMRobertaForMaskedLM, + XLMRobertaForMultipleChoice, + XLMRobertaForQuestionAnswering, + XLMRobertaForSequenceClassification, + XLMRobertaForTokenClassification, + XLMRobertaModel, + XLMRobertaPreTrainedModel, + ) + from .models.xlm_roberta_xl import ( + XLMRobertaXLForCausalLM, + XLMRobertaXLForMaskedLM, + XLMRobertaXLForMultipleChoice, + XLMRobertaXLForQuestionAnswering, + XLMRobertaXLForSequenceClassification, + XLMRobertaXLForTokenClassification, + XLMRobertaXLModel, + XLMRobertaXLPreTrainedModel, + ) + from .models.xlnet import ( + XLNetForMultipleChoice, + XLNetForQuestionAnswering, + XLNetForQuestionAnsweringSimple, + XLNetForSequenceClassification, + XLNetForTokenClassification, + XLNetLMHeadModel, + XLNetModel, + XLNetPreTrainedModel, + load_tf_weights_in_xlnet, + ) + from .models.xmod import ( + XmodForCausalLM, + XmodForMaskedLM, + XmodForMultipleChoice, + XmodForQuestionAnswering, + XmodForSequenceClassification, + XmodForTokenClassification, + XmodModel, + XmodPreTrainedModel, + ) + from .models.yolos import ( + YolosForObjectDetection, + YolosModel, + YolosPreTrainedModel, + ) + from .models.yoso import ( + YosoForMaskedLM, + YosoForMultipleChoice, + YosoForQuestionAnswering, + YosoForSequenceClassification, + YosoForTokenClassification, + YosoModel, + YosoPreTrainedModel, + ) + from .models.zamba import ( + ZambaForCausalLM, + ZambaForSequenceClassification, + ZambaModel, + ZambaPreTrainedModel, + ) + from .models.zamba2 import ( + Zamba2ForCausalLM, + Zamba2ForSequenceClassification, + Zamba2Model, + Zamba2PreTrainedModel, + ) + from .models.zoedepth import ( + ZoeDepthForDepthEstimation, + ZoeDepthPreTrainedModel, + ) + + # Optimization + from .optimization import ( + Adafactor, + get_constant_schedule, + get_constant_schedule_with_warmup, + get_cosine_schedule_with_warmup, + get_cosine_with_hard_restarts_schedule_with_warmup, + get_inverse_sqrt_schedule, + get_linear_schedule_with_warmup, + get_polynomial_decay_schedule_with_warmup, + get_scheduler, + get_wsd_schedule, + ) + from .pytorch_utils import Conv1D, apply_chunking_to_forward, prune_layer + + # Trainer + from .trainer import Trainer + from .trainer_pt_utils import torch_distributed_zero_first + from .trainer_seq2seq import Seq2SeqTrainer + + # TensorFlow + try: + if not is_tf_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + # Import the same objects as dummies to get them in the namespace. + # They will raise an import error if the user tries to instantiate / use them. + from .utils.dummy_tf_objects import * + else: + from .generation import ( + TFForcedBOSTokenLogitsProcessor, + TFForcedEOSTokenLogitsProcessor, + TFForceTokensLogitsProcessor, + TFGenerationMixin, + TFLogitsProcessor, + TFLogitsProcessorList, + TFLogitsWarper, + TFMinLengthLogitsProcessor, + TFNoBadWordsLogitsProcessor, + TFNoRepeatNGramLogitsProcessor, + TFRepetitionPenaltyLogitsProcessor, + TFSuppressTokensAtBeginLogitsProcessor, + TFSuppressTokensLogitsProcessor, + TFTemperatureLogitsWarper, + TFTopKLogitsWarper, + TFTopPLogitsWarper, + ) + from .keras_callbacks import KerasMetricCallback, PushToHubCallback + from .modeling_tf_utils import ( + TFPreTrainedModel, + TFSequenceSummary, + TFSharedEmbeddings, + shape_list, + ) + + # TensorFlow model imports + from .models.albert import ( + TFAlbertForMaskedLM, + TFAlbertForMultipleChoice, + TFAlbertForPreTraining, + TFAlbertForQuestionAnswering, + TFAlbertForSequenceClassification, + TFAlbertForTokenClassification, + TFAlbertMainLayer, + TFAlbertModel, + TFAlbertPreTrainedModel, + ) + from .models.auto import ( + TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING, + TF_MODEL_FOR_CAUSAL_LM_MAPPING, + TF_MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING, + TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING, + TF_MODEL_FOR_MASK_GENERATION_MAPPING, + TF_MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING, + TF_MODEL_FOR_MASKED_LM_MAPPING, + TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING, + TF_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING, + TF_MODEL_FOR_PRETRAINING_MAPPING, + TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING, + TF_MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING, + TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, + TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING, + TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, + TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING, + TF_MODEL_FOR_TEXT_ENCODING_MAPPING, + TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, + TF_MODEL_FOR_VISION_2_SEQ_MAPPING, + TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING, + TF_MODEL_MAPPING, + TF_MODEL_WITH_LM_HEAD_MAPPING, + TFAutoModel, + TFAutoModelForAudioClassification, + TFAutoModelForCausalLM, + TFAutoModelForDocumentQuestionAnswering, + TFAutoModelForImageClassification, + TFAutoModelForMaskedImageModeling, + TFAutoModelForMaskedLM, + TFAutoModelForMaskGeneration, + TFAutoModelForMultipleChoice, + TFAutoModelForNextSentencePrediction, + TFAutoModelForPreTraining, + TFAutoModelForQuestionAnswering, + TFAutoModelForSemanticSegmentation, + TFAutoModelForSeq2SeqLM, + TFAutoModelForSequenceClassification, + TFAutoModelForSpeechSeq2Seq, + TFAutoModelForTableQuestionAnswering, + TFAutoModelForTextEncoding, + TFAutoModelForTokenClassification, + TFAutoModelForVision2Seq, + TFAutoModelForZeroShotImageClassification, + TFAutoModelWithLMHead, + ) + from .models.bart import ( + TFBartForConditionalGeneration, + TFBartForSequenceClassification, + TFBartModel, + TFBartPretrainedModel, + ) + from .models.bert import ( + TFBertForMaskedLM, + TFBertForMultipleChoice, + TFBertForNextSentencePrediction, + TFBertForPreTraining, + TFBertForQuestionAnswering, + TFBertForSequenceClassification, + TFBertForTokenClassification, + TFBertLMHeadModel, + TFBertMainLayer, + TFBertModel, + TFBertPreTrainedModel, + ) + from .models.blenderbot import ( + TFBlenderbotForConditionalGeneration, + TFBlenderbotModel, + TFBlenderbotPreTrainedModel, + ) + from .models.blenderbot_small import ( + TFBlenderbotSmallForConditionalGeneration, + TFBlenderbotSmallModel, + TFBlenderbotSmallPreTrainedModel, + ) + from .models.blip import ( + TFBlipForConditionalGeneration, + TFBlipForImageTextRetrieval, + TFBlipForQuestionAnswering, + TFBlipModel, + TFBlipPreTrainedModel, + TFBlipTextModel, + TFBlipVisionModel, + ) + from .models.camembert import ( + TFCamembertForCausalLM, + TFCamembertForMaskedLM, + TFCamembertForMultipleChoice, + TFCamembertForQuestionAnswering, + TFCamembertForSequenceClassification, + TFCamembertForTokenClassification, + TFCamembertModel, + TFCamembertPreTrainedModel, + ) + from .models.clip import ( + TFCLIPModel, + TFCLIPPreTrainedModel, + TFCLIPTextModel, + TFCLIPVisionModel, + ) + from .models.convbert import ( + TFConvBertForMaskedLM, + TFConvBertForMultipleChoice, + TFConvBertForQuestionAnswering, + TFConvBertForSequenceClassification, + TFConvBertForTokenClassification, + TFConvBertModel, + TFConvBertPreTrainedModel, + ) + from .models.convnext import ( + TFConvNextForImageClassification, + TFConvNextModel, + TFConvNextPreTrainedModel, + ) + from .models.convnextv2 import ( + TFConvNextV2ForImageClassification, + TFConvNextV2Model, + TFConvNextV2PreTrainedModel, + ) + from .models.ctrl import ( + TFCTRLForSequenceClassification, + TFCTRLLMHeadModel, + TFCTRLModel, + TFCTRLPreTrainedModel, + ) + from .models.cvt import ( + TFCvtForImageClassification, + TFCvtModel, + TFCvtPreTrainedModel, + ) + from .models.data2vec import ( + TFData2VecVisionForImageClassification, + TFData2VecVisionForSemanticSegmentation, + TFData2VecVisionModel, + TFData2VecVisionPreTrainedModel, + ) + from .models.deberta import ( + TFDebertaForMaskedLM, + TFDebertaForQuestionAnswering, + TFDebertaForSequenceClassification, + TFDebertaForTokenClassification, + TFDebertaModel, + TFDebertaPreTrainedModel, + ) + from .models.deberta_v2 import ( + TFDebertaV2ForMaskedLM, + TFDebertaV2ForMultipleChoice, + TFDebertaV2ForQuestionAnswering, + TFDebertaV2ForSequenceClassification, + TFDebertaV2ForTokenClassification, + TFDebertaV2Model, + TFDebertaV2PreTrainedModel, + ) + from .models.deit import ( + TFDeiTForImageClassification, + TFDeiTForImageClassificationWithTeacher, + TFDeiTForMaskedImageModeling, + TFDeiTModel, + TFDeiTPreTrainedModel, + ) + from .models.deprecated.efficientformer import ( + TFEfficientFormerForImageClassification, + TFEfficientFormerForImageClassificationWithTeacher, + TFEfficientFormerModel, + TFEfficientFormerPreTrainedModel, + ) + from .models.deprecated.transfo_xl import ( + TFAdaptiveEmbedding, + TFTransfoXLForSequenceClassification, + TFTransfoXLLMHeadModel, + TFTransfoXLMainLayer, + TFTransfoXLModel, + TFTransfoXLPreTrainedModel, + ) + from .models.distilbert import ( + TFDistilBertForMaskedLM, + TFDistilBertForMultipleChoice, + TFDistilBertForQuestionAnswering, + TFDistilBertForSequenceClassification, + TFDistilBertForTokenClassification, + TFDistilBertMainLayer, + TFDistilBertModel, + TFDistilBertPreTrainedModel, + ) + from .models.dpr import ( + TFDPRContextEncoder, + TFDPRPretrainedContextEncoder, + TFDPRPretrainedQuestionEncoder, + TFDPRPretrainedReader, + TFDPRQuestionEncoder, + TFDPRReader, + ) + from .models.electra import ( + TFElectraForMaskedLM, + TFElectraForMultipleChoice, + TFElectraForPreTraining, + TFElectraForQuestionAnswering, + TFElectraForSequenceClassification, + TFElectraForTokenClassification, + TFElectraModel, + TFElectraPreTrainedModel, + ) + from .models.encoder_decoder import TFEncoderDecoderModel + from .models.esm import ( + TFEsmForMaskedLM, + TFEsmForSequenceClassification, + TFEsmForTokenClassification, + TFEsmModel, + TFEsmPreTrainedModel, + ) + from .models.flaubert import ( + TFFlaubertForMultipleChoice, + TFFlaubertForQuestionAnsweringSimple, + TFFlaubertForSequenceClassification, + TFFlaubertForTokenClassification, + TFFlaubertModel, + TFFlaubertPreTrainedModel, + TFFlaubertWithLMHeadModel, + ) + from .models.funnel import ( + TFFunnelBaseModel, + TFFunnelForMaskedLM, + TFFunnelForMultipleChoice, + TFFunnelForPreTraining, + TFFunnelForQuestionAnswering, + TFFunnelForSequenceClassification, + TFFunnelForTokenClassification, + TFFunnelModel, + TFFunnelPreTrainedModel, + ) + from .models.gpt2 import ( + TFGPT2DoubleHeadsModel, + TFGPT2ForSequenceClassification, + TFGPT2LMHeadModel, + TFGPT2MainLayer, + TFGPT2Model, + TFGPT2PreTrainedModel, + ) + from .models.gptj import ( + TFGPTJForCausalLM, + TFGPTJForQuestionAnswering, + TFGPTJForSequenceClassification, + TFGPTJModel, + TFGPTJPreTrainedModel, + ) + from .models.groupvit import ( + TFGroupViTModel, + TFGroupViTPreTrainedModel, + TFGroupViTTextModel, + TFGroupViTVisionModel, + ) + from .models.hubert import ( + TFHubertForCTC, + TFHubertModel, + TFHubertPreTrainedModel, + ) + from .models.idefics import ( + TFIdeficsForVisionText2Text, + TFIdeficsModel, + TFIdeficsPreTrainedModel, + ) + from .models.layoutlm import ( + TFLayoutLMForMaskedLM, + TFLayoutLMForQuestionAnswering, + TFLayoutLMForSequenceClassification, + TFLayoutLMForTokenClassification, + TFLayoutLMMainLayer, + TFLayoutLMModel, + TFLayoutLMPreTrainedModel, + ) + from .models.layoutlmv3 import ( + TFLayoutLMv3ForQuestionAnswering, + TFLayoutLMv3ForSequenceClassification, + TFLayoutLMv3ForTokenClassification, + TFLayoutLMv3Model, + TFLayoutLMv3PreTrainedModel, + ) + from .models.led import ( + TFLEDForConditionalGeneration, + TFLEDModel, + TFLEDPreTrainedModel, + ) + from .models.longformer import ( + TFLongformerForMaskedLM, + TFLongformerForMultipleChoice, + TFLongformerForQuestionAnswering, + TFLongformerForSequenceClassification, + TFLongformerForTokenClassification, + TFLongformerModel, + TFLongformerPreTrainedModel, + ) + from .models.lxmert import ( + TFLxmertForPreTraining, + TFLxmertMainLayer, + TFLxmertModel, + TFLxmertPreTrainedModel, + TFLxmertVisualFeatureEncoder, + ) + from .models.marian import ( + TFMarianModel, + TFMarianMTModel, + TFMarianPreTrainedModel, + ) + from .models.mbart import ( + TFMBartForConditionalGeneration, + TFMBartModel, + TFMBartPreTrainedModel, + ) + from .models.mistral import ( + TFMistralForCausalLM, + TFMistralForSequenceClassification, + TFMistralModel, + TFMistralPreTrainedModel, + ) + from .models.mobilebert import ( + TFMobileBertForMaskedLM, + TFMobileBertForMultipleChoice, + TFMobileBertForNextSentencePrediction, + TFMobileBertForPreTraining, + TFMobileBertForQuestionAnswering, + TFMobileBertForSequenceClassification, + TFMobileBertForTokenClassification, + TFMobileBertMainLayer, + TFMobileBertModel, + TFMobileBertPreTrainedModel, + ) + from .models.mobilevit import ( + TFMobileViTForImageClassification, + TFMobileViTForSemanticSegmentation, + TFMobileViTModel, + TFMobileViTPreTrainedModel, + ) + from .models.mpnet import ( + TFMPNetForMaskedLM, + TFMPNetForMultipleChoice, + TFMPNetForQuestionAnswering, + TFMPNetForSequenceClassification, + TFMPNetForTokenClassification, + TFMPNetMainLayer, + TFMPNetModel, + TFMPNetPreTrainedModel, + ) + from .models.mt5 import ( + TFMT5EncoderModel, + TFMT5ForConditionalGeneration, + TFMT5Model, + ) + from .models.openai import ( + TFOpenAIGPTDoubleHeadsModel, + TFOpenAIGPTForSequenceClassification, + TFOpenAIGPTLMHeadModel, + TFOpenAIGPTMainLayer, + TFOpenAIGPTModel, + TFOpenAIGPTPreTrainedModel, + ) + from .models.opt import TFOPTForCausalLM, TFOPTModel, TFOPTPreTrainedModel + from .models.pegasus import ( + TFPegasusForConditionalGeneration, + TFPegasusModel, + TFPegasusPreTrainedModel, + ) + from .models.rag import ( + TFRagModel, + TFRagPreTrainedModel, + TFRagSequenceForGeneration, + TFRagTokenForGeneration, + ) + from .models.regnet import ( + TFRegNetForImageClassification, + TFRegNetModel, + TFRegNetPreTrainedModel, + ) + from .models.rembert import ( + TFRemBertForCausalLM, + TFRemBertForMaskedLM, + TFRemBertForMultipleChoice, + TFRemBertForQuestionAnswering, + TFRemBertForSequenceClassification, + TFRemBertForTokenClassification, + TFRemBertModel, + TFRemBertPreTrainedModel, + ) + from .models.resnet import ( + TFResNetForImageClassification, + TFResNetModel, + TFResNetPreTrainedModel, + ) + from .models.roberta import ( + TFRobertaForCausalLM, + TFRobertaForMaskedLM, + TFRobertaForMultipleChoice, + TFRobertaForQuestionAnswering, + TFRobertaForSequenceClassification, + TFRobertaForTokenClassification, + TFRobertaMainLayer, + TFRobertaModel, + TFRobertaPreTrainedModel, + ) + from .models.roberta_prelayernorm import ( + TFRobertaPreLayerNormForCausalLM, + TFRobertaPreLayerNormForMaskedLM, + TFRobertaPreLayerNormForMultipleChoice, + TFRobertaPreLayerNormForQuestionAnswering, + TFRobertaPreLayerNormForSequenceClassification, + TFRobertaPreLayerNormForTokenClassification, + TFRobertaPreLayerNormMainLayer, + TFRobertaPreLayerNormModel, + TFRobertaPreLayerNormPreTrainedModel, + ) + from .models.roformer import ( + TFRoFormerForCausalLM, + TFRoFormerForMaskedLM, + TFRoFormerForMultipleChoice, + TFRoFormerForQuestionAnswering, + TFRoFormerForSequenceClassification, + TFRoFormerForTokenClassification, + TFRoFormerModel, + TFRoFormerPreTrainedModel, + ) + from .models.sam import ( + TFSamModel, + TFSamPreTrainedModel, + TFSamVisionModel, + ) + from .models.segformer import ( + TFSegformerDecodeHead, + TFSegformerForImageClassification, + TFSegformerForSemanticSegmentation, + TFSegformerModel, + TFSegformerPreTrainedModel, + ) + from .models.speech_to_text import ( + TFSpeech2TextForConditionalGeneration, + TFSpeech2TextModel, + TFSpeech2TextPreTrainedModel, + ) + from .models.swiftformer import ( + TFSwiftFormerForImageClassification, + TFSwiftFormerModel, + TFSwiftFormerPreTrainedModel, + ) + from .models.swin import ( + TFSwinForImageClassification, + TFSwinForMaskedImageModeling, + TFSwinModel, + TFSwinPreTrainedModel, + ) + from .models.t5 import ( + TFT5EncoderModel, + TFT5ForConditionalGeneration, + TFT5Model, + TFT5PreTrainedModel, + ) + from .models.tapas import ( + TFTapasForMaskedLM, + TFTapasForQuestionAnswering, + TFTapasForSequenceClassification, + TFTapasModel, + TFTapasPreTrainedModel, + ) + from .models.vision_encoder_decoder import TFVisionEncoderDecoderModel + from .models.vision_text_dual_encoder import TFVisionTextDualEncoderModel + from .models.vit import ( + TFViTForImageClassification, + TFViTModel, + TFViTPreTrainedModel, + ) + from .models.vit_mae import ( + TFViTMAEForPreTraining, + TFViTMAEModel, + TFViTMAEPreTrainedModel, + ) + from .models.wav2vec2 import ( + TFWav2Vec2ForCTC, + TFWav2Vec2ForSequenceClassification, + TFWav2Vec2Model, + TFWav2Vec2PreTrainedModel, + ) + from .models.whisper import ( + TFWhisperForConditionalGeneration, + TFWhisperModel, + TFWhisperPreTrainedModel, + ) + from .models.xglm import ( + TFXGLMForCausalLM, + TFXGLMModel, + TFXGLMPreTrainedModel, + ) + from .models.xlm import ( + TFXLMForMultipleChoice, + TFXLMForQuestionAnsweringSimple, + TFXLMForSequenceClassification, + TFXLMForTokenClassification, + TFXLMMainLayer, + TFXLMModel, + TFXLMPreTrainedModel, + TFXLMWithLMHeadModel, + ) + from .models.xlm_roberta import ( + TFXLMRobertaForCausalLM, + TFXLMRobertaForMaskedLM, + TFXLMRobertaForMultipleChoice, + TFXLMRobertaForQuestionAnswering, + TFXLMRobertaForSequenceClassification, + TFXLMRobertaForTokenClassification, + TFXLMRobertaModel, + TFXLMRobertaPreTrainedModel, + ) + from .models.xlnet import ( + TFXLNetForMultipleChoice, + TFXLNetForQuestionAnsweringSimple, + TFXLNetForSequenceClassification, + TFXLNetForTokenClassification, + TFXLNetLMHeadModel, + TFXLNetMainLayer, + TFXLNetModel, + TFXLNetPreTrainedModel, + ) + + # Optimization + from .optimization_tf import ( + AdamWeightDecay, + GradientAccumulator, + WarmUp, + create_optimizer, + ) + + try: + if not ( + is_librosa_available() + and is_essentia_available() + and is_scipy_available() + and is_torch_available() + and is_pretty_midi_available() + ): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_essentia_and_librosa_and_pretty_midi_and_scipy_and_torch_objects import * + else: + from .models.pop2piano import ( + Pop2PianoFeatureExtractor, + Pop2PianoProcessor, + Pop2PianoTokenizer, + ) + + try: + if not is_torchaudio_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_torchaudio_objects import * + else: + from .models.musicgen_melody import MusicgenMelodyFeatureExtractor, MusicgenMelodyProcessor + try: + if not is_flax_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + # Import the same objects as dummies to get them in the namespace. + # They will raise an import error if the user tries to instantiate / use them. + from .utils.dummy_flax_objects import * + else: + from .generation import ( + FlaxForcedBOSTokenLogitsProcessor, + FlaxForcedEOSTokenLogitsProcessor, + FlaxForceTokensLogitsProcessor, + FlaxGenerationMixin, + FlaxLogitsProcessor, + FlaxLogitsProcessorList, + FlaxLogitsWarper, + FlaxMinLengthLogitsProcessor, + FlaxSuppressTokensAtBeginLogitsProcessor, + FlaxSuppressTokensLogitsProcessor, + FlaxTemperatureLogitsWarper, + FlaxTopKLogitsWarper, + FlaxTopPLogitsWarper, + FlaxWhisperTimeStampLogitsProcessor, + ) + from .modeling_flax_utils import FlaxPreTrainedModel + + # Flax model imports + from .models.albert import ( + FlaxAlbertForMaskedLM, + FlaxAlbertForMultipleChoice, + FlaxAlbertForPreTraining, + FlaxAlbertForQuestionAnswering, + FlaxAlbertForSequenceClassification, + FlaxAlbertForTokenClassification, + FlaxAlbertModel, + FlaxAlbertPreTrainedModel, + ) + from .models.auto import ( + FLAX_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING, + FLAX_MODEL_FOR_CAUSAL_LM_MAPPING, + FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING, + FLAX_MODEL_FOR_MASKED_LM_MAPPING, + FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING, + FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING, + FLAX_MODEL_FOR_PRETRAINING_MAPPING, + FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING, + FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, + FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING, + FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, + FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING, + FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING, + FLAX_MODEL_MAPPING, + FlaxAutoModel, + FlaxAutoModelForCausalLM, + FlaxAutoModelForImageClassification, + FlaxAutoModelForMaskedLM, + FlaxAutoModelForMultipleChoice, + FlaxAutoModelForNextSentencePrediction, + FlaxAutoModelForPreTraining, + FlaxAutoModelForQuestionAnswering, + FlaxAutoModelForSeq2SeqLM, + FlaxAutoModelForSequenceClassification, + FlaxAutoModelForSpeechSeq2Seq, + FlaxAutoModelForTokenClassification, + FlaxAutoModelForVision2Seq, + ) + from .models.bart import ( + FlaxBartDecoderPreTrainedModel, + FlaxBartForCausalLM, + FlaxBartForConditionalGeneration, + FlaxBartForQuestionAnswering, + FlaxBartForSequenceClassification, + FlaxBartModel, + FlaxBartPreTrainedModel, + ) + from .models.beit import ( + FlaxBeitForImageClassification, + FlaxBeitForMaskedImageModeling, + FlaxBeitModel, + FlaxBeitPreTrainedModel, + ) + from .models.bert import ( + FlaxBertForCausalLM, + FlaxBertForMaskedLM, + FlaxBertForMultipleChoice, + FlaxBertForNextSentencePrediction, + FlaxBertForPreTraining, + FlaxBertForQuestionAnswering, + FlaxBertForSequenceClassification, + FlaxBertForTokenClassification, + FlaxBertModel, + FlaxBertPreTrainedModel, + ) + from .models.big_bird import ( + FlaxBigBirdForCausalLM, + FlaxBigBirdForMaskedLM, + FlaxBigBirdForMultipleChoice, + FlaxBigBirdForPreTraining, + FlaxBigBirdForQuestionAnswering, + FlaxBigBirdForSequenceClassification, + FlaxBigBirdForTokenClassification, + FlaxBigBirdModel, + FlaxBigBirdPreTrainedModel, + ) + from .models.blenderbot import ( + FlaxBlenderbotForConditionalGeneration, + FlaxBlenderbotModel, + FlaxBlenderbotPreTrainedModel, + ) + from .models.blenderbot_small import ( + FlaxBlenderbotSmallForConditionalGeneration, + FlaxBlenderbotSmallModel, + FlaxBlenderbotSmallPreTrainedModel, + ) + from .models.bloom import ( + FlaxBloomForCausalLM, + FlaxBloomModel, + FlaxBloomPreTrainedModel, + ) + from .models.clip import ( + FlaxCLIPModel, + FlaxCLIPPreTrainedModel, + FlaxCLIPTextModel, + FlaxCLIPTextModelWithProjection, + FlaxCLIPTextPreTrainedModel, + FlaxCLIPVisionModel, + FlaxCLIPVisionPreTrainedModel, + ) + from .models.dinov2 import ( + FlaxDinov2ForImageClassification, + FlaxDinov2Model, + FlaxDinov2PreTrainedModel, + ) + from .models.distilbert import ( + FlaxDistilBertForMaskedLM, + FlaxDistilBertForMultipleChoice, + FlaxDistilBertForQuestionAnswering, + FlaxDistilBertForSequenceClassification, + FlaxDistilBertForTokenClassification, + FlaxDistilBertModel, + FlaxDistilBertPreTrainedModel, + ) + from .models.electra import ( + FlaxElectraForCausalLM, + FlaxElectraForMaskedLM, + FlaxElectraForMultipleChoice, + FlaxElectraForPreTraining, + FlaxElectraForQuestionAnswering, + FlaxElectraForSequenceClassification, + FlaxElectraForTokenClassification, + FlaxElectraModel, + FlaxElectraPreTrainedModel, + ) + from .models.encoder_decoder import FlaxEncoderDecoderModel + from .models.gemma import ( + FlaxGemmaForCausalLM, + FlaxGemmaModel, + FlaxGemmaPreTrainedModel, + ) + from .models.gpt2 import ( + FlaxGPT2LMHeadModel, + FlaxGPT2Model, + FlaxGPT2PreTrainedModel, + ) + from .models.gpt_neo import ( + FlaxGPTNeoForCausalLM, + FlaxGPTNeoModel, + FlaxGPTNeoPreTrainedModel, + ) + from .models.gptj import ( + FlaxGPTJForCausalLM, + FlaxGPTJModel, + FlaxGPTJPreTrainedModel, + ) + from .models.llama import ( + FlaxLlamaForCausalLM, + FlaxLlamaModel, + FlaxLlamaPreTrainedModel, + ) + from .models.longt5 import ( + FlaxLongT5ForConditionalGeneration, + FlaxLongT5Model, + FlaxLongT5PreTrainedModel, + ) + from .models.marian import ( + FlaxMarianModel, + FlaxMarianMTModel, + FlaxMarianPreTrainedModel, + ) + from .models.mbart import ( + FlaxMBartForConditionalGeneration, + FlaxMBartForQuestionAnswering, + FlaxMBartForSequenceClassification, + FlaxMBartModel, + FlaxMBartPreTrainedModel, + ) + from .models.mistral import ( + FlaxMistralForCausalLM, + FlaxMistralModel, + FlaxMistralPreTrainedModel, + ) + from .models.mt5 import ( + FlaxMT5EncoderModel, + FlaxMT5ForConditionalGeneration, + FlaxMT5Model, + ) + from .models.opt import FlaxOPTForCausalLM, FlaxOPTModel, FlaxOPTPreTrainedModel + from .models.pegasus import ( + FlaxPegasusForConditionalGeneration, + FlaxPegasusModel, + FlaxPegasusPreTrainedModel, + ) + from .models.regnet import ( + FlaxRegNetForImageClassification, + FlaxRegNetModel, + FlaxRegNetPreTrainedModel, + ) + from .models.resnet import ( + FlaxResNetForImageClassification, + FlaxResNetModel, + FlaxResNetPreTrainedModel, + ) + from .models.roberta import ( + FlaxRobertaForCausalLM, + FlaxRobertaForMaskedLM, + FlaxRobertaForMultipleChoice, + FlaxRobertaForQuestionAnswering, + FlaxRobertaForSequenceClassification, + FlaxRobertaForTokenClassification, + FlaxRobertaModel, + FlaxRobertaPreTrainedModel, + ) + from .models.roberta_prelayernorm import ( + FlaxRobertaPreLayerNormForCausalLM, + FlaxRobertaPreLayerNormForMaskedLM, + FlaxRobertaPreLayerNormForMultipleChoice, + FlaxRobertaPreLayerNormForQuestionAnswering, + FlaxRobertaPreLayerNormForSequenceClassification, + FlaxRobertaPreLayerNormForTokenClassification, + FlaxRobertaPreLayerNormModel, + FlaxRobertaPreLayerNormPreTrainedModel, + ) + from .models.roformer import ( + FlaxRoFormerForMaskedLM, + FlaxRoFormerForMultipleChoice, + FlaxRoFormerForQuestionAnswering, + FlaxRoFormerForSequenceClassification, + FlaxRoFormerForTokenClassification, + FlaxRoFormerModel, + FlaxRoFormerPreTrainedModel, + ) + from .models.speech_encoder_decoder import FlaxSpeechEncoderDecoderModel + from .models.t5 import ( + FlaxT5EncoderModel, + FlaxT5ForConditionalGeneration, + FlaxT5Model, + FlaxT5PreTrainedModel, + ) + from .models.vision_encoder_decoder import FlaxVisionEncoderDecoderModel + from .models.vision_text_dual_encoder import FlaxVisionTextDualEncoderModel + from .models.vit import ( + FlaxViTForImageClassification, + FlaxViTModel, + FlaxViTPreTrainedModel, + ) + from .models.wav2vec2 import ( + FlaxWav2Vec2ForCTC, + FlaxWav2Vec2ForPreTraining, + FlaxWav2Vec2Model, + FlaxWav2Vec2PreTrainedModel, + ) + from .models.whisper import ( + FlaxWhisperForAudioClassification, + FlaxWhisperForConditionalGeneration, + FlaxWhisperModel, + FlaxWhisperPreTrainedModel, + ) + from .models.xglm import ( + FlaxXGLMForCausalLM, + FlaxXGLMModel, + FlaxXGLMPreTrainedModel, + ) + from .models.xlm_roberta import ( + FlaxXLMRobertaForCausalLM, + FlaxXLMRobertaForMaskedLM, + FlaxXLMRobertaForMultipleChoice, + FlaxXLMRobertaForQuestionAnswering, + FlaxXLMRobertaForSequenceClassification, + FlaxXLMRobertaForTokenClassification, + FlaxXLMRobertaModel, + FlaxXLMRobertaPreTrainedModel, + ) + + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + extra_objects={"__version__": __version__}, + ) + + +if not is_tf_available() and not is_torch_available() and not is_flax_available(): + logger.warning_advice( + "None of PyTorch, TensorFlow >= 2.0, or Flax have been found. " + "Models won't be available and only tokenizers, configuration " + "and file/data utilities can be used." + ) diff --git a/lib/python3.12/site-packages/transformers/activations.py b/lib/python3.12/site-packages/transformers/activations.py new file mode 100644 index 0000000000000000000000000000000000000000..15f0397535e8c6df24cbd2148471ed71742d0636 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/activations.py @@ -0,0 +1,240 @@ +# Copyright 2020 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. + +import math +from collections import OrderedDict + +import torch +from packaging import version +from torch import Tensor, nn + +from .utils import logging + + +logger = logging.get_logger(__name__) + + +class PytorchGELUTanh(nn.Module): + """ + A fast C implementation of the tanh approximation of the GeLU activation function. See + https://arxiv.org/abs/1606.08415. + + This implementation is equivalent to NewGELU and FastGELU but much faster. However, it is not an exact numerical + match due to rounding errors. + """ + + def __init__(self): + super().__init__() + if version.parse(torch.__version__) < version.parse("1.12.0"): + raise ImportError( + f"You are using torch=={torch.__version__}, but torch>=1.12.0 is required to use " + "PytorchGELUTanh. Please upgrade torch." + ) + + def forward(self, input: Tensor) -> Tensor: + return nn.functional.gelu(input, approximate="tanh") + + +class NewGELUActivation(nn.Module): + """ + Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT). Also see + the Gaussian Error Linear Units paper: https://arxiv.org/abs/1606.08415 + """ + + def forward(self, input: Tensor) -> Tensor: + return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0)))) + + +class GELUActivation(nn.Module): + """ + Original Implementation of the GELU activation function in Google BERT repo when initially created. For + information: OpenAI GPT's GELU is slightly different (and gives slightly different results): 0.5 * x * (1 + + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))) This is now written in C in nn.functional + Also see the Gaussian Error Linear Units paper: https://arxiv.org/abs/1606.08415 + """ + + def __init__(self, use_gelu_python: bool = False): + super().__init__() + if use_gelu_python: + self.act = self._gelu_python + else: + self.act = nn.functional.gelu + + def _gelu_python(self, input: Tensor) -> Tensor: + return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0))) + + def forward(self, input: Tensor) -> Tensor: + return self.act(input) + + +class FastGELUActivation(nn.Module): + """ + Applies GELU approximation that is slower than QuickGELU but more accurate. See: https://github.com/hendrycks/GELUs + """ + + def forward(self, input: Tensor) -> Tensor: + return 0.5 * input * (1.0 + torch.tanh(input * 0.7978845608 * (1.0 + 0.044715 * input * input))) + + +class QuickGELUActivation(nn.Module): + """ + Applies GELU approximation that is fast but somewhat inaccurate. See: https://github.com/hendrycks/GELUs + """ + + def forward(self, input: Tensor) -> Tensor: + return input * torch.sigmoid(1.702 * input) + + +class ClippedGELUActivation(nn.Module): + """ + Clip the range of possible GeLU outputs between [min, max]. This is especially useful for quantization purpose, as + it allows mapping negatives values in the GeLU spectrum. For more information on this trick, please refer to + https://arxiv.org/abs/2004.09602. + + Gaussian Error Linear Unit. Original Implementation of the gelu activation function in Google Bert repo when + initially created. + + For information: OpenAI GPT's gelu is slightly different (and gives slightly different results): 0.5 * x * (1 + + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))). See https://arxiv.org/abs/1606.08415 + """ + + def __init__(self, min: float, max: float): + if min > max: + raise ValueError(f"min should be < max (got min: {min}, max: {max})") + + super().__init__() + self.min = min + self.max = max + + def forward(self, x: Tensor) -> Tensor: + return torch.clip(gelu(x), self.min, self.max) + + +class AccurateGELUActivation(nn.Module): + """ + Applies GELU approximation that is faster than default and more accurate than QuickGELU. See: + https://github.com/hendrycks/GELUs + + Implemented along with MEGA (Moving Average Equipped Gated Attention) + """ + + def __init__(self): + super().__init__() + self.precomputed_constant = math.sqrt(2 / math.pi) + + def forward(self, input: Tensor) -> Tensor: + return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.044715 * torch.pow(input, 3)))) + + +class MishActivation(nn.Module): + """ + See Mish: A Self-Regularized Non-Monotonic Activation Function (Misra., https://arxiv.org/abs/1908.08681). Also + visit the official repository for the paper: https://github.com/digantamisra98/Mish + """ + + def __init__(self): + super().__init__() + if version.parse(torch.__version__) < version.parse("1.9.0"): + self.act = self._mish_python + else: + self.act = nn.functional.mish + + def _mish_python(self, input: Tensor) -> Tensor: + return input * torch.tanh(nn.functional.softplus(input)) + + def forward(self, input: Tensor) -> Tensor: + return self.act(input) + + +class LinearActivation(nn.Module): + """ + Applies the linear activation function, i.e. forwarding input directly to output. + """ + + def forward(self, input: Tensor) -> Tensor: + return input + + +class LaplaceActivation(nn.Module): + """ + Applies elementwise activation based on Laplace function, introduced in MEGA as an attention activation. See + https://arxiv.org/abs/2209.10655 + + Inspired by squared relu, but with bounded range and gradient for better stability + """ + + def forward(self, input, mu=0.707107, sigma=0.282095): + input = (input - mu).div(sigma * math.sqrt(2.0)) + return 0.5 * (1.0 + torch.erf(input)) + + +class ReLUSquaredActivation(nn.Module): + """ + Applies the relu^2 activation introduced in https://arxiv.org/abs/2109.08668v2 + """ + + def forward(self, input): + relu_applied = nn.functional.relu(input) + squared = torch.square(relu_applied) + return squared + + +class ClassInstantier(OrderedDict): + def __getitem__(self, key): + content = super().__getitem__(key) + cls, kwargs = content if isinstance(content, tuple) else (content, {}) + return cls(**kwargs) + + +ACT2CLS = { + "gelu": GELUActivation, + "gelu_10": (ClippedGELUActivation, {"min": -10, "max": 10}), + "gelu_fast": FastGELUActivation, + "gelu_new": NewGELUActivation, + "gelu_python": (GELUActivation, {"use_gelu_python": True}), + "gelu_pytorch_tanh": PytorchGELUTanh, + "gelu_accurate": AccurateGELUActivation, + "laplace": LaplaceActivation, + "leaky_relu": nn.LeakyReLU, + "linear": LinearActivation, + "mish": MishActivation, + "quick_gelu": QuickGELUActivation, + "relu": nn.ReLU, + "relu2": ReLUSquaredActivation, + "relu6": nn.ReLU6, + "sigmoid": nn.Sigmoid, + "silu": nn.SiLU, + "swish": nn.SiLU, + "tanh": nn.Tanh, + "prelu": nn.PReLU, +} +ACT2FN = ClassInstantier(ACT2CLS) + + +def get_activation(activation_string): + if activation_string in ACT2FN: + return ACT2FN[activation_string] + else: + raise KeyError(f"function {activation_string} not found in ACT2FN mapping {list(ACT2FN.keys())}") + + +# For backwards compatibility with: from activations import gelu_python +gelu_python = get_activation("gelu_python") +gelu_new = get_activation("gelu_new") +gelu = get_activation("gelu") +gelu_fast = get_activation("gelu_fast") +quick_gelu = get_activation("quick_gelu") +silu = get_activation("silu") +mish = get_activation("mish") +linear_act = get_activation("linear") diff --git a/lib/python3.12/site-packages/transformers/activations_tf.py b/lib/python3.12/site-packages/transformers/activations_tf.py new file mode 100644 index 0000000000000000000000000000000000000000..d12b73ea45176f3a4bc42cdabe8b73078a3b90f2 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/activations_tf.py @@ -0,0 +1,147 @@ +# Copyright 2020 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. + +import math + +import tensorflow as tf +from packaging.version import parse + + +try: + import tf_keras as keras +except (ModuleNotFoundError, ImportError): + import keras + + if parse(keras.__version__).major > 2: + raise ValueError( + "Your currently installed version of Keras is Keras 3, but this is not yet supported in " + "Transformers. Please install the backwards-compatible tf-keras package with " + "`pip install tf-keras`." + ) + + +def _gelu(x): + """ + Gaussian Error Linear Unit. Original Implementation of the gelu activation function in Google Bert repo when + initially created. For information: OpenAI GPT's gelu is slightly different (and gives slightly different results): + 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))) Also see + https://arxiv.org/abs/1606.08415 + """ + x = tf.convert_to_tensor(x) + cdf = 0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0), x.dtype))) + + return x * cdf + + +def _gelu_new(x): + """ + Gaussian Error Linear Unit. This is a smoother version of the GELU. Original paper: https://arxiv.org/abs/1606.0841 + + Args: + x: float Tensor to perform activation + + Returns: + `x` with the GELU activation applied. + """ + x = tf.convert_to_tensor(x) + pi = tf.cast(math.pi, x.dtype) + coeff = tf.cast(0.044715, x.dtype) + cdf = 0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi) * (x + coeff * tf.pow(x, 3)))) + + return x * cdf + + +def mish(x): + x = tf.convert_to_tensor(x) + + return x * tf.tanh(tf.math.softplus(x)) + + +def gelu_fast(x): + x = tf.convert_to_tensor(x) + coeff1 = tf.cast(0.044715, x.dtype) + coeff2 = tf.cast(0.7978845608, x.dtype) + + return 0.5 * x * (1.0 + tf.tanh(x * coeff2 * (1.0 + coeff1 * x * x))) + + +def quick_gelu(x): + x = tf.convert_to_tensor(x) + coeff = tf.cast(1.702, x.dtype) + return x * tf.math.sigmoid(coeff * x) + + +def gelu_10(x): + """ + Clip the range of possible GeLU outputs between [-10, 10]. This is especially useful for quantization purpose, as + it allows mapping 2 negatives values in the GeLU spectrum. For more information on this trick, please refer to + https://arxiv.org/abs/2004.09602 + + Gaussian Error Linear Unit. Original Implementation of the gelu activation function in Google Bert repo when + initially created. For information: OpenAI GPT's gelu is slightly different (and gives slightly different results): + 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3)))) Also see + https://arxiv.org/abs/1606.08415 :param x: :return: + """ + return tf.clip_by_value(_gelu(x), -10, 10) + + +def glu(x, axis=-1): + """ + Gated Linear Unit. Implementation as defined in the original paper (see https://arxiv.org/abs/1612.08083), where + the input `x` is split in two halves across a dimension (`axis`), A and B, returning A * sigmoid(B). + + Args: + `x`: float Tensor to perform activation + `axis`: dimension across which `x` be split in half + + Returns: + `x` with the GLU activation applied (with its size halved across the dimension `axis`). + """ + a, b = tf.split(x, 2, axis=axis) + return a * tf.math.sigmoid(b) + + +if parse(tf.version.VERSION) >= parse("2.4"): + + def approximate_gelu_wrap(x): + return keras.activations.gelu(x, approximate=True) + + gelu = keras.activations.gelu + gelu_new = approximate_gelu_wrap +else: + gelu = _gelu + gelu_new = _gelu_new + + +ACT2FN = { + "gelu": gelu, + "gelu_10": gelu_10, + "gelu_fast": gelu_fast, + "gelu_new": gelu_new, + "glu": glu, + "mish": mish, + "quick_gelu": quick_gelu, + "relu": keras.activations.relu, + "sigmoid": keras.activations.sigmoid, + "silu": keras.activations.swish, + "swish": keras.activations.swish, + "tanh": keras.activations.tanh, +} + + +def get_tf_activation(activation_string): + if activation_string in ACT2FN: + return ACT2FN[activation_string] + else: + raise KeyError(f"function {activation_string} not found in ACT2FN mapping {list(ACT2FN.keys())}") diff --git a/lib/python3.12/site-packages/transformers/agents/__init__.py b/lib/python3.12/site-packages/transformers/agents/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70762c252a832882a53c1baf515c58356dffe1a1 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/__init__.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2023 The HuggingFace Inc. 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. +from typing import TYPE_CHECKING + +from ..utils import ( + OptionalDependencyNotAvailable, + _LazyModule, + is_torch_available, +) + + +_import_structure = { + "agents": ["Agent", "CodeAgent", "ManagedAgent", "ReactAgent", "ReactCodeAgent", "ReactJsonAgent", "Toolbox"], + "llm_engine": ["HfApiEngine", "TransformersEngine"], + "monitoring": ["stream_to_gradio"], + "tools": ["PipelineTool", "Tool", "ToolCollection", "launch_gradio_demo", "load_tool", "tool"], +} + +try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["default_tools"] = ["FinalAnswerTool", "PythonInterpreterTool"] + _import_structure["document_question_answering"] = ["DocumentQuestionAnsweringTool"] + _import_structure["image_question_answering"] = ["ImageQuestionAnsweringTool"] + _import_structure["search"] = ["DuckDuckGoSearchTool", "VisitWebpageTool"] + _import_structure["speech_to_text"] = ["SpeechToTextTool"] + _import_structure["text_to_speech"] = ["TextToSpeechTool"] + _import_structure["translation"] = ["TranslationTool"] + +if TYPE_CHECKING: + from .agents import Agent, CodeAgent, ManagedAgent, ReactAgent, ReactCodeAgent, ReactJsonAgent, Toolbox + from .llm_engine import HfApiEngine, TransformersEngine + from .monitoring import stream_to_gradio + from .tools import PipelineTool, Tool, ToolCollection, launch_gradio_demo, load_tool, tool + + try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .default_tools import FinalAnswerTool, PythonInterpreterTool + from .document_question_answering import DocumentQuestionAnsweringTool + from .image_question_answering import ImageQuestionAnsweringTool + from .search import DuckDuckGoSearchTool, VisitWebpageTool + from .speech_to_text import SpeechToTextTool + from .text_to_speech import TextToSpeechTool + from .translation import TranslationTool +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..215d8449012e13ebe2f207a056662f9ed77985d8 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/agent_types.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/agent_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e8b4355ab29de6c7c3a7bc410367094ff1d770a Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/agent_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/agents.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/agents.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..acc6ec6babd8578a77a748b717b2148811d2b7b0 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/agents.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/default_tools.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/default_tools.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..218ac8315e5852f49b88ede6af91b2e2873f757c Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/default_tools.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/document_question_answering.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/document_question_answering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1311ef354a20f4982d9e5d04aa258557fef076c1 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/document_question_answering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/evaluate_agent.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/evaluate_agent.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8373a4cd9a1d9e5e7686b18130a4996f8790615a Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/evaluate_agent.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/image_question_answering.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/image_question_answering.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ddb5a1bd8a3f9521cca7a190f606266c2485184c Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/image_question_answering.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/llm_engine.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/llm_engine.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..548d7b1892d25f213de611350341636baa61d6a7 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/llm_engine.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/monitoring.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/monitoring.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e17f0493929ceed6546f7796a24dd7b59e464ef2 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/monitoring.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/prompts.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/prompts.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5dae28cf6fbde4db70d1107e20e29d42e71df6c5 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/prompts.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/python_interpreter.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/python_interpreter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e0d25680d03d5805469c5d3aabac43fd827a027 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/python_interpreter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/search.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eab1f26ecd1e2e30fea8f56771c8bd9d45bcc0c1 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/speech_to_text.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/speech_to_text.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f47a15beeb3eadf02db65eda7e9aa3e60fdbe97f Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/speech_to_text.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/text_to_speech.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/text_to_speech.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..968abbd00c96140dc569ec3e0ca7175d27798a1b Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/text_to_speech.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/tools.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/tools.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aee57f2ff4fd591121cbf4daa45753a810d7e945 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/tools.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/__pycache__/translation.cpython-312.pyc b/lib/python3.12/site-packages/transformers/agents/__pycache__/translation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d01468a7d4fbca8cf0e18df0cf90ce4c0b7eec09 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/agents/__pycache__/translation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/agents/agent_types.py b/lib/python3.12/site-packages/transformers/agents/agent_types.py new file mode 100644 index 0000000000000000000000000000000000000000..2da60745fae091efaad4cbbf322040a6ecd54af2 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/agent_types.py @@ -0,0 +1,260 @@ +# coding=utf-8 +# Copyright 2024 HuggingFace Inc. +# +# 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. +import os +import pathlib +import tempfile +import uuid + +import numpy as np + +from ..utils import is_soundfile_available, is_torch_available, is_vision_available, logging + + +logger = logging.get_logger(__name__) + +if is_vision_available(): + from PIL import Image + from PIL.Image import Image as ImageType +else: + ImageType = object + +if is_torch_available(): + import torch + from torch import Tensor +else: + Tensor = object + +if is_soundfile_available(): + import soundfile as sf + + +class AgentType: + """ + Abstract class to be reimplemented to define types that can be returned by agents. + + These objects serve three purposes: + + - They behave as they were the type they're meant to be, e.g., a string for text, a PIL.Image for images + - They can be stringified: str(object) in order to return a string defining the object + - They should be displayed correctly in ipython notebooks/colab/jupyter + """ + + def __init__(self, value): + self._value = value + + def __str__(self): + return self.to_string() + + def to_raw(self): + logger.error( + "This is a raw AgentType of unknown type. Display in notebooks and string conversion will be unreliable" + ) + return self._value + + def to_string(self) -> str: + logger.error( + "This is a raw AgentType of unknown type. Display in notebooks and string conversion will be unreliable" + ) + return str(self._value) + + +class AgentText(AgentType, str): + """ + Text type returned by the agent. Behaves as a string. + """ + + def to_raw(self): + return self._value + + def to_string(self): + return str(self._value) + + +class AgentImage(AgentType, ImageType): + """ + Image type returned by the agent. Behaves as a PIL.Image. + """ + + def __init__(self, value): + AgentType.__init__(self, value) + ImageType.__init__(self) + + if not is_vision_available(): + raise ImportError("PIL must be installed in order to handle images.") + + self._path = None + self._raw = None + self._tensor = None + + if isinstance(value, ImageType): + self._raw = value + elif isinstance(value, (str, pathlib.Path)): + self._path = value + elif isinstance(value, torch.Tensor): + self._tensor = value + elif isinstance(value, np.ndarray): + self._tensor = torch.from_numpy(value) + else: + raise TypeError(f"Unsupported type for {self.__class__.__name__}: {type(value)}") + + def _ipython_display_(self, include=None, exclude=None): + """ + Displays correctly this type in an ipython notebook (ipython, colab, jupyter, ...) + """ + from IPython.display import Image, display + + display(Image(self.to_string())) + + def to_raw(self): + """ + Returns the "raw" version of that object. In the case of an AgentImage, it is a PIL.Image. + """ + if self._raw is not None: + return self._raw + + if self._path is not None: + self._raw = Image.open(self._path) + return self._raw + + if self._tensor is not None: + array = self._tensor.detach().cpu().numpy() + return Image.fromarray((255 - array * 255).astype(np.uint8)) + + def to_string(self): + """ + Returns the stringified version of that object. In the case of an AgentImage, it is a path to the serialized + version of the image. + """ + if self._path is not None: + return self._path + + if self._raw is not None: + directory = tempfile.mkdtemp() + self._path = os.path.join(directory, str(uuid.uuid4()) + ".png") + self._raw.save(self._path) + return self._path + + if self._tensor is not None: + array = self._tensor.detach().cpu().numpy() + + # There is likely simpler than load into image into save + img = Image.fromarray((255 - array * 255).astype(np.uint8)) + + directory = tempfile.mkdtemp() + self._path = os.path.join(directory, str(uuid.uuid4()) + ".png") + + img.save(self._path) + + return self._path + + def save(self, output_bytes, format, **params): + """ + Saves the image to a file. + Args: + output_bytes (bytes): The output bytes to save the image to. + format (str): The format to use for the output image. The format is the same as in PIL.Image.save. + **params: Additional parameters to pass to PIL.Image.save. + """ + img = self.to_raw() + img.save(output_bytes, format, **params) + + +class AgentAudio(AgentType, str): + """ + Audio type returned by the agent. + """ + + def __init__(self, value, samplerate=16_000): + super().__init__(value) + + if not is_soundfile_available(): + raise ImportError("soundfile must be installed in order to handle audio.") + + self._path = None + self._tensor = None + + self.samplerate = samplerate + if isinstance(value, (str, pathlib.Path)): + self._path = value + elif is_torch_available() and isinstance(value, torch.Tensor): + self._tensor = value + elif isinstance(value, tuple): + self.samplerate = value[0] + if isinstance(value[1], np.ndarray): + self._tensor = torch.from_numpy(value[1]) + else: + self._tensor = torch.tensor(value[1]) + else: + raise ValueError(f"Unsupported audio type: {type(value)}") + + def _ipython_display_(self, include=None, exclude=None): + """ + Displays correctly this type in an ipython notebook (ipython, colab, jupyter, ...) + """ + from IPython.display import Audio, display + + display(Audio(self.to_string(), rate=self.samplerate)) + + def to_raw(self): + """ + Returns the "raw" version of that object. It is a `torch.Tensor` object. + """ + if self._tensor is not None: + return self._tensor + + if self._path is not None: + tensor, self.samplerate = sf.read(self._path) + self._tensor = torch.tensor(tensor) + return self._tensor + + def to_string(self): + """ + Returns the stringified version of that object. In the case of an AgentAudio, it is a path to the serialized + version of the audio. + """ + if self._path is not None: + return self._path + + if self._tensor is not None: + directory = tempfile.mkdtemp() + self._path = os.path.join(directory, str(uuid.uuid4()) + ".wav") + sf.write(self._path, self._tensor, samplerate=self.samplerate) + return self._path + + +AGENT_TYPE_MAPPING = {"string": AgentText, "image": AgentImage, "audio": AgentAudio} +INSTANCE_TYPE_MAPPING = {str: AgentText, ImageType: AgentImage} + +if is_torch_available(): + INSTANCE_TYPE_MAPPING[Tensor] = AgentAudio + + +def handle_agent_inputs(*args, **kwargs): + args = [(arg.to_raw() if isinstance(arg, AgentType) else arg) for arg in args] + kwargs = {k: (v.to_raw() if isinstance(v, AgentType) else v) for k, v in kwargs.items()} + return args, kwargs + + +def handle_agent_outputs(output, output_type=None): + if output_type in AGENT_TYPE_MAPPING: + # If the class has defined outputs, we can map directly according to the class definition + decoded_outputs = AGENT_TYPE_MAPPING[output_type](output) + return decoded_outputs + else: + # If the class does not have defined output, then we map according to the type + for _k, _v in INSTANCE_TYPE_MAPPING.items(): + if isinstance(output, _k): + return _v(output) + return output diff --git a/lib/python3.12/site-packages/transformers/agents/agents.py b/lib/python3.12/site-packages/transformers/agents/agents.py new file mode 100644 index 0000000000000000000000000000000000000000..f1c1d58ff217026e9e58b51d2b4cac29c92bb322 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/agents.py @@ -0,0 +1,1292 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2024 The HuggingFace Inc. 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. +import json +import logging +import re +import time +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +from huggingface_hub.utils._deprecation import _deprecate_method + +from .. import is_torch_available +from ..utils import logging as transformers_logging +from ..utils.import_utils import is_pygments_available +from .agent_types import AgentAudio, AgentImage +from .default_tools import BASE_PYTHON_TOOLS, FinalAnswerTool, setup_default_tools +from .llm_engine import HfApiEngine, MessageRole +from .monitoring import Monitor +from .prompts import ( + DEFAULT_CODE_SYSTEM_PROMPT, + DEFAULT_REACT_CODE_SYSTEM_PROMPT, + DEFAULT_REACT_JSON_SYSTEM_PROMPT, + PLAN_UPDATE_FINAL_PLAN_REDACTION, + PROMPTS_FOR_INITIAL_PLAN, + PROMPTS_FOR_PLAN_UPDATE, + SUPPORTED_PLAN_TYPES, + SYSTEM_PROMPT_FACTS, + SYSTEM_PROMPT_FACTS_UPDATE, + USER_PROMPT_FACTS_UPDATE, +) +from .python_interpreter import LIST_SAFE_MODULES, evaluate_python_code +from .tools import ( + DEFAULT_TOOL_DESCRIPTION_TEMPLATE, + Tool, + get_tool_description_with_args, + load_tool, +) + + +if is_pygments_available(): + from pygments import highlight + from pygments.formatters import Terminal256Formatter + from pygments.lexers import PythonLexer + + +class CustomFormatter(logging.Formatter): + grey = "\x1b[38;20m" + bold_yellow = "\x1b[33;1m" + red = "\x1b[31;20m" + green = "\x1b[32;20m" + bold_green = "\x1b[32;20;1m" + bold_red = "\x1b[31;1m" + bold_white = "\x1b[37;1m" + orange = "\x1b[38;5;214m" + bold_orange = "\x1b[38;5;214;1m" + reset = "\x1b[0m" + format = "%(message)s" + + FORMATS = { + logging.DEBUG: grey + format + reset, + logging.INFO: format, + logging.WARNING: bold_yellow + format + reset, + logging.ERROR: red + format + reset, + logging.CRITICAL: bold_red + format + reset, + 31: reset + format + reset, + 32: green + format + reset, + 33: bold_green + format + reset, + 34: bold_white + format + reset, + 35: orange + format + reset, + 36: bold_orange + format + reset, + } + + def format(self, record): + log_fmt = self.FORMATS.get(record.levelno) + formatter = logging.Formatter(log_fmt) + return formatter.format(record) + + +logger = transformers_logging.get_logger(__name__) +logger.propagate = False +ch = logging.StreamHandler() +ch.setFormatter(CustomFormatter()) +logger.addHandler(ch) + + +def parse_json_blob(json_blob: str) -> Dict[str, str]: + try: + first_accolade_index = json_blob.find("{") + last_accolade_index = [a.start() for a in list(re.finditer("}", json_blob))][-1] + json_blob = json_blob[first_accolade_index : last_accolade_index + 1].replace('\\"', "'") + json_data = json.loads(json_blob, strict=False) + return json_data + except json.JSONDecodeError as e: + place = e.pos + if json_blob[place - 1 : place + 2] == "},\n": + raise ValueError( + "JSON is invalid: you probably tried to provide multiple tool calls in one action. PROVIDE ONLY ONE TOOL CALL." + ) + raise ValueError( + f"The JSON blob you used is invalid due to the following error: {e}.\n" + f"JSON blob was: {json_blob}, decoding failed on that specific part of the blob:\n" + f"'{json_blob[place - 4 : place + 5]}'." + ) + except Exception as e: + raise ValueError(f"Error in parsing the JSON blob: {e}") + + +def parse_code_blob(code_blob: str) -> str: + try: + pattern = r"```(?:py|python)?\n(.*?)\n```" + match = re.search(pattern, code_blob, re.DOTALL) + return match.group(1).strip() + except Exception as e: + raise ValueError( + f""" +The code blob you used is invalid: due to the following error: {e} +This means that the regex pattern {pattern} was not respected: make sure to include code with the correct pattern, for instance: +Thoughts: Your thoughts +Code: +```py +# Your python code here +```""" + ) + + +def parse_json_tool_call(json_blob: str) -> Tuple[str, Dict[str, str]]: + json_blob = json_blob.replace("```json", "").replace("```", "") + tool_call = parse_json_blob(json_blob) + if "action" in tool_call and "action_input" in tool_call: + return tool_call["action"], tool_call["action_input"] + elif "action" in tool_call: + return tool_call["action"], None + else: + raise ValueError( + f"Missing keys: {[key for key in ['action', 'action_input'] if key not in tool_call]} in blob {tool_call}" + ) + + +def parse_text_tool_call(text: str) -> Tuple[str, Union[str, Dict[str, str]]]: + """ + Expects a text in the format: 'Action:', 'Action input:', 'Observation:'. 'Action input:' contains a json string with input arguments. + """ + try: + if "Observation:" in text: + text = text.split("Observation:")[0] + if "Action:" in text: + text = text.split("Action:")[1] + tool_name, tool_input = text.split("Action input:") + if "{" in tool_input: + tool_input = parse_json_blob(tool_input) + else: + tool_input = tool_input.strip().replace('"', "") + return tool_name.strip().replace('"', "").replace("\\", ""), tool_input + except Exception as e: + raise ValueError( + f"Error in parsing the text tool call: {e}. Be sure to provide the correct format. DO NOT repeat your previous incorrect tool call." + ) + + +def to_text(input: Union[List[Dict[str, str]], Dict[str, str], str]) -> str: + if isinstance(input, list): + return "\n".join([m["content"] for m in input]) + elif isinstance(input, dict): + return input["content"] + else: + return input + + +HUGGINGFACE_DEFAULT_TOOLS = {} +_tools_are_initialized = False + + +class Toolbox: + """ + The toolbox contains all tools that the agent can perform operations with, as well as a few methods to + manage them. + + Args: + tools (`List[Tool]`): + The list of tools to instantiate the toolbox with + add_base_tools (`bool`, defaults to `False`, *optional*, defaults to `False`): + Whether to add the tools available within `transformers` to the toolbox. + """ + + def __init__(self, tools: List[Tool], add_base_tools: bool = False): + self._tools = {tool.name: tool for tool in tools} + if add_base_tools: + self.add_base_tools() + self._load_tools_if_needed() + + def add_base_tools(self, add_python_interpreter: bool = False): + global _tools_are_initialized + global HUGGINGFACE_DEFAULT_TOOLS + if not _tools_are_initialized: + HUGGINGFACE_DEFAULT_TOOLS = setup_default_tools(logger) + _tools_are_initialized = True + for tool in HUGGINGFACE_DEFAULT_TOOLS.values(): + if tool.name != "python_interpreter" or add_python_interpreter: + self.add_tool(tool) + self._load_tools_if_needed() + + @property + def tools(self) -> Dict[str, Tool]: + """Get all tools currently in the toolbox""" + return self._tools + + def show_tool_descriptions(self, tool_description_template: Optional[str] = None) -> str: + """ + Returns the description of all tools in the toolbox + + Args: + tool_description_template (`str`, *optional*): + The template to use to describe the tools. If not provided, the default template will be used. + """ + return "\n".join( + [get_tool_description_with_args(tool, tool_description_template) for tool in self._tools.values()] + ) + + def add_tool(self, tool: Tool): + """ + Adds a tool to the toolbox + + Args: + tool (`Tool`): + The tool to add to the toolbox. + """ + if tool.name in self._tools: + raise KeyError(f"Error: tool '{tool.name}' already exists in the toolbox.") + self._tools[tool.name] = tool + + def remove_tool(self, tool_name: str): + """ + Removes a tool from the toolbox + + Args: + tool_name (`str`): + The tool to remove from the toolbox. + """ + if tool_name not in self._tools: + raise KeyError( + f"Error: tool {tool_name} not found in toolbox for removal, should be instead one of {list(self._tools.keys())}." + ) + del self._tools[tool_name] + + def update_tool(self, tool: Tool): + """ + Updates a tool in the toolbox according to its name. + + Args: + tool (`Tool`): + The tool to update to the toolbox. + """ + if tool.name not in self._tools: + raise KeyError( + f"Error: tool {tool.name} not found in toolbox for update, should be instead one of {list(self._tools.keys())}." + ) + self._tools[tool.name] = tool + + def clear_toolbox(self): + """Clears the toolbox""" + self._tools = {} + + def _load_tools_if_needed(self): + for name, tool in self._tools.items(): + if not isinstance(tool, Tool): + task_or_repo_id = tool.task if tool.repo_id is None else tool.repo_id + self._tools[name] = load_tool(task_or_repo_id) + + def __repr__(self): + toolbox_description = "Toolbox contents:\n" + for tool in self._tools.values(): + toolbox_description += f"\t{tool.name}: {tool.description}\n" + return toolbox_description + + +class AgentError(Exception): + """Base class for other agent-related exceptions""" + + def __init__(self, message): + super().__init__(message) + self.message = message + + +class AgentParsingError(AgentError): + """Exception raised for errors in parsing in the agent""" + + pass + + +class AgentExecutionError(AgentError): + """Exception raised for errors in execution in the agent""" + + pass + + +class AgentMaxIterationsError(AgentError): + """Exception raised for errors in execution in the agent""" + + pass + + +class AgentGenerationError(AgentError): + """Exception raised for errors in generation in the agent""" + + pass + + +def format_prompt_with_tools(toolbox: Toolbox, prompt_template: str, tool_description_template: str) -> str: + tool_descriptions = toolbox.show_tool_descriptions(tool_description_template) + prompt = prompt_template.replace("<>", tool_descriptions) + + if "<>" in prompt: + tool_names = [f"'{tool_name}'" for tool_name in toolbox.tools.keys()] + prompt = prompt.replace("<>", ", ".join(tool_names)) + + return prompt + + +def show_agents_descriptions(managed_agents: list): + managed_agents_descriptions = """ +You can also give requests to team members. +Calling a team member works the same as for calling a tool: simply, the only argument you can give in the call is 'request', a long string explaining your request. +Given that this team member is a real human, you should be very verbose in your request. +Here is a list of the team members that you can call:""" + for agent in managed_agents.values(): + managed_agents_descriptions += f"\n- {agent.name}: {agent.description}" + return managed_agents_descriptions + + +def format_prompt_with_managed_agents_descriptions(prompt_template, managed_agents=None) -> str: + if managed_agents is not None: + return prompt_template.replace("<>", show_agents_descriptions(managed_agents)) + else: + return prompt_template.replace("<>", "") + + +def format_prompt_with_imports(prompt_template: str, authorized_imports: List[str]) -> str: + if "<>" not in prompt_template: + raise AgentError("Tag '<>' should be provided in the prompt.") + return prompt_template.replace("<>", str(authorized_imports)) + + +class Agent: + def __init__( + self, + tools: Union[List[Tool], Toolbox], + llm_engine: Callable = None, + system_prompt: Optional[str] = None, + tool_description_template: Optional[str] = None, + additional_args: Dict = {}, + max_iterations: int = 6, + tool_parser: Optional[Callable] = None, + add_base_tools: bool = False, + verbose: int = 0, + grammar: Optional[Dict[str, str]] = None, + managed_agents: Optional[List] = None, + step_callbacks: Optional[List[Callable]] = None, + monitor_metrics: bool = True, + ): + if system_prompt is None: + system_prompt = DEFAULT_REACT_CODE_SYSTEM_PROMPT + if tool_parser is None: + tool_parser = parse_json_tool_call + self.agent_name = self.__class__.__name__ + self.llm_engine = llm_engine + self.system_prompt_template = system_prompt + self.tool_description_template = ( + tool_description_template if tool_description_template else DEFAULT_TOOL_DESCRIPTION_TEMPLATE + ) + self.additional_args = additional_args + self.max_iterations = max_iterations + self.logger = logger + self.tool_parser = tool_parser + self.grammar = grammar + + self.managed_agents = None + if managed_agents is not None: + self.managed_agents = {agent.name: agent for agent in managed_agents} + + if isinstance(tools, Toolbox): + self._toolbox = tools + if add_base_tools: + if not is_torch_available(): + raise ImportError("Using the base tools requires torch to be installed.") + + self._toolbox.add_base_tools(add_python_interpreter=(self.__class__ == ReactJsonAgent)) + else: + self._toolbox = Toolbox(tools, add_base_tools=add_base_tools) + self._toolbox.add_tool(FinalAnswerTool()) + + self.system_prompt = format_prompt_with_tools( + self._toolbox, self.system_prompt_template, self.tool_description_template + ) + self.system_prompt = format_prompt_with_managed_agents_descriptions(self.system_prompt, self.managed_agents) + self.prompt = None + self.logs = [] + self.task = None + + if verbose == 0: + logger.setLevel(logging.WARNING) + elif verbose == 1: + logger.setLevel(logging.INFO) + elif verbose == 2: + logger.setLevel(logging.DEBUG) + + # Initialize step callbacks + self.step_callbacks = step_callbacks if step_callbacks is not None else [] + + # Initialize Monitor if monitor_metrics is True + self.monitor = None + if monitor_metrics: + self.monitor = Monitor(self.llm_engine) + self.step_callbacks.append(self.monitor.update_metrics) + + @property + def toolbox(self) -> Toolbox: + """Get the toolbox currently available to the agent""" + return self._toolbox + + def initialize_for_run(self): + self.token_count = 0 + self.system_prompt = format_prompt_with_tools( + self._toolbox, + self.system_prompt_template, + self.tool_description_template, + ) + self.system_prompt = format_prompt_with_managed_agents_descriptions(self.system_prompt, self.managed_agents) + if hasattr(self, "authorized_imports"): + self.system_prompt = format_prompt_with_imports( + self.system_prompt, list(set(LIST_SAFE_MODULES) | set(self.authorized_imports)) + ) + self.logs = [{"system_prompt": self.system_prompt, "task": self.task}] + self.logger.log(33, "======== New task ========") + self.logger.log(34, self.task) + self.logger.debug("System prompt is as follows:") + self.logger.debug(self.system_prompt) + + def write_inner_memory_from_logs(self, summary_mode: Optional[bool] = False) -> List[Dict[str, str]]: + """ + Reads past llm_outputs, actions, and observations or errors from the logs into a series of messages + that can be used as input to the LLM. + """ + prompt_message = {"role": MessageRole.SYSTEM, "content": self.logs[0]["system_prompt"]} + task_message = { + "role": MessageRole.USER, + "content": "Task: " + self.logs[0]["task"], + } + if summary_mode: + memory = [task_message] + else: + memory = [prompt_message, task_message] + for i, step_log in enumerate(self.logs[1:]): + if "llm_output" in step_log and not summary_mode: + thought_message = {"role": MessageRole.ASSISTANT, "content": step_log["llm_output"].strip()} + memory.append(thought_message) + if "facts" in step_log: + thought_message = { + "role": MessageRole.ASSISTANT, + "content": "[FACTS LIST]:\n" + step_log["facts"].strip(), + } + memory.append(thought_message) + + if "plan" in step_log and not summary_mode: + thought_message = {"role": MessageRole.ASSISTANT, "content": "[PLAN]:\n" + step_log["plan"].strip()} + memory.append(thought_message) + + if "tool_call" in step_log and summary_mode: + tool_call_message = { + "role": MessageRole.ASSISTANT, + "content": f"[STEP {i} TOOL CALL]: " + str(step_log["tool_call"]).strip(), + } + memory.append(tool_call_message) + + if "task" in step_log: + tool_call_message = { + "role": MessageRole.USER, + "content": "New task:\n" + step_log["task"], + } + memory.append(tool_call_message) + + if "error" in step_log or "observation" in step_log: + if "error" in step_log: + message_content = ( + f"[OUTPUT OF STEP {i}] -> Error:\n" + + str(step_log["error"]) + + "\nNow let's retry: take care not to repeat previous errors! If you have retried several times, try a completely different approach.\n" + ) + elif "observation" in step_log: + message_content = f"[OUTPUT OF STEP {i}] -> Observation:\n{step_log['observation']}" + tool_response_message = {"role": MessageRole.TOOL_RESPONSE, "content": message_content} + memory.append(tool_response_message) + + return memory + + def get_succinct_logs(self): + return [{key: value for key, value in log.items() if key != "agent_memory"} for log in self.logs] + + def extract_action(self, llm_output: str, split_token: str) -> str: + """ + Parse action from the LLM output + + Args: + llm_output (`str`): Output of the LLM + split_token (`str`): Separator for the action. Should match the example in the system prompt. + """ + try: + split = llm_output.split(split_token) + rationale, action = ( + split[-2], + split[-1], + ) # NOTE: using indexes starting from the end solves for when you have more than one split_token in the output + except Exception as e: + self.logger.error(e, exc_info=1) + raise AgentParsingError( + f"Error: No '{split_token}' token provided in your output.\nYour output:\n{llm_output}\n. Be sure to include an action, prefaced with '{split_token}'!" + ) + return rationale.strip(), action.strip() + + def execute_tool_call(self, tool_name: str, arguments: Dict[str, str]) -> Any: + """ + Execute tool with the provided input and returns the result. + This method replaces arguments with the actual values from the state if they refer to state variables. + + Args: + tool_name (`str`): Name of the Tool to execute (should be one from self.toolbox). + arguments (Dict[str, str]): Arguments passed to the Tool. + """ + available_tools = self.toolbox.tools + if self.managed_agents is not None: + available_tools = {**available_tools, **self.managed_agents} + if tool_name not in available_tools: + error_msg = f"Error: unknown tool {tool_name}, should be instead one of {list(available_tools.keys())}." + self.logger.error(error_msg, exc_info=1) + raise AgentExecutionError(error_msg) + + try: + if isinstance(arguments, str): + observation = available_tools[tool_name](arguments) + elif isinstance(arguments, dict): + for key, value in arguments.items(): + # if the value is the name of a state variable like "image.png", replace it with the actual value + if isinstance(value, str) and value in self.state: + arguments[key] = self.state[value] + observation = available_tools[tool_name](**arguments) + else: + raise AgentExecutionError( + f"Arguments passed to tool should be a dict or string: got a {type(arguments)}." + ) + return observation + except Exception as e: + if tool_name in self.toolbox.tools: + raise AgentExecutionError( + f"Error in tool call execution: {e}\nYou should only use this tool with a correct input.\n" + f"As a reminder, this tool's description is the following:\n{get_tool_description_with_args(available_tools[tool_name])}" + ) + elif tool_name in self.managed_agents: + raise AgentExecutionError( + f"Error in calling team member: {e}\nYou should only ask this team member with a correct request.\n" + f"As a reminder, this team member's description is the following:\n{available_tools[tool_name]}" + ) + + def log_rationale_code_action(self, rationale: str, code_action: str) -> None: + self.logger.warning("=== Agent thoughts:") + self.logger.log(31, rationale) + self.logger.warning(">>> Agent is executing the code below:") + if is_pygments_available(): + self.logger.log( + 31, highlight(code_action, PythonLexer(ensurenl=False), Terminal256Formatter(style="nord")) + ) + else: + self.logger.log(31, code_action) + self.logger.warning("====") + + def run(self, **kwargs): + """To be implemented in the child class""" + raise NotImplementedError + + +class CodeAgent(Agent): + """ + A class for an agent that solves the given task using a single block of code. It plans all its actions, then executes all in one shot. + """ + + def __init__( + self, + tools: List[Tool], + llm_engine: Optional[Callable] = None, + system_prompt: Optional[str] = None, + tool_description_template: Optional[str] = None, + grammar: Optional[Dict[str, str]] = None, + additional_authorized_imports: Optional[List[str]] = None, + **kwargs, + ): + if llm_engine is None: + llm_engine = HfApiEngine() + if system_prompt is None: + system_prompt = DEFAULT_CODE_SYSTEM_PROMPT + if tool_description_template is None: + tool_description_template = DEFAULT_TOOL_DESCRIPTION_TEMPLATE + super().__init__( + tools=tools, + llm_engine=llm_engine, + system_prompt=system_prompt, + tool_description_template=tool_description_template, + grammar=grammar, + **kwargs, + ) + + if not is_pygments_available(): + transformers_logging.warning_once( + logger, + "pygments isn't installed. Installing pygments will enable color syntax highlighting in the " + "CodeAgent.", + ) + + self.python_evaluator = evaluate_python_code + self.additional_authorized_imports = additional_authorized_imports if additional_authorized_imports else [] + self.authorized_imports = list(set(LIST_SAFE_MODULES) | set(self.additional_authorized_imports)) + self.system_prompt = self.system_prompt.replace("<>", str(self.authorized_imports)) + + def parse_code_blob(self, result: str) -> str: + """ + Override this method if you want to change the way the code is + cleaned in the `run` method. + """ + return parse_code_blob(result) + + def run(self, task: str, return_generated_code: bool = False, **kwargs): + """ + Runs the agent for the given task. + + Args: + task (`str`): The task to perform + return_generated_code (`bool`, *optional*, defaults to `False`): Whether to return the generated code instead of running it + kwargs (additional keyword arguments, *optional*): + Any keyword argument to send to the agent when evaluating the code. + + Example: + + ```py + from transformers.agents import CodeAgent + + agent = CodeAgent(tools=[]) + agent.run("What is the result of 2 power 3.7384?") + ``` + """ + self.task = task + if len(kwargs) > 0: + self.task += f"\nYou have been provided with these initial arguments: {str(kwargs)}." + self.state = kwargs.copy() + self.initialize_for_run() + + # Run LLM + prompt_message = {"role": MessageRole.SYSTEM, "content": self.system_prompt} + task_message = { + "role": MessageRole.USER, + "content": "Task: " + self.task, + } + + self.prompt = [prompt_message, task_message] + self.logger.info("====Executing with this prompt====") + self.logger.info(self.prompt) + + additional_args = {"grammar": self.grammar} if self.grammar is not None else {} + llm_output = self.llm_engine(self.prompt, stop_sequences=[""], **additional_args) + + if return_generated_code: + return llm_output + + # Parse + try: + rationale, code_action = self.extract_action(llm_output=llm_output, split_token="Code:") + except Exception as e: + self.logger.debug( + f"Error in extracting action, trying to parse the whole output as code. Error trace: {e}" + ) + rationale, code_action = "", llm_output + + try: + code_action = self.parse_code_blob(code_action) + except Exception as e: + error_msg = f"Error in code parsing: {e}. Be sure to provide correct code" + self.logger.error(error_msg, exc_info=1) + return error_msg + + # Execute + self.log_rationale_code_action(rationale, code_action) + try: + available_tools = {**BASE_PYTHON_TOOLS.copy(), **self.toolbox.tools} + output = self.python_evaluator( + code_action, + static_tools=available_tools, + custom_tools={}, + state=self.state, + authorized_imports=self.authorized_imports, + ) + self.logger.info(self.state["print_outputs"]) + return output + except Exception as e: + error_msg = f"Error in execution: {e}. Be sure to provide correct code." + self.logger.error(error_msg, exc_info=1) + return error_msg + + +class ReactAgent(Agent): + """ + This agent that solves the given task step by step, using the ReAct framework: + While the objective is not reached, the agent will perform a cycle of thinking and acting. + The action will be parsed from the LLM output: it consists in calls to tools from the toolbox, with arguments chosen by the LLM engine. + """ + + @_deprecate_method( + version="4.51.0", + message="Switch to smolagents instead, with the same functionalities and similar API (https://huggingface.co/docs/smolagents/index)", + ) + def __init__( + self, + tools: List[Tool], + llm_engine: Optional[Callable] = None, + system_prompt: Optional[str] = None, + tool_description_template: Optional[str] = None, + grammar: Optional[Dict[str, str]] = None, + plan_type: Optional[str] = None, + planning_interval: Optional[int] = None, + **kwargs, + ): + if llm_engine is None: + llm_engine = HfApiEngine() + if system_prompt is None: + system_prompt = DEFAULT_REACT_CODE_SYSTEM_PROMPT + if tool_description_template is None: + tool_description_template = DEFAULT_TOOL_DESCRIPTION_TEMPLATE + if plan_type is None: + plan_type = SUPPORTED_PLAN_TYPES[0] + else: + assert plan_type in SUPPORTED_PLAN_TYPES, f"plan type {plan_type} is not supported" + super().__init__( + tools=tools, + llm_engine=llm_engine, + system_prompt=system_prompt, + tool_description_template=tool_description_template, + grammar=grammar, + **kwargs, + ) + self.planning_interval = planning_interval + self.plan_type = plan_type + + def provide_final_answer(self, task) -> str: + """ + This method provides a final answer to the task, based on the logs of the agent's interactions. + """ + self.prompt = [ + { + "role": MessageRole.SYSTEM, + "content": "An agent tried to answer an user query but it got stuck and failed to do so. You are tasked with providing an answer instead. Here is the agent's memory:", + } + ] + self.prompt += self.write_inner_memory_from_logs()[1:] + self.prompt += [ + { + "role": MessageRole.USER, + "content": f"Based on the above, please provide an answer to the following user request:\n{task}", + } + ] + try: + return self.llm_engine(self.prompt) + except Exception as e: + return f"Error in generating final llm output: {e}." + + @_deprecate_method( + version="4.51.0", + message="Switch to smolagents instead, with the same functionalities and similar API (https://huggingface.co/docs/smolagents/index)", + ) + def run(self, task: str, stream: bool = False, reset: bool = True, **kwargs): + """ + Runs the agent for the given task. + + Args: + task (`str`): The task to perform + + Example: + ```py + from transformers.agents import ReactCodeAgent + agent = ReactCodeAgent(tools=[]) + agent.run("What is the result of 2 power 3.7384?") + ``` + """ + self.task = task + if len(kwargs) > 0: + self.task += f"\nYou have been provided with these initial arguments: {str(kwargs)}." + self.state = kwargs.copy() + if reset: + self.initialize_for_run() + else: + self.logs.append({"task": task}) + if stream: + return self.stream_run(task) + else: + return self.direct_run(task) + + def stream_run(self, task: str): + """ + Runs the agent in streaming mode, yielding steps as they are executed: should be launched only in the `run` method. + """ + final_answer = None + iteration = 0 + while final_answer is None and iteration < self.max_iterations: + step_start_time = time.time() + step_log_entry = {"iteration": iteration, "start_time": step_start_time} + try: + self.step(step_log_entry) + if "final_answer" in step_log_entry: + final_answer = step_log_entry["final_answer"] + except AgentError as e: + self.logger.error(e, exc_info=1) + step_log_entry["error"] = e + finally: + step_end_time = time.time() + step_log_entry["step_end_time"] = step_end_time + step_log_entry["step_duration"] = step_end_time - step_start_time + self.logs.append(step_log_entry) + for callback in self.step_callbacks: + callback(step_log_entry) + iteration += 1 + yield step_log_entry + + if final_answer is None and iteration == self.max_iterations: + error_message = "Reached max iterations." + final_step_log = {"error": AgentMaxIterationsError(error_message)} + self.logs.append(final_step_log) + self.logger.error(error_message, exc_info=1) + final_answer = self.provide_final_answer(task) + final_step_log["final_answer"] = final_answer + final_step_log["step_duration"] = 0 + for callback in self.step_callbacks: + callback(final_step_log) + yield final_step_log + + yield final_answer + + def direct_run(self, task: str): + """ + Runs the agent in direct mode, returning outputs only at the end: should be launched only in the `run` method. + """ + final_answer = None + iteration = 0 + while final_answer is None and iteration < self.max_iterations: + step_start_time = time.time() + step_log_entry = {"iteration": iteration, "start_time": step_start_time} + try: + if self.planning_interval is not None and iteration % self.planning_interval == 0: + self.planning_step(task, is_first_step=(iteration == 0), iteration=iteration) + self.step(step_log_entry) + if "final_answer" in step_log_entry: + final_answer = step_log_entry["final_answer"] + except AgentError as e: + self.logger.error(e, exc_info=1) + step_log_entry["error"] = e + finally: + step_end_time = time.time() + step_log_entry["step_end_time"] = step_end_time + step_log_entry["step_duration"] = step_end_time - step_start_time + self.logs.append(step_log_entry) + for callback in self.step_callbacks: + callback(step_log_entry) + iteration += 1 + + if final_answer is None and iteration == self.max_iterations: + error_message = "Reached max iterations." + final_step_log = {"error": AgentMaxIterationsError(error_message)} + self.logs.append(final_step_log) + self.logger.error(error_message, exc_info=1) + final_answer = self.provide_final_answer(task) + final_step_log["final_answer"] = final_answer + final_step_log["step_duration"] = 0 + for callback in self.step_callbacks: + callback(final_step_log) + + return final_answer + + def planning_step(self, task, is_first_step: bool = False, iteration: Optional[int] = None): + """ + Used periodically by the agent to plan the next steps to reach the objective. + + Args: + task (`str`): The task to perform + is_first_step (`bool`): If this step is not the first one, the plan should be an update over a previous plan. + iteration (`int`): The number of the current step, used as an indication for the LLM. + """ + if is_first_step: + message_prompt_facts = {"role": MessageRole.SYSTEM, "content": SYSTEM_PROMPT_FACTS} + message_prompt_task = { + "role": MessageRole.USER, + "content": f"""Here is the task: +``` +{task} +``` +Now begin!""", + } + + answer_facts = self.llm_engine([message_prompt_facts, message_prompt_task]) + + message_system_prompt_plan = { + "role": MessageRole.SYSTEM, + "content": PROMPTS_FOR_INITIAL_PLAN[self.plan_type]["system"], + } + message_user_prompt_plan = { + "role": MessageRole.USER, + "content": PROMPTS_FOR_INITIAL_PLAN[self.plan_type]["user"].format( + task=task, + tool_descriptions=self._toolbox.show_tool_descriptions(self.tool_description_template), + managed_agents_descriptions=( + show_agents_descriptions(self.managed_agents) if self.managed_agents is not None else "" + ), + answer_facts=answer_facts, + ), + } + answer_plan = self.llm_engine( + [message_system_prompt_plan, message_user_prompt_plan], stop_sequences=[""] + ) + + final_plan_redaction = f"""Here is the plan of action that I will follow to solve the task: +``` +{answer_plan} +```""" + final_facts_redaction = f"""Here are the facts that I know so far: +``` +{answer_facts} +```""".strip() + self.logs.append({"plan": final_plan_redaction, "facts": final_facts_redaction}) + self.logger.log(36, "===== Initial plan =====") + self.logger.log(35, final_plan_redaction) + else: # update plan + agent_memory = self.write_inner_memory_from_logs( + summary_mode=False + ) # This will not log the plan but will log facts + + # Redact updated facts + facts_update_system_prompt = { + "role": MessageRole.SYSTEM, + "content": SYSTEM_PROMPT_FACTS_UPDATE, + } + facts_update_message = { + "role": MessageRole.USER, + "content": USER_PROMPT_FACTS_UPDATE, + } + facts_update = self.llm_engine([facts_update_system_prompt] + agent_memory + [facts_update_message]) + + # Redact updated plan + plan_update_message = { + "role": MessageRole.SYSTEM, + "content": PROMPTS_FOR_PLAN_UPDATE[self.plan_type]["system"].format(task=task), + } + plan_update_message_user = { + "role": MessageRole.USER, + "content": PROMPTS_FOR_PLAN_UPDATE[self.plan_type]["user"].format( + task=task, + tool_descriptions=self._toolbox.show_tool_descriptions(self.tool_description_template), + managed_agents_descriptions=( + show_agents_descriptions(self.managed_agents) if self.managed_agents is not None else "" + ), + facts_update=facts_update, + remaining_steps=(self.max_iterations - iteration), + ), + } + plan_update = self.llm_engine( + [plan_update_message] + agent_memory + [plan_update_message_user], stop_sequences=[""] + ) + + # Log final facts and plan + final_plan_redaction = PLAN_UPDATE_FINAL_PLAN_REDACTION.format(task=task, plan_update=plan_update) + final_facts_redaction = f"""Here is the updated list of the facts that I know: +``` +{facts_update} +```""" + self.logs.append({"plan": final_plan_redaction, "facts": final_facts_redaction}) + self.logger.log(36, "===== Updated plan =====") + self.logger.log(35, final_plan_redaction) + + +class ReactJsonAgent(ReactAgent): + """ + This agent that solves the given task step by step, using the ReAct framework: + While the objective is not reached, the agent will perform a cycle of thinking and acting. + The tool calls will be formulated by the LLM in JSON format, then parsed and executed. + """ + + def __init__( + self, + tools: List[Tool], + llm_engine: Optional[Callable] = None, + system_prompt: Optional[str] = None, + tool_description_template: Optional[str] = None, + grammar: Optional[Dict[str, str]] = None, + planning_interval: Optional[int] = None, + **kwargs, + ): + if llm_engine is None: + llm_engine = HfApiEngine() + if system_prompt is None: + system_prompt = DEFAULT_REACT_JSON_SYSTEM_PROMPT + if tool_description_template is None: + tool_description_template = DEFAULT_TOOL_DESCRIPTION_TEMPLATE + super().__init__( + tools=tools, + llm_engine=llm_engine, + system_prompt=system_prompt, + tool_description_template=tool_description_template, + grammar=grammar, + planning_interval=planning_interval, + **kwargs, + ) + + def step(self, log_entry: Dict[str, Any]): + """ + Perform one step in the ReAct framework: the agent thinks, acts, and observes the result. + The errors are raised here, they are caught and logged in the run() method. + """ + agent_memory = self.write_inner_memory_from_logs() + + self.prompt = agent_memory + self.logger.debug("===== New step =====") + + # Add new step in logs + log_entry["agent_memory"] = agent_memory.copy() + + self.logger.info("===== Calling LLM with this last message: =====") + self.logger.info(self.prompt[-1]) + + try: + additional_args = {"grammar": self.grammar} if self.grammar is not None else {} + llm_output = self.llm_engine( + self.prompt, stop_sequences=["", "Observation:"], **additional_args + ) + except Exception as e: + raise AgentGenerationError(f"Error in generating llm output: {e}.") + self.logger.debug("===== Output message of the LLM: =====") + self.logger.debug(llm_output) + log_entry["llm_output"] = llm_output + + # Parse + self.logger.debug("===== Extracting action =====") + rationale, action = self.extract_action(llm_output=llm_output, split_token="Action:") + + try: + tool_name, arguments = self.tool_parser(action) + except Exception as e: + raise AgentParsingError(f"Could not parse the given action: {e}.") + + log_entry["rationale"] = rationale + log_entry["tool_call"] = {"tool_name": tool_name, "tool_arguments": arguments} + + # Execute + self.logger.warning("=== Agent thoughts:") + self.logger.log(31, rationale) + self.logger.warning(f">>> Calling tool: '{tool_name}' with arguments: {arguments}") + if tool_name == "final_answer": + if isinstance(arguments, dict): + if "answer" in arguments: + answer = arguments["answer"] + if ( + isinstance(answer, str) and answer in self.state.keys() + ): # if the answer is a state variable, return the value + answer = self.state[answer] + else: + answer = arguments + else: + answer = arguments + log_entry["final_answer"] = answer + return answer + else: + if arguments is None: + arguments = {} + observation = self.execute_tool_call(tool_name, arguments) + observation_type = type(observation) + if observation_type in [AgentImage, AgentAudio]: + if observation_type == AgentImage: + observation_name = "image.png" + elif observation_type == AgentAudio: + observation_name = "audio.mp3" + # TODO: observation naming could allow for different names of same type + + self.state[observation_name] = observation + updated_information = f"Stored '{observation_name}' in memory." + else: + updated_information = str(observation).strip() + self.logger.info(updated_information) + log_entry["observation"] = updated_information + return log_entry + + +class ReactCodeAgent(ReactAgent): + """ + This agent that solves the given task step by step, using the ReAct framework: + While the objective is not reached, the agent will perform a cycle of thinking and acting. + The tool calls will be formulated by the LLM in code format, then parsed and executed. + """ + + def __init__( + self, + tools: List[Tool], + llm_engine: Optional[Callable] = None, + system_prompt: Optional[str] = None, + tool_description_template: Optional[str] = None, + grammar: Optional[Dict[str, str]] = None, + additional_authorized_imports: Optional[List[str]] = None, + planning_interval: Optional[int] = None, + **kwargs, + ): + if llm_engine is None: + llm_engine = HfApiEngine() + if system_prompt is None: + system_prompt = DEFAULT_REACT_CODE_SYSTEM_PROMPT + if tool_description_template is None: + tool_description_template = DEFAULT_TOOL_DESCRIPTION_TEMPLATE + super().__init__( + tools=tools, + llm_engine=llm_engine, + system_prompt=system_prompt, + tool_description_template=tool_description_template, + grammar=grammar, + planning_interval=planning_interval, + **kwargs, + ) + + if not is_pygments_available(): + transformers_logging.warning_once( + logger, + "pygments isn't installed. Installing pygments will enable color syntax highlighting in the " + "ReactCodeAgent.", + ) + + self.python_evaluator = evaluate_python_code + self.additional_authorized_imports = additional_authorized_imports if additional_authorized_imports else [] + self.authorized_imports = list(set(LIST_SAFE_MODULES) | set(self.additional_authorized_imports)) + self.system_prompt = self.system_prompt.replace("<>", str(self.authorized_imports)) + self.custom_tools = {} + + def step(self, log_entry: Dict[str, Any]): + """ + Perform one step in the ReAct framework: the agent thinks, acts, and observes the result. + The errors are raised here, they are caught and logged in the run() method. + """ + agent_memory = self.write_inner_memory_from_logs() + + self.prompt = agent_memory.copy() + self.logger.debug("===== New step =====") + + # Add new step in logs + log_entry["agent_memory"] = agent_memory.copy() + + self.logger.info("===== Calling LLM with these last messages: =====") + self.logger.info(self.prompt[-2:]) + + try: + additional_args = {"grammar": self.grammar} if self.grammar is not None else {} + llm_output = self.llm_engine( + self.prompt, stop_sequences=["", "Observation:"], **additional_args + ) + except Exception as e: + raise AgentGenerationError(f"Error in generating llm output: {e}.") + + self.logger.debug("=== Output message of the LLM:") + self.logger.debug(llm_output) + log_entry["llm_output"] = llm_output + + # Parse + self.logger.debug("=== Extracting action ===") + try: + rationale, raw_code_action = self.extract_action(llm_output=llm_output, split_token="Code:") + except Exception as e: + self.logger.debug(f"Error in extracting action, trying to parse the whole output. Error trace: {e}") + rationale, raw_code_action = llm_output, llm_output + + try: + code_action = parse_code_blob(raw_code_action) + except Exception as e: + error_msg = f"Error in code parsing: {e}. Make sure to provide correct code" + raise AgentParsingError(error_msg) + + log_entry["rationale"] = rationale + log_entry["tool_call"] = {"tool_name": "code interpreter", "tool_arguments": code_action} + + # Execute + self.log_rationale_code_action(rationale, code_action) + try: + static_tools = { + **BASE_PYTHON_TOOLS.copy(), + **self.toolbox.tools, + } + if self.managed_agents is not None: + static_tools = {**static_tools, **self.managed_agents} + result = self.python_evaluator( + code_action, + static_tools=static_tools, + custom_tools=self.custom_tools, + state=self.state, + authorized_imports=self.authorized_imports, + ) + self.logger.warning("Print outputs:") + self.logger.log(32, self.state["print_outputs"]) + observation = "Print outputs:\n" + self.state["print_outputs"] + if result is not None: + self.logger.warning("Last output from code snippet:") + self.logger.log(32, str(result)) + observation += "Last output from code snippet:\n" + str(result)[:100000] + log_entry["observation"] = observation + except Exception as e: + error_msg = f"Code execution failed due to the following error:\n{str(e)}" + if "'dict' object has no attribute 'read'" in str(e): + error_msg += "\nYou get this error because you passed a dict as input for one of the arguments instead of a string." + raise AgentExecutionError(error_msg) + for line in code_action.split("\n"): + if line[: len("final_answer")] == "final_answer": + self.logger.log(33, "Final answer:") + self.logger.log(32, result) + log_entry["final_answer"] = result + return result + + +LENGTH_TRUNCATE_REPORTS = 1000 + + +class ManagedAgent: + @_deprecate_method( + version="4.51.0", + message="Switch to smolagents instead, with the same functionalities and similar API (https://huggingface.co/docs/smolagents/index)", + ) + def __init__(self, agent, name, description, additional_prompting=None, provide_run_summary=False): + self.agent = agent + self.name = name + self.description = description + self.additional_prompting = additional_prompting + self.provide_run_summary = provide_run_summary + + def write_full_task(self, task): + full_task = f"""You're a helpful agent named '{self.name}'. +You have been submitted this task by your manager. +--- +Task: +{task} +--- +You're helping your manager solve a wider task: so make sure to not provide a one-line answer, but give as much information as possible so that they have a clear understanding of the answer. + +Your final_answer WILL HAVE to contain these parts: +### 1. Task outcome (short version): +### 2. Task outcome (extremely detailed version): +### 3. Additional context (if relevant): + +Put all these in your final_answer tool, everything that you do not pass as an argument to final_answer will be lost. +And even if your task resolution is not successful, please return as much context as possible, so that your manager can act upon this feedback. +<>""" + if self.additional_prompting: + full_task = full_task.replace("\n<>", self.additional_prompting).strip() + else: + full_task = full_task.replace("\n<>", "").strip() + return full_task + + def __call__(self, request, **kwargs): + full_task = self.write_full_task(request) + output = self.agent.run(full_task, **kwargs) + if self.provide_run_summary: + answer = f"Here is the final answer from your managed agent '{self.name}':\n" + answer += str(output) + answer += f"\n\nFor more detail, find below a summary of this agent's work:\nSUMMARY OF WORK FROM AGENT '{self.name}':\n" + for message in self.agent.write_inner_memory_from_logs(summary_mode=True): + content = message["content"] + if len(str(content)) < LENGTH_TRUNCATE_REPORTS or "[FACTS LIST]" in str(content): + answer += "\n" + str(content) + "\n---" + else: + answer += ( + "\n" + + str(content)[:LENGTH_TRUNCATE_REPORTS] + + "\n(...Step was truncated because too long)...\n---" + ) + answer += f"\nEND OF SUMMARY OF WORK FROM AGENT '{self.name}'." + return answer + else: + return output diff --git a/lib/python3.12/site-packages/transformers/agents/default_tools.py b/lib/python3.12/site-packages/transformers/agents/default_tools.py new file mode 100644 index 0000000000000000000000000000000000000000..3946aa9f873503c9a9564da2d5e0dbba7773db89 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/default_tools.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2024 The HuggingFace Inc. 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. +import importlib.util +import json +import math +from dataclasses import dataclass +from math import sqrt +from typing import Dict + +from huggingface_hub import hf_hub_download, list_spaces + +from ..utils import is_offline_mode +from .python_interpreter import LIST_SAFE_MODULES, evaluate_python_code +from .tools import TOOL_CONFIG_FILE, TOOL_MAPPING, Tool + + +def custom_print(*args): + return None + + +BASE_PYTHON_TOOLS = { + "print": custom_print, + "isinstance": isinstance, + "range": range, + "float": float, + "int": int, + "bool": bool, + "str": str, + "set": set, + "list": list, + "dict": dict, + "tuple": tuple, + "round": round, + "ceil": math.ceil, + "floor": math.floor, + "log": math.log, + "exp": math.exp, + "sin": math.sin, + "cos": math.cos, + "tan": math.tan, + "asin": math.asin, + "acos": math.acos, + "atan": math.atan, + "atan2": math.atan2, + "degrees": math.degrees, + "radians": math.radians, + "pow": math.pow, + "sqrt": sqrt, + "len": len, + "sum": sum, + "max": max, + "min": min, + "abs": abs, + "enumerate": enumerate, + "zip": zip, + "reversed": reversed, + "sorted": sorted, + "all": all, + "any": any, + "map": map, + "filter": filter, + "ord": ord, + "chr": chr, + "next": next, + "iter": iter, + "divmod": divmod, + "callable": callable, + "getattr": getattr, + "hasattr": hasattr, + "setattr": setattr, + "issubclass": issubclass, + "type": type, +} + + +@dataclass +class PreTool: + name: str + inputs: Dict[str, str] + output_type: type + task: str + description: str + repo_id: str + + +HUGGINGFACE_DEFAULT_TOOLS_FROM_HUB = [ + "image-transformation", + "text-to-image", +] + + +def get_remote_tools(logger, organization="huggingface-tools"): + if is_offline_mode(): + logger.info("You are in offline mode, so remote tools are not available.") + return {} + + spaces = list_spaces(author=organization) + tools = {} + for space_info in spaces: + repo_id = space_info.id + resolved_config_file = hf_hub_download(repo_id, TOOL_CONFIG_FILE, repo_type="space") + with open(resolved_config_file, encoding="utf-8") as reader: + config = json.load(reader) + task = repo_id.split("/")[-1] + tools[config["name"]] = PreTool( + task=task, + description=config["description"], + repo_id=repo_id, + name=task, + inputs=config["inputs"], + output_type=config["output_type"], + ) + + return tools + + +def setup_default_tools(logger): + default_tools = {} + main_module = importlib.import_module("transformers") + tools_module = main_module.agents + + for task_name, tool_class_name in TOOL_MAPPING.items(): + tool_class = getattr(tools_module, tool_class_name) + tool_instance = tool_class() + default_tools[tool_class.name] = PreTool( + name=tool_instance.name, + inputs=tool_instance.inputs, + output_type=tool_instance.output_type, + task=task_name, + description=tool_instance.description, + repo_id=None, + ) + + return default_tools + + +class PythonInterpreterTool(Tool): + name = "python_interpreter" + description = "This is a tool that evaluates python code. It can be used to perform calculations." + + output_type = "string" + + def __init__(self, *args, authorized_imports=None, **kwargs): + if authorized_imports is None: + self.authorized_imports = list(set(LIST_SAFE_MODULES)) + else: + self.authorized_imports = list(set(LIST_SAFE_MODULES) | set(authorized_imports)) + self.inputs = { + "code": { + "type": "string", + "description": ( + "The code snippet to evaluate. All variables used in this snippet must be defined in this same snippet, " + f"else you will get an error. This code can only import the following python libraries: {authorized_imports}." + ), + } + } + super().__init__(*args, **kwargs) + + def forward(self, code): + output = str( + evaluate_python_code(code, static_tools=BASE_PYTHON_TOOLS, authorized_imports=self.authorized_imports) + ) + return output + + +class FinalAnswerTool(Tool): + name = "final_answer" + description = "Provides a final answer to the given problem." + inputs = {"answer": {"type": "any", "description": "The final answer to the problem"}} + output_type = "any" + + def forward(self, answer): + return answer diff --git a/lib/python3.12/site-packages/transformers/agents/document_question_answering.py b/lib/python3.12/site-packages/transformers/agents/document_question_answering.py new file mode 100644 index 0000000000000000000000000000000000000000..23ae5b0429120dd6dd9c2eb7f454959c92625012 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/document_question_answering.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2023 The HuggingFace Inc. 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. +import re + +import numpy as np +import torch + +from ..models.auto import AutoProcessor +from ..models.vision_encoder_decoder import VisionEncoderDecoderModel +from ..utils import is_vision_available +from .tools import PipelineTool + + +if is_vision_available(): + from PIL import Image + + +class DocumentQuestionAnsweringTool(PipelineTool): + default_checkpoint = "naver-clova-ix/donut-base-finetuned-docvqa" + description = "This is a tool that answers a question about an document (pdf). It returns a string that contains the answer to the question." + name = "document_qa" + pre_processor_class = AutoProcessor + model_class = VisionEncoderDecoderModel + + inputs = { + "document": { + "type": "image", + "description": "The image containing the information. Can be a PIL Image or a string path to the image.", + }, + "question": {"type": "string", "description": "The question in English"}, + } + output_type = "string" + + def __init__(self, *args, **kwargs): + if not is_vision_available(): + raise ValueError("Pillow must be installed to use the DocumentQuestionAnsweringTool.") + + super().__init__(*args, **kwargs) + + def encode(self, document: "Image", question: str): + task_prompt = "{user_input}" + prompt = task_prompt.replace("{user_input}", question) + decoder_input_ids = self.pre_processor.tokenizer( + prompt, add_special_tokens=False, return_tensors="pt" + ).input_ids + if isinstance(document, str): + img = Image.open(document).convert("RGB") + img_array = np.array(img).transpose(2, 0, 1) + document = torch.from_numpy(img_array) + pixel_values = self.pre_processor(document, return_tensors="pt").pixel_values + + return {"decoder_input_ids": decoder_input_ids, "pixel_values": pixel_values} + + def forward(self, inputs): + return self.model.generate( + inputs["pixel_values"].to(self.device), + decoder_input_ids=inputs["decoder_input_ids"].to(self.device), + max_length=self.model.decoder.config.max_position_embeddings, + early_stopping=True, + pad_token_id=self.pre_processor.tokenizer.pad_token_id, + eos_token_id=self.pre_processor.tokenizer.eos_token_id, + use_cache=True, + num_beams=1, + bad_words_ids=[[self.pre_processor.tokenizer.unk_token_id]], + return_dict_in_generate=True, + ).sequences + + def decode(self, outputs): + sequence = self.pre_processor.batch_decode(outputs)[0] + sequence = sequence.replace(self.pre_processor.tokenizer.eos_token, "") + sequence = sequence.replace(self.pre_processor.tokenizer.pad_token, "") + sequence = re.sub(r"<.*?>", "", sequence, count=1).strip() # remove first task start token + sequence = self.pre_processor.token2json(sequence) + + return sequence["answer"] diff --git a/lib/python3.12/site-packages/transformers/agents/evaluate_agent.py b/lib/python3.12/site-packages/transformers/agents/evaluate_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..555acef28c8514f0d06c35284aa86e43329fac94 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/evaluate_agent.py @@ -0,0 +1,414 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2023 The HuggingFace Inc. 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. +from .agents import BASE_PYTHON_TOOLS +from .python_interpreter import InterpreterError, evaluate + + +### Fake tools for test +def classifier(text, labels): + return f"This is the classification of {text} along {labels}." + + +def translator(text, src_lang, tgt_lang): + return f"This is the translation of {text} from {src_lang} to {tgt_lang}." + + +def speaker(text): + return f"This is actually a sound reading {text}." + + +def transcriber(audio): + if "sound" not in audio: + raise ValueError(f"`audio` ({audio}) is not a sound.") + return f"This is the transcribed text from {audio}." + + +def image_generator(prompt): + return f"This is actually an image representing {prompt}." + + +def image_captioner(image): + if "image" not in image: + raise ValueError(f"`image` ({image}) is not an image.") + return f"This is a description of {image}." + + +def image_transformer(image, prompt): + if "image" not in image: + raise ValueError(f"`image` ({image}) is not an image.") + return f"This is a transformation of {image} according to {prompt}." + + +def question_answerer(text, question): + return f"This is the answer to {question} from {text}." + + +def image_qa(image, question): + if "image" not in image: + raise ValueError(f"`image` ({image}) is not an image.") + return f"This is the answer to {question} from {image}." + + +def text_downloader(url): + return f"This is the content of {url}." + + +def summarizer(text): + return f"This is a summary of {text}." + + +def video_generator(prompt, seconds=2): + return f"A video of {prompt}" + + +def document_qa(image, question): + return f"This is the answer to {question} from the document {image}." + + +def image_segmenter(image, prompt): + return f"This is the mask of {prompt} in {image}" + + +TEST_TOOLS = { + "text_classifier": classifier, + "translator": translator, + "text_reader": speaker, + "summarizer": summarizer, + "transcriber": transcriber, + "image_generator": image_generator, + "image_captioner": image_captioner, + "image_transformer": image_transformer, + "text_qa": question_answerer, + "text_downloader": text_downloader, + "image_qa": image_qa, + "video_generator": video_generator, + "document_qa": document_qa, + "image_segmenter": image_segmenter, +} + + +class Problem: + """ + A class regrouping all the information to solve a problem on which we will evaluate agents. + + Args: + task (`str` ou `list[str]`): + One or several descriptions of the task to perform. If a list, it should contain variations on the + phrasing, but for the same task. + inputs (`list[str]` or `dict[str, str]`): + The inputs that will be fed to the tools. For this testing environment, only strings are accepted as + values. Pass along a dictionary when you want to specify the values of each inputs, or just the list of + inputs expected (the value used will be `<>` in this case). + answer (`str` or `list[str]`): + The theoretical answer (or list of possible valid answers) to the problem, as code. + """ + + def __init__(self, task, inputs, answer): + self.task = task + self.inputs = inputs + self.answer = answer + + +### The list of problems the agent will be evaluated on. +EVALUATION_TASKS = [ + Problem( + task=[ + "Is the following `text` (in Spanish) positive or negative?", + "Is the text in the variable `text` (in Spanish) positive or negative?", + "Translate the following `text` from Spanish to English then tell me if its positive or negative.", + ], + inputs=["text"], + answer="""text_classifier(translator(text, src_lang="Spanish", tgt_lang="English"), labels=["positive", "negative"])""", + ), + Problem( + task=[ + "Tell me out loud what the `image` contains.", + "Describe the following `image` out loud.", + "Find what is in the picture stored in `image` then read it out loud.", + ], + inputs=["image"], + answer=[ + "text_reader(image_captioner(image))", + "text_reader(image_qa(image, question='What is in the image?'))", + ], + ), + Problem( + task=[ + "Generate an image from the text given in `text_input`. Then transform it according to the text in `prompt`.", + "Use the following `text_input` to generate an image, then transform it by using the text in `prompt`.", + ], + inputs=["text_input", "prompt"], + answer="image_transformer(image_generator(text_input), prompt)", + ), + Problem( + task=[ + "Download the content of `url`, summarize it then generate an image from its content.", + "Use a summary of the web page at `url` to generate an image.", + "Summarize the content of the web page at `url`, and use the result to generate an image.", + ], + inputs=["url"], + answer="image_generator(summarizer(text_downloader(url)))", + ), + Problem( + task=[ + "Transform the following `image` using the prompt in `text`. The prompt is in Spanish.", + "Use the text prompt in `text` (in Spanish) to transform the following `image`.", + "Translate the `text` from Spanish to English then use it to transform the picture in `image`.", + ], + inputs=["text", "image"], + answer="image_transformer(image, translator(text, src_lang='Spanish', tgt_lang='English'))", + ), + Problem( + task=[ + "Download the content of `url`, summarize it then read it out loud to me.", + "Read me a summary of the web page at `url`.", + ], + inputs=["url"], + answer="text_reader(summarizer(text_downloader(url)))", + ), + Problem( + task=[ + "Generate an image from the text given in `text_input`.", + ], + inputs=["text_input"], + answer="image_generator(text_input)", + ), + Problem( + task=[ + "Replace the beaver in the `image` by the `prompt`.", + "Transform the `image` so that it contains the `prompt`.", + "Use `prompt` to transform this `image`.", + ], + inputs=["image", "prompt"], + answer="image_transformer(image, prompt)", + ), + Problem( + task=[ + "Provide me the summary of the `text`, then read it to me before transcribing it and translating it in French.", + "Summarize `text`, read it out loud then transcribe the audio and translate it in French.", + "Read me a summary of the `text` out loud. Transcribe this and translate it in French.", + ], + inputs=["text"], + answer="translator(transcriber(text_reader(summarizer(text))), src_lang='English', tgt_lang='French')", + ), + Problem( + task=["Generate a video of the `prompt`", "Animate a `prompt`", "Make me a short video using `prompt`."], + inputs={"prompt": "A lobster swimming"}, + answer="video_generator('A lobster swimming')", + ), + Problem( + task=[ + "Download the following file `url`, summarize it in a few words and generate a video from it." + "Fetch the file at this `url`, summarize it, and create an animation out of it." + ], + inputs=["url"], + answer="video_generator(summarizer(text_downloader(url)))", + ), +] + + +def get_theoretical_tools(agent_answer, theoretical_answer, code_answer): + if not isinstance(theoretical_answer, list): + return {name for name in TEST_TOOLS if name in code_answer} + + if isinstance(agent_answer, dict): + for one_answer, one_code in zip(theoretical_answer, code_answer): + if one_answer in agent_answer.values(): + return {name for name in TEST_TOOLS if name in one_code} + + for one_answer, one_code in zip(theoretical_answer, code_answer): + if agent_answer == one_answer: + return {name for name in TEST_TOOLS if name in one_code} + + return {name for name in TEST_TOOLS if name in code_answer[0]} + + +def evaluate_code(code, inputs=None, state=None, verbose=False, return_interpretor_error=False): + tools = BASE_PYTHON_TOOLS.copy() + for name, tool in TEST_TOOLS.items(): + if name not in code: + continue + tools[name] = tool + + if isinstance(inputs, dict): + inputs = inputs.copy() + elif inputs is not None: + inputs = {inp: f"<<{inp}>>" for inp in inputs} + + if state is not None: + state.update(inputs) + else: + state = inputs + + try: + return evaluate(code, tools, state) + except InterpreterError as e: + return str(e) + except Exception as e: + if verbose: + print(e) + return None + + +def score_code(agent_answer, theoretical_answer, verbose: bool = False): + if verbose: + print(agent_answer, theoretical_answer) + theoretical_answer = theoretical_answer if isinstance(theoretical_answer, list) else [theoretical_answer] + + if agent_answer in theoretical_answer: + if verbose: + print("Perfect!") + return 1 + elif isinstance(agent_answer, dict) and any(v in theoretical_answer for v in agent_answer.values()): + if verbose: + print("Almost perfect, result in state!") + return 0.75 + else: + if verbose: + print("Result is not the right one but code executed.") + return 0.3 + + +def evaluate_one_result(code, agent_answer, theoretical_answer, answer, verbose=False): + tools_in_code = {name for name in TEST_TOOLS if f"`{name}`" in code} + theoretical_tools = get_theoretical_tools(agent_answer, theoretical_answer, answer) + if tools_in_code == theoretical_tools: + tool_selection_score = 1.0 + tool_selection_errors = None + else: + missing_tools = len(theoretical_tools - tools_in_code) + unexpected_tools = len(tools_in_code - theoretical_tools) + tool_selection_score = max(0, 1.0 - 0.25 * missing_tools - 0.25 * unexpected_tools) + + tool_selection_errors = { + "selected_tools": tools_in_code, + "theoretical_tools": theoretical_tools, + } + + tools_in_code = {name for name in TEST_TOOLS if name in code} + if tools_in_code == theoretical_tools: + tool_used_score = 1.0 + tool_used_errors = None + else: + missing_tools = len(theoretical_tools - tools_in_code) + unexpected_tools = len(tools_in_code - theoretical_tools) + tool_used_score = max(0, 1.0 - 0.25 * missing_tools - 0.25 * unexpected_tools) + + tool_used_errors = { + "selected_tools": tools_in_code, + "theoretical_tools": theoretical_tools, + } + + score = score_code(agent_answer, theoretical_answer, verbose=verbose) + if score < 1.0: + code_errors = { + "code_produced": code, + "evaluation": agent_answer, + "theoretical_answer": theoretical_answer, + } + else: + code_errors = None + + return (tool_selection_score, tool_used_score, score), (tool_selection_errors, tool_used_errors, code_errors) + + +def evaluate_agent(agent, batch_size=8, verbose=False, return_errors=False): + """ + Evaluates a new agent on all `EVALUATION_TASKS`. + + Example: + + ```py + agent = NewOpenAiAgent(model="text-davinci-003", api_key=your_api_key) + bads = new_evaluate_agent(agent) + for bad in bads: + print(bad) + ``` + """ + # Sanity check + agent_tools = set(agent.toolbox.keys()) + if agent_tools != set(TEST_TOOLS): + missing_tools = set(TEST_TOOLS) - agent_tools + unexpected_tools = set(agent_tools) - TEST_TOOLS + raise ValueError( + f"Fix the test tools in the evaluate_agent module. Tools missing: {missing_tools}. Extra tools: {unexpected_tools}." + ) + + eval_tasks = [] + eval_idx = [] + for idx, pb in enumerate(EVALUATION_TASKS): + if isinstance(pb.task, list): + eval_tasks.extend(pb.task) + eval_idx.extend([idx] * len(pb.task)) + else: + eval_tasks.append(pb.task) + eval_idx.append(idx) + + tool_selection_score = 0 + tool_used_score = 0 + code_score = 0 + + if return_errors: + tool_selection_errors = {} + tool_used_errors = {} + code_errors = {} + + for start_idx in range(0, len(eval_tasks), batch_size): + end_idx = min(start_idx + batch_size, len(eval_tasks)) + batch_tasks = eval_tasks[start_idx:end_idx] + + results = [agent.run(task, return_generated_code=True) for task in batch_tasks] + + for idx, result in enumerate(results): + problem = EVALUATION_TASKS[eval_idx[start_idx + idx]] + if verbose: + print(f"====Task {start_idx + idx}====\n{batch_tasks[idx]}\n") + code = agent.extract_action(result, split_token="Answer:") + + # Evaluate agent answer and code answer + agent_answer = evaluate_code(code, problem.inputs, verbose=verbose) + if isinstance(problem.answer, list): + theoretical_answer = [evaluate_code(answer, problem.inputs) for answer in problem.answer] + else: + theoretical_answer = evaluate_code(problem.answer, problem.inputs) + + scores, errors = evaluate_one_result( + code, agent_answer, theoretical_answer, problem.answer, verbose=verbose + ) + + tool_selection_score += scores[0] + tool_used_score += scores[1] + code_score += scores[2] + + if return_errors: + if errors[0] is not None: + tool_selection_errors[batch_tasks[idx]] = errors[0] + if errors[1] is not None: + tool_used_errors[batch_tasks[idx]] = errors[1] + if errors[2] is not None: + code_errors[batch_tasks[idx]] = errors[2] + + scores = { + "tool selection score": 100 * (tool_selection_score / len(eval_tasks)), + "tool used score": 100 * (tool_used_score / len(eval_tasks)), + "code score": 100 * (code_score / len(eval_tasks)), + } + + if return_errors: + return scores, tool_selection_errors, tool_used_errors, code_errors + else: + return scores diff --git a/lib/python3.12/site-packages/transformers/agents/image_question_answering.py b/lib/python3.12/site-packages/transformers/agents/image_question_answering.py new file mode 100644 index 0000000000000000000000000000000000000000..832f77a4ebe72597543a52be8b38e63508bab310 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/image_question_answering.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2023 The HuggingFace Inc. 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. + +import torch +from PIL import Image + +from ..models.auto import AutoModelForVisualQuestionAnswering, AutoProcessor +from ..utils import requires_backends +from .tools import PipelineTool + + +class ImageQuestionAnsweringTool(PipelineTool): + default_checkpoint = "dandelin/vilt-b32-finetuned-vqa" + description = ( + "This is a tool that answers a question about an image. It returns a text that is the answer to the question." + ) + name = "image_qa" + pre_processor_class = AutoProcessor + model_class = AutoModelForVisualQuestionAnswering + + inputs = { + "image": { + "type": "image", + "description": "The image containing the information. Can be a PIL Image or a string path to the image.", + }, + "question": {"type": "string", "description": "The question in English"}, + } + output_type = "string" + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + super().__init__(*args, **kwargs) + + def encode(self, image: "Image", question: str): + return self.pre_processor(image, question, return_tensors="pt") + + def forward(self, inputs): + with torch.no_grad(): + return self.model(**inputs).logits + + def decode(self, outputs): + idx = outputs.argmax(-1).item() + return self.model.config.id2label[idx] diff --git a/lib/python3.12/site-packages/transformers/agents/llm_engine.py b/lib/python3.12/site-packages/transformers/agents/llm_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..49c96b95797144b683779a47ef4722593e093a38 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/llm_engine.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2024 The HuggingFace Inc. 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. +from copy import deepcopy +from enum import Enum +from typing import Dict, List, Optional + +from huggingface_hub import InferenceClient +from huggingface_hub.utils._deprecation import _deprecate_method + +from .. import AutoTokenizer +from ..pipelines.base import Pipeline +from ..utils import logging + + +logger = logging.get_logger(__name__) + + +class MessageRole(str, Enum): + USER = "user" + ASSISTANT = "assistant" + SYSTEM = "system" + TOOL_CALL = "tool-call" + TOOL_RESPONSE = "tool-response" + + @classmethod + def roles(cls): + return [r.value for r in cls] + + +def get_clean_message_list(message_list: List[Dict[str, str]], role_conversions: Dict[str, str] = {}): + """ + Subsequent messages with the same role will be concatenated to a single message. + + Args: + message_list (`List[Dict[str, str]]`): List of chat messages. + """ + final_message_list = [] + message_list = deepcopy(message_list) # Avoid modifying the original list + for message in message_list: + if not set(message.keys()) == {"role", "content"}: + raise ValueError("Message should contain only 'role' and 'content' keys!") + + role = message["role"] + if role not in MessageRole.roles(): + raise ValueError(f"Incorrect role {role}, only {MessageRole.roles()} are supported for now.") + + if role in role_conversions: + message["role"] = role_conversions[role] + + if len(final_message_list) > 0 and message["role"] == final_message_list[-1]["role"]: + final_message_list[-1]["content"] += "\n=======\n" + message["content"] + else: + final_message_list.append(message) + return final_message_list + + +llama_role_conversions = { + MessageRole.TOOL_RESPONSE: MessageRole.USER, +} + + +class HfEngine: + @_deprecate_method( + version="4.51.0", + message="Switch to smolagents instead, with the same functionalities and similar API (https://huggingface.co/docs/smolagents/index)", + ) + def __init__(self, model_id: Optional[str] = None): + self.last_input_token_count = None + self.last_output_token_count = None + if model_id is None: + model_id = "HuggingFaceTB/SmolLM2-1.7B-Instruct" + logger.warning(f"Using default model for token counting: '{model_id}'") + try: + self.tokenizer = AutoTokenizer.from_pretrained(model_id) + except Exception as e: + logger.warning(f"Failed to load tokenizer for model {model_id}: {e}. Loading default tokenizer instead.") + self.tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-1.7B-Instruct") + + def get_token_counts(self): + return { + "input_token_count": self.last_input_token_count, + "output_token_count": self.last_output_token_count, + } + + def generate( + self, messages: List[Dict[str, str]], stop_sequences: Optional[List[str]] = None, grammar: Optional[str] = None + ): + raise NotImplementedError + + def __call__( + self, messages: List[Dict[str, str]], stop_sequences: Optional[List[str]] = None, grammar: Optional[str] = None + ) -> str: + """Process the input messages and return the model's response. + + This method sends a list of messages to the Hugging Face Inference API, optionally with stop sequences and grammar customization. + + Parameters: + messages (`List[Dict[str, str]]`): + A list of message dictionaries to be processed. Each dictionary should have the structure `{"role": "user/system", "content": "message content"}`. + stop_sequences (`List[str]`, *optional*): + A list of strings that will stop the generation if encountered in the model's output. + grammar (`str`, *optional*): + The grammar or formatting structure to use in the model's response. + + Returns: + `str`: The text content of the model's response. + + Example: + ```python + >>> engine = HfApiEngine( + ... model="meta-llama/Meta-Llama-3.1-8B-Instruct", + ... token="your_hf_token_here", + ... max_tokens=2000 + ... ) + >>> messages = [{"role": "user", "content": "Explain quantum mechanics in simple terms."}] + >>> response = engine(messages, stop_sequences=["END"]) + >>> print(response) + "Quantum mechanics is the branch of physics that studies..." + ``` + """ + if not isinstance(messages, List): + raise ValueError("Messages should be a list of dictionaries with 'role' and 'content' keys.") + if stop_sequences is None: + stop_sequences = [] + response = self.generate(messages, stop_sequences, grammar) + self.last_input_token_count = len(self.tokenizer.apply_chat_template(messages, tokenize=True)) + self.last_output_token_count = len(self.tokenizer.encode(response)) + + # Remove stop sequences from LLM output + for stop_seq in stop_sequences: + if response[-len(stop_seq) :] == stop_seq: + response = response[: -len(stop_seq)] + return response + + +class HfApiEngine(HfEngine): + """A class to interact with Hugging Face's Inference API for language model interaction. + + This engine allows you to communicate with Hugging Face's models using the Inference API. It can be used in both serverless mode or with a dedicated endpoint, supporting features like stop sequences and grammar customization. + + Parameters: + model (`str`, *optional*, defaults to `"meta-llama/Meta-Llama-3.1-8B-Instruct"`): + The Hugging Face model ID to be used for inference. This can be a path or model identifier from the Hugging Face model hub. + token (`str`, *optional*): + Token used by the Hugging Face API for authentication. + If not provided, the class will use the token stored in the Hugging Face CLI configuration. + max_tokens (`int`, *optional*, defaults to 1500): + The maximum number of tokens allowed in the output. + timeout (`int`, *optional*, defaults to 120): + Timeout for the API request, in seconds. + + Raises: + ValueError: + If the model name is not provided. + """ + + def __init__( + self, + model: str = "meta-llama/Meta-Llama-3.1-8B-Instruct", + token: Optional[str] = None, + max_tokens: Optional[int] = 1500, + timeout: Optional[int] = 120, + ): + super().__init__(model_id=model) + self.model = model + self.client = InferenceClient(self.model, token=token, timeout=timeout) + self.max_tokens = max_tokens + + def generate( + self, messages: List[Dict[str, str]], stop_sequences: Optional[List[str]] = None, grammar: Optional[str] = None + ) -> str: + # Get clean message list + messages = get_clean_message_list(messages, role_conversions=llama_role_conversions) + + # Send messages to the Hugging Face Inference API + if grammar is not None: + response = self.client.chat_completion( + messages, stop=stop_sequences, max_tokens=self.max_tokens, response_format=grammar + ) + else: + response = self.client.chat_completion(messages, stop=stop_sequences, max_tokens=self.max_tokens) + + response = response.choices[0].message.content + return response + + +class TransformersEngine(HfEngine): + """This engine uses a pre-initialized local text-generation pipeline.""" + + def __init__(self, pipeline: Pipeline, model_id: Optional[str] = None): + super().__init__(model_id) + self.pipeline = pipeline + + def generate( + self, + messages: List[Dict[str, str]], + stop_sequences: Optional[List[str]] = None, + grammar: Optional[str] = None, + max_length: int = 1500, + ) -> str: + # Get clean message list + messages = get_clean_message_list(messages, role_conversions=llama_role_conversions) + + # Get LLM output + if stop_sequences is not None and len(stop_sequences) > 0: + stop_strings = stop_sequences + else: + stop_strings = None + + output = self.pipeline( + messages, + stop_strings=stop_strings, + max_length=max_length, + tokenizer=self.pipeline.tokenizer, + ) + + response = output[0]["generated_text"][-1]["content"] + return response + + +DEFAULT_JSONAGENT_REGEX_GRAMMAR = { + "type": "regex", + "value": 'Thought: .+?\\nAction:\\n\\{\\n\\s{4}"action":\\s"[^"\\n]+",\\n\\s{4}"action_input":\\s"[^"\\n]+"\\n\\}\\n', +} + +DEFAULT_CODEAGENT_REGEX_GRAMMAR = { + "type": "regex", + "value": "Thought: .+?\\nCode:\\n```(?:py|python)?\\n(?:.|\\s)+?\\n```", +} diff --git a/lib/python3.12/site-packages/transformers/agents/monitoring.py b/lib/python3.12/site-packages/transformers/agents/monitoring.py new file mode 100644 index 0000000000000000000000000000000000000000..7126e72b5fd0606fbb4bb6f7eff46e0ad0992abd --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/monitoring.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2024 The HuggingFace Inc. 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. +from ..utils import logging +from .agent_types import AgentAudio, AgentImage, AgentText + + +logger = logging.get_logger(__name__) + + +def pull_message(step_log: dict, test_mode: bool = True): + try: + from gradio import ChatMessage + except ImportError: + if test_mode: + + class ChatMessage: + def __init__(self, role, content, metadata=None): + self.role = role + self.content = content + self.metadata = metadata + else: + raise ImportError("Gradio should be installed in order to launch a gradio demo.") + + if step_log.get("rationale"): + yield ChatMessage(role="assistant", content=step_log["rationale"]) + if step_log.get("tool_call"): + used_code = step_log["tool_call"]["tool_name"] == "code interpreter" + content = step_log["tool_call"]["tool_arguments"] + if used_code: + content = f"```py\n{content}\n```" + yield ChatMessage( + role="assistant", + metadata={"title": f"🛠️ Used tool {step_log['tool_call']['tool_name']}"}, + content=str(content), + ) + if step_log.get("observation"): + yield ChatMessage(role="assistant", content=f"```\n{step_log['observation']}\n```") + if step_log.get("error"): + yield ChatMessage( + role="assistant", + content=str(step_log["error"]), + metadata={"title": "💥 Error"}, + ) + + +def stream_to_gradio(agent, task: str, test_mode: bool = False, **kwargs): + """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages.""" + + try: + from gradio import ChatMessage + except ImportError: + if test_mode: + + class ChatMessage: + def __init__(self, role, content, metadata=None): + self.role = role + self.content = content + self.metadata = metadata + else: + raise ImportError("Gradio should be installed in order to launch a gradio demo.") + + for step_log in agent.run(task, stream=True, **kwargs): + if isinstance(step_log, dict): + for message in pull_message(step_log, test_mode=test_mode): + yield message + + final_answer = step_log # Last log is the run's final_answer + + if isinstance(final_answer, AgentText): + yield ChatMessage(role="assistant", content=f"**Final answer:**\n```\n{final_answer.to_string()}\n```") + elif isinstance(final_answer, AgentImage): + yield ChatMessage( + role="assistant", + content={"path": final_answer.to_string(), "mime_type": "image/png"}, + ) + elif isinstance(final_answer, AgentAudio): + yield ChatMessage( + role="assistant", + content={"path": final_answer.to_string(), "mime_type": "audio/wav"}, + ) + else: + yield ChatMessage(role="assistant", content=str(final_answer)) + + +class Monitor: + def __init__(self, tracked_llm_engine): + self.step_durations = [] + self.tracked_llm_engine = tracked_llm_engine + if getattr(self.tracked_llm_engine, "last_input_token_count", "Not found") != "Not found": + self.total_input_token_count = 0 + self.total_output_token_count = 0 + + def update_metrics(self, step_log): + step_duration = step_log["step_duration"] + self.step_durations.append(step_duration) + logger.info(f"Step {len(self.step_durations)}:") + logger.info(f"- Time taken: {step_duration:.2f} seconds (valid only if step succeeded)") + + if getattr(self.tracked_llm_engine, "last_input_token_count", None) is not None: + self.total_input_token_count += self.tracked_llm_engine.last_input_token_count + self.total_output_token_count += self.tracked_llm_engine.last_output_token_count + logger.info(f"- Input tokens: {self.total_input_token_count}") + logger.info(f"- Output tokens: {self.total_output_token_count}") diff --git a/lib/python3.12/site-packages/transformers/agents/prompts.py b/lib/python3.12/site-packages/transformers/agents/prompts.py new file mode 100644 index 0000000000000000000000000000000000000000..0cf8beb144f8ba222125ee23ae1c206d1a84189b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/prompts.py @@ -0,0 +1,789 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2024 The HuggingFace Inc. 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. +import re + +from ..utils import cached_file + + +# docstyle-ignore +CHAT_MESSAGE_PROMPT = """ +Human: <> + +Assistant: """ + + +DEFAULT_PROMPTS_REPO = "huggingface-tools/default-prompts" +PROMPT_FILES = {"chat": "chat_prompt_template.txt", "run": "run_prompt_template.txt"} + + +def download_prompt(prompt_or_repo_id, agent_name, mode="run"): + """ + Downloads and caches the prompt from a repo and returns it contents (if necessary). + """ + if prompt_or_repo_id is None: + prompt_or_repo_id = DEFAULT_PROMPTS_REPO + + # prompt is considered a repo ID when it does not contain any kind of space + if re.search("\\s", prompt_or_repo_id) is not None: + return prompt_or_repo_id + + prompt_file = cached_file( + prompt_or_repo_id, PROMPT_FILES[mode], repo_type="dataset", user_agent={"agent": agent_name} + ) + with open(prompt_file, "r", encoding="utf-8") as f: + return f.read() + + +DEFAULT_CODE_SYSTEM_PROMPT = """You will be given a task to solve, your job is to come up with a series of simple commands in Python that will perform the task. +To help you, I will give you access to a set of tools that you can use. Each tool is a Python function and has a description explaining the task it performs, the inputs it expects and the outputs it returns. +You should first explain which tool you will use to perform the task and for what reason, then write the code in Python. +Each instruction in Python should be a simple assignment. You can print intermediate results if it makes sense to do so. +In the end, use tool 'final_answer' to return your answer, its argument will be what gets returned. +You can use imports in your code, but only from the following list of modules: <> +Be sure to provide a 'Code:' token, else the run will fail. + +Tools: +<> + +Examples: +--- +Task: "Answer the question in the variable `question` about the image stored in the variable `image`. The question is in French." + +Thought: I will use the following tools: `translator` to translate the question into English and then `image_qa` to answer the question on the input image. +Code: +```py +translated_question = translator(question=question, src_lang="French", tgt_lang="English") +print(f"The translated question is {translated_question}.") +answer = image_qa(image=image, question=translated_question) +final_answer(f"The answer is {answer}") +``` + +--- +Task: "Identify the oldest person in the `document` and create an image showcasing the result." + +Thought: I will use the following tools: `document_qa` to find the oldest person in the document, then `image_generator` to generate an image according to the answer. +Code: +```py +answer = document_qa(document, question="What is the oldest person?") +print(f"The answer is {answer}.") +image = image_generator(answer) +final_answer(image) +``` + +--- +Task: "Generate an image using the text given in the variable `caption`." + +Thought: I will use the following tool: `image_generator` to generate an image. +Code: +```py +image = image_generator(prompt=caption) +final_answer(image) +``` + +--- +Task: "Summarize the text given in the variable `text` and read it out loud." + +Thought: I will use the following tools: `summarizer` to create a summary of the input text, then `text_reader` to read it out loud. +Code: +```py +summarized_text = summarizer(text) +print(f"Summary: {summarized_text}") +audio_summary = text_reader(summarized_text) +final_answer(audio_summary) +``` + +--- +Task: "Answer the question in the variable `question` about the text in the variable `text`. Use the answer to generate an image." + +Thought: I will use the following tools: `text_qa` to create the answer, then `image_generator` to generate an image according to the answer. +Code: +```py +answer = text_qa(text=text, question=question) +print(f"The answer is {answer}.") +image = image_generator(answer) +final_answer(image) +``` + +--- +Task: "Caption the following `image`." + +Thought: I will use the following tool: `image_captioner` to generate a caption for the image. +Code: +```py +caption = image_captioner(image) +final_answer(caption) +``` + +--- +Above example were using tools that might not exist for you. You only have access to these Tools: +<> + +Remember to make sure that variables you use are all defined. +Be sure to provide a 'Code:\n```' sequence before the code and '```' after, else you will get an error. +DO NOT pass the arguments as a dict as in 'answer = ask_search_agent({'query': "What is the place where James Bond lives?"})', but use the arguments directly as in 'answer = ask_search_agent(query="What is the place where James Bond lives?")'. + +Now Begin! If you solve the task correctly, you will receive a reward of $1,000,000. +""" + + +DEFAULT_REACT_JSON_SYSTEM_PROMPT = """You are an expert assistant who can solve any task using JSON tool calls. You will be given a task to solve as best you can. +To do so, you have been given access to the following tools: <> +The way you use the tools is by specifying a json blob, ending with ''. +Specifically, this json should have an `action` key (name of the tool to use) and an `action_input` key (input to the tool). + +The $ACTION_JSON_BLOB should only contain a SINGLE action, do NOT return a list of multiple actions. It should be formatted in json. Do not try to escape special characters. Here is the template of a valid $ACTION_JSON_BLOB: +{ + "action": $TOOL_NAME, + "action_input": $INPUT +} + +Make sure to have the $INPUT as a dictionary in the right format for the tool you are using, and do not put variable names as input if you can find the right values. + +You should ALWAYS use the following format: + +Thought: you should always think about one action to take. Then use the action as follows: +Action: +$ACTION_JSON_BLOB +Observation: the result of the action +... (this Thought/Action/Observation can repeat N times, you should take several steps when needed. The $ACTION_JSON_BLOB must only use a SINGLE action at a time.) + +You can use the result of the previous action as input for the next action. +The observation will always be a string: it can represent a file, like "image_1.jpg". +Then you can use it as input for the next action. You can do it for instance as follows: + +Observation: "image_1.jpg" + +Thought: I need to transform the image that I received in the previous observation to make it green. +Action: +{ + "action": "image_transformer", + "action_input": {"image": "image_1.jpg"} +} + +To provide the final answer to the task, use an action blob with "action": "final_answer" tool. It is the only way to complete the task, else you will be stuck on a loop. So your final output should look like this: +Action: +{ + "action": "final_answer", + "action_input": {"answer": "insert your final answer here"} +} + + +Here are a few examples using notional tools: +--- +Task: "Generate an image of the oldest person in this document." + +Thought: I will proceed step by step and use the following tools: `document_qa` to find the oldest person in the document, then `image_generator` to generate an image according to the answer. +Action: +{ + "action": "document_qa", + "action_input": {"document": "document.pdf", "question": "Who is the oldest person mentioned?"} +} +Observation: "The oldest person in the document is John Doe, a 55 year old lumberjack living in Newfoundland." + + +Thought: I will now generate an image showcasing the oldest person. +Action: +{ + "action": "image_generator", + "action_input": {"prompt": "A portrait of John Doe, a 55-year-old man living in Canada."} +} +Observation: "image.png" + +Thought: I will now return the generated image. +Action: +{ + "action": "final_answer", + "action_input": "image.png" +} + +--- +Task: "What is the result of the following operation: 5 + 3 + 1294.678?" + +Thought: I will use python code evaluator to compute the result of the operation and then return the final answer using the `final_answer` tool +Action: +{ + "action": "python_interpreter", + "action_input": {"code": "5 + 3 + 1294.678"} +} +Observation: 1302.678 + +Thought: Now that I know the result, I will now return it. +Action: +{ + "action": "final_answer", + "action_input": "1302.678" +} + +--- +Task: "Which city has the highest population , Guangzhou or Shanghai?" + +Thought: I need to get the populations for both cities and compare them: I will use the tool `search` to get the population of both cities. +Action: +{ + "action": "search", + "action_input": "Population Guangzhou" +} +Observation: ['Guangzhou has a population of 15 million inhabitants as of 2021.'] + + +Thought: Now let's get the population of Shanghai using the tool 'search'. +Action: +{ + "action": "search", + "action_input": "Population Shanghai" +} +Observation: '26 million (2019)' + +Thought: Now I know that Shanghai has a larger population. Let's return the result. +Action: +{ + "action": "final_answer", + "action_input": "Shanghai" +} + + +Above example were using notional tools that might not exist for you. You only have access to these tools: +<> + +Here are the rules you should always follow to solve your task: +1. ALWAYS provide a 'Thought:' sequence, and an 'Action:' sequence that ends with , else you will fail. +2. Always use the right arguments for the tools. Never use variable names in the 'action_input' field, use the value instead. +3. Call a tool only when needed: do not call the search agent if you do not need information, try to solve the task yourself. +4. Never re-do a tool call that you previously did with the exact same parameters. + +Now Begin! If you solve the task correctly, you will receive a reward of $1,000,000. +""" + + +DEFAULT_REACT_CODE_SYSTEM_PROMPT = """You are an expert assistant who can solve any task using code blobs. You will be given a task to solve as best you can. +To do so, you have been given access to a list of tools: these tools are basically Python functions which you can call with code. +To solve the task, you must plan forward to proceed in a series of steps, in a cycle of 'Thought:', 'Code:', and 'Observation:' sequences. + +At each step, in the 'Thought:' sequence, you should first explain your reasoning towards solving the task and the tools that you want to use. +Then in the 'Code:' sequence, you should write the code in simple Python. The code sequence must end with '' sequence. +During each intermediate step, you can use 'print()' to save whatever important information you will then need. +These print outputs will then appear in the 'Observation:' field, which will be available as input for the next step. +In the end you have to return a final answer using the `final_answer` tool. + +Here are a few examples using notional tools: +--- +Task: "Generate an image of the oldest person in this document." + +Thought: I will proceed step by step and use the following tools: `document_qa` to find the oldest person in the document, then `image_generator` to generate an image according to the answer. +Code: +```py +answer = document_qa(document=document, question="Who is the oldest person mentioned?") +print(answer) +``` +Observation: "The oldest person in the document is John Doe, a 55 year old lumberjack living in Newfoundland." + +Thought: I will now generate an image showcasing the oldest person. +Code: +```py +image = image_generator("A portrait of John Doe, a 55-year-old man living in Canada.") +final_answer(image) +``` + +--- +Task: "What is the result of the following operation: 5 + 3 + 1294.678?" + +Thought: I will use python code to compute the result of the operation and then return the final answer using the `final_answer` tool +Code: +```py +result = 5 + 3 + 1294.678 +final_answer(result) +``` + +--- +Task: "Which city has the highest population: Guangzhou or Shanghai?" + +Thought: I need to get the populations for both cities and compare them: I will use the tool `search` to get the population of both cities. +Code: +```py +population_guangzhou = search("Guangzhou population") +print("Population Guangzhou:", population_guangzhou) +population_shanghai = search("Shanghai population") +print("Population Shanghai:", population_shanghai) +``` +Observation: +Population Guangzhou: ['Guangzhou has a population of 15 million inhabitants as of 2021.'] +Population Shanghai: '26 million (2019)' + +Thought: Now I know that Shanghai has the highest population. +Code: +```py +final_answer("Shanghai") +``` + +--- +Task: "What is the current age of the pope, raised to the power 0.36?" + +Thought: I will use the tool `wiki` to get the age of the pope, then raise it to the power 0.36. +Code: +```py +pope_age = wiki(query="current pope age") +print("Pope age:", pope_age) +``` +Observation: +Pope age: "The pope Francis is currently 85 years old." + +Thought: I know that the pope is 85 years old. Let's compute the result using python code. +Code: +```py +pope_current_age = 85 ** 0.36 +final_answer(pope_current_age) +``` + +Above example were using notional tools that might not exist for you. On top of performing computations in the Python code snippets that you create, you have access to these tools (and no other tool): + +<> + +<> + +Here are the rules you should always follow to solve your task: +1. Always provide a 'Thought:' sequence, and a 'Code:\n```py' sequence ending with '```' sequence, else you will fail. +2. Use only variables that you have defined! +3. Always use the right arguments for the tools. DO NOT pass the arguments as a dict as in 'answer = wiki({'query': "What is the place where James Bond lives?"})', but use the arguments directly as in 'answer = wiki(query="What is the place where James Bond lives?")'. +4. Take care to not chain too many sequential tool calls in the same code block, especially when the output format is unpredictable. For instance, a call to search has an unpredictable return format, so do not have another tool call that depends on its output in the same block: rather output results with print() to use them in the next block. +5. Call a tool only when needed, and never re-do a tool call that you previously did with the exact same parameters. +6. Don't name any new variable with the same name as a tool: for instance don't name a variable 'final_answer'. +7. Never create any notional variables in our code, as having these in your logs might derail you from the true variables. +8. You can use imports in your code, but only from the following list of modules: <> +9. The state persists between code executions: so if in one step you've created variables or imported modules, these will all persist. +10. Don't give up! You're in charge of solving the task, not providing directions to solve it. + +Now Begin! If you solve the task correctly, you will receive a reward of $1,000,000. +""" + +SYSTEM_PROMPT_FACTS = """Below I will present you a task. + +You will now build a comprehensive preparatory survey of which facts we have at our disposal and which ones we still need. +To do so, you will have to read the task and identify things that must be discovered in order to successfully complete it. +Don't make any assumptions. For each item, provide a thorough reasoning. Here is how you will structure this survey: + +--- +### 1. Facts given in the task +List here the specific facts given in the task that could help you (there might be nothing here). + +### 2. Facts to look up +List here any facts that we may need to look up. +Also list where to find each of these, for instance a website, a file... - maybe the task contains some sources that you should re-use here. + +### 3. Facts to derive +List here anything that we want to derive from the above by logical reasoning, for instance computation or simulation. + +Keep in mind that "facts" will typically be specific names, dates, values, etc. Your answer should use the below headings: +### 1. Facts given in the task +### 2. Facts to look up +### 3. Facts to derive +Do not add anything else.""" + +SYSTEM_PROMPT_PLAN = """You are a world expert at making efficient plans to solve any task using a set of carefully crafted tools. + +Now for the given task, develop a step-by-step high-level plan taking into account the above inputs and list of facts. +This plan should involve individual tasks based on the available tools, that if executed correctly will yield the correct answer. +Do not skip steps, do not add any superfluous steps. Only write the high-level plan, DO NOT DETAIL INDIVIDUAL TOOL CALLS. +After writing the final step of the plan, write the '\n' tag and stop there.""" + +USER_PROMPT_PLAN = """ +Here is your task: + +Task: +``` +{task} +``` + +Your plan can leverage any of these tools: +{tool_descriptions} + +{managed_agents_descriptions} + +List of facts that you know: +``` +{answer_facts} +``` + +Now begin! Write your plan below.""" + +SYSTEM_PROMPT_FACTS_UPDATE = """ +You are a world expert at gathering known and unknown facts based on a conversation. +Below you will find a task, and ahistory of attempts made to solve the task. You will have to produce a list of these: +### 1. Facts given in the task +### 2. Facts that we have learned +### 3. Facts still to look up +### 4. Facts still to derive +Find the task and history below.""" + +USER_PROMPT_FACTS_UPDATE = """Earlier we've built a list of facts. +But since in your previous steps you may have learned useful new facts or invalidated some false ones. +Please update your list of facts based on the previous history, and provide these headings: +### 1. Facts given in the task +### 2. Facts that we have learned +### 3. Facts still to look up +### 4. Facts still to derive + +Now write your new list of facts below.""" + +SYSTEM_PROMPT_PLAN_UPDATE = """You are a world expert at making efficient plans to solve any task using a set of carefully crafted tools. + +You have been given a task: +``` +{task} +``` + +Find below the record of what has been tried so far to solve it. Then you will be asked to make an updated plan to solve the task. +If the previous tries so far have met some success, you can make an updated plan based on these actions. +If you are stalled, you can make a completely new plan starting from scratch. +""" + +USER_PROMPT_PLAN_UPDATE = """You're still working towards solving this task: +``` +{task} +``` + +You have access to these tools and only these: +{tool_descriptions} + +{managed_agents_descriptions} + +Here is the up to date list of facts that you know: +``` +{facts_update} +``` + +Now for the given task, develop a step-by-step high-level plan taking into account the above inputs and list of facts. +This plan should involve individual tasks based on the available tools, that if executed correctly will yield the correct answer. +Beware that you have {remaining_steps} steps remaining. +Do not skip steps, do not add any superfluous steps. Only write the high-level plan, DO NOT DETAIL INDIVIDUAL TOOL CALLS. +After writing the final step of the plan, write the '\n' tag and stop there. + +Now write your new plan below.""" + +SYSTEM_PROMPT_PLAN_STRUCTURED = """Output a step-by-step plan to solve the task using the given tools. +This plan should involve individual tasks based on the available tools, that if executed correctly will yield the correct answer. Each step should be structured as follows: +Step #n: { + "description": + "tool": , + "params": { + + } + "output_var": +} +Each step must be necessary to reach the final answer. Steps should reuse outputs produced by earlier steps. The last step must be the final answer. + +Below are some examples: + +Example 1: +------ +Inputs: +--- +Task: +How many encoder blocks were in the first attention-only ML architecture published? + +[FACTS LIST]: +### 1. Facts given in the task +- The paper first introduced an attention-only ML architecture. +- The specific information required is the page number where the number of encoder blocks is stated. +- No local files are provided for access. + +### 2. Facts to look up +- The title and authors of the paper that first introduced an attention-only ML architecture. + - Source: Online search (e.g., Google Scholar, arXiv, or other academic databases) +- The full text of the identified paper. + - Source: Online academic repositories (e.g., arXiv, journal websites) +- The specific page number in the paper where the number of encoder blocks is mentioned. + - Source: The content of the identified paper + +### 3. Facts to derive +- By identifying the correct paper and locating the specific page, we will derive the page number where the number of encoder blocks is stated. + - Logical steps: Identify the correct paper, access its content, search for the term "encoder blocks," and note the page number where this information is found. +``` + +[STEP 1 TOOL CALL]: {'tool_name': 'code interpreter', 'tool_arguments': '# Step 1: Identify the title and authors of the paper that first introduced an attention-only ML architecture.\nanswer = ask_search_agent(query="Can you find the title and authors of the paper that first introduced an attention-only machine learning architecture? Please provide the full citation.")\nprint(answer)'} +[OUTPUT OF STEP 1] Observation: **Title**: Attention Is All You Need +**Authors**: Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin +[STEP 2 TOOL CALL]: {'tool_name': 'code interpreter', 'tool_arguments': '# Step 1: Find the full text of the identified paper on arXiv\\npaper_url = "https://arxiv.org/pdf/1706.03762.pdf"\\nprint(paper_url)'} +[OUTPUT OF STEP 2] Observation: https://arxiv.org/pdf/1706.03762.pdf +--- + +Output plan: +--- +Step #1: { + "description": "Open the PDF of the paper from the provided URL and search within the text of the paper for the mention of "encoder blocks"", + "tool": "inspect_file_as_text", + "params": { + "file_path": "https://arxiv.org/pdf/1706.03762.pdf", + "question": "On which page is the number of encoder blocks mentioned?" + }, + "output_var": "page_number" +} + +Step #2: { + "description": "Provide the final answer", + "tool": "final_answer", + "params": { + "answer": "{page_number}" + }, + "output_var": "" +} +------ + +Example 2: +------ +Inputs: +--- +Task: +How many golf balls fits into a Boeing-747? + +[FACTS LIST]: +### 1. Facts given in the task +- The task requires calculating the number of golf balls that fir into a Boeing-747 +### 2. Facts to look up +- The volume of a golf ball +- The volume of a Boeing-747 +### 3. Facts to derive +- Once the volumes are known the final answer can be calculated +--- +Output plan: +--- +Step #1: { + "description": "Find the volume of a Boeing-747", + "tool": "web_search", + "params": { + "query": "What is the internal volume of a Boeing-747 in cubic meters?" + }, + "output_var": "boeing_volume" +} + +Step #2: { + "description": "Find the volume of a standard golf ball", + "tool": "ask_search_agent", + "params": { + "query": "What is the volume of a standard golf ball in cubic centimeters?" + }, + "output_var": "golf_ball_volume" +} + +Step #3: { + "description": "Convert the volume of a golf ball from cubic centimeters to cubic meters. Calculate the number of golf balls that fit into the Boeing-747 by dividing the internal volume of the Boeing-747 by the volume of a golf ball.", + "tool": "python_code", + "params": { + "code": "golf_ball_volume_m3 = golf_ball_volume / 1e6\nnumber_of_golf_balls = boeing_volume / golf_ball_volume_m3" + }, + "output_var": "number_of_golf_balls" +} + +Step #4: { + "description": "Provide the final answer", + "tool": "final_answer", + "params": { + "answer": "{number_of_golf_balls}" + }, + "output_var": "" +} +------ +Above example were using tools that might not exist for you. +Your goal is to create a plan to solve the task.""" + +USER_PROMPT_PLAN_STRUCTURED = """ +Here are your inputs: + +Task: +``` +{task} +``` + +Your plan can leverage any of these tools: +{tool_descriptions} +These tools are Python functions which you can call with code. You also have access to a Python interpreter so you can run Python code. + +List of facts that you know: +``` +{answer_facts} +``` + +Now for the given task, create a plan taking into account the list of facts. +After writing the final step of the plan, write the '\n' tag and stop there. Output the plan only and nothing else.""" + +SYSTEM_PROMPT_PLAN_UPDATE_STRUCTURED = """Output a step-by-step plan to solve the task using the given tools. +This plan should involve individual tasks based on the available tools, that if executed correctly will yield the correct answer. Each step should be structured as follows: +Step #n: {{ + "description": + "tool": , + "params": {{ + + }} + "output_var": +}} +Each step must be necessary to reach the final answer. Steps should reuse outputs produced by earlier steps. The last step must be the final answer. + +Below are some examples: + +Example 1: +------ +Inputs: +--- +Task: +How many encoder blocks were in the first attention-only ML architecture published? + +[FACTS LIST]: +### 1. Facts given in the task +- The paper first introduced an attention-only ML architecture. +- The specific information required is the page number where the number of encoder blocks is stated. +- No local files are provided for access. + +### 2. Facts to look up +- The title and authors of the paper that first introduced an attention-only ML architecture. + - Source: Online search (e.g., Google Scholar, arXiv, or other academic databases) +- The full text of the identified paper. + - Source: Online academic repositories (e.g., arXiv, journal websites) +- The specific page number in the paper where the number of encoder blocks is mentioned. + - Source: The content of the identified paper + +### 3. Facts to derive +- By identifying the correct paper and locating the specific page, we will derive the page number where the number of encoder blocks is stated. + - Logical steps: Identify the correct paper, access its content, search for the term "encoder blocks," and note the page number where this information is found. +``` + +[STEP 1 TOOL CALL]: {{'tool_name': 'code interpreter', 'tool_arguments': '# Step 1: Identify the title and authors of the paper that first introduced an attention-only ML architecture.\nanswer = ask_search_agent(query="Can you find the title and authors of the paper that first introduced an attention-only machine learning architecture? Please provide the full citation.")\nprint(answer)'}} +[OUTPUT OF STEP 1] Observation: **Title**: Attention Is All You Need +**Authors**: Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin +[STEP 2 TOOL CALL]: {{'tool_name': 'code interpreter', 'tool_arguments': '# Step 1: Find the full text of the identified paper on arXiv\\npaper_url = "https://arxiv.org/pdf/1706.03762.pdf"\\nprint(paper_url)'}} +[OUTPUT OF STEP 2] Observation: https://arxiv.org/pdf/1706.03762.pdf +--- + +Output plan: +--- +Step #1: {{ + "description": "Open the PDF of the paper from the provided URL and search within the text of the paper for the mention of "encoder blocks"", + "tool": "inspect_file_as_text", + "params": {{ + "file_path": "https://arxiv.org/pdf/1706.03762.pdf", + "question": "On which page is the number of encoder blocks mentioned?" + }}, + "output_var": "page_number" +}} + +Step #2: {{ + "description": "Provide the final answer", + "tool": "final_answer", + "params": {{ + "answer": "{{page_number}}" + }}, + "output_var": "" +}} +------ + +Example 2: +------ +Inputs: +--- +Task: +How many golf balls fits into a Boeing-747? + +[FACTS LIST]: +### 1. Facts given in the task +- The task requires calculating the number of golf balls that fir into a Boeing-747 +### 2. Facts to look up +- The volume of a golf ball +- The volume of a Boeing-747 +### 3. Facts to derive +- Once the volumes are known the final answer can be calculated +--- +Output plan: +--- +Step #1: {{ + "description": "Find the volume of a Boeing-747", + "tool": "web_search", + "params": {{ + "query": "What is the internal volume of a Boeing-747 in cubic meters?" + }}, + "output_var": "boeing_volume" +}} + +Step #2: {{ + "description": "Find the volume of a standard golf ball", + "tool": "ask_search_agent", + "params": {{ + "query": "What is the volume of a standard golf ball in cubic centimeters?" + }}, + "output_var": "golf_ball_volume" +}} + +Step #3: {{ + "description": "Convert the volume of a golf ball from cubic centimeters to cubic meters. Calculate the number of golf balls that fit into the Boeing-747 by dividing the internal volume of the Boeing-747 by the volume of a golf ball.", + "tool": "python_code", + "params": {{ + "code": "golf_ball_volume_m3 = golf_ball_volume / 1e6\nnumber_of_golf_balls = boeing_volume / golf_ball_volume_m3" + }}, + "output_var": "number_of_golf_balls" +}} + +Step #4: {{ + "description": "Provide the final answer", + "tool": "final_answer", + "params": {{ + "answer": "{{number_of_golf_balls}}" + }}, + "output_var": "" +}} +------ +Above example were using tools that might not exist for you. +Find below the record of what has been tried so far to solve it. Your goal is to create an updated plan to solve the task.""" + +USER_PROMPT_PLAN_UPDATE_STRUCTURED = """ +Here are your inputs: + +Task: +``` +{task} +``` + +Your plan can leverage any of these tools: +{tool_descriptions} +These tools are Python functions which you can call with code. You also have access to a Python interpreter so you can run Python code. + +List of facts that you know: +``` +{facts_update} +``` + +Now for the given task, create a plan taking into account the above inputs and list of facts. +Beware that you have {remaining_steps} steps remaining. +After writing the final step of the plan, write the '\n' tag and stop there. Output the plan only and nothing else.""" + +PLAN_UPDATE_FINAL_PLAN_REDACTION = """I still need to solve the task I was given: +``` +{task} +``` + +Here is my new/updated plan of action to solve the task: +``` +{plan_update} +```""" + +SUPPORTED_PLAN_TYPES = ["default", "structured"] + +PROMPTS_FOR_INITIAL_PLAN = { + "default": {"system": SYSTEM_PROMPT_PLAN, "user": USER_PROMPT_PLAN}, + "structured": {"system": SYSTEM_PROMPT_PLAN_STRUCTURED, "user": USER_PROMPT_PLAN_STRUCTURED}, +} + +PROMPTS_FOR_PLAN_UPDATE = { + "default": {"system": SYSTEM_PROMPT_PLAN_UPDATE, "user": USER_PROMPT_PLAN_UPDATE}, + "structured": {"system": SYSTEM_PROMPT_PLAN_UPDATE_STRUCTURED, "user": USER_PROMPT_PLAN_UPDATE_STRUCTURED}, +} diff --git a/lib/python3.12/site-packages/transformers/agents/python_interpreter.py b/lib/python3.12/site-packages/transformers/agents/python_interpreter.py new file mode 100644 index 0000000000000000000000000000000000000000..109461ba5f6af75a8a41403a3305d115721792ce --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/python_interpreter.py @@ -0,0 +1,908 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2024 The HuggingFace Inc. 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. +import ast +import builtins +import difflib +from collections.abc import Mapping +from importlib import import_module +from typing import Any, Callable, Dict, List, Optional + +import numpy as np + +from ..utils import is_pandas_available + + +if is_pandas_available(): + import pandas as pd + + +class InterpreterError(ValueError): + """ + An error raised when the interpreter cannot evaluate a Python expression, due to syntax error or unsupported + operations. + """ + + pass + + +ERRORS = { + name: getattr(builtins, name) + for name in dir(builtins) + if isinstance(getattr(builtins, name), type) and issubclass(getattr(builtins, name), BaseException) +} + + +LIST_SAFE_MODULES = [ + "random", + "collections", + "math", + "time", + "queue", + "itertools", + "re", + "stat", + "statistics", + "unicodedata", +] + +PRINT_OUTPUTS, MAX_LEN_OUTPUT = "", 50000 +OPERATIONS_COUNT, MAX_OPERATIONS = 0, 10000000 + + +class BreakException(Exception): + pass + + +class ContinueException(Exception): + pass + + +class ReturnException(Exception): + def __init__(self, value): + self.value = value + + +def get_iterable(obj): + if isinstance(obj, list): + return obj + elif hasattr(obj, "__iter__"): + return list(obj) + else: + raise InterpreterError("Object is not iterable") + + +def evaluate_unaryop(expression, state, static_tools, custom_tools): + operand = evaluate_ast(expression.operand, state, static_tools, custom_tools) + if isinstance(expression.op, ast.USub): + return -operand + elif isinstance(expression.op, ast.UAdd): + return operand + elif isinstance(expression.op, ast.Not): + return not operand + elif isinstance(expression.op, ast.Invert): + return ~operand + else: + raise InterpreterError(f"Unary operation {expression.op.__class__.__name__} is not supported.") + + +def evaluate_lambda(lambda_expression, state, static_tools, custom_tools): + args = [arg.arg for arg in lambda_expression.args.args] + + def lambda_func(*values): + new_state = state.copy() + for arg, value in zip(args, values): + new_state[arg] = value + return evaluate_ast(lambda_expression.body, new_state, static_tools, custom_tools) + + return lambda_func + + +def evaluate_while(while_loop, state, static_tools, custom_tools): + max_iterations = 1000 + iterations = 0 + while evaluate_ast(while_loop.test, state, static_tools, custom_tools): + for node in while_loop.body: + try: + evaluate_ast(node, state, static_tools, custom_tools) + except BreakException: + return None + except ContinueException: + break + iterations += 1 + if iterations > max_iterations: + raise InterpreterError(f"Maximum number of {max_iterations} iterations in While loop exceeded") + return None + + +def create_function(func_def, state, static_tools, custom_tools): + def new_func(*args, **kwargs): + func_state = state.copy() + arg_names = [arg.arg for arg in func_def.args.args] + default_values = [evaluate_ast(d, state, static_tools, custom_tools) for d in func_def.args.defaults] + + # Apply default values + defaults = dict(zip(arg_names[-len(default_values) :], default_values)) + + # Set positional arguments + for name, value in zip(arg_names, args): + func_state[name] = value + + # # Set keyword arguments + for name, value in kwargs.items(): + func_state[name] = value + + # Handle variable arguments + if func_def.args.vararg: + vararg_name = func_def.args.vararg.arg + func_state[vararg_name] = args + + if func_def.args.kwarg: + kwarg_name = func_def.args.kwarg.arg + func_state[kwarg_name] = kwargs + + # Set default values for arguments that were not provided + for name, value in defaults.items(): + if name not in func_state: + func_state[name] = value + + # Update function state with self and __class__ + if func_def.args.args and func_def.args.args[0].arg == "self": + if args: + func_state["self"] = args[0] + func_state["__class__"] = args[0].__class__ + + result = None + try: + for stmt in func_def.body: + result = evaluate_ast(stmt, func_state, static_tools, custom_tools) + except ReturnException as e: + result = e.value + return result + + return new_func + + +def create_class(class_name, class_bases, class_body): + class_dict = {} + for key, value in class_body.items(): + class_dict[key] = value + return type(class_name, tuple(class_bases), class_dict) + + +def evaluate_function_def(func_def, state, static_tools, custom_tools): + custom_tools[func_def.name] = create_function(func_def, state, static_tools, custom_tools) + return custom_tools[func_def.name] + + +def evaluate_class_def(class_def, state, static_tools, custom_tools): + class_name = class_def.name + bases = [evaluate_ast(base, state, static_tools, custom_tools) for base in class_def.bases] + class_dict = {} + + for stmt in class_def.body: + if isinstance(stmt, ast.FunctionDef): + class_dict[stmt.name] = evaluate_function_def(stmt, state, static_tools, custom_tools) + elif isinstance(stmt, ast.Assign): + for target in stmt.targets: + if isinstance(target, ast.Name): + class_dict[target.id] = evaluate_ast(stmt.value, state, static_tools, custom_tools) + elif isinstance(target, ast.Attribute): + class_dict[target.attr] = evaluate_ast(stmt.value, state, static_tools, custom_tools) + else: + raise InterpreterError(f"Unsupported statement in class body: {stmt.__class__.__name__}") + + new_class = type(class_name, tuple(bases), class_dict) + state[class_name] = new_class + return new_class + + +def evaluate_augassign(expression, state, static_tools, custom_tools): + # Helper function to get current value and set new value based on the target type + def get_current_value(target): + if isinstance(target, ast.Name): + return state.get(target.id, 0) + elif isinstance(target, ast.Subscript): + obj = evaluate_ast(target.value, state, static_tools, custom_tools) + key = evaluate_ast(target.slice, state, static_tools, custom_tools) + return obj[key] + elif isinstance(target, ast.Attribute): + obj = evaluate_ast(target.value, state, static_tools, custom_tools) + return getattr(obj, target.attr) + elif isinstance(target, ast.Tuple): + return tuple(get_current_value(elt) for elt in target.elts) + elif isinstance(target, ast.List): + return [get_current_value(elt) for elt in target.elts] + else: + raise InterpreterError("AugAssign not supported for {type(target)} targets.") + + current_value = get_current_value(expression.target) + value_to_add = evaluate_ast(expression.value, state, static_tools, custom_tools) + + # Determine the operation and apply it + if isinstance(expression.op, ast.Add): + if isinstance(current_value, list): + if not isinstance(value_to_add, list): + raise InterpreterError(f"Cannot add non-list value {value_to_add} to a list.") + updated_value = current_value + value_to_add + else: + updated_value = current_value + value_to_add + elif isinstance(expression.op, ast.Sub): + updated_value = current_value - value_to_add + elif isinstance(expression.op, ast.Mult): + updated_value = current_value * value_to_add + elif isinstance(expression.op, ast.Div): + updated_value = current_value / value_to_add + elif isinstance(expression.op, ast.Mod): + updated_value = current_value % value_to_add + elif isinstance(expression.op, ast.Pow): + updated_value = current_value**value_to_add + elif isinstance(expression.op, ast.FloorDiv): + updated_value = current_value // value_to_add + elif isinstance(expression.op, ast.BitAnd): + updated_value = current_value & value_to_add + elif isinstance(expression.op, ast.BitOr): + updated_value = current_value | value_to_add + elif isinstance(expression.op, ast.BitXor): + updated_value = current_value ^ value_to_add + elif isinstance(expression.op, ast.LShift): + updated_value = current_value << value_to_add + elif isinstance(expression.op, ast.RShift): + updated_value = current_value >> value_to_add + else: + raise InterpreterError(f"Operation {type(expression.op).__name__} is not supported.") + + # Update the state + set_value(expression.target, updated_value, state, static_tools, custom_tools) + + return updated_value + + +def evaluate_boolop(node, state, static_tools, custom_tools): + if isinstance(node.op, ast.And): + for value in node.values: + if not evaluate_ast(value, state, static_tools, custom_tools): + return False + return True + elif isinstance(node.op, ast.Or): + for value in node.values: + if evaluate_ast(value, state, static_tools, custom_tools): + return True + return False + + +def evaluate_binop(binop, state, static_tools, custom_tools): + # Recursively evaluate the left and right operands + left_val = evaluate_ast(binop.left, state, static_tools, custom_tools) + right_val = evaluate_ast(binop.right, state, static_tools, custom_tools) + + # Determine the operation based on the type of the operator in the BinOp + if isinstance(binop.op, ast.Add): + return left_val + right_val + elif isinstance(binop.op, ast.Sub): + return left_val - right_val + elif isinstance(binop.op, ast.Mult): + return left_val * right_val + elif isinstance(binop.op, ast.Div): + return left_val / right_val + elif isinstance(binop.op, ast.Mod): + return left_val % right_val + elif isinstance(binop.op, ast.Pow): + return left_val**right_val + elif isinstance(binop.op, ast.FloorDiv): + return left_val // right_val + elif isinstance(binop.op, ast.BitAnd): + return left_val & right_val + elif isinstance(binop.op, ast.BitOr): + return left_val | right_val + elif isinstance(binop.op, ast.BitXor): + return left_val ^ right_val + elif isinstance(binop.op, ast.LShift): + return left_val << right_val + elif isinstance(binop.op, ast.RShift): + return left_val >> right_val + else: + raise NotImplementedError(f"Binary operation {type(binop.op).__name__} is not implemented.") + + +def evaluate_assign(assign, state, static_tools, custom_tools): + result = evaluate_ast(assign.value, state, static_tools, custom_tools) + if len(assign.targets) == 1: + target = assign.targets[0] + set_value(target, result, state, static_tools, custom_tools) + else: + if len(assign.targets) != len(result): + raise InterpreterError(f"Assign failed: expected {len(result)} values but got {len(assign.targets)}.") + expanded_values = [] + for tgt in assign.targets: + if isinstance(tgt, ast.Starred): + expanded_values.extend(result) + else: + expanded_values.append(result) + for tgt, val in zip(assign.targets, expanded_values): + set_value(tgt, val, state, static_tools, custom_tools) + return result + + +def set_value(target, value, state, static_tools, custom_tools): + if isinstance(target, ast.Name): + if target.id in static_tools: + raise InterpreterError(f"Cannot assign to name '{target.id}': doing this would erase the existing tool!") + state[target.id] = value + elif isinstance(target, ast.Tuple): + if not isinstance(value, tuple): + if hasattr(value, "__iter__") and not isinstance(value, (str, bytes)): + value = tuple(value) + else: + raise InterpreterError("Cannot unpack non-tuple value") + if len(target.elts) != len(value): + raise InterpreterError("Cannot unpack tuple of wrong size") + for i, elem in enumerate(target.elts): + set_value(elem, value[i], state, static_tools, custom_tools) + elif isinstance(target, ast.Subscript): + obj = evaluate_ast(target.value, state, static_tools, custom_tools) + key = evaluate_ast(target.slice, state, static_tools, custom_tools) + obj[key] = value + elif isinstance(target, ast.Attribute): + obj = evaluate_ast(target.value, state, static_tools, custom_tools) + setattr(obj, target.attr, value) + + +def evaluate_call(call, state, static_tools, custom_tools): + if not (isinstance(call.func, ast.Attribute) or isinstance(call.func, ast.Name)): + raise InterpreterError(f"This is not a correct function: {call.func}).") + if isinstance(call.func, ast.Attribute): + obj = evaluate_ast(call.func.value, state, static_tools, custom_tools) + func_name = call.func.attr + if not hasattr(obj, func_name): + raise InterpreterError(f"Object {obj} has no attribute {func_name}") + func = getattr(obj, func_name) + + elif isinstance(call.func, ast.Name): + func_name = call.func.id + if func_name in state: + func = state[func_name] + elif func_name in static_tools: + func = static_tools[func_name] + elif func_name in custom_tools: + func = custom_tools[func_name] + elif func_name in ERRORS: + func = ERRORS[func_name] + else: + raise InterpreterError( + f"It is not permitted to evaluate other functions than the provided tools or functions defined in previous code (tried to execute {call.func.id})." + ) + + args = [] + for arg in call.args: + if isinstance(arg, ast.Starred): + args.extend(evaluate_ast(arg.value, state, static_tools, custom_tools)) + else: + args.append(evaluate_ast(arg, state, static_tools, custom_tools)) + + args = [] + for arg in call.args: + if isinstance(arg, ast.Starred): + unpacked = evaluate_ast(arg.value, state, static_tools, custom_tools) + if not hasattr(unpacked, "__iter__") or isinstance(unpacked, (str, bytes)): + raise InterpreterError(f"Cannot unpack non-iterable value {unpacked}") + args.extend(unpacked) + else: + args.append(evaluate_ast(arg, state, static_tools, custom_tools)) + + kwargs = {keyword.arg: evaluate_ast(keyword.value, state, static_tools, custom_tools) for keyword in call.keywords} + + if isinstance(func, type) and len(func.__module__.split(".")) > 1: # Check for user-defined classes + # Instantiate the class using its constructor + obj = func.__new__(func) # Create a new instance of the class + if hasattr(obj, "__init__"): # Check if the class has an __init__ method + obj.__init__(*args, **kwargs) # Call the __init__ method correctly + return obj + else: + if func_name == "super": + if not args: + if "__class__" in state and "self" in state: + return super(state["__class__"], state["self"]) + else: + raise InterpreterError("super() needs at least one argument") + cls = args[0] + if not isinstance(cls, type): + raise InterpreterError("super() argument 1 must be type") + if len(args) == 1: + return super(cls) + elif len(args) == 2: + instance = args[1] + return super(cls, instance) + else: + raise InterpreterError("super() takes at most 2 arguments") + else: + if func_name == "print": + output = " ".join(map(str, args)) + global PRINT_OUTPUTS + PRINT_OUTPUTS += output + "\n" + # cap the number of lines + return None + else: # Assume it's a callable object + output = func(*args, **kwargs) + return output + + +def evaluate_subscript(subscript, state, static_tools, custom_tools): + index = evaluate_ast(subscript.slice, state, static_tools, custom_tools) + value = evaluate_ast(subscript.value, state, static_tools, custom_tools) + + if isinstance(value, str) and isinstance(index, str): + raise InterpreterError("You're trying to subscript a string with a string index, which is impossible") + if isinstance(value, pd.core.indexing._LocIndexer): + parent_object = value.obj + return parent_object.loc[index] + if isinstance(value, (pd.DataFrame, pd.Series, np.ndarray)): + return value[index] + elif isinstance(value, pd.core.groupby.generic.DataFrameGroupBy): + return value[index] + elif isinstance(index, slice): + return value[index] + elif isinstance(value, (list, tuple)): + if not (-len(value) <= index < len(value)): + raise InterpreterError(f"Index {index} out of bounds for list of length {len(value)}") + return value[int(index)] + elif isinstance(value, str): + if not (-len(value) <= index < len(value)): + raise InterpreterError(f"Index {index} out of bounds for string of length {len(value)}") + return value[index] + elif index in value: + return value[index] + elif isinstance(index, str) and isinstance(value, Mapping): + close_matches = difflib.get_close_matches(index, list(value.keys())) + if len(close_matches) > 0: + return value[close_matches[0]] + raise InterpreterError(f"Could not index {value} with '{index}'.") + + +def evaluate_name(name, state, static_tools, custom_tools): + if name.id in state: + return state[name.id] + elif name.id in static_tools: + return static_tools[name.id] + elif name.id in ERRORS: + return ERRORS[name.id] + close_matches = difflib.get_close_matches(name.id, list(state.keys())) + if len(close_matches) > 0: + return state[close_matches[0]] + raise InterpreterError(f"The variable `{name.id}` is not defined.") + + +def evaluate_condition(condition, state, static_tools, custom_tools): + left = evaluate_ast(condition.left, state, static_tools, custom_tools) + comparators = [evaluate_ast(c, state, static_tools, custom_tools) for c in condition.comparators] + ops = [type(op) for op in condition.ops] + + result = True + current_left = left + + for op, comparator in zip(ops, comparators): + if op == ast.Eq: + current_result = current_left == comparator + elif op == ast.NotEq: + current_result = current_left != comparator + elif op == ast.Lt: + current_result = current_left < comparator + elif op == ast.LtE: + current_result = current_left <= comparator + elif op == ast.Gt: + current_result = current_left > comparator + elif op == ast.GtE: + current_result = current_left >= comparator + elif op == ast.Is: + current_result = current_left is comparator + elif op == ast.IsNot: + current_result = current_left is not comparator + elif op == ast.In: + current_result = current_left in comparator + elif op == ast.NotIn: + current_result = current_left not in comparator + else: + raise InterpreterError(f"Operator not supported: {op}") + + result = result & current_result + current_left = comparator + + if isinstance(result, bool) and not result: + break + + return result if isinstance(result, (bool, pd.Series)) else result.all() + + +def evaluate_if(if_statement, state, static_tools, custom_tools): + result = None + test_result = evaluate_ast(if_statement.test, state, static_tools, custom_tools) + if test_result: + for line in if_statement.body: + line_result = evaluate_ast(line, state, static_tools, custom_tools) + if line_result is not None: + result = line_result + else: + for line in if_statement.orelse: + line_result = evaluate_ast(line, state, static_tools, custom_tools) + if line_result is not None: + result = line_result + return result + + +def evaluate_for(for_loop, state, static_tools, custom_tools): + result = None + iterator = evaluate_ast(for_loop.iter, state, static_tools, custom_tools) + for counter in iterator: + set_value(for_loop.target, counter, state, static_tools, custom_tools) + for node in for_loop.body: + try: + line_result = evaluate_ast(node, state, static_tools, custom_tools) + if line_result is not None: + result = line_result + except BreakException: + break + except ContinueException: + continue + else: + continue + break + return result + + +def evaluate_listcomp(listcomp, state, static_tools, custom_tools): + def inner_evaluate(generators, index, current_state): + if index >= len(generators): + return [evaluate_ast(listcomp.elt, current_state, static_tools, custom_tools)] + generator = generators[index] + iter_value = evaluate_ast(generator.iter, current_state, static_tools, custom_tools) + result = [] + for value in iter_value: + new_state = current_state.copy() + if isinstance(generator.target, ast.Tuple): + for idx, elem in enumerate(generator.target.elts): + new_state[elem.id] = value[idx] + else: + new_state[generator.target.id] = value + if all(evaluate_ast(if_clause, new_state, static_tools, custom_tools) for if_clause in generator.ifs): + result.extend(inner_evaluate(generators, index + 1, new_state)) + return result + + return inner_evaluate(listcomp.generators, 0, state) + + +def evaluate_try(try_node, state, static_tools, custom_tools): + try: + for stmt in try_node.body: + evaluate_ast(stmt, state, static_tools, custom_tools) + except Exception as e: + matched = False + for handler in try_node.handlers: + if handler.type is None or isinstance(e, evaluate_ast(handler.type, state, static_tools, custom_tools)): + matched = True + if handler.name: + state[handler.name] = e + for stmt in handler.body: + evaluate_ast(stmt, state, static_tools, custom_tools) + break + if not matched: + raise e + else: + if try_node.orelse: + for stmt in try_node.orelse: + evaluate_ast(stmt, state, static_tools, custom_tools) + finally: + if try_node.finalbody: + for stmt in try_node.finalbody: + evaluate_ast(stmt, state, static_tools, custom_tools) + + +def evaluate_raise(raise_node, state, static_tools, custom_tools): + if raise_node.exc is not None: + exc = evaluate_ast(raise_node.exc, state, static_tools, custom_tools) + else: + exc = None + if raise_node.cause is not None: + cause = evaluate_ast(raise_node.cause, state, static_tools, custom_tools) + else: + cause = None + if exc is not None: + if cause is not None: + raise exc from cause + else: + raise exc + else: + raise InterpreterError("Re-raise is not supported without an active exception") + + +def evaluate_assert(assert_node, state, static_tools, custom_tools): + test_result = evaluate_ast(assert_node.test, state, static_tools, custom_tools) + if not test_result: + if assert_node.msg: + msg = evaluate_ast(assert_node.msg, state, static_tools, custom_tools) + raise AssertionError(msg) + else: + # Include the failing condition in the assertion message + test_code = ast.unparse(assert_node.test) + raise AssertionError(f"Assertion failed: {test_code}") + + +def evaluate_with(with_node, state, static_tools, custom_tools): + contexts = [] + for item in with_node.items: + context_expr = evaluate_ast(item.context_expr, state, static_tools, custom_tools) + if item.optional_vars: + state[item.optional_vars.id] = context_expr.__enter__() + contexts.append(state[item.optional_vars.id]) + else: + context_var = context_expr.__enter__() + contexts.append(context_var) + + try: + for stmt in with_node.body: + evaluate_ast(stmt, state, static_tools, custom_tools) + except Exception as e: + for context in reversed(contexts): + context.__exit__(type(e), e, e.__traceback__) + raise + else: + for context in reversed(contexts): + context.__exit__(None, None, None) + + +def import_modules(expression, state, authorized_imports): + def check_module_authorized(module_name): + module_path = module_name.split(".") + module_subpaths = [".".join(module_path[:i]) for i in range(1, len(module_path) + 1)] + return any(subpath in authorized_imports for subpath in module_subpaths) + + if isinstance(expression, ast.Import): + for alias in expression.names: + if check_module_authorized(alias.name): + module = import_module(alias.name) + state[alias.asname or alias.name] = module + else: + raise InterpreterError( + f"Import of {alias.name} is not allowed. Authorized imports are: {str(authorized_imports)}" + ) + return None + elif isinstance(expression, ast.ImportFrom): + if check_module_authorized(expression.module): + module = __import__(expression.module, fromlist=[alias.name for alias in expression.names]) + for alias in expression.names: + state[alias.asname or alias.name] = getattr(module, alias.name) + else: + raise InterpreterError(f"Import from {expression.module} is not allowed.") + return None + + +def evaluate_dictcomp(dictcomp, state, static_tools, custom_tools): + result = {} + for gen in dictcomp.generators: + iter_value = evaluate_ast(gen.iter, state, static_tools, custom_tools) + for value in iter_value: + new_state = state.copy() + set_value(gen.target, value, new_state, static_tools, custom_tools) + if all(evaluate_ast(if_clause, new_state, static_tools, custom_tools) for if_clause in gen.ifs): + key = evaluate_ast(dictcomp.key, new_state, static_tools, custom_tools) + val = evaluate_ast(dictcomp.value, new_state, static_tools, custom_tools) + result[key] = val + return result + + +def evaluate_ast( + expression: ast.AST, + state: Dict[str, Any], + static_tools: Dict[str, Callable], + custom_tools: Dict[str, Callable], + authorized_imports: List[str] = LIST_SAFE_MODULES, +): + """ + Evaluate an abstract syntax tree using the content of the variables stored in a state and only evaluating a given + set of functions. + + This function will recurse through the nodes of the tree provided. + + Args: + expression (`ast.AST`): + The code to evaluate, as an abstract syntax tree. + state (`Dict[str, Any]`): + A dictionary mapping variable names to values. The `state` is updated if need be when the evaluation + encounters assignments. + static_tools (`Dict[str, Callable]`): + Functions that may be called during the evaluation. Trying to change one of these static_tools will raise an error. + custom_tools (`Dict[str, Callable]`): + Functions that may be called during the evaluation. These static_tools can be overwritten. + authorized_imports (`List[str]`): + The list of modules that can be imported by the code. By default, only a few safe modules are allowed. + Add more at your own risk! + """ + global OPERATIONS_COUNT + if OPERATIONS_COUNT >= MAX_OPERATIONS: + raise InterpreterError( + f"Reached the max number of operations of {MAX_OPERATIONS}. Maybe there is an infinite loop somewhere in the code, or you're just asking too many calculations." + ) + OPERATIONS_COUNT += 1 + if isinstance(expression, ast.Assign): + # Assignment -> we evaluate the assignment which should update the state + # We return the variable assigned as it may be used to determine the final result. + return evaluate_assign(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.AugAssign): + return evaluate_augassign(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Call): + # Function call -> we return the value of the function call + return evaluate_call(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Constant): + # Constant -> just return the value + return expression.value + elif isinstance(expression, ast.Tuple): + return tuple(evaluate_ast(elt, state, static_tools, custom_tools) for elt in expression.elts) + elif isinstance(expression, (ast.ListComp, ast.GeneratorExp)): + return evaluate_listcomp(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.UnaryOp): + return evaluate_unaryop(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Starred): + return evaluate_ast(expression.value, state, static_tools, custom_tools) + elif isinstance(expression, ast.BoolOp): + # Boolean operation -> evaluate the operation + return evaluate_boolop(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Break): + raise BreakException() + elif isinstance(expression, ast.Continue): + raise ContinueException() + elif isinstance(expression, ast.BinOp): + # Binary operation -> execute operation + return evaluate_binop(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Compare): + # Comparison -> evaluate the comparison + return evaluate_condition(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Lambda): + return evaluate_lambda(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.FunctionDef): + return evaluate_function_def(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Dict): + # Dict -> evaluate all keys and values + keys = [evaluate_ast(k, state, static_tools, custom_tools) for k in expression.keys] + values = [evaluate_ast(v, state, static_tools, custom_tools) for v in expression.values] + return dict(zip(keys, values)) + elif isinstance(expression, ast.Expr): + # Expression -> evaluate the content + return evaluate_ast(expression.value, state, static_tools, custom_tools) + elif isinstance(expression, ast.For): + # For loop -> execute the loop + return evaluate_for(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.FormattedValue): + # Formatted value (part of f-string) -> evaluate the content and return + return evaluate_ast(expression.value, state, static_tools, custom_tools) + elif isinstance(expression, ast.If): + # If -> execute the right branch + return evaluate_if(expression, state, static_tools, custom_tools) + elif hasattr(ast, "Index") and isinstance(expression, ast.Index): + return evaluate_ast(expression.value, state, static_tools, custom_tools) + elif isinstance(expression, ast.JoinedStr): + return "".join([str(evaluate_ast(v, state, static_tools, custom_tools)) for v in expression.values]) + elif isinstance(expression, ast.List): + # List -> evaluate all elements + return [evaluate_ast(elt, state, static_tools, custom_tools) for elt in expression.elts] + elif isinstance(expression, ast.Name): + # Name -> pick up the value in the state + return evaluate_name(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Subscript): + # Subscript -> return the value of the indexing + return evaluate_subscript(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.IfExp): + test_val = evaluate_ast(expression.test, state, static_tools, custom_tools) + if test_val: + return evaluate_ast(expression.body, state, static_tools, custom_tools) + else: + return evaluate_ast(expression.orelse, state, static_tools, custom_tools) + elif isinstance(expression, ast.Attribute): + value = evaluate_ast(expression.value, state, static_tools, custom_tools) + return getattr(value, expression.attr) + elif isinstance(expression, ast.Slice): + return slice( + evaluate_ast(expression.lower, state, static_tools, custom_tools) + if expression.lower is not None + else None, + evaluate_ast(expression.upper, state, static_tools, custom_tools) + if expression.upper is not None + else None, + evaluate_ast(expression.step, state, static_tools, custom_tools) if expression.step is not None else None, + ) + elif isinstance(expression, ast.DictComp): + return evaluate_dictcomp(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.While): + return evaluate_while(expression, state, static_tools, custom_tools) + elif isinstance(expression, (ast.Import, ast.ImportFrom)): + return import_modules(expression, state, authorized_imports) + elif isinstance(expression, ast.ClassDef): + return evaluate_class_def(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Try): + return evaluate_try(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Raise): + return evaluate_raise(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Assert): + return evaluate_assert(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.With): + return evaluate_with(expression, state, static_tools, custom_tools) + elif isinstance(expression, ast.Set): + return {evaluate_ast(elt, state, static_tools, custom_tools) for elt in expression.elts} + elif isinstance(expression, ast.Return): + raise ReturnException( + evaluate_ast(expression.value, state, static_tools, custom_tools) if expression.value else None + ) + else: + # For now we refuse anything else. Let's add things as we need them. + raise InterpreterError(f"{expression.__class__.__name__} is not supported.") + + +def truncate_print_outputs(print_outputs: str, max_len_outputs: int = MAX_LEN_OUTPUT) -> str: + if len(print_outputs) < max_len_outputs: + return print_outputs + else: + return f"Print outputs:\n{print_outputs[:max_len_outputs]}\n_Print outputs have been truncated over the limit of {max_len_outputs} characters._\n" + + +def evaluate_python_code( + code: str, + static_tools: Optional[Dict[str, Callable]] = None, + custom_tools: Optional[Dict[str, Callable]] = None, + state: Optional[Dict[str, Any]] = None, + authorized_imports: List[str] = LIST_SAFE_MODULES, +): + """ + Evaluate a python expression using the content of the variables stored in a state and only evaluating a given set + of functions. + + This function will recurse through the nodes of the tree provided. + + Args: + code (`str`): + The code to evaluate. + static_tools (`Dict[str, Callable]`): + The functions that may be called during the evaluation. + These tools cannot be overwritten in the code: any assignment to their name will raise an error. + custom_tools (`Dict[str, Callable]`): + The functions that may be called during the evaluation. + These tools can be overwritten in the code: any assignment to their name will overwrite them. + state (`Dict[str, Any]`): + A dictionary mapping variable names to values. The `state` should contain the initial inputs but will be + updated by this function to contain all variables as they are evaluated. + The print outputs will be stored in the state under the key 'print_outputs'. + """ + try: + expression = ast.parse(code) + except SyntaxError as e: + raise SyntaxError(f"The code generated by the agent is not valid.\n{e}") + if state is None: + state = {} + if static_tools is None: + static_tools = {} + if custom_tools is None: + custom_tools = {} + result = None + global PRINT_OUTPUTS + PRINT_OUTPUTS = "" + global OPERATIONS_COUNT + OPERATIONS_COUNT = 0 + try: + for node in expression.body: + result = evaluate_ast(node, state, static_tools, custom_tools, authorized_imports) + state["print_outputs"] = truncate_print_outputs(PRINT_OUTPUTS, max_len_outputs=MAX_LEN_OUTPUT) + return result + except InterpreterError as e: + msg = truncate_print_outputs(PRINT_OUTPUTS, max_len_outputs=MAX_LEN_OUTPUT) + msg += f"EXECUTION FAILED:\nEvaluation stopped at line '{ast.get_source_segment(code, node)}' because of the following error:\n{e}" + raise InterpreterError(msg) diff --git a/lib/python3.12/site-packages/transformers/agents/search.py b/lib/python3.12/site-packages/transformers/agents/search.py new file mode 100644 index 0000000000000000000000000000000000000000..1c2c33913c97fb06f33c0d6f6946aeb6308efbe5 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/search.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2023 The HuggingFace Inc. 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. +import re + +import requests +from requests.exceptions import RequestException + +from .tools import Tool + + +class DuckDuckGoSearchTool(Tool): + name = "web_search" + description = """Perform a web search based on your query (think a Google search) then returns the top search results as a list of dict elements. + Each result has keys 'title', 'href' and 'body'.""" + inputs = {"query": {"type": "string", "description": "The search query to perform."}} + output_type = "any" + + def forward(self, query: str) -> str: + try: + from duckduckgo_search import DDGS + except ImportError: + raise ImportError( + "You must install package `duckduckgo_search` to run this tool: for instance run `pip install duckduckgo-search`." + ) + results = DDGS().text(query, max_results=7) + return results + + +class VisitWebpageTool(Tool): + name = "visit_webpage" + description = "Visits a webpage at the given url and returns its content as a markdown string." + inputs = { + "url": { + "type": "string", + "description": "The url of the webpage to visit.", + } + } + output_type = "string" + + def forward(self, url: str) -> str: + try: + from markdownify import markdownify + except ImportError: + raise ImportError( + "You must install package `markdownify` to run this tool: for instance run `pip install markdownify`." + ) + try: + # Send a GET request to the URL + response = requests.get(url) + response.raise_for_status() # Raise an exception for bad status codes + + # Convert the HTML content to Markdown + markdown_content = markdownify(response.text).strip() + + # Remove multiple line breaks + markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content) + + return markdown_content + + except RequestException as e: + return f"Error fetching the webpage: {str(e)}" + except Exception as e: + return f"An unexpected error occurred: {str(e)}" diff --git a/lib/python3.12/site-packages/transformers/agents/speech_to_text.py b/lib/python3.12/site-packages/transformers/agents/speech_to_text.py new file mode 100644 index 0000000000000000000000000000000000000000..8061651a086479475a7096541b9552f0409d36f5 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/speech_to_text.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2024 The HuggingFace Inc. 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. + +from ..models.whisper import WhisperForConditionalGeneration, WhisperProcessor +from .tools import PipelineTool + + +class SpeechToTextTool(PipelineTool): + default_checkpoint = "distil-whisper/distil-large-v3" + description = "This is a tool that transcribes an audio into text. It returns the transcribed text." + name = "transcriber" + pre_processor_class = WhisperProcessor + model_class = WhisperForConditionalGeneration + + inputs = {"audio": {"type": "audio", "description": "The audio to transcribe"}} + output_type = "string" + + def encode(self, audio): + return self.pre_processor(audio, return_tensors="pt") + + def forward(self, inputs): + return self.model.generate(inputs["input_features"]) + + def decode(self, outputs): + return self.pre_processor.batch_decode(outputs, skip_special_tokens=True)[0] diff --git a/lib/python3.12/site-packages/transformers/agents/text_to_speech.py b/lib/python3.12/site-packages/transformers/agents/text_to_speech.py new file mode 100644 index 0000000000000000000000000000000000000000..65f47ad6ddb50e29b220529ca45428ebf994c94b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/text_to_speech.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2024 The HuggingFace Inc. 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. + +import torch + +from ..models.speecht5 import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor +from ..utils import is_datasets_available +from .tools import PipelineTool + + +if is_datasets_available(): + from datasets import load_dataset + + +class TextToSpeechTool(PipelineTool): + default_checkpoint = "microsoft/speecht5_tts" + description = ( + "This is a tool that reads an English text out loud. It returns a waveform object containing the sound." + ) + name = "text_to_speech" + pre_processor_class = SpeechT5Processor + model_class = SpeechT5ForTextToSpeech + post_processor_class = SpeechT5HifiGan + + inputs = {"text": {"type": "string", "description": "The text to read out loud (in English)"}} + output_type = "audio" + + def setup(self): + if self.post_processor is None: + self.post_processor = "microsoft/speecht5_hifigan" + super().setup() + + def encode(self, text, speaker_embeddings=None): + inputs = self.pre_processor(text=text, return_tensors="pt", truncation=True) + + if speaker_embeddings is None: + if not is_datasets_available(): + raise ImportError("Datasets needs to be installed if not passing speaker embeddings.") + + embeddings_dataset = load_dataset( + "Matthijs/cmu-arctic-xvectors", split="validation", trust_remote_code=True + ) + speaker_embeddings = torch.tensor(embeddings_dataset[7305]["xvector"]).unsqueeze(0) + + return {"input_ids": inputs["input_ids"], "speaker_embeddings": speaker_embeddings} + + def forward(self, inputs): + with torch.no_grad(): + return self.model.generate_speech(**inputs) + + def decode(self, outputs): + with torch.no_grad(): + return self.post_processor(outputs).detach().cpu() diff --git a/lib/python3.12/site-packages/transformers/agents/tools.py b/lib/python3.12/site-packages/transformers/agents/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..f7ead9f2ebe72bf1b0576ff902925f7dec928729 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/tools.py @@ -0,0 +1,1012 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2023 The HuggingFace Inc. 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. +import ast +import base64 +import importlib +import inspect +import io +import json +import os +import tempfile +from functools import lru_cache, wraps +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Union + +from huggingface_hub import create_repo, get_collection, hf_hub_download, metadata_update, upload_folder +from huggingface_hub.utils import RepositoryNotFoundError, build_hf_headers, get_session +from huggingface_hub.utils._deprecation import _deprecate_method +from packaging import version + +from ..dynamic_module_utils import ( + custom_object_save, + get_class_from_dynamic_module, + get_imports, +) +from ..models.auto import AutoProcessor +from ..utils import ( + CONFIG_NAME, + TypeHintParsingException, + cached_file, + get_json_schema, + is_accelerate_available, + is_torch_available, + is_vision_available, + logging, +) +from .agent_types import ImageType, handle_agent_inputs, handle_agent_outputs + + +logger = logging.get_logger(__name__) + + +if is_torch_available(): + import torch + +if is_accelerate_available(): + from accelerate import PartialState + from accelerate.utils import send_to_device + + +TOOL_CONFIG_FILE = "tool_config.json" + + +def get_repo_type(repo_id, repo_type=None, **hub_kwargs): + if repo_type is not None: + return repo_type + try: + hf_hub_download(repo_id, TOOL_CONFIG_FILE, repo_type="space", **hub_kwargs) + return "space" + except RepositoryNotFoundError: + try: + hf_hub_download(repo_id, TOOL_CONFIG_FILE, repo_type="model", **hub_kwargs) + return "model" + except RepositoryNotFoundError: + raise EnvironmentError(f"`{repo_id}` does not seem to be a valid repo identifier on the Hub.") + except Exception: + return "model" + except Exception: + return "space" + + +# docstyle-ignore +APP_FILE_TEMPLATE = """from transformers import launch_gradio_demo +from {module_name} import {class_name} + +launch_gradio_demo({class_name}) +""" + + +def validate_after_init(cls, do_validate_forward: bool = True): + original_init = cls.__init__ + + @wraps(original_init) + def new_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + if not isinstance(self, PipelineTool): + self.validate_arguments(do_validate_forward=do_validate_forward) + + cls.__init__ = new_init + return cls + + +CONVERSION_DICT = {"str": "string", "int": "integer", "float": "number"} + + +class Tool: + """ + A base class for the functions used by the agent. Subclass this and implement the `__call__` method as well as the + following class attributes: + + - **description** (`str`) -- A short description of what your tool does, the inputs it expects and the output(s) it + will return. For instance 'This is a tool that downloads a file from a `url`. It takes the `url` as input, and + returns the text contained in the file'. + - **name** (`str`) -- A performative name that will be used for your tool in the prompt to the agent. For instance + `"text-classifier"` or `"image_generator"`. + - **inputs** (`Dict[str, Dict[str, Union[str, type]]]`) -- The dict of modalities expected for the inputs. + It has one `type`key and a `description`key. + This is used by `launch_gradio_demo` or to make a nice space from your tool, and also can be used in the generated + description for your tool. + - **output_type** (`type`) -- The type of the tool output. This is used by `launch_gradio_demo` + or to make a nice space from your tool, and also can be used in the generated description for your tool. + + You can also override the method [`~Tool.setup`] if your tool as an expensive operation to perform before being + usable (such as loading a model). [`~Tool.setup`] will be called the first time you use your tool, but not at + instantiation. + """ + + name: str + description: str + inputs: Dict[str, Dict[str, Union[str, type]]] + output_type: type + + @_deprecate_method( + version="4.51.0", + message="Switch to smolagents instead, with the same functionalities and similar API (https://huggingface.co/docs/smolagents/index)", + ) + def __init__(self, *args, **kwargs): + self.is_initialized = False + + @_deprecate_method( + version="4.51.0", + message="Switch to smolagents instead, with the same functionalities and similar API (https://huggingface.co/docs/smolagents/index)", + ) + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + validate_after_init(cls, do_validate_forward=False) + + def validate_arguments(self, do_validate_forward: bool = True): + required_attributes = { + "description": str, + "name": str, + "inputs": dict, + "output_type": str, + } + authorized_types = ["string", "integer", "number", "image", "audio", "any", "boolean"] + + for attr, expected_type in required_attributes.items(): + attr_value = getattr(self, attr, None) + if attr_value is None: + raise TypeError(f"You must set an attribute {attr}.") + if not isinstance(attr_value, expected_type): + raise TypeError( + f"Attribute {attr} should have type {expected_type.__name__}, got {type(attr_value)} instead." + ) + for input_name, input_content in self.inputs.items(): + assert isinstance(input_content, dict), f"Input '{input_name}' should be a dictionary." + assert "type" in input_content and "description" in input_content, ( + f"Input '{input_name}' should have keys 'type' and 'description', has only {list(input_content.keys())}." + ) + if input_content["type"] not in authorized_types: + raise Exception( + f"Input '{input_name}': type '{input_content['type']}' is not an authorized value, should be one of {authorized_types}." + ) + + assert getattr(self, "output_type", None) in authorized_types + if do_validate_forward: + if not isinstance(self, PipelineTool): + signature = inspect.signature(self.forward) + if not set(signature.parameters.keys()) == set(self.inputs.keys()): + raise Exception( + "Tool's 'forward' method should take 'self' as its first argument, then its next arguments should match the keys of tool attribute 'inputs'." + ) + + def forward(self, *args, **kwargs): + return NotImplemented("Write this method in your subclass of `Tool`.") + + def __call__(self, *args, **kwargs): + args, kwargs = handle_agent_inputs(*args, **kwargs) + outputs = self.forward(*args, **kwargs) + return handle_agent_outputs(outputs, self.output_type) + + def setup(self): + """ + Overwrite this method here for any operation that is expensive and needs to be executed before you start using + your tool. Such as loading a big model. + """ + self.is_initialized = True + + def save(self, output_dir): + """ + Saves the relevant code files for your tool so it can be pushed to the Hub. This will copy the code of your + tool in `output_dir` as well as autogenerate: + + - a config file named `tool_config.json` + - an `app.py` file so that your tool can be converted to a space + - a `requirements.txt` containing the names of the module used by your tool (as detected when inspecting its + code) + + You should only use this method to save tools that are defined in a separate module (not `__main__`). + + Args: + output_dir (`str`): The folder in which you want to save your tool. + """ + os.makedirs(output_dir, exist_ok=True) + # Save module file + if self.__module__ == "__main__": + raise ValueError( + f"We can't save the code defining {self} in {output_dir} as it's been defined in __main__. You " + "have to put this code in a separate module so we can include it in the saved folder." + ) + module_files = custom_object_save(self, output_dir) + + module_name = self.__class__.__module__ + last_module = module_name.split(".")[-1] + full_name = f"{last_module}.{self.__class__.__name__}" + + # Save config file + config_file = os.path.join(output_dir, "tool_config.json") + if os.path.isfile(config_file): + with open(config_file, "r", encoding="utf-8") as f: + tool_config = json.load(f) + else: + tool_config = {} + + tool_config = { + "tool_class": full_name, + "description": self.description, + "name": self.name, + "inputs": self.inputs, + "output_type": str(self.output_type), + } + with open(config_file, "w", encoding="utf-8") as f: + f.write(json.dumps(tool_config, indent=2, sort_keys=True) + "\n") + + # Save app file + app_file = os.path.join(output_dir, "app.py") + with open(app_file, "w", encoding="utf-8") as f: + f.write(APP_FILE_TEMPLATE.format(module_name=last_module, class_name=self.__class__.__name__)) + + # Save requirements file + requirements_file = os.path.join(output_dir, "requirements.txt") + imports = [] + for module in module_files: + imports.extend(get_imports(module)) + imports = list(set(imports)) + with open(requirements_file, "w", encoding="utf-8") as f: + f.write("\n".join(imports) + "\n") + + @classmethod + def from_hub( + cls, + repo_id: str, + token: Optional[str] = None, + **kwargs, + ): + """ + Loads a tool defined on the Hub. + + + + Loading a tool from the Hub means that you'll download the tool and execute it locally. + ALWAYS inspect the tool you're downloading before loading it within your runtime, as you would do when + installing a package using pip/npm/apt. + + + + Args: + repo_id (`str`): + The name of the repo on the Hub where your tool is defined. + token (`str`, *optional*): + The token to identify you on hf.co. If unset, will use the token generated when running + `huggingface-cli login` (stored in `~/.huggingface`). + kwargs (additional keyword arguments, *optional*): + Additional keyword arguments that will be split in two: all arguments relevant to the Hub (such as + `cache_dir`, `revision`, `subfolder`) will be used when downloading the files for your tool, and the + others will be passed along to its init. + """ + hub_kwargs_names = [ + "cache_dir", + "force_download", + "resume_download", + "proxies", + "revision", + "repo_type", + "subfolder", + "local_files_only", + ] + hub_kwargs = {k: v for k, v in kwargs.items() if k in hub_kwargs_names} + + # Try to get the tool config first. + hub_kwargs["repo_type"] = get_repo_type(repo_id, **hub_kwargs) + resolved_config_file = cached_file( + repo_id, + TOOL_CONFIG_FILE, + token=token, + **hub_kwargs, + _raise_exceptions_for_gated_repo=False, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + ) + is_tool_config = resolved_config_file is not None + if resolved_config_file is None: + resolved_config_file = cached_file( + repo_id, + CONFIG_NAME, + token=token, + **hub_kwargs, + _raise_exceptions_for_gated_repo=False, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + ) + if resolved_config_file is None: + raise EnvironmentError( + f"{repo_id} does not appear to provide a valid configuration in `tool_config.json` or `config.json`." + ) + + with open(resolved_config_file, encoding="utf-8") as reader: + config = json.load(reader) + + if not is_tool_config: + if "custom_tool" not in config: + raise EnvironmentError( + f"{repo_id} does not provide a mapping to custom tools in its configuration `config.json`." + ) + custom_tool = config["custom_tool"] + else: + custom_tool = config + + tool_class = custom_tool["tool_class"] + tool_class = get_class_from_dynamic_module(tool_class, repo_id, token=token, **hub_kwargs) + + if len(tool_class.name) == 0: + tool_class.name = custom_tool["name"] + if tool_class.name != custom_tool["name"]: + logger.warning( + f"{tool_class.__name__} implements a different name in its configuration and class. Using the tool " + "configuration name." + ) + tool_class.name = custom_tool["name"] + + if len(tool_class.description) == 0: + tool_class.description = custom_tool["description"] + if tool_class.description != custom_tool["description"]: + logger.warning( + f"{tool_class.__name__} implements a different description in its configuration and class. Using the " + "tool configuration description." + ) + tool_class.description = custom_tool["description"] + + if tool_class.inputs != custom_tool["inputs"]: + tool_class.inputs = custom_tool["inputs"] + if tool_class.output_type != custom_tool["output_type"]: + tool_class.output_type = custom_tool["output_type"] + + if not isinstance(tool_class.inputs, dict): + tool_class.inputs = ast.literal_eval(tool_class.inputs) + + return tool_class(**kwargs) + + def push_to_hub( + self, + repo_id: str, + commit_message: str = "Upload tool", + private: Optional[bool] = None, + token: Optional[Union[bool, str]] = None, + create_pr: bool = False, + ) -> str: + """ + Upload the tool to the Hub. + + For this method to work properly, your tool must have been defined in a separate module (not `__main__`). + For instance: + ``` + from my_tool_module import MyTool + my_tool = MyTool() + my_tool.push_to_hub("my-username/my-space") + ``` + + Parameters: + repo_id (`str`): + The name of the repository you want to push your tool to. It should contain your organization name when + pushing to a given organization. + commit_message (`str`, *optional*, defaults to `"Upload tool"`): + Message to commit while pushing. + private (`bool`, *optional*): + Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. + token (`bool` or `str`, *optional*): + The token to use as HTTP bearer authorization for remote files. If unset, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). + create_pr (`bool`, *optional*, defaults to `False`): + Whether or not to create a PR with the uploaded files or directly commit. + """ + repo_url = create_repo( + repo_id=repo_id, + token=token, + private=private, + exist_ok=True, + repo_type="space", + space_sdk="gradio", + ) + repo_id = repo_url.repo_id + metadata_update(repo_id, {"tags": ["tool"]}, repo_type="space") + + with tempfile.TemporaryDirectory() as work_dir: + # Save all files. + self.save(work_dir) + logger.info(f"Uploading the following files to {repo_id}: {','.join(os.listdir(work_dir))}") + return upload_folder( + repo_id=repo_id, + commit_message=commit_message, + folder_path=work_dir, + token=token, + create_pr=create_pr, + repo_type="space", + ) + + @staticmethod + def from_space( + space_id: str, name: str, description: str, api_name: Optional[str] = None, token: Optional[str] = None + ): + """ + Creates a [`Tool`] from a Space given its id on the Hub. + + Args: + space_id (`str`): + The id of the Space on the Hub. + name (`str`): + The name of the tool. + description (`str`): + The description of the tool. + api_name (`str`, *optional*): + The specific api_name to use, if the space has several tabs. If not precised, will default to the first available api. + token (`str`, *optional*): + Add your token to access private spaces or increase your GPU quotas. + Returns: + [`Tool`]: + The Space, as a tool. + + Examples: + ``` + image_generator = Tool.from_space( + space_id="black-forest-labs/FLUX.1-schnell", + name="image-generator", + description="Generate an image from a prompt" + ) + image = image_generator("Generate an image of a cool surfer in Tahiti") + ``` + ``` + face_swapper = Tool.from_space( + "tuan2308/face-swap", + "face_swapper", + "Tool that puts the face shown on the first image on the second image. You can give it paths to images.", + ) + image = face_swapper('./aymeric.jpeg', './ruth.jpg') + ``` + """ + from gradio_client import Client, handle_file + from gradio_client.utils import is_http_url_like + + class SpaceToolWrapper(Tool): + def __init__( + self, + space_id: str, + name: str, + description: str, + api_name: Optional[str] = None, + token: Optional[str] = None, + ): + self.client = Client(space_id, hf_token=token) + self.name = name + self.description = description + space_description = self.client.view_api(return_format="dict", print_info=False)["named_endpoints"] + + # If api_name is not defined, take the first of the available APIs for this space + if api_name is None: + api_name = list(space_description.keys())[0] + logger.warning( + f"Since `api_name` was not defined, it was automatically set to the first available API: `{api_name}`." + ) + self.api_name = api_name + + try: + space_description_api = space_description[api_name] + except KeyError: + raise KeyError(f"Could not find specified {api_name=} among available api names.") + + self.inputs = {} + for parameter in space_description_api["parameters"]: + if not parameter["parameter_has_default"]: + parameter_type = parameter["type"]["type"] + if parameter_type == "object": + parameter_type = "any" + self.inputs[parameter["parameter_name"]] = { + "type": parameter_type, + "description": parameter["python_type"]["description"], + } + output_component = space_description_api["returns"][0]["component"] + if output_component == "Image": + self.output_type = "image" + elif output_component == "Audio": + self.output_type = "audio" + else: + self.output_type = "any" + + def sanitize_argument_for_prediction(self, arg): + if isinstance(arg, ImageType): + temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False) + arg.save(temp_file.name) + arg = temp_file.name + if (isinstance(arg, (str, Path)) and Path(arg).exists() and Path(arg).is_file()) or is_http_url_like( + arg + ): + arg = handle_file(arg) + return arg + + def forward(self, *args, **kwargs): + # Preprocess args and kwargs: + args = list(args) + for i, arg in enumerate(args): + args[i] = self.sanitize_argument_for_prediction(arg) + for arg_name, arg in kwargs.items(): + kwargs[arg_name] = self.sanitize_argument_for_prediction(arg) + + output = self.client.predict(*args, api_name=self.api_name, **kwargs) + if isinstance(output, tuple) or isinstance(output, list): + return output[ + 0 + ] # Sometime the space also returns the generation seed, in which case the result is at index 0 + return output + + return SpaceToolWrapper(space_id, name, description, api_name=api_name, token=token) + + @staticmethod + def from_gradio(gradio_tool): + """ + Creates a [`Tool`] from a gradio tool. + """ + import inspect + + class GradioToolWrapper(Tool): + def __init__(self, _gradio_tool): + self.name = _gradio_tool.name + self.description = _gradio_tool.description + self.output_type = "string" + self._gradio_tool = _gradio_tool + func_args = list(inspect.signature(_gradio_tool.run).parameters.items()) + self.inputs = { + key: {"type": CONVERSION_DICT[value.annotation], "description": ""} for key, value in func_args + } + self.forward = self._gradio_tool.run + + return GradioToolWrapper(gradio_tool) + + @staticmethod + def from_langchain(langchain_tool): + """ + Creates a [`Tool`] from a langchain tool. + """ + + class LangChainToolWrapper(Tool): + def __init__(self, _langchain_tool): + self.name = _langchain_tool.name.lower() + self.description = _langchain_tool.description + self.inputs = _langchain_tool.args.copy() + for input_content in self.inputs.values(): + if "title" in input_content: + input_content.pop("title") + input_content["description"] = "" + self.output_type = "string" + self.langchain_tool = _langchain_tool + + def forward(self, *args, **kwargs): + tool_input = kwargs.copy() + for index, argument in enumerate(args): + if index < len(self.inputs): + input_key = next(iter(self.inputs)) + tool_input[input_key] = argument + return self.langchain_tool.run(tool_input) + + return LangChainToolWrapper(langchain_tool) + + +DEFAULT_TOOL_DESCRIPTION_TEMPLATE = """ +- {{ tool.name }}: {{ tool.description }} + Takes inputs: {{tool.inputs}} + Returns an output of type: {{tool.output_type}} +""" + + +def get_tool_description_with_args(tool: Tool, description_template: str = DEFAULT_TOOL_DESCRIPTION_TEMPLATE) -> str: + compiled_template = compile_jinja_template(description_template) + rendered = compiled_template.render( + tool=tool, + ) + return rendered + + +@lru_cache +def compile_jinja_template(template): + try: + import jinja2 + from jinja2.exceptions import TemplateError + from jinja2.sandbox import ImmutableSandboxedEnvironment + except ImportError: + raise ImportError("template requires jinja2 to be installed.") + + if version.parse(jinja2.__version__) < version.parse("3.1.0"): + raise ImportError(f"template requires jinja2>=3.1.0 to be installed. Your version is {jinja2.__version__}.") + + def raise_exception(message): + raise TemplateError(message) + + jinja_env = ImmutableSandboxedEnvironment(trim_blocks=True, lstrip_blocks=True) + jinja_env.globals["raise_exception"] = raise_exception + return jinja_env.from_string(template) + + +class PipelineTool(Tool): + """ + A [`Tool`] tailored towards Transformer models. On top of the class attributes of the base class [`Tool`], you will + need to specify: + + - **model_class** (`type`) -- The class to use to load the model in this tool. + - **default_checkpoint** (`str`) -- The default checkpoint that should be used when the user doesn't specify one. + - **pre_processor_class** (`type`, *optional*, defaults to [`AutoProcessor`]) -- The class to use to load the + pre-processor + - **post_processor_class** (`type`, *optional*, defaults to [`AutoProcessor`]) -- The class to use to load the + post-processor (when different from the pre-processor). + + Args: + model (`str` or [`PreTrainedModel`], *optional*): + The name of the checkpoint to use for the model, or the instantiated model. If unset, will default to the + value of the class attribute `default_checkpoint`. + pre_processor (`str` or `Any`, *optional*): + The name of the checkpoint to use for the pre-processor, or the instantiated pre-processor (can be a + tokenizer, an image processor, a feature extractor or a processor). Will default to the value of `model` if + unset. + post_processor (`str` or `Any`, *optional*): + The name of the checkpoint to use for the post-processor, or the instantiated pre-processor (can be a + tokenizer, an image processor, a feature extractor or a processor). Will default to the `pre_processor` if + unset. + device (`int`, `str` or `torch.device`, *optional*): + The device on which to execute the model. Will default to any accelerator available (GPU, MPS etc...), the + CPU otherwise. + device_map (`str` or `dict`, *optional*): + If passed along, will be used to instantiate the model. + model_kwargs (`dict`, *optional*): + Any keyword argument to send to the model instantiation. + token (`str`, *optional*): + The token to use as HTTP bearer authorization for remote files. If unset, will use the token generated when + running `huggingface-cli login` (stored in `~/.huggingface`). + hub_kwargs (additional keyword arguments, *optional*): + Any additional keyword argument to send to the methods that will load the data from the Hub. + """ + + pre_processor_class = AutoProcessor + model_class = None + post_processor_class = AutoProcessor + default_checkpoint = None + description = "This is a pipeline tool" + name = "pipeline" + inputs = {"prompt": str} + output_type = str + + def __init__( + self, + model=None, + pre_processor=None, + post_processor=None, + device=None, + device_map=None, + model_kwargs=None, + token=None, + **hub_kwargs, + ): + if not is_torch_available(): + raise ImportError("Please install torch in order to use this tool.") + + if not is_accelerate_available(): + raise ImportError("Please install accelerate in order to use this tool.") + + if model is None: + if self.default_checkpoint is None: + raise ValueError("This tool does not implement a default checkpoint, you need to pass one.") + model = self.default_checkpoint + if pre_processor is None: + pre_processor = model + + self.model = model + self.pre_processor = pre_processor + self.post_processor = post_processor + self.device = device + self.device_map = device_map + self.model_kwargs = {} if model_kwargs is None else model_kwargs + if device_map is not None: + self.model_kwargs["device_map"] = device_map + self.hub_kwargs = hub_kwargs + self.hub_kwargs["token"] = token + + super().__init__() + + def setup(self): + """ + Instantiates the `pre_processor`, `model` and `post_processor` if necessary. + """ + if isinstance(self.pre_processor, str): + self.pre_processor = self.pre_processor_class.from_pretrained(self.pre_processor, **self.hub_kwargs) + + if isinstance(self.model, str): + self.model = self.model_class.from_pretrained(self.model, **self.model_kwargs, **self.hub_kwargs) + + if self.post_processor is None: + self.post_processor = self.pre_processor + elif isinstance(self.post_processor, str): + self.post_processor = self.post_processor_class.from_pretrained(self.post_processor, **self.hub_kwargs) + + if self.device is None: + if self.device_map is not None: + self.device = list(self.model.hf_device_map.values())[0] + else: + self.device = PartialState().default_device + + if self.device_map is None: + self.model.to(self.device) + + super().setup() + + def encode(self, raw_inputs): + """ + Uses the `pre_processor` to prepare the inputs for the `model`. + """ + return self.pre_processor(raw_inputs) + + def forward(self, inputs): + """ + Sends the inputs through the `model`. + """ + with torch.no_grad(): + return self.model(**inputs) + + def decode(self, outputs): + """ + Uses the `post_processor` to decode the model output. + """ + return self.post_processor(outputs) + + def __call__(self, *args, **kwargs): + args, kwargs = handle_agent_inputs(*args, **kwargs) + + if not self.is_initialized: + self.setup() + + encoded_inputs = self.encode(*args, **kwargs) + + tensor_inputs = {k: v for k, v in encoded_inputs.items() if isinstance(v, torch.Tensor)} + non_tensor_inputs = {k: v for k, v in encoded_inputs.items() if not isinstance(v, torch.Tensor)} + + encoded_inputs = send_to_device(tensor_inputs, self.device) + outputs = self.forward({**encoded_inputs, **non_tensor_inputs}) + outputs = send_to_device(outputs, "cpu") + decoded_outputs = self.decode(outputs) + + return handle_agent_outputs(decoded_outputs, self.output_type) + + +def launch_gradio_demo(tool_class: Tool): + """ + Launches a gradio demo for a tool. The corresponding tool class needs to properly implement the class attributes + `inputs` and `output_type`. + + Args: + tool_class (`type`): The class of the tool for which to launch the demo. + """ + try: + import gradio as gr + except ImportError: + raise ImportError("Gradio should be installed in order to launch a gradio demo.") + + tool = tool_class() + + def fn(*args, **kwargs): + return tool(*args, **kwargs) + + TYPE_TO_COMPONENT_CLASS_MAPPING = { + "image": gr.Image, + "audio": gr.Audio, + "string": gr.Textbox, + "integer": gr.Textbox, + "number": gr.Textbox, + } + + gradio_inputs = [] + for input_name, input_details in tool_class.inputs.items(): + input_gradio_component_class = TYPE_TO_COMPONENT_CLASS_MAPPING[input_details["type"]] + new_component = input_gradio_component_class(label=input_name) + gradio_inputs.append(new_component) + + output_gradio_componentclass = TYPE_TO_COMPONENT_CLASS_MAPPING[tool_class.output_type] + gradio_output = output_gradio_componentclass(label=input_name) + + gr.Interface( + fn=fn, + inputs=gradio_inputs, + outputs=gradio_output, + title=tool_class.__name__, + article=tool.description, + ).launch() + + +TOOL_MAPPING = { + "document_question_answering": "DocumentQuestionAnsweringTool", + "image_question_answering": "ImageQuestionAnsweringTool", + "speech_to_text": "SpeechToTextTool", + "text_to_speech": "TextToSpeechTool", + "translation": "TranslationTool", + "python_interpreter": "PythonInterpreterTool", + "web_search": "DuckDuckGoSearchTool", +} + + +def load_tool(task_or_repo_id, model_repo_id=None, token=None, **kwargs): + """ + Main function to quickly load a tool, be it on the Hub or in the Transformers library. + + + + Loading a tool means that you'll download the tool and execute it locally. + ALWAYS inspect the tool you're downloading before loading it within your runtime, as you would do when + installing a package using pip/npm/apt. + + + + Args: + task_or_repo_id (`str`): + The task for which to load the tool or a repo ID of a tool on the Hub. Tasks implemented in Transformers + are: + + - `"document_question_answering"` + - `"image_question_answering"` + - `"speech_to_text"` + - `"text_to_speech"` + - `"translation"` + + model_repo_id (`str`, *optional*): + Use this argument to use a different model than the default one for the tool you selected. + token (`str`, *optional*): + The token to identify you on hf.co. If unset, will use the token generated when running `huggingface-cli + login` (stored in `~/.huggingface`). + kwargs (additional keyword arguments, *optional*): + Additional keyword arguments that will be split in two: all arguments relevant to the Hub (such as + `cache_dir`, `revision`, `subfolder`) will be used when downloading the files for your tool, and the others + will be passed along to its init. + """ + if task_or_repo_id in TOOL_MAPPING: + tool_class_name = TOOL_MAPPING[task_or_repo_id] + main_module = importlib.import_module("transformers") + tools_module = main_module.agents + tool_class = getattr(tools_module, tool_class_name) + return tool_class(model_repo_id, token=token, **kwargs) + else: + logger.warning_once( + f"You're loading a tool from the Hub from {model_repo_id}. Please make sure this is a source that you " + f"trust as the code within that tool will be executed on your machine. Always verify the code of " + f"the tools that you load. We recommend specifying a `revision` to ensure you're loading the " + f"code that you have checked." + ) + return Tool.from_hub(task_or_repo_id, model_repo_id=model_repo_id, token=token, **kwargs) + + +def add_description(description): + """ + A decorator that adds a description to a function. + """ + + def inner(func): + func.description = description + func.name = func.__name__ + return func + + return inner + + +## Will move to the Hub +class EndpointClient: + def __init__(self, endpoint_url: str, token: Optional[str] = None): + self.headers = { + **build_hf_headers(token=token), + "Content-Type": "application/json", + } + self.endpoint_url = endpoint_url + + @staticmethod + def encode_image(image): + _bytes = io.BytesIO() + image.save(_bytes, format="PNG") + b64 = base64.b64encode(_bytes.getvalue()) + return b64.decode("utf-8") + + @staticmethod + def decode_image(raw_image): + if not is_vision_available(): + raise ImportError( + "This tool returned an image but Pillow is not installed. Please install it (`pip install Pillow`)." + ) + + from PIL import Image + + b64 = base64.b64decode(raw_image) + _bytes = io.BytesIO(b64) + return Image.open(_bytes) + + def __call__( + self, + inputs: Optional[Union[str, Dict, List[str], List[List[str]]]] = None, + params: Optional[Dict] = None, + data: Optional[bytes] = None, + output_image: bool = False, + ) -> Any: + # Build payload + payload = {} + if inputs: + payload["inputs"] = inputs + if params: + payload["parameters"] = params + + # Make API call + response = get_session().post(self.endpoint_url, headers=self.headers, json=payload, data=data) + + # By default, parse the response for the user. + if output_image: + return self.decode_image(response.content) + else: + return response.json() + + +class ToolCollection: + """ + Tool collections enable loading all Spaces from a collection in order to be added to the agent's toolbox. + + > [!NOTE] + > Only Spaces will be fetched, so you can feel free to add models and datasets to your collection if you'd + > like for this collection to showcase them. + + Args: + collection_slug (str): + The collection slug referencing the collection. + token (str, *optional*): + The authentication token if the collection is private. + + Example: + + ```py + >>> from transformers import ToolCollection, ReactCodeAgent + + >>> image_tool_collection = ToolCollection(collection_slug="huggingface-tools/diffusion-tools-6630bb19a942c2306a2cdb6f") + >>> agent = ReactCodeAgent(tools=[*image_tool_collection.tools], add_base_tools=True) + + >>> agent.run("Please draw me a picture of rivers and lakes.") + ``` + """ + + def __init__(self, collection_slug: str, token: Optional[str] = None): + self._collection = get_collection(collection_slug, token=token) + self._hub_repo_ids = {item.item_id for item in self._collection.items if item.item_type == "space"} + self.tools = {Tool.from_hub(repo_id) for repo_id in self._hub_repo_ids} + + +def tool(tool_function: Callable) -> Tool: + """ + Converts a function into an instance of a Tool subclass. + + Args: + tool_function: Your function. Should have type hints for each input and a type hint for the output. + Should also have a docstring description including an 'Args:' part where each argument is described. + """ + parameters = get_json_schema(tool_function)["function"] + if "return" not in parameters: + raise TypeHintParsingException("Tool return type not found: make sure your function has a return type hint!") + class_name = f"{parameters['name'].capitalize()}Tool" + + class SpecificTool(Tool): + name = parameters["name"] + description = parameters["description"] + inputs = parameters["parameters"]["properties"] + output_type = parameters["return"]["type"] + + @wraps(tool_function) + def forward(self, *args, **kwargs): + return tool_function(*args, **kwargs) + + original_signature = inspect.signature(tool_function) + new_parameters = [inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD)] + list( + original_signature.parameters.values() + ) + new_signature = original_signature.replace(parameters=new_parameters) + SpecificTool.forward.__signature__ = new_signature + + SpecificTool.__name__ = class_name + return SpecificTool() diff --git a/lib/python3.12/site-packages/transformers/agents/translation.py b/lib/python3.12/site-packages/transformers/agents/translation.py new file mode 100644 index 0000000000000000000000000000000000000000..6875aac4bc64f8fa5609c759d788866567591a8c --- /dev/null +++ b/lib/python3.12/site-packages/transformers/agents/translation.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2023 The HuggingFace Inc. 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. +from ..models.auto import AutoModelForSeq2SeqLM, AutoTokenizer +from .tools import PipelineTool + + +LANGUAGE_CODES = { + "Acehnese Arabic": "ace_Arab", + "Acehnese Latin": "ace_Latn", + "Mesopotamian Arabic": "acm_Arab", + "Ta'izzi-Adeni Arabic": "acq_Arab", + "Tunisian Arabic": "aeb_Arab", + "Afrikaans": "afr_Latn", + "South Levantine Arabic": "ajp_Arab", + "Akan": "aka_Latn", + "Amharic": "amh_Ethi", + "North Levantine Arabic": "apc_Arab", + "Modern Standard Arabic": "arb_Arab", + "Modern Standard Arabic Romanized": "arb_Latn", + "Najdi Arabic": "ars_Arab", + "Moroccan Arabic": "ary_Arab", + "Egyptian Arabic": "arz_Arab", + "Assamese": "asm_Beng", + "Asturian": "ast_Latn", + "Awadhi": "awa_Deva", + "Central Aymara": "ayr_Latn", + "South Azerbaijani": "azb_Arab", + "North Azerbaijani": "azj_Latn", + "Bashkir": "bak_Cyrl", + "Bambara": "bam_Latn", + "Balinese": "ban_Latn", + "Belarusian": "bel_Cyrl", + "Bemba": "bem_Latn", + "Bengali": "ben_Beng", + "Bhojpuri": "bho_Deva", + "Banjar Arabic": "bjn_Arab", + "Banjar Latin": "bjn_Latn", + "Standard Tibetan": "bod_Tibt", + "Bosnian": "bos_Latn", + "Buginese": "bug_Latn", + "Bulgarian": "bul_Cyrl", + "Catalan": "cat_Latn", + "Cebuano": "ceb_Latn", + "Czech": "ces_Latn", + "Chokwe": "cjk_Latn", + "Central Kurdish": "ckb_Arab", + "Crimean Tatar": "crh_Latn", + "Welsh": "cym_Latn", + "Danish": "dan_Latn", + "German": "deu_Latn", + "Southwestern Dinka": "dik_Latn", + "Dyula": "dyu_Latn", + "Dzongkha": "dzo_Tibt", + "Greek": "ell_Grek", + "English": "eng_Latn", + "Esperanto": "epo_Latn", + "Estonian": "est_Latn", + "Basque": "eus_Latn", + "Ewe": "ewe_Latn", + "Faroese": "fao_Latn", + "Fijian": "fij_Latn", + "Finnish": "fin_Latn", + "Fon": "fon_Latn", + "French": "fra_Latn", + "Friulian": "fur_Latn", + "Nigerian Fulfulde": "fuv_Latn", + "Scottish Gaelic": "gla_Latn", + "Irish": "gle_Latn", + "Galician": "glg_Latn", + "Guarani": "grn_Latn", + "Gujarati": "guj_Gujr", + "Haitian Creole": "hat_Latn", + "Hausa": "hau_Latn", + "Hebrew": "heb_Hebr", + "Hindi": "hin_Deva", + "Chhattisgarhi": "hne_Deva", + "Croatian": "hrv_Latn", + "Hungarian": "hun_Latn", + "Armenian": "hye_Armn", + "Igbo": "ibo_Latn", + "Ilocano": "ilo_Latn", + "Indonesian": "ind_Latn", + "Icelandic": "isl_Latn", + "Italian": "ita_Latn", + "Javanese": "jav_Latn", + "Japanese": "jpn_Jpan", + "Kabyle": "kab_Latn", + "Jingpho": "kac_Latn", + "Kamba": "kam_Latn", + "Kannada": "kan_Knda", + "Kashmiri Arabic": "kas_Arab", + "Kashmiri Devanagari": "kas_Deva", + "Georgian": "kat_Geor", + "Central Kanuri Arabic": "knc_Arab", + "Central Kanuri Latin": "knc_Latn", + "Kazakh": "kaz_Cyrl", + "Kabiyè": "kbp_Latn", + "Kabuverdianu": "kea_Latn", + "Khmer": "khm_Khmr", + "Kikuyu": "kik_Latn", + "Kinyarwanda": "kin_Latn", + "Kyrgyz": "kir_Cyrl", + "Kimbundu": "kmb_Latn", + "Northern Kurdish": "kmr_Latn", + "Kikongo": "kon_Latn", + "Korean": "kor_Hang", + "Lao": "lao_Laoo", + "Ligurian": "lij_Latn", + "Limburgish": "lim_Latn", + "Lingala": "lin_Latn", + "Lithuanian": "lit_Latn", + "Lombard": "lmo_Latn", + "Latgalian": "ltg_Latn", + "Luxembourgish": "ltz_Latn", + "Luba-Kasai": "lua_Latn", + "Ganda": "lug_Latn", + "Luo": "luo_Latn", + "Mizo": "lus_Latn", + "Standard Latvian": "lvs_Latn", + "Magahi": "mag_Deva", + "Maithili": "mai_Deva", + "Malayalam": "mal_Mlym", + "Marathi": "mar_Deva", + "Minangkabau Arabic ": "min_Arab", + "Minangkabau Latin": "min_Latn", + "Macedonian": "mkd_Cyrl", + "Plateau Malagasy": "plt_Latn", + "Maltese": "mlt_Latn", + "Meitei Bengali": "mni_Beng", + "Halh Mongolian": "khk_Cyrl", + "Mossi": "mos_Latn", + "Maori": "mri_Latn", + "Burmese": "mya_Mymr", + "Dutch": "nld_Latn", + "Norwegian Nynorsk": "nno_Latn", + "Norwegian Bokmål": "nob_Latn", + "Nepali": "npi_Deva", + "Northern Sotho": "nso_Latn", + "Nuer": "nus_Latn", + "Nyanja": "nya_Latn", + "Occitan": "oci_Latn", + "West Central Oromo": "gaz_Latn", + "Odia": "ory_Orya", + "Pangasinan": "pag_Latn", + "Eastern Panjabi": "pan_Guru", + "Papiamento": "pap_Latn", + "Western Persian": "pes_Arab", + "Polish": "pol_Latn", + "Portuguese": "por_Latn", + "Dari": "prs_Arab", + "Southern Pashto": "pbt_Arab", + "Ayacucho Quechua": "quy_Latn", + "Romanian": "ron_Latn", + "Rundi": "run_Latn", + "Russian": "rus_Cyrl", + "Sango": "sag_Latn", + "Sanskrit": "san_Deva", + "Santali": "sat_Olck", + "Sicilian": "scn_Latn", + "Shan": "shn_Mymr", + "Sinhala": "sin_Sinh", + "Slovak": "slk_Latn", + "Slovenian": "slv_Latn", + "Samoan": "smo_Latn", + "Shona": "sna_Latn", + "Sindhi": "snd_Arab", + "Somali": "som_Latn", + "Southern Sotho": "sot_Latn", + "Spanish": "spa_Latn", + "Tosk Albanian": "als_Latn", + "Sardinian": "srd_Latn", + "Serbian": "srp_Cyrl", + "Swati": "ssw_Latn", + "Sundanese": "sun_Latn", + "Swedish": "swe_Latn", + "Swahili": "swh_Latn", + "Silesian": "szl_Latn", + "Tamil": "tam_Taml", + "Tatar": "tat_Cyrl", + "Telugu": "tel_Telu", + "Tajik": "tgk_Cyrl", + "Tagalog": "tgl_Latn", + "Thai": "tha_Thai", + "Tigrinya": "tir_Ethi", + "Tamasheq Latin": "taq_Latn", + "Tamasheq Tifinagh": "taq_Tfng", + "Tok Pisin": "tpi_Latn", + "Tswana": "tsn_Latn", + "Tsonga": "tso_Latn", + "Turkmen": "tuk_Latn", + "Tumbuka": "tum_Latn", + "Turkish": "tur_Latn", + "Twi": "twi_Latn", + "Central Atlas Tamazight": "tzm_Tfng", + "Uyghur": "uig_Arab", + "Ukrainian": "ukr_Cyrl", + "Umbundu": "umb_Latn", + "Urdu": "urd_Arab", + "Northern Uzbek": "uzn_Latn", + "Venetian": "vec_Latn", + "Vietnamese": "vie_Latn", + "Waray": "war_Latn", + "Wolof": "wol_Latn", + "Xhosa": "xho_Latn", + "Eastern Yiddish": "ydd_Hebr", + "Yoruba": "yor_Latn", + "Yue Chinese": "yue_Hant", + "Chinese Simplified": "zho_Hans", + "Chinese Traditional": "zho_Hant", + "Standard Malay": "zsm_Latn", + "Zulu": "zul_Latn", +} + + +class TranslationTool(PipelineTool): + """ + Example: + + ```py + from transformers.agents import TranslationTool + + translator = TranslationTool() + translator("This is a super nice API!", src_lang="English", tgt_lang="French") + ``` + """ + + lang_to_code = LANGUAGE_CODES + default_checkpoint = "facebook/nllb-200-distilled-600M" + description = ( + "This is a tool that translates text from a language to another." + f"Both `src_lang`and `tgt_lang` should belong to this list of languages: {list(lang_to_code.keys())}." + ) + name = "translator" + pre_processor_class = AutoTokenizer + model_class = AutoModelForSeq2SeqLM + + inputs = { + "text": {"type": "string", "description": "The text to translate"}, + "src_lang": { + "type": "string", + "description": "The language of the text to translate. Written in plain English, such as 'Romanian', or 'Albanian'", + }, + "tgt_lang": { + "type": "string", + "description": "The language for the desired output language. Written in plain English, such as 'Romanian', or 'Albanian'", + }, + } + output_type = "string" + + def encode(self, text, src_lang, tgt_lang): + if src_lang not in self.lang_to_code: + raise ValueError(f"{src_lang} is not a supported language.") + if tgt_lang not in self.lang_to_code: + raise ValueError(f"{tgt_lang} is not a supported language.") + src_lang = self.lang_to_code[src_lang] + tgt_lang = self.lang_to_code[tgt_lang] + return self.pre_processor._build_translation_inputs( + text, return_tensors="pt", src_lang=src_lang, tgt_lang=tgt_lang + ) + + def forward(self, inputs): + return self.model.generate(**inputs) + + def decode(self, outputs): + return self.post_processor.decode(outputs[0].tolist(), skip_special_tokens=True) diff --git a/lib/python3.12/site-packages/transformers/audio_utils.py b/lib/python3.12/site-packages/transformers/audio_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8420a84e089e03be9a80fb63c237e34203ea28a0 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/audio_utils.py @@ -0,0 +1,1193 @@ +# Copyright 2023 The HuggingFace Inc. team and the librosa & torchaudio authors. +# +# 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. +""" +Audio processing functions to extract features from audio waveforms. This code is pure numpy to support all frameworks +and remove unnecessary dependencies. +""" + +import os +import warnings +from io import BytesIO +from typing import List, Optional, Tuple, Union + +import numpy as np +import requests + +from .utils import is_librosa_available, requires_backends + + +if is_librosa_available(): + import librosa + + +def load_audio(audio: Union[str, np.ndarray], sampling_rate=16000, timeout=None) -> np.ndarray: + """ + Loads `audio` to an np.ndarray object. + + Args: + audio (`str` or `np.ndarray`): + The audio to be laoded to the numpy array format. + sampling_rate (`int`, *optional*, defaults to 16000): + The samlping rate to be used when loading the audio. It should be same as the + sampling rate the model you will be using further was trained with. + timeout (`float`, *optional*): + The timeout value in seconds for the URL request. + + Returns: + `np.ndarray`: A numpy artay representing the audio. + """ + requires_backends(load_audio, ["librosa"]) + + if isinstance(audio, str): + # Load audio from URL (e.g https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2-Audio/audio/translate_to_chinese.wav) + if audio.startswith("http://") or audio.startswith("https://"): + audio = librosa.load(BytesIO(requests.get(audio, timeout=timeout).content), sr=sampling_rate)[0] + elif os.path.isfile(audio): + audio = librosa.load(audio, sr=sampling_rate)[0] + elif isinstance(audio, np.ndarray): + audio = audio + else: + raise TypeError( + "Incorrect format used for `audio`. Should be an url linking to an audio, a local path, or numpy array." + ) + return audio + + +AudioInput = Union[ + np.ndarray, "torch.Tensor", List[np.ndarray], Tuple[np.ndarray], List["torch.Tensor"], Tuple["torch.Tensor"] # noqa: F821 +] + + +def hertz_to_mel(freq: Union[float, np.ndarray], mel_scale: str = "htk") -> Union[float, np.ndarray]: + """ + Convert frequency from hertz to mels. + + Args: + freq (`float` or `np.ndarray`): + The frequency, or multiple frequencies, in hertz (Hz). + mel_scale (`str`, *optional*, defaults to `"htk"`): + The mel frequency scale to use, `"htk"`, `"kaldi"` or `"slaney"`. + + Returns: + `float` or `np.ndarray`: The frequencies on the mel scale. + """ + + if mel_scale not in ["slaney", "htk", "kaldi"]: + raise ValueError('mel_scale should be one of "htk", "slaney" or "kaldi".') + + if mel_scale == "htk": + return 2595.0 * np.log10(1.0 + (freq / 700.0)) + elif mel_scale == "kaldi": + return 1127.0 * np.log(1.0 + (freq / 700.0)) + + min_log_hertz = 1000.0 + min_log_mel = 15.0 + logstep = 27.0 / np.log(6.4) + mels = 3.0 * freq / 200.0 + + if isinstance(freq, np.ndarray): + log_region = freq >= min_log_hertz + mels[log_region] = min_log_mel + np.log(freq[log_region] / min_log_hertz) * logstep + elif freq >= min_log_hertz: + mels = min_log_mel + np.log(freq / min_log_hertz) * logstep + + return mels + + +def mel_to_hertz(mels: Union[float, np.ndarray], mel_scale: str = "htk") -> Union[float, np.ndarray]: + """ + Convert frequency from mels to hertz. + + Args: + mels (`float` or `np.ndarray`): + The frequency, or multiple frequencies, in mels. + mel_scale (`str`, *optional*, `"htk"`): + The mel frequency scale to use, `"htk"`, `"kaldi"` or `"slaney"`. + + Returns: + `float` or `np.ndarray`: The frequencies in hertz. + """ + + if mel_scale not in ["slaney", "htk", "kaldi"]: + raise ValueError('mel_scale should be one of "htk", "slaney" or "kaldi".') + + if mel_scale == "htk": + return 700.0 * (np.power(10, mels / 2595.0) - 1.0) + elif mel_scale == "kaldi": + return 700.0 * (np.exp(mels / 1127.0) - 1.0) + + min_log_hertz = 1000.0 + min_log_mel = 15.0 + logstep = np.log(6.4) / 27.0 + freq = 200.0 * mels / 3.0 + + if isinstance(mels, np.ndarray): + log_region = mels >= min_log_mel + freq[log_region] = min_log_hertz * np.exp(logstep * (mels[log_region] - min_log_mel)) + elif mels >= min_log_mel: + freq = min_log_hertz * np.exp(logstep * (mels - min_log_mel)) + + return freq + + +def hertz_to_octave( + freq: Union[float, np.ndarray], tuning: Optional[float] = 0.0, bins_per_octave: Optional[int] = 12 +): + """ + Convert frequency from hertz to fractional octave numbers. + Adapted from *librosa*. + + Args: + freq (`float` or `np.ndarray`): + The frequency, or multiple frequencies, in hertz (Hz). + tuning (`float`, defaults to `0.`): + Tuning deviation from the Stuttgart pitch (A440) in (fractional) bins per octave. + bins_per_octave (`int`, defaults to `12`): + Number of bins per octave. + + Returns: + `float` or `np.ndarray`: The frequencies on the octave scale. + """ + stuttgart_pitch = 440.0 * 2.0 ** (tuning / bins_per_octave) + octave = np.log2(freq / (float(stuttgart_pitch) / 16)) + return octave + + +def _create_triangular_filter_bank(fft_freqs: np.ndarray, filter_freqs: np.ndarray) -> np.ndarray: + """ + Creates a triangular filter bank. + + Adapted from *torchaudio* and *librosa*. + + Args: + fft_freqs (`np.ndarray` of shape `(num_frequency_bins,)`): + Discrete frequencies of the FFT bins in Hz. + filter_freqs (`np.ndarray` of shape `(num_mel_filters,)`): + Center frequencies of the triangular filters to create, in Hz. + + Returns: + `np.ndarray` of shape `(num_frequency_bins, num_mel_filters)` + """ + filter_diff = np.diff(filter_freqs) + slopes = np.expand_dims(filter_freqs, 0) - np.expand_dims(fft_freqs, 1) + down_slopes = -slopes[:, :-2] / filter_diff[:-1] + up_slopes = slopes[:, 2:] / filter_diff[1:] + return np.maximum(np.zeros(1), np.minimum(down_slopes, up_slopes)) + + +def chroma_filter_bank( + num_frequency_bins: int, + num_chroma: int, + sampling_rate: int, + tuning: float = 0.0, + power: Optional[float] = 2.0, + weighting_parameters: Optional[tuple[float, float]] = (5.0, 2.0), + start_at_c_chroma: Optional[bool] = True, +): + """ + Creates a chroma filter bank, i.e a linear transformation to project spectrogram bins onto chroma bins. + + Adapted from *librosa*. + + Args: + num_frequency_bins (`int`): + Number of frequencies used to compute the spectrogram (should be the same as in `stft`). + num_chroma (`int`): + Number of chroma bins (i.e pitch classes). + sampling_rate (`float`): + Sample rate of the audio waveform. + tuning (`float`): + Tuning deviation from A440 in fractions of a chroma bin. + power (`float`, *optional*, defaults to 2.0): + If 12.0, normalizes each column with their L2 norm. If 1.0, normalizes each column with their L1 norm. + weighting_parameters (`Tuple[float, float]`, *optional*, defaults to `(5., 2.)`): + If specified, apply a Gaussian weighting parameterized by the first element of the tuple being the center and + the second element being the Gaussian half-width. + start_at_c_chroma (`float`, *optional*, defaults to `True`): + If True, the filter bank will start at the 'C' pitch class. Otherwise, it will start at 'A'. + Returns: + `np.ndarray` of shape `(num_frequency_bins, num_chroma)` + """ + # Get the FFT bins, not counting the DC component + frequencies = np.linspace(0, sampling_rate, num_frequency_bins, endpoint=False)[1:] + + freq_bins = num_chroma * hertz_to_octave(frequencies, tuning=tuning, bins_per_octave=num_chroma) + + # make up a value for the 0 Hz bin = 1.5 octaves below bin 1 + # (so chroma is 50% rotated from bin 1, and bin width is broad) + freq_bins = np.concatenate(([freq_bins[0] - 1.5 * num_chroma], freq_bins)) + + bins_width = np.concatenate((np.maximum(freq_bins[1:] - freq_bins[:-1], 1.0), [1])) + + chroma_filters = np.subtract.outer(freq_bins, np.arange(0, num_chroma, dtype="d")).T + + num_chroma2 = np.round(float(num_chroma) / 2) + + # Project into range -num_chroma/2 .. num_chroma/2 + # add on fixed offset of 10*num_chroma to ensure all values passed to + # rem are positive + chroma_filters = np.remainder(chroma_filters + num_chroma2 + 10 * num_chroma, num_chroma) - num_chroma2 + + # Gaussian bumps - 2*D to make them narrower + chroma_filters = np.exp(-0.5 * (2 * chroma_filters / np.tile(bins_width, (num_chroma, 1))) ** 2) + + # normalize each column + if power is not None: + chroma_filters = chroma_filters / np.sum(chroma_filters**power, axis=0, keepdims=True) ** (1.0 / power) + + # Maybe apply scaling for fft bins + if weighting_parameters is not None: + center, half_width = weighting_parameters + chroma_filters *= np.tile( + np.exp(-0.5 * (((freq_bins / num_chroma - center) / half_width) ** 2)), + (num_chroma, 1), + ) + + if start_at_c_chroma: + chroma_filters = np.roll(chroma_filters, -3 * (num_chroma // 12), axis=0) + + # remove aliasing columns, copy to ensure row-contiguity + return np.ascontiguousarray(chroma_filters[:, : int(1 + num_frequency_bins / 2)]) + + +def mel_filter_bank( + num_frequency_bins: int, + num_mel_filters: int, + min_frequency: float, + max_frequency: float, + sampling_rate: int, + norm: Optional[str] = None, + mel_scale: str = "htk", + triangularize_in_mel_space: bool = False, +) -> np.ndarray: + """ + Creates a frequency bin conversion matrix used to obtain a mel spectrogram. This is called a *mel filter bank*, and + various implementation exist, which differ in the number of filters, the shape of the filters, the way the filters + are spaced, the bandwidth of the filters, and the manner in which the spectrum is warped. The goal of these + features is to approximate the non-linear human perception of the variation in pitch with respect to the frequency. + + Different banks of mel filters were introduced in the literature. The following variations are supported: + + - MFCC FB-20: introduced in 1980 by Davis and Mermelstein, it assumes a sampling frequency of 10 kHz and a speech + bandwidth of `[0, 4600]` Hz. + - MFCC FB-24 HTK: from the Cambridge HMM Toolkit (HTK) (1995) uses a filter bank of 24 filters for a speech + bandwidth of `[0, 8000]` Hz. This assumes sampling rate ≥ 16 kHz. + - MFCC FB-40: from the Auditory Toolbox for MATLAB written by Slaney in 1998, assumes a sampling rate of 16 kHz and + speech bandwidth of `[133, 6854]` Hz. This version also includes area normalization. + - HFCC-E FB-29 (Human Factor Cepstral Coefficients) of Skowronski and Harris (2004), assumes a sampling rate of + 12.5 kHz and speech bandwidth of `[0, 6250]` Hz. + + This code is adapted from *torchaudio* and *librosa*. Note that the default parameters of torchaudio's + `melscale_fbanks` implement the `"htk"` filters while librosa uses the `"slaney"` implementation. + + Args: + num_frequency_bins (`int`): + Number of frequency bins (should be the same as `n_fft // 2 + 1` where `n_fft` is the size of the Fourier Transform used to compute the spectrogram). + num_mel_filters (`int`): + Number of mel filters to generate. + min_frequency (`float`): + Lowest frequency of interest in Hz. + max_frequency (`float`): + Highest frequency of interest in Hz. This should not exceed `sampling_rate / 2`. + sampling_rate (`int`): + Sample rate of the audio waveform. + norm (`str`, *optional*): + If `"slaney"`, divide the triangular mel weights by the width of the mel band (area normalization). + mel_scale (`str`, *optional*, defaults to `"htk"`): + The mel frequency scale to use, `"htk"`, `"kaldi"` or `"slaney"`. + triangularize_in_mel_space (`bool`, *optional*, defaults to `False`): + If this option is enabled, the triangular filter is applied in mel space rather than frequency space. This + should be set to `true` in order to get the same results as `torchaudio` when computing mel filters. + + Returns: + `np.ndarray` of shape (`num_frequency_bins`, `num_mel_filters`): Triangular filter bank matrix. This is a + projection matrix to go from a spectrogram to a mel spectrogram. + """ + if norm is not None and norm != "slaney": + raise ValueError('norm must be one of None or "slaney"') + + if num_frequency_bins < 2: + raise ValueError(f"Require num_frequency_bins: {num_frequency_bins} >= 2") + + if min_frequency > max_frequency: + raise ValueError(f"Require min_frequency: {min_frequency} <= max_frequency: {max_frequency}") + + # center points of the triangular mel filters + mel_min = hertz_to_mel(min_frequency, mel_scale=mel_scale) + mel_max = hertz_to_mel(max_frequency, mel_scale=mel_scale) + mel_freqs = np.linspace(mel_min, mel_max, num_mel_filters + 2) + filter_freqs = mel_to_hertz(mel_freqs, mel_scale=mel_scale) + + if triangularize_in_mel_space: + # frequencies of FFT bins in Hz, but filters triangularized in mel space + fft_bin_width = sampling_rate / ((num_frequency_bins - 1) * 2) + fft_freqs = hertz_to_mel(fft_bin_width * np.arange(num_frequency_bins), mel_scale=mel_scale) + filter_freqs = mel_freqs + else: + # frequencies of FFT bins in Hz + fft_freqs = np.linspace(0, sampling_rate // 2, num_frequency_bins) + + mel_filters = _create_triangular_filter_bank(fft_freqs, filter_freqs) + + if norm is not None and norm == "slaney": + # Slaney-style mel is scaled to be approx constant energy per channel + enorm = 2.0 / (filter_freqs[2 : num_mel_filters + 2] - filter_freqs[:num_mel_filters]) + mel_filters *= np.expand_dims(enorm, 0) + + if (mel_filters.max(axis=0) == 0.0).any(): + warnings.warn( + "At least one mel filter has all zero values. " + f"The value for `num_mel_filters` ({num_mel_filters}) may be set too high. " + f"Or, the value for `num_frequency_bins` ({num_frequency_bins}) may be set too low." + ) + + return mel_filters + + +def optimal_fft_length(window_length: int) -> int: + """ + Finds the best FFT input size for a given `window_length`. This function takes a given window length and, if not + already a power of two, rounds it up to the next power or two. + + The FFT algorithm works fastest when the length of the input is a power of two, which may be larger than the size + of the window or analysis frame. For example, if the window is 400 samples, using an FFT input size of 512 samples + is more optimal than an FFT size of 400 samples. Using a larger FFT size does not affect the detected frequencies, + it simply gives a higher frequency resolution (i.e. the frequency bins are smaller). + """ + return 2 ** int(np.ceil(np.log2(window_length))) + + +def window_function( + window_length: int, + name: str = "hann", + periodic: bool = True, + frame_length: Optional[int] = None, + center: bool = True, +) -> np.ndarray: + """ + Returns an array containing the specified window. This window is intended to be used with `stft`. + + The following window types are supported: + + - `"boxcar"`: a rectangular window + - `"hamming"`: the Hamming window + - `"hann"`: the Hann window + - `"povey"`: the Povey window + + Args: + window_length (`int`): + The length of the window in samples. + name (`str`, *optional*, defaults to `"hann"`): + The name of the window function. + periodic (`bool`, *optional*, defaults to `True`): + Whether the window is periodic or symmetric. + frame_length (`int`, *optional*): + The length of the analysis frames in samples. Provide a value for `frame_length` if the window is smaller + than the frame length, so that it will be zero-padded. + center (`bool`, *optional*, defaults to `True`): + Whether to center the window inside the FFT buffer. Only used when `frame_length` is provided. + + Returns: + `np.ndarray` of shape `(window_length,)` or `(frame_length,)` containing the window. + """ + length = window_length + 1 if periodic else window_length + + if name == "boxcar": + window = np.ones(length) + elif name in ["hamming", "hamming_window"]: + window = np.hamming(length) + elif name in ["hann", "hann_window"]: + window = np.hanning(length) + elif name in ["povey"]: + window = np.power(np.hanning(length), 0.85) + else: + raise ValueError(f"Unknown window function '{name}'") + + if periodic: + window = window[:-1] + + if frame_length is None: + return window + + if window_length > frame_length: + raise ValueError( + f"Length of the window ({window_length}) may not be larger than frame_length ({frame_length})" + ) + + padded_window = np.zeros(frame_length) + offset = (frame_length - window_length) // 2 if center else 0 + padded_window[offset : offset + window_length] = window + return padded_window + + +# TODO This method does not support batching yet as we are mainly focused on inference. +def spectrogram( + waveform: np.ndarray, + window: np.ndarray, + frame_length: int, + hop_length: int, + fft_length: Optional[int] = None, + power: Optional[float] = 1.0, + center: bool = True, + pad_mode: str = "reflect", + onesided: bool = True, + dither: float = 0.0, + preemphasis: Optional[float] = None, + mel_filters: Optional[np.ndarray] = None, + mel_floor: float = 1e-10, + log_mel: Optional[str] = None, + reference: float = 1.0, + min_value: float = 1e-10, + db_range: Optional[float] = None, + remove_dc_offset: Optional[bool] = None, + dtype: np.dtype = np.float32, +) -> np.ndarray: + """ + Calculates a spectrogram over one waveform using the Short-Time Fourier Transform. + + This function can create the following kinds of spectrograms: + + - amplitude spectrogram (`power = 1.0`) + - power spectrogram (`power = 2.0`) + - complex-valued spectrogram (`power = None`) + - log spectrogram (use `log_mel` argument) + - mel spectrogram (provide `mel_filters`) + - log-mel spectrogram (provide `mel_filters` and `log_mel`) + + How this works: + + 1. The input waveform is split into frames of size `frame_length` that are partially overlapping by `frame_length + - hop_length` samples. + 2. Each frame is multiplied by the window and placed into a buffer of size `fft_length`. + 3. The DFT is taken of each windowed frame. + 4. The results are stacked into a spectrogram. + + We make a distinction between the following "blocks" of sample data, each of which may have a different lengths: + + - The analysis frame. This is the size of the time slices that the input waveform is split into. + - The window. Each analysis frame is multiplied by the window to avoid spectral leakage. + - The FFT input buffer. The length of this determines how many frequency bins are in the spectrogram. + + In this implementation, the window is assumed to be zero-padded to have the same size as the analysis frame. A + padded window can be obtained from `window_function()`. The FFT input buffer may be larger than the analysis frame, + typically the next power of two. + + Note: This function is not optimized for speed yet. It should be mostly compatible with `librosa.stft` and + `torchaudio.functional.transforms.Spectrogram`, although it is more flexible due to the different ways spectrograms + can be constructed. + + Args: + waveform (`np.ndarray` of shape `(length,)`): + The input waveform. This must be a single real-valued, mono waveform. + window (`np.ndarray` of shape `(frame_length,)`): + The windowing function to apply, including zero-padding if necessary. The actual window length may be + shorter than `frame_length`, but we're assuming the array has already been zero-padded. + frame_length (`int`): + The length of the analysis frames in samples. With librosa this is always equal to `fft_length` but we also + allow smaller sizes. + hop_length (`int`): + The stride between successive analysis frames in samples. + fft_length (`int`, *optional*): + The size of the FFT buffer in samples. This determines how many frequency bins the spectrogram will have. + For optimal speed, this should be a power of two. If `None`, uses `frame_length`. + power (`float`, *optional*, defaults to 1.0): + If 1.0, returns the amplitude spectrogram. If 2.0, returns the power spectrogram. If `None`, returns + complex numbers. + center (`bool`, *optional*, defaults to `True`): + Whether to pad the waveform so that frame `t` is centered around time `t * hop_length`. If `False`, frame + `t` will start at time `t * hop_length`. + pad_mode (`str`, *optional*, defaults to `"reflect"`): + Padding mode used when `center` is `True`. Possible values are: `"constant"` (pad with zeros), `"edge"` + (pad with edge values), `"reflect"` (pads with mirrored values). + onesided (`bool`, *optional*, defaults to `True`): + If True, only computes the positive frequencies and returns a spectrogram containing `fft_length // 2 + 1` + frequency bins. If False, also computes the negative frequencies and returns `fft_length` frequency bins. + dither (`float`, *optional*, defaults to 0.0): + Adds dithering. In other words, adds a small Gaussian noise to each frame. + E.g. use 4.0 to add dithering with a normal distribution centered + around 0.0 with standard deviation 4.0, 0.0 means no dithering. + Dithering has similar effect as `mel_floor`. It reduces the high log_mel_fbank + values for signals with hard-zero sections, when VAD cutoff is present in the signal. + preemphasis (`float`, *optional*) + Coefficient for a low-pass filter that applies pre-emphasis before the DFT. + mel_filters (`np.ndarray` of shape `(num_freq_bins, num_mel_filters)`, *optional*): + The mel filter bank. If supplied, applies a this filter bank to create a mel spectrogram. + mel_floor (`float`, *optional*, defaults to 1e-10): + Minimum value of mel frequency banks. + log_mel (`str`, *optional*): + How to convert the spectrogram to log scale. Possible options are: `None` (don't convert), `"log"` (take + the natural logarithm) `"log10"` (take the base-10 logarithm), `"dB"` (convert to decibels). Can only be + used when `power` is not `None`. + reference (`float`, *optional*, defaults to 1.0): + Sets the input spectrogram value that corresponds to 0 dB. For example, use `np.max(spectrogram)` to set + the loudest part to 0 dB. Must be greater than zero. + min_value (`float`, *optional*, defaults to `1e-10`): + The spectrogram will be clipped to this minimum value before conversion to decibels, to avoid taking + `log(0)`. For a power spectrogram, the default of `1e-10` corresponds to a minimum of -100 dB. For an + amplitude spectrogram, the value `1e-5` corresponds to -100 dB. Must be greater than zero. + db_range (`float`, *optional*): + Sets the maximum dynamic range in decibels. For example, if `db_range = 80`, the difference between the + peak value and the smallest value will never be more than 80 dB. Must be greater than zero. + remove_dc_offset (`bool`, *optional*): + Subtract mean from waveform on each frame, applied before pre-emphasis. This should be set to `true` in + order to get the same results as `torchaudio.compliance.kaldi.fbank` when computing mel filters. + dtype (`np.dtype`, *optional*, defaults to `np.float32`): + Data type of the spectrogram tensor. If `power` is None, this argument is ignored and the dtype will be + `np.complex64`. + + Returns: + `nd.array` containing a spectrogram of shape `(num_frequency_bins, length)` for a regular spectrogram or shape + `(num_mel_filters, length)` for a mel spectrogram. + """ + window_length = len(window) + + if fft_length is None: + fft_length = frame_length + + if frame_length > fft_length: + raise ValueError(f"frame_length ({frame_length}) may not be larger than fft_length ({fft_length})") + + if window_length != frame_length: + raise ValueError(f"Length of the window ({window_length}) must equal frame_length ({frame_length})") + + if hop_length <= 0: + raise ValueError("hop_length must be greater than zero") + + if waveform.ndim != 1: + raise ValueError(f"Input waveform must have only one dimension, shape is {waveform.shape}") + + if np.iscomplexobj(waveform): + raise ValueError("Complex-valued input waveforms are not currently supported") + + if power is None and mel_filters is not None: + raise ValueError( + "You have provided `mel_filters` but `power` is `None`. Mel spectrogram computation is not yet supported for complex-valued spectrogram." + "Specify `power` to fix this issue." + ) + + # center pad the waveform + if center: + padding = [(int(frame_length // 2), int(frame_length // 2))] + waveform = np.pad(waveform, padding, mode=pad_mode) + + # promote to float64, since np.fft uses float64 internally + waveform = waveform.astype(np.float64) + window = window.astype(np.float64) + + # split waveform into frames of frame_length size + num_frames = int(1 + np.floor((waveform.size - frame_length) / hop_length)) + + num_frequency_bins = (fft_length // 2) + 1 if onesided else fft_length + spectrogram = np.empty((num_frames, num_frequency_bins), dtype=np.complex64) + + # rfft is faster than fft + fft_func = np.fft.rfft if onesided else np.fft.fft + buffer = np.zeros(fft_length) + + timestep = 0 + for frame_idx in range(num_frames): + buffer[:frame_length] = waveform[timestep : timestep + frame_length] + + if dither != 0.0: + buffer[:frame_length] += dither * np.random.randn(frame_length) + + if remove_dc_offset: + buffer[:frame_length] = buffer[:frame_length] - buffer[:frame_length].mean() + + if preemphasis is not None: + buffer[1:frame_length] -= preemphasis * buffer[: frame_length - 1] + buffer[0] *= 1 - preemphasis + + buffer[:frame_length] *= window + + spectrogram[frame_idx] = fft_func(buffer) + timestep += hop_length + + # note: ** is much faster than np.power + if power is not None: + spectrogram = np.abs(spectrogram, dtype=np.float64) ** power + + spectrogram = spectrogram.T + + if mel_filters is not None: + spectrogram = np.maximum(mel_floor, np.dot(mel_filters.T, spectrogram)) + + if power is not None and log_mel is not None: + if log_mel == "log": + spectrogram = np.log(spectrogram) + elif log_mel == "log10": + spectrogram = np.log10(spectrogram) + elif log_mel == "dB": + if power == 1.0: + spectrogram = amplitude_to_db(spectrogram, reference, min_value, db_range) + elif power == 2.0: + spectrogram = power_to_db(spectrogram, reference, min_value, db_range) + else: + raise ValueError(f"Cannot use log_mel option '{log_mel}' with power {power}") + else: + raise ValueError(f"Unknown log_mel option: {log_mel}") + + spectrogram = np.asarray(spectrogram, dtype) + + return spectrogram + + +def spectrogram_batch( + waveform_list: list[np.ndarray], + window: np.ndarray, + frame_length: int, + hop_length: int, + fft_length: Optional[int] = None, + power: Optional[float] = 1.0, + center: bool = True, + pad_mode: str = "reflect", + onesided: bool = True, + dither: float = 0.0, + preemphasis: Optional[float] = None, + mel_filters: Optional[np.ndarray] = None, + mel_floor: float = 1e-10, + log_mel: Optional[str] = None, + reference: float = 1.0, + min_value: float = 1e-10, + db_range: Optional[float] = None, + remove_dc_offset: Optional[bool] = None, + dtype: np.dtype = np.float32, +) -> list[np.ndarray]: + """ + Calculates spectrograms for a list of waveforms using the Short-Time Fourier Transform, optimized for batch processing. + This function extends the capabilities of the `spectrogram` function to handle multiple waveforms efficiently by leveraging broadcasting. + + It supports generating various types of spectrograms: + + - amplitude spectrogram (`power = 1.0`) + - power spectrogram (`power = 2.0`) + - complex-valued spectrogram (`power = None`) + - log spectrogram (use `log_mel` argument) + - mel spectrogram (provide `mel_filters`) + - log-mel spectrogram (provide `mel_filters` and `log_mel`) + + How this works: + + 1. The input waveform is split into frames of size `frame_length` that are partially overlapping by `frame_length + - hop_length` samples. + 2. Each frame is multiplied by the window and placed into a buffer of size `fft_length`. + 3. The DFT is taken of each windowed frame. + 4. The results are stacked into a spectrogram. + + We make a distinction between the following "blocks" of sample data, each of which may have a different lengths: + + - The analysis frame. This is the size of the time slices that the input waveform is split into. + - The window. Each analysis frame is multiplied by the window to avoid spectral leakage. + - The FFT input buffer. The length of this determines how many frequency bins are in the spectrogram. + + In this implementation, the window is assumed to be zero-padded to have the same size as the analysis frame. A + padded window can be obtained from `window_function()`. The FFT input buffer may be larger than the analysis frame, + typically the next power of two. + + Note: This function is designed for efficient batch processing of multiple waveforms but retains compatibility with individual waveform processing methods like `librosa.stft`. + + Args: + waveform_list (`List[np.ndarray]` with arrays of shape `(length,)`): + The list of input waveforms, each a single-channel (mono) signal. + window (`np.ndarray` of shape `(frame_length,)`): + The windowing function to apply, including zero-padding if necessary. + frame_length (`int`): + The length of each frame for analysis. + hop_length (`int`): + The step size between successive frames. + fft_length (`int`, *optional*): + The size of the FFT buffer, defining frequency bin resolution. + power (`float`, *optional*, defaults to 1.0): + Determines the type of spectrogram: 1.0 for amplitude, 2.0 for power, None for complex. + center (`bool`, *optional*, defaults to `True`): + Whether to center-pad the waveform frames. + pad_mode (`str`, *optional*, defaults to `"reflect"`): + The padding strategy when `center` is `True`. + onesided (`bool`, *optional*, defaults to `True`): + If True, returns a one-sided spectrogram for real input signals. + dither (`float`, *optional*, defaults to 0.0): + Adds dithering. In other words, adds a small Gaussian noise to each frame. + E.g. use 4.0 to add dithering with a normal distribution centered + around 0.0 with standard deviation 4.0, 0.0 means no dithering. + preemphasis (`float`, *optional*): + Applies a pre-emphasis filter to each frame. + mel_filters (`np.ndarray`, *optional*): + Mel filter bank for converting to mel spectrogram. + mel_floor (`float`, *optional*, defaults to 1e-10): + Floor value for mel spectrogram to avoid log(0). + log_mel (`str`, *optional*): + Specifies log scaling strategy; options are None, "log", "log10", "dB". + reference (`float`, *optional*, defaults to 1.0): + Reference value for dB conversion in log_mel. + min_value (`float`, *optional*, defaults to 1e-10): + Minimum floor value for log scale conversions. + db_range (`float`, *optional*): + Dynamic range for dB scale spectrograms. + remove_dc_offset (`bool`, *optional*): + Whether to remove the DC offset from each frame. + dtype (`np.dtype`, *optional*, defaults to `np.float32`): + Data type of the output spectrogram. + + Returns: + List[`np.ndarray`]: A list of spectrogram arrays, one for each input waveform. + """ + window_length = len(window) + + if fft_length is None: + fft_length = frame_length + + if frame_length > fft_length: + raise ValueError(f"frame_length ({frame_length}) may not be larger than fft_length ({fft_length})") + + if window_length != frame_length: + raise ValueError(f"Length of the window ({window_length}) must equal frame_length ({frame_length})") + + if hop_length <= 0: + raise ValueError("hop_length must be greater than zero") + + # Check the dimensions of the waveform , and if waveform is complex + for waveform in waveform_list: + if waveform.ndim != 1: + raise ValueError(f"Input waveform must have only one dimension, shape is {waveform.shape}") + if np.iscomplexobj(waveform): + raise ValueError("Complex-valued input waveforms are not currently supported") + # Center pad the waveform + if center: + padding = [(int(frame_length // 2), int(frame_length // 2))] + waveform_list = [ + np.pad( + waveform, + padding, + mode=pad_mode, + ) + for waveform in waveform_list + ] + original_waveform_lengths = [ + len(waveform) for waveform in waveform_list + ] # these lengths will be used to remove padding later + + # Batch pad the waveform + max_length = max(original_waveform_lengths) + padded_waveform_batch = np.array( + [ + np.pad(waveform, (0, max_length - len(waveform)), mode="constant", constant_values=0) + for waveform in waveform_list + ], + dtype=dtype, + ) + + # Promote to float64, since np.fft uses float64 internally + padded_waveform_batch = padded_waveform_batch.astype(np.float64) + window = window.astype(np.float64) + + # Split waveform into frames of frame_length size + num_frames = int(1 + np.floor((padded_waveform_batch.shape[1] - frame_length) / hop_length)) + # these lengths will be used to remove padding later + true_num_frames = [int(1 + np.floor((length - frame_length) / hop_length)) for length in original_waveform_lengths] + num_batches = padded_waveform_batch.shape[0] + + num_frequency_bins = (fft_length // 2) + 1 if onesided else fft_length + spectrogram = np.empty((num_batches, num_frames, num_frequency_bins), dtype=np.complex64) + + # rfft is faster than fft + fft_func = np.fft.rfft if onesided else np.fft.fft + buffer = np.zeros((num_batches, fft_length)) + + for frame_idx in range(num_frames): + timestep = frame_idx * hop_length + buffer[:, :frame_length] = padded_waveform_batch[:, timestep : timestep + frame_length] + + if dither != 0.0: + buffer[:, :frame_length] += dither * np.random.randn(*buffer[:, :frame_length].shape) + + if remove_dc_offset: + buffer[:, :frame_length] -= buffer[:, :frame_length].mean(axis=1, keepdims=True) + + if preemphasis is not None: + buffer[:, 1:frame_length] -= preemphasis * buffer[:, : frame_length - 1] + buffer[:, 0] *= 1 - preemphasis + + buffer[:, :frame_length] *= window + + spectrogram[:, frame_idx] = fft_func(buffer) + + # Note: ** is much faster than np.power + if power is not None: + spectrogram = np.abs(spectrogram, dtype=np.float64) ** power + + # Apply mel filters if provided + if mel_filters is not None: + result = np.tensordot(spectrogram, mel_filters.T, axes=([2], [1])) + spectrogram = np.maximum(mel_floor, result) + + # Convert to log scale if specified + if power is not None and log_mel is not None: + if log_mel == "log": + spectrogram = np.log(spectrogram) + elif log_mel == "log10": + spectrogram = np.log10(spectrogram) + elif log_mel == "dB": + if power == 1.0: + spectrogram = amplitude_to_db_batch(spectrogram, reference, min_value, db_range) + elif power == 2.0: + spectrogram = power_to_db_batch(spectrogram, reference, min_value, db_range) + else: + raise ValueError(f"Cannot use log_mel option '{log_mel}' with power {power}") + else: + raise ValueError(f"Unknown log_mel option: {log_mel}") + + spectrogram = np.asarray(spectrogram, dtype) + + spectrogram_list = [spectrogram[i, : true_num_frames[i], :].T for i in range(len(true_num_frames))] + + return spectrogram_list + + +def power_to_db( + spectrogram: np.ndarray, + reference: float = 1.0, + min_value: float = 1e-10, + db_range: Optional[float] = None, +) -> np.ndarray: + """ + Converts a power spectrogram to the decibel scale. This computes `10 * log10(spectrogram / reference)`, using basic + logarithm properties for numerical stability. + + The motivation behind applying the log function on the (mel) spectrogram is that humans do not hear loudness on a + linear scale. Generally to double the perceived volume of a sound we need to put 8 times as much energy into it. + This means that large variations in energy may not sound all that different if the sound is loud to begin with. + This compression operation makes the (mel) spectrogram features match more closely what humans actually hear. + + Based on the implementation of `librosa.power_to_db`. + + Args: + spectrogram (`np.ndarray`): + The input power (mel) spectrogram. Note that a power spectrogram has the amplitudes squared! + reference (`float`, *optional*, defaults to 1.0): + Sets the input spectrogram value that corresponds to 0 dB. For example, use `np.max(spectrogram)` to set + the loudest part to 0 dB. Must be greater than zero. + min_value (`float`, *optional*, defaults to `1e-10`): + The spectrogram will be clipped to this minimum value before conversion to decibels, to avoid taking + `log(0)`. The default of `1e-10` corresponds to a minimum of -100 dB. Must be greater than zero. + db_range (`float`, *optional*): + Sets the maximum dynamic range in decibels. For example, if `db_range = 80`, the difference between the + peak value and the smallest value will never be more than 80 dB. Must be greater than zero. + + Returns: + `np.ndarray`: the spectrogram in decibels + """ + if reference <= 0.0: + raise ValueError("reference must be greater than zero") + if min_value <= 0.0: + raise ValueError("min_value must be greater than zero") + + reference = max(min_value, reference) + + spectrogram = np.clip(spectrogram, a_min=min_value, a_max=None) + spectrogram = 10.0 * (np.log10(spectrogram) - np.log10(reference)) + + if db_range is not None: + if db_range <= 0.0: + raise ValueError("db_range must be greater than zero") + spectrogram = np.clip(spectrogram, a_min=spectrogram.max() - db_range, a_max=None) + + return spectrogram + + +def power_to_db_batch( + spectrogram: np.ndarray, + reference: float = 1.0, + min_value: float = 1e-10, + db_range: Optional[float] = None, +) -> np.ndarray: + """ + Converts a batch of power spectrograms to the decibel scale. This computes `10 * log10(spectrogram / reference)`, + using basic logarithm properties for numerical stability. + + This function supports batch processing, where each item in the batch is an individual power (mel) spectrogram. + + Args: + spectrogram (`np.ndarray`): + The input batch of power (mel) spectrograms. Expected shape is (batch_size, *spectrogram_shape). + Note that a power spectrogram has the amplitudes squared! + reference (`float`, *optional*, defaults to 1.0): + Sets the input spectrogram value that corresponds to 0 dB. For example, use `np.max(spectrogram)` to set + the loudest part to 0 dB. Must be greater than zero. + min_value (`float`, *optional*, defaults to `1e-10`): + The spectrogram will be clipped to this minimum value before conversion to decibels, to avoid taking + `log(0)`. The default of `1e-10` corresponds to a minimum of -100 dB. Must be greater than zero. + db_range (`float`, *optional*): + Sets the maximum dynamic range in decibels. For example, if `db_range = 80`, the difference between the + peak value and the smallest value will never be more than 80 dB. Must be greater than zero. + + Returns: + `np.ndarray`: the batch of spectrograms in decibels + """ + if reference <= 0.0: + raise ValueError("reference must be greater than zero") + if min_value <= 0.0: + raise ValueError("min_value must be greater than zero") + + reference = max(min_value, reference) + + spectrogram = np.clip(spectrogram, a_min=min_value, a_max=None) + spectrogram = 10.0 * (np.log10(spectrogram) - np.log10(reference)) + + if db_range is not None: + if db_range <= 0.0: + raise ValueError("db_range must be greater than zero") + # Apply db_range clipping per batch item + max_values = spectrogram.max(axis=(1, 2), keepdims=True) + spectrogram = np.clip(spectrogram, a_min=max_values - db_range, a_max=None) + + return spectrogram + + +def amplitude_to_db( + spectrogram: np.ndarray, + reference: float = 1.0, + min_value: float = 1e-5, + db_range: Optional[float] = None, +) -> np.ndarray: + """ + Converts an amplitude spectrogram to the decibel scale. This computes `20 * log10(spectrogram / reference)`, using + basic logarithm properties for numerical stability. + + The motivation behind applying the log function on the (mel) spectrogram is that humans do not hear loudness on a + linear scale. Generally to double the perceived volume of a sound we need to put 8 times as much energy into it. + This means that large variations in energy may not sound all that different if the sound is loud to begin with. + This compression operation makes the (mel) spectrogram features match more closely what humans actually hear. + + Args: + spectrogram (`np.ndarray`): + The input amplitude (mel) spectrogram. + reference (`float`, *optional*, defaults to 1.0): + Sets the input spectrogram value that corresponds to 0 dB. For example, use `np.max(spectrogram)` to set + the loudest part to 0 dB. Must be greater than zero. + min_value (`float`, *optional*, defaults to `1e-5`): + The spectrogram will be clipped to this minimum value before conversion to decibels, to avoid taking + `log(0)`. The default of `1e-5` corresponds to a minimum of -100 dB. Must be greater than zero. + db_range (`float`, *optional*): + Sets the maximum dynamic range in decibels. For example, if `db_range = 80`, the difference between the + peak value and the smallest value will never be more than 80 dB. Must be greater than zero. + + Returns: + `np.ndarray`: the spectrogram in decibels + """ + if reference <= 0.0: + raise ValueError("reference must be greater than zero") + if min_value <= 0.0: + raise ValueError("min_value must be greater than zero") + + reference = max(min_value, reference) + + spectrogram = np.clip(spectrogram, a_min=min_value, a_max=None) + spectrogram = 20.0 * (np.log10(spectrogram) - np.log10(reference)) + + if db_range is not None: + if db_range <= 0.0: + raise ValueError("db_range must be greater than zero") + spectrogram = np.clip(spectrogram, a_min=spectrogram.max() - db_range, a_max=None) + + return spectrogram + + +def amplitude_to_db_batch( + spectrogram: np.ndarray, reference: float = 1.0, min_value: float = 1e-5, db_range: Optional[float] = None +) -> np.ndarray: + """ + Converts a batch of amplitude spectrograms to the decibel scale. This computes `20 * log10(spectrogram / reference)`, + using basic logarithm properties for numerical stability. + + The function supports batch processing, where each item in the batch is an individual amplitude (mel) spectrogram. + + Args: + spectrogram (`np.ndarray`): + The input batch of amplitude (mel) spectrograms. Expected shape is (batch_size, *spectrogram_shape). + reference (`float`, *optional*, defaults to 1.0): + Sets the input spectrogram value that corresponds to 0 dB. For example, use `np.max(spectrogram)` to set + the loudest part to 0 dB. Must be greater than zero. + min_value (`float`, *optional*, defaults to `1e-5`): + The spectrogram will be clipped to this minimum value before conversion to decibels, to avoid taking + `log(0)`. The default of `1e-5` corresponds to a minimum of -100 dB. Must be greater than zero. + db_range (`float`, *optional*): + Sets the maximum dynamic range in decibels. For example, if `db_range = 80`, the difference between the + peak value and the smallest value will never be more than 80 dB. Must be greater than zero. + + Returns: + `np.ndarray`: the batch of spectrograms in decibels + """ + if reference <= 0.0: + raise ValueError("reference must be greater than zero") + if min_value <= 0.0: + raise ValueError("min_value must be greater than zero") + + reference = max(min_value, reference) + + spectrogram = np.clip(spectrogram, a_min=min_value, a_max=None) + spectrogram = 20.0 * (np.log10(spectrogram) - np.log10(reference)) + + if db_range is not None: + if db_range <= 0.0: + raise ValueError("db_range must be greater than zero") + # Apply db_range clipping per batch item + max_values = spectrogram.max(axis=(1, 2), keepdims=True) + spectrogram = np.clip(spectrogram, a_min=max_values - db_range, a_max=None) + + return spectrogram + + +### deprecated functions below this line ### + + +def get_mel_filter_banks( + nb_frequency_bins: int, + nb_mel_filters: int, + frequency_min: float, + frequency_max: float, + sample_rate: int, + norm: Optional[str] = None, + mel_scale: str = "htk", +) -> np.array: + warnings.warn( + "The function `get_mel_filter_banks` is deprecated and will be removed in version 4.31.0 of Transformers", + FutureWarning, + ) + return mel_filter_bank( + num_frequency_bins=nb_frequency_bins, + num_mel_filters=nb_mel_filters, + min_frequency=frequency_min, + max_frequency=frequency_max, + sampling_rate=sample_rate, + norm=norm, + mel_scale=mel_scale, + ) + + +def fram_wave(waveform: np.array, hop_length: int = 160, fft_window_size: int = 400, center: bool = True): + """ + In order to compute the short time fourier transform, the waveform needs to be split in overlapping windowed + segments called `frames`. + + The window length (window_length) defines how much of the signal is contained in each frame, while the hop length + defines the step between the beginning of each new frame. + + + Args: + waveform (`np.array` of shape `(sample_length,)`): + The raw waveform which will be split into smaller chunks. + hop_length (`int`, *optional*, defaults to 160): + Step between each window of the waveform. + fft_window_size (`int`, *optional*, defaults to 400): + Defines the size of the window. + center (`bool`, defaults to `True`): + Whether or not to center each frame around the middle of the frame. Centering is done by reflecting the + waveform on the left and on the right. + + Return: + framed_waveform (`np.array` of shape `(waveform.shape // hop_length , fft_window_size)`): + The framed waveforms that can be fed to `np.fft`. + """ + warnings.warn( + "The function `fram_wave` is deprecated and will be removed in version 4.31.0 of Transformers", + FutureWarning, + ) + frames = [] + for i in range(0, waveform.shape[0] + 1, hop_length): + if center: + half_window = (fft_window_size - 1) // 2 + 1 + start = i - half_window if i > half_window else 0 + end = i + half_window if i < waveform.shape[0] - half_window else waveform.shape[0] + frame = waveform[start:end] + if start == 0: + padd_width = (-i + half_window, 0) + frame = np.pad(frame, pad_width=padd_width, mode="reflect") + + elif end == waveform.shape[0]: + padd_width = (0, (i - waveform.shape[0] + half_window)) + frame = np.pad(frame, pad_width=padd_width, mode="reflect") + + else: + frame = waveform[i : i + fft_window_size] + frame_width = frame.shape[0] + if frame_width < waveform.shape[0]: + frame = np.lib.pad( + frame, pad_width=(0, fft_window_size - frame_width), mode="constant", constant_values=0 + ) + frames.append(frame) + + frames = np.stack(frames, 0) + return frames + + +def stft(frames: np.array, windowing_function: np.array, fft_window_size: Optional[int] = None): + """ + Calculates the complex Short-Time Fourier Transform (STFT) of the given framed signal. Should give the same results + as `torch.stft`. + + Args: + frames (`np.array` of dimension `(num_frames, fft_window_size)`): + A framed audio signal obtained using `audio_utils.fram_wav`. + windowing_function (`np.array` of dimension `(nb_frequency_bins, nb_mel_filters)`: + A array representing the function that will be used to reduces the amplitude of the discontinuities at the + boundaries of each frame when computing the STFT. Each frame will be multiplied by the windowing_function. + For more information on the discontinuities, called *Spectral leakage*, refer to [this + tutorial]https://download.ni.com/evaluation/pxi/Understanding%20FFTs%20and%20Windowing.pdf + fft_window_size (`int`, *optional*): + Size of the window om which the Fourier transform is applied. This controls the frequency resolution of the + spectrogram. 400 means that the fourrier transform is computed on windows of 400 samples. The number of + frequency bins (`nb_frequency_bins`) used to divide the window into equal strips is equal to + `(1+fft_window_size)//2`. An increase of the fft_window_size slows the calculus time proportionnally. + + Example: + + ```python + >>> from transformers.audio_utils import stft, fram_wave + >>> import numpy as np + + >>> audio = np.random.rand(50) + >>> fft_window_size = 10 + >>> hop_length = 2 + >>> framed_audio = fram_wave(audio, hop_length, fft_window_size) + >>> spectrogram = stft(framed_audio, np.hanning(fft_window_size + 1)) + ``` + + Returns: + spectrogram (`np.ndarray`): + A spectrogram of shape `(num_frames, nb_frequency_bins)` obtained using the STFT algorithm + """ + warnings.warn( + "The function `stft` is deprecated and will be removed in version 4.31.0 of Transformers", + FutureWarning, + ) + frame_size = frames.shape[1] + + if fft_window_size is None: + fft_window_size = frame_size + + if fft_window_size < frame_size: + raise ValueError("FFT size must greater or equal the frame size") + # number of FFT bins to store + nb_frequency_bins = (fft_window_size >> 1) + 1 + + spectrogram = np.empty((len(frames), nb_frequency_bins), dtype=np.complex64) + fft_signal = np.zeros(fft_window_size) + + for f, frame in enumerate(frames): + if windowing_function is not None: + np.multiply(frame, windowing_function, out=fft_signal[:frame_size]) + else: + fft_signal[:frame_size] = frame + spectrogram[f] = np.fft.fft(fft_signal, axis=0)[:nb_frequency_bins] + return spectrogram.T diff --git a/lib/python3.12/site-packages/transformers/cache_utils.py b/lib/python3.12/site-packages/transformers/cache_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..760e676d96b89ddf69c0372d165376d227429e8c --- /dev/null +++ b/lib/python3.12/site-packages/transformers/cache_utils.py @@ -0,0 +1,2373 @@ +import copy +import importlib.metadata +import json +import os +from dataclasses import dataclass +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union + +import torch +from packaging import version + +from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_6 + +from .configuration_utils import PretrainedConfig +from .utils import is_hqq_available, is_optimum_quanto_available, is_torch_greater_or_equal, logging + + +if is_hqq_available(): + from hqq.core.quantize import Quantizer as HQQQuantizer + +logger = logging.get_logger(__name__) + + +class Cache: + """ + Base, abstract class for all caches. The actual data structure is specific to each subclass. + """ + + is_compileable = False + + def __init__(self): + super().__init__() + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. These are specific to each subclass and allow new types of + cache to be created. + + Return: + A tuple containing the updated key and value states. + """ + raise NotImplementedError("Make sure to implement `update` in a subclass.") + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # TODO: deprecate this function in favor of `cache_position` + raise NotImplementedError("Make sure to implement `get_seq_length` in a subclass.") + + def get_max_cache_shape(self) -> Optional[int]: + """Returns the maximum sequence length (i.e. max capacity) of the cache object""" + raise NotImplementedError("Make sure to implement `get_max_cache_shape` in a subclass.") + + def get_usable_length(self, new_seq_length: int, layer_idx: Optional[int] = 0) -> int: + """Given the sequence length of the new inputs, returns the usable length of the cache.""" + # Cache without size limit -> all cache is usable + # Cache with size limit -> if the length cache plus the length of the new inputs is larger the maximum cache + # length, we will need to evict part of the cache (and thus not all cache is usable) + max_length = self.get_max_cache_shape() + previous_seq_length = self.get_seq_length(layer_idx) + if max_length is not None and previous_seq_length + new_seq_length > max_length: + return max_length - new_seq_length + return previous_seq_length + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorders the cache for beam search, given the selected beam indices.""" + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx].numel(): + device = self.key_cache[layer_idx].device + self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device)) + if self.value_cache[layer_idx].numel(): + device = self.value_cache[layer_idx].device + self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device)) + + @property + def seen_tokens(self): + logger.warning_once( + "The `seen_tokens` attribute is deprecated and will be removed in v4.41. Use the `cache_position` " + "model input instead." + ) + if hasattr(self, "_seen_tokens"): + return self._seen_tokens + else: + return None + + +@dataclass +class CacheConfig: + """ + Base class for cache configs + """ + + cache_implementation: None + + @classmethod + def from_dict(cls, config_dict, **kwargs): + """ + Constructs a CacheConfig instance from a dictionary of parameters. + Args: + config_dict (Dict[str, Any]): Dictionary containing configuration parameters. + **kwargs: Additional keyword arguments to override dictionary values. + + Returns: + CacheConfig: Instance of CacheConfig constructed from the dictionary. + """ + config = cls(**config_dict) + to_remove = [] + for key, value in kwargs.items(): + if hasattr(config, key): + setattr(config, key, value) + to_remove.append(key) + for key in to_remove: + kwargs.pop(key, None) + return config + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_json_file + def to_json_file(self, json_file_path: Union[str, os.PathLike]): + """ + Save this instance to a JSON file. + + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file in which this configuration instance's parameters will be saved. + use_diff (`bool`, *optional*, defaults to `True`): + If set to `True`, only the difference between the config instance and the default + `QuantizationConfig()` is serialized to JSON file. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + config_dict = self.to_dict() + json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n" + + writer.write(json_string) + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_dict + def to_dict(self) -> Dict[str, Any]: + """ + Serializes this instance to a Python dictionary. Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + return copy.deepcopy(self.__dict__) + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__iter__ + def __iter__(self): + """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin""" + for attr, value in copy.deepcopy(self.__dict__).items(): + yield attr, value + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__repr__ + def __repr__(self): + return f"{self.__class__.__name__} {self.to_json_string()}" + + def to_json_string(self): + """ + Serializes this instance to a JSON formatted string. + Returns: + str: JSON formatted string representing the configuration instance. + """ + return json.dumps(self.__dict__, indent=2) + "\n" + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.update + def update(self, **kwargs): + """ + Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes, + returning all the unused kwargs. + + Args: + kwargs (`Dict[str, Any]`): + Dictionary of attributes to tentatively update this class. + + Returns: + `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance. + """ + to_remove = [] + for key, value in kwargs.items(): + if hasattr(self, key): + setattr(self, key, value) + to_remove.append(key) + + # Remove all the attributes that were updated, without modifying the input dict + unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove} + return unused_kwargs + + +@dataclass +class QuantizedCacheConfig(CacheConfig): + """ + Configuration class for quantized cache settings. + + Attributes: + backend (`str`, *optional*, defaults to `"quanto"`): + Backend to use when performing quantization, Can be one of [`quanto`, `HQQ`] + nbits (`Optional[int]`, *optional*, defaults to 4): + Number of bits, can be 2 or 4 for the `quanto` backend and one of [1, 2, 3, 4, 8] for the `HQQ` backend. Defaults to 2. + axis_key (`int`, *optional*, defaults to 0): + Axis over which to perform grouping for the key tensors. Can be [0, -1] for `quanto` backend and [0, 1] for `HQQ` backend. + axis_value (`int`, *optional*, defaults to 0): + Axis over which to perform grouping for the value tensors. Can be [0, -1] for `quanto` backend and [0, 1] for `HQQ` backend. + q_group_size (`Optional[int]`, *optional*, defaults to 64): + Size of the quantization group, should be a divisor of the model's hidden dimension. + Defaults to 64. + residual_length (`Optional[int]`, *optional*, defaults to 128): + Length of the residual cache which will always be stored in original precision. + Defaults to 128. + compute_dtype (`torch.dtype`, *optional*, defaults to `torch.float16`): + The default dtype used for computations in the model. Keys and Values will be cast to this dtype after dequantization. + device (`str`, *optional*, defaults to `"cpu"`): + Device on which to perform computations, should be same as the model's device. + """ + + def __init__( + self, + backend: str = "quanto", + nbits: Optional[int] = 4, + axis_key: Optional[int] = 0, + axis_value: Optional[int] = 0, + q_group_size: Optional[int] = 64, + residual_length: Optional[int] = 128, + compute_dtype: Optional[torch.dtype] = torch.float16, + device: Optional[str] = "cpu", + ): + self.backend = backend + self.nbits = nbits + self.axis_key = axis_key + self.axis_value = axis_value + self.q_group_size = q_group_size + self.residual_length = residual_length + self.compute_dtype = compute_dtype + self.device = device + + def validate(self): + """Validates if the arguments passed are correct""" + + incorrect_arg_msg = ( + "Some of the keys in `cache_config` are defined incorrectly. `{key}` should be {correct_value}` " + "but found {found_value}" + ) + # Check that the values are reasonable in general (nbits, axis) + # Later in QuantizedCache init we check if they are supported for that particular backend + if self.nbits not in [1, 2, 3, 4, 8]: + raise ValueError( + incorrect_arg_msg.format( + key="nbits", + correct_value="2 or 4 or 8", + found_value=self.nbits, + ), + ) + if self.q_group_size <= 0: + raise ValueError( + incorrect_arg_msg.format( + key="q_group_size", + correct_value="a positive integer", + found_value=self.q_group_size, + ), + ) + if self.residual_length < 0: + raise ValueError( + incorrect_arg_msg.format( + key="residual_length", + correct_value="a positive integer", + found_value=self.residual_length, + ), + ) + + if self.axis_key not in [0, 1, -1]: + raise ValueError( + incorrect_arg_msg.format( + key="axis_key", + correct_value="`1` or `0`, `-1`", + found_value=self.axis_key, + ), + ) + + if self.axis_value not in [0, 1, -1]: + raise ValueError( + incorrect_arg_msg.format( + key="axis_value", + correct_value="`1` or `0` or `-1`", + found_value=self.axis_value, + ), + ) + + +@dataclass +class StaticCacheConfig(CacheConfig): + """ + Configuration class for static cache settings. + """ + + cache_implementation = "static" + + def __init__(self, batch_size: int, max_cache_len: int, device="cpu"): + self.batch_size = batch_size + self.max_cache_len = max_cache_len + self.device = device + + def validate(self): + """Validates if the arguments passed are correct""" + + incorrect_arg_msg = ( + "Some of the keys in `cache_config` are defined incorrectly. `{key}` should be {correct_value}` " + "but found {found_value}" + ) + + if self.batch_size <= 0: + raise ValueError( + incorrect_arg_msg.format( + key="batch_size", + correct_value="> 0", + found_value=self.batch_size, + ), + ) + + if self.max_cache_len <= 0: + raise ValueError( + incorrect_arg_msg.format( + key="max_cache_len", + correct_value="> 0", + found_value=self.max_cache_len, + ), + ) + + +class DynamicCache(Cache): + """ + A cache that grows dynamically as more tokens are generated. This is the default for generative models. + + It stores the Key and Value states as a list of tensors, one for each layer. The expected shape for each tensor is + `[batch_size, num_heads, seq_len, head_dim]`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + + >>> inputs = tokenizer(text="My name is Qwen2", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> past_key_values = DynamicCache() + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + DynamicCache() + ``` + """ + + def __init__(self, _distributed_cache_data: Iterable = None) -> None: + super().__init__() + self._seen_tokens = 0 # Used in `generate` to keep tally of how many tokens the cache has seen + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + + # `_distributed_cache_data` was originally added for compatibility with `torch.distributed` (DDP). See #36121 + # and #36373 for more information. In a nutshell, it is `map(gather_map, zip(*caches))`, i.e. each item in the + # iterable contains the key and value states for a layer gathered across replicas by torch.distributed + # (shape=[global batch size, num_heads, seq_len, head_dim]). + # WARNING: `_distributed_cache_data` must be the first argument in `__init__`, otherwise we'll break + # compatibility. The name of the argument doesn't matter. + if _distributed_cache_data is not None: + for key_states, value_states in _distributed_cache_data: + self.key_cache.append(key_states) + self.value_cache.append(value_states) + + def __getitem__(self, layer_idx: int) -> List[Tuple[torch.Tensor]]: + """ + Support for backwards-compatible `past_key_value` indexing, e.g. `past_key_value[0][0].shape[2]` to get the + sequence length. + """ + if layer_idx < len(self): + return (self.key_cache[layer_idx], self.value_cache[layer_idx]) + else: + raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}") + + def __iter__(self): + """ + Support for backwards-compatible `past_key_value` iteration, e.g. `for x in past_key_value:` to iterate over + keys and values + """ + for layer_idx in range(len(self)): + yield (self.key_cache[layer_idx], self.value_cache[layer_idx]) + + def __len__(self): + """ + Support for backwards-compatible `past_key_value` length, e.g. `len(past_key_value)`. This value corresponds + to the number of layers in the model. + """ + return len(self.key_cache) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. No additional arguments are used in `DynamicCache`. + + Return: + A tuple containing the updated key and value states. + """ + # Update the number of seen tokens + if layer_idx == 0: + self._seen_tokens += key_states.shape[-2] + + # Update the cache + if key_states is not None: + if len(self.key_cache) <= layer_idx: + # There may be skipped layers, fill them with empty lists + for _ in range(len(self.key_cache), layer_idx): + self.key_cache.append(torch.tensor([])) + self.value_cache.append(torch.tensor([])) + self.key_cache.append(key_states) + self.value_cache.append(value_states) + elif ( + not self.key_cache[layer_idx].numel() # prefers not t.numel() to len(t) == 0 to export the model + ): # fills previously skipped layers; checking for tensor causes errors + self.key_cache[layer_idx] = key_states + self.value_cache[layer_idx] = value_states + else: + self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2) + self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # TODO: deprecate this function in favor of `cache_position` + is_empty_layer = ( + len(self.key_cache) == 0 # no cache in any layer + or len(self.key_cache) <= layer_idx # skipped `layer_idx` and hasn't run a layer with cache after it + or not self.key_cache[layer_idx].numel() # the layer has no cache + ) + layer_seq_length = self.key_cache[layer_idx].shape[-2] if not is_empty_layer else 0 + return layer_seq_length + + def get_max_cache_shape(self) -> Optional[int]: + """Returns the maximum sequence length of the cache object. DynamicCache does not have a maximum length.""" + return None + + def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor]]: + """Converts the `DynamicCache` instance into the its equivalent in the legacy cache format. Used for + backward compatibility.""" + legacy_cache = () + for layer_idx in range(len(self)): + legacy_cache += ((self.key_cache[layer_idx], self.value_cache[layer_idx]),) + return legacy_cache + + @classmethod + def from_legacy_cache(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None) -> "DynamicCache": + """Converts a cache in the legacy cache format into an equivalent `DynamicCache`. Used for + backward compatibility.""" + cache = cls() + if past_key_values is not None: + for layer_idx in range(len(past_key_values)): + key_states, value_states = past_key_values[layer_idx] + cache.update(key_states, value_states, layer_idx) + return cache + + def crop(self, max_length: int): + """Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be + negative to remove `max_length` tokens. This is used in assisted decoding and contrastive search.""" + # In case it is negative + if max_length < 0: + max_length = self.get_seq_length() - abs(max_length) + + if self.get_seq_length() <= max_length: + return + + self._seen_tokens = max_length + for idx in range(len(self.key_cache)): + if self.key_cache[idx].numel(): + self.key_cache[idx] = self.key_cache[idx][..., :max_length, :] + self.value_cache[idx] = self.value_cache[idx][..., :max_length, :] + + def batch_split(self, full_batch_size: int, split_size: int) -> List["DynamicCache"]: + """Split the current instance into a list of `DynamicCache` by the batch size. This will be used by + `_split_model_inputs()` in `generation.utils`""" + out = [] + for i in range(0, full_batch_size, split_size): + current_split = DynamicCache() + current_split._seen_tokens = self._seen_tokens + current_split.key_cache = [tensor[i : i + split_size] for tensor in self.key_cache] + current_split.value_cache = [tensor[i : i + split_size] for tensor in self.value_cache] + out.append(current_split) + return out + + @classmethod + def from_batch_splits(cls, splits: List["DynamicCache"]) -> "DynamicCache": + """This is the opposite of the above `batch_split()` method. This will be used by `stack_model_outputs` in + `generation.utils`""" + cache = cls() + for idx in range(len(splits[0])): + key_cache = [current.key_cache[idx] for current in splits if current.key_cache[idx].numel()] + value_cache = [current.value_cache[idx] for current in splits if current.value_cache[idx].numel()] + if key_cache != []: + layer_keys = torch.cat(key_cache, dim=0) + layer_values = torch.cat(value_cache, dim=0) + cache.update(layer_keys, layer_values, idx) + return cache + + def batch_repeat_interleave(self, repeats: int): + """Repeat the cache `repeats` times in the batch dimension. Used in contrastive search.""" + for layer_idx in range(len(self)): + self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave(repeats, dim=0) + self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave(repeats, dim=0) + + def batch_select_indices(self, indices: torch.Tensor): + """Only keep the `indices` in the batch dimension of the cache. Used in contrastive search.""" + for layer_idx in range(len(self)): + self.key_cache[layer_idx] = self.key_cache[layer_idx][indices, ...] + self.value_cache[layer_idx] = self.value_cache[layer_idx][indices, ...] + + +# Utilities for `DynamicCache` <> torch.export support +def _flatten_dynamic_cache( + dynamic_cache: DynamicCache, +): + """Flattens DynamicCache into flat list of tensors for `torch.export.export` to consume""" + if not isinstance(dynamic_cache, DynamicCache): + raise RuntimeError("This pytree flattening function should only be applied to DynamicCache") + + if not is_torch_greater_or_equal_than_2_6: + logger.warning_once( + "DynamicCache + torch.export is tested on torch 2.6.0+ and may not work on earlier versions." + ) + + # NOTE it seems _seen_tokens is deprecated, so probably doesn't need tracking + dictionary = { + "key_cache": getattr(dynamic_cache, "key_cache"), + "value_cache": getattr(dynamic_cache, "value_cache"), + } + return torch.utils._pytree._dict_flatten(dictionary) + + +def _flatten_with_keys_dynamic_cache(dynamic_cache: DynamicCache): + dictionary = { + "key_cache": getattr(dynamic_cache, "key_cache"), + "value_cache": getattr(dynamic_cache, "value_cache"), + } + return torch.utils._pytree._dict_flatten_with_keys(dictionary) + + +def _unflatten_dynamic_cache( + values, + context: torch.utils._pytree.Context, +): + dictionary = torch.utils._pytree._dict_unflatten(values, context) + cache = DynamicCache() + for k, v in dictionary.items(): + setattr(cache, k, v) + return cache + + +def _flatten_dynamic_cache_for_fx(cache, spec): + dictionary = { + "key_cache": getattr(cache, "key_cache"), + "value_cache": getattr(cache, "value_cache"), + } + return torch.utils._pytree.tree_flatten(dictionary)[0] + + +if is_torch_greater_or_equal("2.3"): + torch.utils._pytree.register_pytree_node( + DynamicCache, + _flatten_dynamic_cache, + _unflatten_dynamic_cache, + serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}", + flatten_with_keys_fn=_flatten_with_keys_dynamic_cache, + ) + # TODO (tmanlaibaatar) This won't be needed in torch 2.7. + torch.fx._pytree.register_pytree_flatten_spec(DynamicCache, _flatten_dynamic_cache_for_fx) + + +class OffloadedCache(DynamicCache): + """ + A drop-in replacement for DynamicCache that conserves accelerator(GPU, XPU) memory at the expense of more CPU memory. + Useful for generating from models with very long context. + + In addition to the default accelerator stream, where all forward() computations happen, + this class uses another stream, the prefetch stream, which it creates itself. + Since scheduling of operations on separate streams happens independently, this class uses + the prefetch stream to asynchronously prefetch the KV cache of layer k+1 when layer k is executing. + The movement of the layer k-1 cache to the CPU is handled by the default stream as a simple way to + ensure the eviction is scheduled after all computations on that cache are finished. + """ + + def __init__(self) -> None: + if not ( + torch.cuda.is_available() + or (is_torch_greater_or_equal("2.7", accept_dev=True) and torch.xpu.is_available()) + ): + raise RuntimeError( + "OffloadedCache can only be used with a GPU" + + (" or XPU" if is_torch_greater_or_equal("2.7", accept_dev=True) else "") + ) + + super().__init__() + self.original_device = [] + self.prefetch_stream = None + self.prefetch_stream = ( + torch.Stream() if is_torch_greater_or_equal("2.7", accept_dev=True) else torch.cuda.Stream() + ) + self.beam_idx = None # used to delay beam search operations + + def prefetch_layer(self, layer_idx: int): + "Starts prefetching the next layer cache" + if layer_idx < len(self): + with ( + self.prefetch_stream + if is_torch_greater_or_equal("2.7", accept_dev=True) + else torch.cuda.stream(self.prefetch_stream) + ): + # Prefetch next layer tensors to GPU + device = self.original_device[layer_idx] + self.key_cache[layer_idx] = self.key_cache[layer_idx].to(device, non_blocking=True) + self.value_cache[layer_idx] = self.value_cache[layer_idx].to(device, non_blocking=True) + + def evict_previous_layer(self, layer_idx: int): + "Moves the previous layer cache to the CPU" + if len(self) > 2: + # We do it on the default stream so it occurs after all earlier computations on these tensors are done + prev_layer_idx = (layer_idx - 1) % len(self) + self.key_cache[prev_layer_idx] = self.key_cache[prev_layer_idx].to("cpu", non_blocking=True) + self.value_cache[prev_layer_idx] = self.value_cache[prev_layer_idx].to("cpu", non_blocking=True) + + def __getitem__(self, layer_idx: int) -> List[Tuple[torch.Tensor]]: + "Gets the cache for this layer to the device. Prefetches the next and evicts the previous layer." + if layer_idx < len(self): + # Evict the previous layer if necessary + if is_torch_greater_or_equal("2.7", accept_dev=True): + torch.accelerator.current_stream().synchronize() + else: + torch.cuda.current_stream().synchronize() + self.evict_previous_layer(layer_idx) + # Load current layer cache to its original device if not already there + original_device = self.original_device[layer_idx] + self.prefetch_stream.synchronize() + key_tensor = self.key_cache[layer_idx] + value_tensor = self.value_cache[layer_idx] + # Now deal with beam search ops which were delayed + if self.beam_idx is not None: + self.beam_idx = self.beam_idx.to(original_device) + key_tensor = key_tensor.index_select(0, self.beam_idx) + value_tensor = value_tensor.index_select(0, self.beam_idx) + # Prefetch the next layer + self.prefetch_layer((layer_idx + 1) % len(self)) + return (key_tensor, value_tensor) + else: + raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}") + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Saves the beam indices and reorders the cache when the tensor is back to its device.""" + # We delay this operation until the tensors are back to their original + # device because performing torch.index_select on the CPU is very slow + del self.beam_idx + self.beam_idx = beam_idx.clone() + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. No additional arguments are used in `OffloadedCache`. + Return: + A tuple containing the updated key and value states. + """ + # Update the number of seen tokens + if layer_idx == 0: + self._seen_tokens += key_states.shape[-2] + + # Update the cache + if len(self.key_cache) < layer_idx: + raise ValueError("OffloadedCache does not support model usage where layers are skipped. Use DynamicCache.") + elif len(self.key_cache) == layer_idx: + self.key_cache.append(key_states) + self.value_cache.append(value_states) + self.original_device.append(key_states.device) + self.evict_previous_layer(layer_idx) + else: + key_tensor, value_tensor = self[layer_idx] + self.key_cache[layer_idx] = torch.cat([key_tensor, key_states], dim=-2) + self.value_cache[layer_idx] = torch.cat([value_tensor, value_states], dim=-2) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + # According to https://docs.python.org/3/library/exceptions.html#NotImplementedError + # if a method is not supposed to be supported in a subclass we should set it to None + from_legacy_cache = None + + to_legacy_cache = None + + +class QuantizedCache(DynamicCache): + """ + A quantizer cache similar to what is described in the [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache paper](https://arxiv.org/abs/2402.02750). + It allows the model to generate longer sequence length without allocating too much memory for Key and Value cache by applying quantization. + + The cache has two types of storage, one for original precision and one for the quantized cache. A `residual length` is set as a maximum capacity for the + original precision cache. When the length goes beyond maximum capacity, the original precision cache is discarded and moved into the quantized cache. The + quantization is done per-channel with a set `q_group_size` for both Keys and Values, in contrast to what was described in the paper. + + It stores Keys and Values a list of quantized tensors (tuples in case we need to store metadata), one for each layer. Additionally, it stores the Key and + Value in original precision states as a list of tensors, one for each layer. The size of each tensor + is `[batch_size, num_heads, seq_len - residual_length, head_dim]` + """ + + def __init__(self, cache_config: QuantizedCacheConfig) -> None: + super().__init__() + self._quantized_key_cache: List[torch.Tensor] = [] + self._quantized_value_cache: List[torch.Tensor] = [] + + self.nbits = cache_config.nbits + self.residual_length = cache_config.residual_length + self.q_group_size = cache_config.q_group_size + self.axis_key = cache_config.axis_key + self.axis_value = cache_config.axis_value + self.compute_dtype = cache_config.compute_dtype + self.device = cache_config.device + + super().__init__() + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Update the number of seen tokens + if layer_idx == 0: + self._seen_tokens += key_states.shape[-2] + + if len(self.key_cache) < layer_idx: + raise ValueError("QuantizedCache does not support model usage where layers are skipped. Use DynamicCache.") + elif len(self.key_cache) == layer_idx: + self._quantized_key_cache.append(self._quantize(key_states.contiguous(), axis=self.axis_key)) + self._quantized_value_cache.append(self._quantize(value_states.contiguous(), axis=self.axis_value)) + self.key_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device)) + self.value_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device)) + keys_to_return, values_to_return = key_states, value_states + else: + dequant_key = self._dequantize(self._quantized_key_cache[layer_idx]) + dequant_value = self._dequantize(self._quantized_value_cache[layer_idx]) + keys_to_return = [dequant_key, self.key_cache[layer_idx], key_states] + values_to_return = [dequant_value, self.value_cache[layer_idx], value_states] + + keys_to_return = torch.cat(keys_to_return, dim=-2) + values_to_return = torch.cat(values_to_return, dim=-2) + if ( + self.key_cache[layer_idx].dim() == 4 + and self.key_cache[layer_idx].shape[-2] + 1 >= self.residual_length + ): + self._quantized_key_cache[layer_idx] = self._quantize(keys_to_return.contiguous(), axis=self.axis_key) + self._quantized_value_cache[layer_idx] = self._quantize( + values_to_return.contiguous(), axis=self.axis_value + ) + self.key_cache[layer_idx] = torch.zeros(0, dtype=key_states.dtype, device=key_states.device) + self.value_cache[layer_idx] = torch.zeros(0, dtype=key_states.dtype, device=key_states.device) + else: + self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2) + self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2) + + return keys_to_return, values_to_return + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + if len(self.key_cache) <= layer_idx: + return 0 + # since we cannot get the seq_length of each layer directly and rely on `_seen_tokens` which is + # updated every "layer_idx" == 0, this is a hack to get the actual seq_length for the given layer_idx + # this part of code otherwise fails when used to verify attn_weight shape in some models + return self._seen_tokens if layer_idx == 0 else self._seen_tokens - 1 + + def _quantize(self, tensor, axis): + """Quantizes a key/value using a defined quantization method.""" + raise NotImplementedError("Make sure to implement `_quantize` in a subclass.") + + def _dequantize(self, q_tensor): + """Dequantizes back the tensor that was quantized by `self._quantize()`""" + raise NotImplementedError("Make sure to implement `_dequantize` in a subclass.") + + +class QuantoQuantizedCache(QuantizedCache): + """ + Quantized Cache class that uses `quanto` as a backend to perform quantization. Current implementation supports `int2` and `int4` dtypes only. + + Parameters: + cache_config (`QuantizedCacheConfig`): + A configuration containing all the arguments to be used by the quantizer, including axis, qtype and group size. + + Example: + + ```python + >>> # Run pip install quanto first if you don't have it yet + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, QuantoQuantizedCache, QuantizedCacheConfig + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + + >>> inputs = tokenizer(text="My name is Qwen2", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> cache_config = QuantizedCacheConfig(nbits=4) + >>> past_key_values = QuantoQuantizedCache(cache_config=cache_config) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + QuantoQuantizedCache() + ``` + """ + + def __init__(self, cache_config: CacheConfig) -> None: + super().__init__(cache_config) + + if is_optimum_quanto_available(): + optimum_quanto_version = version.parse(importlib.metadata.version("optimum-quanto")) + if optimum_quanto_version <= version.parse("0.2.5"): + raise ImportError( + f"You need optimum-quanto package version to be greater or equal than 0.2.5 to use `QuantoQuantizedCache`. Detected version {optimum_quanto_version}." + ) + from optimum.quanto import MaxOptimizer, qint2, qint4 + + if self.nbits not in [2, 4]: + raise ValueError(f"`nbits` for `quanto` backend has to be one of [`2`, `4`] but got {self.nbits}") + + if self.axis_key not in [0, -1]: + raise ValueError(f"`axis_key` for `quanto` backend has to be one of [`0`, `-1`] but got {self.axis_key}") + + if self.axis_value not in [0, -1]: + raise ValueError( + f"`axis_value` for `quanto` backend has to be one of [`0`, `-1`] but got {self.axis_value}" + ) + + self.qtype = qint4 if self.nbits == 4 else qint2 + self.optimizer = MaxOptimizer() # hardcode as it's the only one for per-channel quantization + + def _quantize(self, tensor, axis): + # We have two different API since in optimum-quanto, we don't use AffineQuantizer anymore + if is_optimum_quanto_available(): + from optimum.quanto import quantize_weight + + scale, zeropoint = self.optimizer(tensor, self.qtype, axis, self.q_group_size) + qtensor = quantize_weight(tensor, self.qtype, axis, scale, zeropoint, self.q_group_size) + return qtensor + + def _dequantize(self, qtensor): + return qtensor.dequantize() + + +class HQQQuantizedCache(QuantizedCache): + """ + Quantized Cache class that uses `HQQ` as a backend to perform quantization. Current implementation supports `int2`, `int4`, `int8` dtypes. + + Parameters: + cache_config (`QuantizedCacheConfig`): + A configuration containing all the arguments to be used by the quantizer, including axis, qtype and group size. + + Example: + + ```python + >>> # Run pip install hqq first if you don't have it yet + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, HQQQuantizedCache, QuantizedCacheConfig + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + + >>> inputs = tokenizer(text="My name is Qwen2", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> cache_config = QuantizedCacheConfig(nbits=4, axis_key=1, axis_value=1) + >>> past_key_values = HQQQuantizedCache(cache_config=cache_config) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + HQQQuantizedCache() + ``` + """ + + def __init__(self, cache_config: CacheConfig) -> None: + super().__init__(cache_config) + if self.nbits not in [1, 2, 3, 4, 8]: + raise ValueError( + f"`nbits` for `HQQ` backend has to be one of [`1`, `2`, `3`, `4`, `8`] but got {self.nbits}" + ) + + if self.axis_key not in [0, 1]: + raise ValueError(f"`axis_key` for `HQQ` backend has to be one of [`0`, `1`] but got {self.axis_key}") + + if self.axis_value not in [0, 1]: + raise ValueError(f"`axis_value` for `HQQ` backend has to be one of [`0`, `1`] but got {self.axis_value}") + + self.quantizer = HQQQuantizer + + def _quantize(self, tensor, axis): + qtensor, meta = self.quantizer.quantize( + tensor, + axis=axis, + device=self.device, + compute_dtype=self.compute_dtype, + nbits=self.nbits, + group_size=self.q_group_size, + ) + meta["compute_dtype"] = self.compute_dtype + self.quantizer.cuda(qtensor, meta=meta, device=self.device) # Move to device and cast to dtype + return qtensor, meta + + def _dequantize(self, qtensor): + quant_tensor, meta = qtensor + tensor = self.quantizer.dequantize(quant_tensor, meta) + return tensor + + +class SinkCache(Cache): + """ + A cache that as described in the [Attention Sinks paper](https://arxiv.org/abs/2309.17453). It allows the model to + generate beyond the length of its context window, without losing fluency in the conversation. As it discards past + tokens, the model will lose the ability to generate tokens that depend on the context that was discarded. + + It stores the Key and Value states as a list of tensors, one for each layer. The expected shape for each tensor is + `[batch_size, num_heads, seq_len, head_dim]`. + + Parameters: + window_length (`int`): + The length of the context window. + num_sink_tokens (`int`): + The number of sink tokens. See the original paper for more information. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, SinkCache + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + + >>> inputs = tokenizer(text="My name is Qwen2", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> past_key_values = SinkCache(window_length=256, num_sink_tokens=4) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + SinkCache() + ``` + """ + + is_sliding = True + + def __init__(self, window_length: int, num_sink_tokens: int) -> None: + super().__init__() + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + self.window_length = window_length + self.num_sink_tokens = num_sink_tokens + self.cos_sin_rerotation_cache = {} + self._cos_cache = None + self._sin_cache = None + self._seen_tokens = 0 # Used in `generate` to keep tally of how many tokens the cache has seen + + @staticmethod + def _rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + def _apply_key_rotary_pos_emb( + self, key_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + rotated_key_states = (key_states * cos) + (self._rotate_half(key_states) * sin) + return rotated_key_states + + def _get_rerotation_cos_sin( + self, key_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + if key_states.shape[-2] not in self.cos_sin_rerotation_cache: + # Upcast to float32 temporarily for better accuracy + cos = cos.to(torch.float32) + sin = sin.to(torch.float32) + + # Compute the cos and sin required for back- and forward-rotating to one position earlier in the sequence + original_cos = cos[self.num_sink_tokens + key_states.shape[-2] :] + shifted_cos = cos[self.num_sink_tokens : -key_states.shape[-2]] + original_sin = sin[self.num_sink_tokens + key_states.shape[-2] :] + shifted_sin = sin[self.num_sink_tokens : -key_states.shape[-2]] + rerotation_cos = original_cos * shifted_cos + original_sin * shifted_sin + rerotation_sin = -original_sin * shifted_cos + original_cos * shifted_sin + + self.cos_sin_rerotation_cache[key_states.shape[-2]] = ( + rerotation_cos.to(key_states.dtype).unsqueeze(0), + rerotation_sin.to(key_states.dtype).unsqueeze(0), + ) + return self.cos_sin_rerotation_cache[key_states.shape[-2]] + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # TODO: deprecate this function in favor of `cache_position` + # Workaround to make 'key_states.shape[-2] + past_key_value.get_seq_length(self.layer_idx)' <= window_length + if len(self.key_cache) <= layer_idx: + return 0 + return self.key_cache[layer_idx].shape[-2] + + def get_max_cache_shape(self) -> Optional[int]: + """Returns the maximum sequence length of the cache object, in case of SinkCache it is the window length.""" + return self.window_length + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. The following arguments can be used in `SinkCache`: `sin`, + `cos` and `partial_rotation_size`. These arguments are used with models using RoPE, to recompute the + rotation as the tokens are shifted. + + Return: + A tuple containing the updated key and value states. + """ + # Optional kwargs for `SinkCache` -- needed on models using RoPE. `partial_rotation_size` is used on models + # with partially rotated position embeddings, like Phi or Persimmon. + if cache_kwargs is None: + cache_kwargs = {} + sin = cache_kwargs.get("sin") + cos = cache_kwargs.get("cos") + partial_rotation_size = cache_kwargs.get("partial_rotation_size") + using_rope = cos is not None and sin is not None + + # Update the number of seen tokens + if layer_idx == 0: + self._seen_tokens += key_states.shape[-2] + + # Update the sin/cos cache, which holds sin/cos values for all possible positions + if using_rope and layer_idx == 0: + # BC: some models still pass `sin`/`cos` with 2 dims. In those models, they are the full sin/cos. Remove + # after all RoPE models have a llama-like cache utilization. + if cos.dim() == 2: + self._cos_cache = cos + self._sin_cache = sin + else: + if self._cos_cache is None: + self._cos_cache = cos[0, ...] + self._sin_cache = sin[0, ...] + elif self._cos_cache.shape[0] < self.window_length: + self._cos_cache = torch.cat([self._cos_cache, cos[0, ...]], dim=0) + self._sin_cache = torch.cat([self._sin_cache, sin[0, ...]], dim=0) + + # [bsz, num_heads, seq_len, head_dim] + if len(self.key_cache) <= layer_idx: + # Empty cache + self.key_cache.append(key_states) + self.value_cache.append(value_states) + + elif key_states.shape[-2] + self.get_seq_length(layer_idx) < self.window_length: + # Growing cache + self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2) + self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2) + + else: + # Shifting cache + keys_to_keep = self.key_cache[layer_idx][ + :, :, -self.window_length + self.num_sink_tokens + key_states.shape[-2] : + ] + + # On RoPE models, we need to recompute the Key rotation as the tokens are shifted + if using_rope: + rerotation_cos, rerotation_sin = self._get_rerotation_cos_sin( + key_states, self._cos_cache[: self.window_length], self._sin_cache[: self.window_length] + ) + if partial_rotation_size is not None: + keys_to_keep, keys_pass = ( + keys_to_keep[..., :partial_rotation_size], + keys_to_keep[..., partial_rotation_size:], + ) + keys_to_keep = self._apply_key_rotary_pos_emb(keys_to_keep, rerotation_cos, rerotation_sin) + if partial_rotation_size is not None: + keys_to_keep = torch.cat((keys_to_keep, keys_pass), dim=-1) + + # Concatenate sink tokens, shifted & rotated tokens (if needed), and new tokens + sink_keys = self.key_cache[layer_idx][:, :, : self.num_sink_tokens] + self.key_cache[layer_idx] = torch.cat([sink_keys, keys_to_keep, key_states], dim=-2) + + sink_values = self.value_cache[layer_idx][:, :, : self.num_sink_tokens] + values_to_keep = self.value_cache[layer_idx][ + :, :, -self.window_length + self.num_sink_tokens + value_states.shape[-2] : + ] + self.value_cache[layer_idx] = torch.cat([sink_values, values_to_keep, value_states], dim=-2) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + +class StaticCache(Cache): + """ + Static Cache class to be used with `torch.compile(model)` and `torch.export()`. + + Parameters: + config (`PretrainedConfig`): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a + smaller batch size is used. If you are manually setting the batch size, make sure to take into account the + number of beams if you are running beam search + max_cache_len (`int`, *optional*): + The maximum sequence length with which the model will be used. + device (`torch.device` or `str`, *optional*): + The device on which the cache should be initialized. If you're using more than 1 computation device, you + should pass the `layer_device_map` argument instead. + dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): + The default `dtype` to use when initializing the layer. + layer_device_map (`Optional[Dict[int, Union[str, torch.device, int]]]]`, *optional*): + Mapping between the layers and its device. This is required when you are manually initializing the cache + and the model is split between different gpus. You can know which layers mapped to which device by + checking the associated device_map: `model.hf_device_map`. + + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache + + >>> model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf") + + >>> inputs = tokenizer(text="My name is Llama", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate + >>> max_generated_length = inputs.input_ids.shape[1] + 10 + >>> past_key_values = StaticCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + StaticCache() + ``` + """ + + is_compileable = True + + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int] = None, + device: Union[torch.device, str, None] = None, + dtype: torch.dtype = torch.float32, + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ) -> None: + super().__init__() + self.max_batch_size = max_batch_size + self.max_cache_len = config.max_position_embeddings if max_cache_len is None else max_cache_len + + # Some model define a custom `head_dim` != config.hidden_size // config.num_attention_heads + self.head_dim = ( + config.head_dim if hasattr(config, "head_dim") else config.hidden_size // config.num_attention_heads + ) + + self._dtype = dtype + self.num_key_value_heads = ( + config.num_attention_heads + if getattr(config, "num_key_value_heads", None) is None + else config.num_key_value_heads + ) + + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + # Note: There will be significant perf decrease if switching to use 5D tensors instead. + cache_shape = (self.max_batch_size, self.num_key_value_heads, self.max_cache_len, self.head_dim) + device = torch.device(device) if device is not None else None + for idx in range(config.num_hidden_layers): + if layer_device_map is not None: + layer_device = layer_device_map[idx] + else: + layer_device = device + new_layer_key_cache = torch.zeros(cache_shape, dtype=self._dtype, device=layer_device) + new_layer_value_cache = torch.zeros(cache_shape, dtype=self._dtype, device=layer_device) + # Note: `mark_static_address` is used to tag the cache as a fixed data pointer, + # preventing compiled graph breaks when updating the cache. + torch._dynamo.mark_static_address(new_layer_key_cache) + torch._dynamo.mark_static_address(new_layer_value_cache) + self.key_cache.append(new_layer_key_cache) + self.value_cache.append(new_layer_value_cache) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + It is VERY important to index using a tensor, otherwise you introduce a copy to the device. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. The `StaticCache` needs the `cache_position` input + to know how where to write in the cache. + + Return: + A tuple containing the updated key and value states. + """ + if cache_kwargs is None: + cache_kwargs = {} + cache_position = cache_kwargs.get("cache_position") + k_out = self.key_cache[layer_idx] + v_out = self.value_cache[layer_idx] + key_states = key_states.to(k_out.dtype) + value_states = value_states.to(v_out.dtype) + + if cache_position is None: + k_out.copy_(key_states) + v_out.copy_(value_states) + else: + # Note: here we use `tensor.index_copy_(dim, index, tensor)` that is equivalent to + # `tensor[:, :, index] = tensor`, but the first one is compile-friendly and it does explicitly an in-place + # operation, that avoids copies and uses less memory. + try: + k_out.index_copy_(2, cache_position, key_states) + v_out.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # The operator 'aten::index_copy.out' is not currently implemented for the MPS device. + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + + return k_out, v_out + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states that were seen by the model.""" + # Occupied cache == any slot in the 3rd dim (sequence length) holds a non-zero value. To save on compute, let's + # limit the check to the first batch member and head dimension. + # TODO: deprecate this function in favor of `cache_position` + return (self.key_cache[layer_idx][0, 0].any(dim=-1)).sum() + + def get_max_cache_shape(self) -> Optional[int]: + return self.max_cache_len + + def reset(self): + """Resets the cache values while preserving the objects""" + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + +class SlidingWindowCache(StaticCache): + """ + Sliding Window Cache class to be used with `torch.compile` for models like Mistral that support sliding window attention. + Every time when we try to update the cache, we compute the `indices` based on `cache_position >= self.config.sliding_window - 1`, + if true(which means the cache can not hold all the old key value states and new states together because of the sliding window constraint), + we need to do a cycle shift based on `indices` to replace the oldest states by the new key value states passed in. + + The `to_shift` is only true once we are above sliding_window. Thus with `sliding_window==64`: + + indices = (slicing + to_shift[-1].int()-1) % self.config.sliding_window + tensor([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, + 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, + 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, + 55, 56, 57, 58, 59, 60, 61, 62, 63, 0]) + + We overwrite the cache using these, then we always write at cache_position (clamped to `sliding_window`) + + Parameters: + config (`PretrainedConfig`): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a + smaller batch size is used. + max_cache_len (`int`, *optional*): + The maximum sequence length with which the model will be used. + device (`torch.device` or `str`, *optional*): + The device on which the cache should be initialized. If you're using more than 1 computation device, you + should pass the `layer_device_map` argument instead. + dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): + The default `dtype` to use when initializing the layer. + layer_device_map (`Optional[Dict[int, Union[str, torch.device, int]]]]`, *optional*): + Mapping between the layers and its device. This is required when you are manually initializing the cache + and the model is split between different gpus. You can know which layers mapped to which device by + checking the associated device_map: `model.hf_device_map`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, SlidingWindowCache + + >>> model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3") + >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3") + + >>> inputs = tokenizer(text="My name is Mistral", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate + >>> max_generated_length = inputs.input_ids.shape[1] + 10 + >>> past_key_values = SlidingWindowCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + SlidingWindowCache() + ``` + """ + + is_sliding = True + is_compileable = True + + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int] = None, + device: Union[torch.device, str, None] = None, + dtype: torch.dtype = torch.float32, + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ) -> None: + if not hasattr(config, "sliding_window") or config.sliding_window is None: + raise ValueError( + "Setting `cache_implementation` to 'sliding_window' requires the model config supporting " + "sliding window attention, please check if there is a `sliding_window` field in the model " + "config and it's not set to None." + ) + max_cache_len = min(config.sliding_window, max_cache_len) + super().__init__( + config=config, + max_batch_size=max_batch_size, + max_cache_len=max_cache_len, + device=device, + dtype=dtype, + layer_device_map=layer_device_map, + ) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if cache_kwargs is None: + cache_kwargs = {} + cache_position = cache_kwargs.get("cache_position") + k_out = self.key_cache[layer_idx] + v_out = self.value_cache[layer_idx] + key_states = key_states.to(k_out.dtype) + value_states = value_states.to(v_out.dtype) + + # assume this only happens in prefill phase when prompt length > sliding_window_size (= max_cache_len) + if cache_position.shape[0] > self.max_cache_len: + k_out = key_states[:, :, -self.max_cache_len :, :] + v_out = value_states[:, :, -self.max_cache_len :, :] + # Assumption: caches are all zeros at this point, `+=` is equivalent to `=` but compile-friendly + self.key_cache[layer_idx] += k_out + self.value_cache[layer_idx] += v_out + # we should return the whole states instead of k_out, v_out to take the whole prompt + # into consideration when building kv cache instead of just throwing away tokens outside of the window + return key_states, value_states + + slicing = torch.ones(self.max_cache_len, dtype=torch.long, device=value_states.device).cumsum(0) + cache_position = cache_position.clamp(0, self.max_cache_len - 1) + to_shift = cache_position >= self.max_cache_len - 1 + indices = (slicing + to_shift[-1].int() - 1) % self.max_cache_len + + k_out = k_out[:, :, indices] + v_out = v_out[:, :, indices] + + try: + k_out.index_copy_(2, cache_position, key_states) + v_out.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # The operator 'aten::index_copy.out' is not currently implemented for the MPS device. + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + + # `_.zero()` followed by `+=` is equivalent `=`, but compile-friendly (without graph breaks due to assignment) + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + self.key_cache[layer_idx] += k_out + self.value_cache[layer_idx] += v_out + + return k_out, v_out + + def get_max_cache_shape(self) -> Optional[int]: + return self.max_cache_len + + def reset(self): + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + +class EncoderDecoderCache(Cache): + """ + Base, abstract class for all encoder-decoder caches. Can be used to hold combinations of self-attention and + cross-attention caches. + + Example: + + ```python + >>> from transformers import AutoProcessor, AutoModelForCausalLM, DynamicCache, EncoderDecoderCache + + >>> model = AutoModelForCausalLM.from_pretrained("openai/whisper-small") + >>> processor = AutoProcessor.from_pretrained("openai/whisper-small") + + >>> inputs = processor(audio=YOUR-AUDIO, return_tensors="pt") + + >>> # Prepare cache classes for encoder and decoder and pass it to model's forward + >>> self_attention_cache = DynamicCache() + >>> cross_attention_cache = DynamicCache() + >>> past_key_values = EncoderDecoderCache(self_attention_cache, cross_attention_cache) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + EncoderDecoderCache() + ``` + + """ + + def __init__(self, self_attention_cache: Cache, cross_attention_cache: Cache): + super().__init__() + self.self_attention_cache = self_attention_cache + self.cross_attention_cache = cross_attention_cache + self.is_compileable = getattr(self.self_attention_cache, "is_compileable", False) + + self.is_updated = {} + for layer_idx in range(len(cross_attention_cache.key_cache)): + self.is_updated[layer_idx] = bool(cross_attention_cache.get_seq_length(layer_idx) > 0) + + def __getitem__(self, layer_idx: int) -> List[Tuple[torch.Tensor]]: + """ + Support for backwards-compatible `past_key_value` indexing, e.g. `past_key_value[0][0].shape[2]` to get the + sequence length. + """ + if layer_idx < len(self): + return ( + self.self_attention_cache.key_cache[layer_idx], + self.self_attention_cache.value_cache[layer_idx], + self.cross_attention_cache.key_cache[layer_idx], + self.cross_attention_cache.value_cache[layer_idx], + ) + else: + raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}") + + def __len__(self): + """ + Support for backwards-compatible `past_key_value` length, e.g. `len(past_key_value)`. This value corresponds + to the number of layers in the model. + """ + return len(self.self_attention_cache) + + def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor]]: + """Converts the `EncoderDecoderCache` instance into its equivalent in the legacy cache format.""" + legacy_cache = () + if len(self.cross_attention_cache) > 0: + for self_attn, cross_attn in zip( + self.self_attention_cache.to_legacy_cache(), self.cross_attention_cache.to_legacy_cache() + ): + legacy_cache += (self_attn + cross_attn,) + else: + legacy_cache = self.self_attention_cache.to_legacy_cache() + return legacy_cache + + @classmethod + def from_legacy_cache( + cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + ) -> "EncoderDecoderCache": + """Converts a cache in the legacy cache format into an equivalent `EncoderDecoderCache`.""" + cache = cls( + self_attention_cache=DynamicCache(), + cross_attention_cache=DynamicCache(), + ) + if past_key_values is not None: + for layer_idx in range(len(past_key_values)): + key_states, value_states = past_key_values[layer_idx][:2] + cache.self_attention_cache.update(key_states, value_states, layer_idx) + if len(past_key_values[layer_idx]) > 2: + key_states, value_states = past_key_values[layer_idx][2:] + cache.cross_attention_cache.update(key_states, value_states, layer_idx) + cache.is_updated[layer_idx] = True + return cache + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # check if empty list because in case of static cache it will be a tensors and we can't check `if not torch.Tensor` + return self.self_attention_cache.get_seq_length(layer_idx) + + def reset(self): + if hasattr(self.self_attention_cache, "reset"): + self.self_attention_cache.reset() + if hasattr(self.cross_attention_cache, "reset"): + self.cross_attention_cache.reset() + elif not hasattr(self.self_attention_cache, "reset") and not hasattr(self.cross_attention_cache, "reset"): + raise ValueError( + "Neither self nor cross-attention cache have valid `.reset()` methods. `.reset()` should " + "only be called on compatible cache classes, such as `StaticCache` or `SlidingWindowCache`. " + f"Got {self.self_attention_cache.__str__()} for the self attention cache and " + f"{self.cross_attention_cache.__str__()} for the cross attention cache." + ) + for layer_idx in self.is_updated: + self.is_updated[layer_idx] = False + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorders the cache for beam search, given the selected beam indices.""" + self.self_attention_cache.reorder_cache(beam_idx) + self.cross_attention_cache.reorder_cache(beam_idx) + + def check_dynamic_cache(self, method: str): + if not ( + isinstance(self.self_attention_cache, DynamicCache) + and isinstance(self.cross_attention_cache, DynamicCache) + ): + raise ValueError( + f"`{method}` is only defined for dynamic cache, got {self.self_attention_cache.__str__()} for the self " + f"attention cache and {self.cross_attention_cache.__str__()} for the cross attention cache." + ) + + # TODO(gante, sanchit-gandhi): move following functionality into `.generate` + def crop(self, maximum_length: int): + """Crop the past key values up to a new `maximum_length` in terms of tokens. `maximum_length` can also be + negative to remove `maximum_length` tokens. This is used in assisted decoding and contrastive search.""" + self.check_dynamic_cache(self.crop.__name__) + self.self_attention_cache.crop(maximum_length) + + def batch_split(self, full_batch_size: int, split_size: int) -> "List[EncoderDecoderCache]": + """Split the current instance into a list of `DynamicCache` by the batch size. This will be used by + `_split_model_inputs()` in `generation.utils`""" + self.check_dynamic_cache(self.batch_split.__name__) + self_attention_cache = self.self_attention_cache.batch_split(full_batch_size, split_size) + cross_attention_cache = self.cross_attention_cache.batch_split(full_batch_size, split_size) + + out = [] + for self_attn, cross_attn in zip(self_attention_cache, cross_attention_cache): + out.append(EncoderDecoderCache(self_attn, cross_attn)) + return out + + @classmethod + def from_batch_splits(cls, splits: List["EncoderDecoderCache"]) -> "EncoderDecoderCache": + """This is the opposite of the above `batch_split()` method. This will be used by `stack_model_outputs` in + `generation.utils`""" + self_attention_cache = DynamicCache() + cross_attention_cache = DynamicCache() + for idx in range(len(splits[0])): + layer_keys = torch.cat([current.self_attention_cache.key_cache[idx] for current in splits], dim=0) + layer_values = torch.cat([current.self_attention_cache.value_cache[idx] for current in splits], dim=0) + self_attention_cache.update(layer_keys, layer_values, idx) + + layer_keys = torch.cat([current.cross_attention_cache.key_cache[idx] for current in splits], dim=0) + layer_values = torch.cat([current.cross_attention_cache.value_cache[idx] for current in splits], dim=0) + cross_attention_cache.update(layer_keys, layer_values, idx) + return cls(self_attention_cache, cross_attention_cache) + + def batch_repeat_interleave(self, repeats: int): + """Repeat the cache `repeats` times in the batch dimension. Used in contrastive search.""" + self.check_dynamic_cache(self.batch_repeat_interleave.__name__) + self.self_attention_cache.batch_repeat_interleave(repeats) + self.cross_attention_cache.batch_repeat_interleave(repeats) + + def batch_select_indices(self, indices: torch.Tensor): + """Only keep the `indices` in the batch dimension of the cache. Used in contrastive search.""" + self.check_dynamic_cache(self.batch_select_indices.__name__) + self.self_attention_cache.batch_select_indices(indices) + self.cross_attention_cache.batch_select_indices(indices) + + +class HybridCache(Cache): + """ + Hybrid Cache class to be used with `torch.compile` for Gemma2 models that alternate between a local sliding window attention + and global attention in every other layer. Under the hood, Hybrid Cache leverages ["SlidingWindowCache"] for sliding window attention + and ["StaticCache"] for global attention. For more information, see the documentation of each subcomponeent cache class. + + Parameters: + config (`PretrainedConfig): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a + smaller batch size is used. + max_cache_len (`int`, *optional*): + The maximum sequence length with which the model will be used. + device (`torch.device` or `str`, *optional*): + The device on which the cache should be initialized. If you're using more than 1 computation device, you + should pass the `layer_device_map` argument instead. + dtype (torch.dtype, *optional*, defaults to `torch.float32`): + The default `dtype` to use when initializing the layer. + layer_device_map (`Optional[Dict[int, Union[str, torch.device, int]]]]`, *optional*): + Mapping between the layers and its device. This is required when you are manually initializing the cache + and the model is split between different gpus. You can know which layers mapped to which device by + checking the associated device_map: `model.hf_device_map`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, HybridCache + + >>> model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b") + + >>> inputs = tokenizer(text="My name is Gemma", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate + >>> max_generated_length = inputs.input_ids.shape[1] + 10 + >>> past_key_values = HybridCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + HybridCache() + ``` + """ + + # TODO (joao): dive deeper into gemma2 and paligemma -- there are reports of speed loss with compilation. Revert + # ALL changes from the PR that commented the line below when reactivating it. + # is_compileable = True + + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int] = None, + device: Union[torch.device, str, None] = None, + dtype: torch.dtype = torch.float32, + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ) -> None: + super().__init__() + if not hasattr(config, "sliding_window") or config.sliding_window is None: + raise ValueError( + "Setting `cache_implementation` to 'sliding_window' requires the model config supporting " + "sliding window attention, please check if there is a `sliding_window` field in the model " + "config and it's not set to None." + ) + self.max_cache_len = max_cache_len + self.max_batch_size = max_batch_size + # Some model define a custom `head_dim` != config.hidden_size // config.num_attention_heads + self.head_dim = ( + config.head_dim if hasattr(config, "head_dim") else config.hidden_size // config.num_attention_heads + ) + + self._dtype = dtype + self.num_key_value_heads = ( + config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads + ) + + layer_switch = config.sliding_window_pattern if hasattr(config, "sliding_window_pattern") else 2 # 2 is for BC + self.is_sliding = torch.tensor( + [bool((i + 1) % layer_switch) for i in range(config.num_hidden_layers)], dtype=torch.bool + ) + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + global_cache_shape = (self.max_batch_size, self.num_key_value_heads, max_cache_len, self.head_dim) + sliding_cache_shape = ( + self.max_batch_size, + self.num_key_value_heads, + min(config.sliding_window, max_cache_len), + self.head_dim, + ) + device = torch.device(device) if device is not None and isinstance(device, str) else None + for i in range(config.num_hidden_layers): + if layer_device_map is not None: + layer_device = layer_device_map[i] + else: + layer_device = device + # Note: `mark_static_address` is used to tag the cache as an fixed data pointer, preventing cuda graph + # breaks when updating the cache. + cache_shape = global_cache_shape if not self.is_sliding[i] else sliding_cache_shape + new_layer_key_cache = torch.zeros(cache_shape, dtype=self._dtype, device=layer_device) + new_layer_value_cache = torch.zeros(cache_shape, dtype=self._dtype, device=layer_device) + torch._dynamo.mark_static_address(new_layer_key_cache) + torch._dynamo.mark_static_address(new_layer_value_cache) + self.key_cache.append(new_layer_key_cache) + self.value_cache.append(new_layer_value_cache) + + def _sliding_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len): + if cache_position.shape[0] > max_cache_len: + k_out = key_states[:, :, -max_cache_len:, :] + v_out = value_states[:, :, -max_cache_len:, :] + # Assumption: caches are all zeros at this point, `+=` is equivalent to `=` but compile-friendly + self.key_cache[layer_idx] += k_out + self.value_cache[layer_idx] += v_out + # we should return the whole states instead of k_out, v_out to take the whole prompt + # into consideration when building kv cache instead of just throwing away tokens outside of the window + return key_states, value_states + + slicing = torch.ones(max_cache_len, dtype=torch.long, device=value_states.device).cumsum(0) + cache_position = cache_position.clamp(0, max_cache_len - 1) + to_shift = cache_position >= max_cache_len - 1 + indices = (slicing + to_shift[-1].int() - 1) % max_cache_len + k_out = k_out[:, :, indices] + v_out = v_out[:, :, indices] + + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + # `_.zero()` followed by `+=` is equivalent `=`, but compile-friendly (without graph breaks due to assignment) + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + self.key_cache[layer_idx] += k_out + self.value_cache[layer_idx] += v_out + return k_out, v_out + + def _static_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len): + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + + self.key_cache[layer_idx] = k_out + self.value_cache[layer_idx] = v_out + return k_out, v_out + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if cache_kwargs is None: + cache_kwargs = {} + cache_position = cache_kwargs.get("cache_position") + sliding_window = cache_kwargs.get("sliding_window") + + # These two `if` blocks are only reached in multigpu and if `layer_device_map` is not passed. They are used + # when the cache is initialized in the forward pass (e.g. Gemma2) + if self.key_cache[layer_idx].device != key_states.device: + self.key_cache[layer_idx] = self.key_cache[layer_idx].to(key_states.device) + if self.value_cache[layer_idx].device != value_states.device: + self.value_cache[layer_idx] = self.value_cache[layer_idx].to(value_states.device) + + k_out = self.key_cache[layer_idx] + v_out = self.value_cache[layer_idx] + key_states = key_states.to(k_out.dtype) + value_states = value_states.to(v_out.dtype) + + if sliding_window: + update_fn = self._sliding_update + else: + update_fn = self._static_update + + return update_fn( + cache_position, + layer_idx, + key_states, + value_states, + k_out, + v_out, + k_out.shape[2], + ) + + def get_max_cache_shape(self) -> Optional[int]: + return self.max_cache_len + + def get_seq_length(self, layer_idx: Optional[int] = 0): + # Occupied cache == any slot in the 3rd dim (sequence length) holds a non-zero value. To save on compute, let's + # limit the check to the first batch member and head dimension. + # TODO: deprecate this function in favor of `cache_position` + if layer_idx != 0: + raise ValueError( + "`get_seq_length` on `HybridCache` may get inconsistent results depending on the layer index. " + "Using the `layer_idx` argument is not supported." + ) + return (self.key_cache[layer_idx][0, 0].any(dim=-1)).sum() + + def reset(self): + """Resets the cache values while preserving the objects""" + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + +class HybridChunkedCache(Cache): + """ + Hybrid Cache class to be used with `torch.compile` for Gemma2 models that alternate between a local sliding window attention + and global attention in every other layer. Under the hood, Hybrid Cache leverages ["SlidingWindowCache"] for sliding window attention + and ["StaticCache"] for global attention. For more information, see the documentation of each subcomponeent cache class. + + Parameters: + config (`PretrainedConfig): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a + smaller batch size is used. + max_cache_len (`int`, *optional*): + The maximum sequence length with which the model will be used. + device (`torch.device` or `str`, *optional*): + The device on which the cache should be initialized. If you're using more than 1 computation device, you + should pass the `layer_device_map` argument instead. + dtype (torch.dtype, *optional*, defaults to `torch.bfloat16`): + The default `dtype` to use when initializing the layer. + layer_device_map (`Optional[Dict[int, Union[str, torch.device, int]]]]`, *optional*): + Mapping between the layers and its device. This is required when you are manually initializing the cache + and the model is split between different gpus. You can know which layers mapped to which device by + checking the associated device_map: `model.hf_device_map`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, HybridCache + + >>> model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b") + + >>> inputs = tokenizer(text="My name is Gemma", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate + >>> max_generated_length = inputs.input_ids.shape[1] + 10 + >>> past_key_values = HybridCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + HybridCache() + ``` + """ + + # TODO (joao): dive deeper into gemma2 and paligemma -- there are reports of speed loss with compilation. Revert + # ALL changes from the PR that commented the line below when reactivating it. + is_compileable = True + + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int] = None, + device: Union[torch.device, str, None] = None, + dtype: torch.dtype = torch.bfloat16, + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ) -> None: + super().__init__() + if not hasattr(config, "sliding_window") or config.sliding_window is None: + self.sliding_window = getattr(config.get_text_config(), "attention_chunk_size", 8192) + else: + self.sliding_window = config.sliding_window + self.max_cache_len = max_cache_len + self.max_batch_size = max_batch_size + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self._dtype = dtype + + if hasattr(config.get_text_config(), "no_rope_layers"): + self.is_sliding = config.no_rope_layers + else: + layer_switch = getattr(config, "sliding_window_pattern", 2) + self.is_sliding = [bool((i + 1) % layer_switch) for i in range(config.num_hidden_layers)] + + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + self.cumulative_length = [0 for _ in range(config.num_hidden_layers)] + + def initialise_cache_layer(self, layer_idx, key_states): + if len(self.key_cache) > layer_idx: + return + + num_key_value_heads = key_states.shape[1] + device = key_states.device + global_cache_shape = (self.max_batch_size, num_key_value_heads, self.max_cache_len, self.head_dim) + sliding_cache_shape = ( + self.max_batch_size, + num_key_value_heads, + self.sliding_window, + self.head_dim, + ) + # Note: `mark_static_address` is used to tag the cache as an fixed data pointer, preventing cuda graph + # breaks when updating the cache. + cache_shape = sliding_cache_shape if self.is_sliding[layer_idx] else global_cache_shape + new_layer_key_cache = torch.zeros(cache_shape, dtype=self._dtype, device=device) + new_layer_value_cache = torch.zeros(cache_shape, dtype=self._dtype, device=device) + torch._dynamo.mark_static_address(new_layer_key_cache) + torch._dynamo.mark_static_address(new_layer_value_cache) + self.key_cache.append(new_layer_key_cache) + self.value_cache.append(new_layer_value_cache) + + def _sliding_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len): + if cache_position.shape[0] > max_cache_len: + cache_position = cache_position.clamp(0, max_cache_len - 1) + k_out = key_states[:, :, -max_cache_len:, :] + v_out = value_states[:, :, -max_cache_len:, :] + # Assumption: caches are all zeros at this point, `+=` is equivalent to `=` but compile-friendly + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + self.key_cache[layer_idx] += k_out + self.value_cache[layer_idx] += v_out + # we should return the whole states instead of k_out, v_out to take the whole prompt + # into consideration when building kv cache instead of just throwing away tokens outside of the window + return key_states, value_states + + # otherwise we are decoding. Most efficient way to cat 1 token + slicing = torch.ones(max_cache_len, dtype=torch.long, device=value_states.device).cumsum(0) + cache_position = cache_position.clamp(0, max_cache_len - 1) + to_shift = cache_position >= max_cache_len - 1 + indices = (slicing + to_shift[-1].int() - 1) % max_cache_len + k_out = k_out[:, :, indices] + v_out = v_out[:, :, indices] + + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + # `_.zero()` followed by `+=` is equivalent `=`, but compile-friendly (without graph breaks due to assignment) + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + self.key_cache[layer_idx] += k_out + self.value_cache[layer_idx] += v_out + return k_out, v_out + + def _static_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len): + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + + self.key_cache[layer_idx] = k_out + self.value_cache[layer_idx] = v_out + return k_out, v_out + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if cache_kwargs is None: + cache_kwargs = {} + cache_position = cache_kwargs.get("cache_position") + self.initialise_cache_layer(layer_idx, key_states) + + k_out = self.key_cache[layer_idx] + v_out = self.value_cache[layer_idx] + key_states = key_states.to(k_out.dtype) + value_states = value_states.to(v_out.dtype) + + if self.is_sliding[layer_idx]: + update_fn = self._sliding_update + else: + update_fn = self._static_update + + return update_fn( + cache_position, + layer_idx, + key_states, + value_states, + k_out, + v_out, + k_out.shape[2], + ) + + def get_max_cache_shape(self) -> Optional[int]: + return self.max_cache_len + + def get_seq_length(self, layer_idx: Optional[int] = 0): + # Occupied cache == any slot in the 3rd dim (sequence length) holds a non-zero value. To save on compute, let's + # limit the check to the first batch member and head dimension. + # TODO: deprecate this function in favor of `cache_position` + if layer_idx != 0: + raise ValueError( + "`get_seq_length` on `HybridCache` may get inconsistent results depending on the layer index. " + "Using the `layer_idx` argument is not supported." + ) + if len(self.key_cache) == 0: + return 0 + return (self.key_cache[layer_idx][0, 0].any(dim=-1)).sum() + + def reset(self): + """Resets the cache values while preserving the objects""" + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + self.cumulative_length = [0 for _ in range(len(self.cumulative_length))] + + +class MambaCache: + """ + Cache for mamba model which does not have attention mechanism and key value states. + + Arguments: + config (`PretrainedConfig): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a smaller batch size is used. + dtype (`torch.dtype`, *optional*, defaults to `torch.float16`): + The default `dtype` to use when initializing the layer. + device (`torch.device` or `str`, *optional*): + The device on which the cache should be initialized. Should be the same as the layer. + + Example: + + ```python + >>> from transformers import AutoTokenizer, MambaForCausalLM, MambaCache + + >>> model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf") + + >>> inputs = tokenizer(text="My name is Mamba", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> past_key_values = MambaCache(config=model.config, max_batch_size=1, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values + MambaCache() + ``` + """ + + is_compileable = True + + # TODO (joao): add layer_device_map arg and update code in `generate` accordingly + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + dtype: torch.dtype = torch.float16, + device: Union[torch.device, str, None] = None, + ): + self.max_batch_size = max_batch_size + self._dtype = dtype + self.intermediate_size = config.intermediate_size + self.ssm_state_size = config.state_size + self.conv_kernel_size = config.conv_kernel + + self.conv_states: List[torch.Tensor] = [] + self.ssm_states: List[torch.Tensor] = [] + device = torch.device(device) if device is not None else None + for _ in range(config.num_hidden_layers): + conv_state: torch.Tensor = torch.zeros( + self.max_batch_size, + self.intermediate_size, + self.conv_kernel_size, + device=device, + dtype=self._dtype, + ) + ssm_state: torch.Tensor = torch.zeros( + self.max_batch_size, + self.intermediate_size, + self.ssm_state_size, + device=device, + dtype=self._dtype, + ) + + torch._dynamo.mark_static_address(conv_state) + torch._dynamo.mark_static_address(ssm_state) + self.conv_states.append(conv_state) + self.ssm_states.append(ssm_state) + + def update_conv_state( + self, layer_idx: int, new_conv_state: torch.Tensor, cache_position: torch.LongTensor + ) -> torch.Tensor: + # This `if` blocks is only reached in multigpu and if `layer_device_map` is not passed. It is used + # when the cache is initialized in the forward pass (e.g. Mamba) + if self.conv_states[layer_idx].device != new_conv_state.device: + self.conv_states[layer_idx] = self.conv_states[layer_idx].to(new_conv_state.device) + + conv_state = self.conv_states[layer_idx] + cache_position = cache_position.clamp(0, self.conv_kernel_size - 1) + + conv_state = conv_state.roll(shifts=-1, dims=-1) + conv_state[:, :, cache_position] = new_conv_state.to(device=conv_state.device, dtype=conv_state.dtype) + self.conv_states[layer_idx].zero_() + self.conv_states[layer_idx] += conv_state + return self.conv_states[layer_idx] + + def update_ssm_state(self, layer_idx: int, new_ssm_state: torch.Tensor): + self.ssm_states[layer_idx] = new_ssm_state.to(self.ssm_states[layer_idx].device) + return self.ssm_states[layer_idx] + + def reset(self): + for layer_idx in range(len(self.conv_states)): + # In-place ops prevent breaking the static address + self.conv_states[layer_idx].zero_() + self.ssm_states[layer_idx].zero_() + + +class OffloadedStaticCache(StaticCache): + """ + Static cache class to be used with `torch.compile(model)` that offloads to the CPU or + another device. + + Args: + config (`PretrainedConfig): + The configuration file defining the shape-related attributes required to initialize + the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. + max_cache_len (`int`): + The maximum sequence length with which the model will be used. + device (`Union[str, torch.device]`): + The device on which the cache should be initialized. If you're using more than 1 computation device, you + should pass the `layer_device_map` argument instead. + dtype (`torch.dtype`, *optional*): + The default `dtype` to use when initializing the cache. + offload_device (`Union[str, torch.device]`, *optional*, defaults to `cpu`): + The device to offload to. Defaults to CPU. + layer_device_map (`Dict[int, Union[str, torch.device, int]]`, *optional*): + Mapping between the layers and its device. This is required when you are manually initializing the cache + and the model is splitted between differents gpus. You can know which layers mapped to which device by + checking the associated device_map: `model.hf_device_map`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, OffloadedStaticCache + + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + + >>> inputs = tokenizer(text="My name is GPT2", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate + >>> max_generated_length = inputs.input_ids.shape[1] + 10 + >>> past_key_values = OffloadedStaticCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> past_kv_length = outputs.past_key_values # access cache filled with key/values from generation + ``` + """ + + is_compileable = True + + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int], + device: Union[str, torch.device], + dtype: Optional[torch.dtype] = None, + offload_device: Union[str, torch.device] = torch.device("cpu"), + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ) -> None: + super(Cache, self).__init__() + self.max_batch_size = max_batch_size + self.max_cache_len = config.max_position_embeddings if max_cache_len is None else max_cache_len + self.device = torch.device(device) if layer_device_map is None else torch.device(layer_device_map[0]) + self.offload_device = torch.device(offload_device) + self._dtype = dtype if dtype is not None else torch.float32 + + # Some model define a custom `head_dim` != config.hidden_size // config.num_attention_heads + head_dim = config.head_dim if hasattr(config, "head_dim") else config.hidden_size // config.num_attention_heads + + num_key_value_heads = ( + config.num_attention_heads + if getattr(config, "num_key_value_heads", None) is None + else config.num_key_value_heads + ) + + cache_shape = (max_batch_size, num_key_value_heads, self.max_cache_len, head_dim) + + # Create offloaded CPU tensors. + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + + for i in range(config.num_hidden_layers): + # First layer is always on-device. + device = self.device if i == 0 else self.offload_device + + key_cache, value_cache = self._create_key_value_cache_tensors(cache_shape, device) + + self.key_cache.append(key_cache) + self.value_cache.append(value_cache) + + # Create device tensors. + self._device_key_cache: List[torch.Tensor] = [] + self._device_value_cache: List[torch.Tensor] = [] + + for i in range(2): + key_cache, value_cache = self._create_key_value_cache_tensors(cache_shape, self.device) + + self._device_key_cache.append(key_cache) + self._device_value_cache.append(value_cache) + + # For backwards compatibility. + # TODO(gante): Remove this. + self._seen_tokens = 0 + + # Create new CUDA stream for parallel prefetching. + self._prefetch_stream = torch.cuda.Stream() if self.device.type == "cuda" else None + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + It is VERY important to index using a tensor, otherwise you introduce a copy to the device. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, *optional*): + Additional arguments for the cache subclass. The `OffloadedStaticCache` needs the + `cache_position` input to know how where to write in the cache. + + Return: + A tuple containing the updated key and value states. + """ + + if layer_idx == 0: + # Update seen tokens. + # TODO(gante): Remove this. + self._seen_tokens += key_states.shape[-2] + + # Always there. + k_out = self.key_cache[0] + v_out = self.value_cache[0] + else: + # Wait for prefetch stream. + if self._prefetch_stream is not None: + torch.cuda.default_stream(self.device).wait_stream(self._prefetch_stream) + + k_out = self._device_key_cache[layer_idx & 1] + v_out = self._device_value_cache[layer_idx & 1] + + self._prefetch_layer(layer_idx + 1) + + cache_position = cache_kwargs.get("cache_position") if cache_kwargs is not None else None + if cache_position is None: + k_out.copy_(key_states) + v_out.copy_(value_states) + + # Copy the values to the offloaded device as well. + if layer_idx == 0: + self.key_cache[layer_idx].copy_(key_states.to(self.offload_device)) + self.value_cache[layer_idx].copy_(value_states.to(self.offload_device)) + else: + # Note: here we use `tensor.index_copy_(dim, index, tensor)` that is equivalent to + # `tensor[:, :, index] = tensor`, but the first one is compile-friendly and it does + # explicitly an in-place operation, that avoids copies and uses less memory. + try: + k_out.index_copy_(2, cache_position, key_states) + v_out.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # The operator 'aten::index_copy.out' is not currently implemented for the MPS + # device. + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + + # Copy the values to the offloaded device as well. + if layer_idx != 0: + cache_position = cache_position.to(self.offload_device) + key_states = key_states.to(self.offload_device) + value_states = value_states.to(self.offload_device) + + try: + self.key_cache[layer_idx].index_copy_(2, cache_position, key_states) + self.value_cache[layer_idx].index_copy_(2, cache_position, value_states) + except NotImplementedError: + # The operator 'aten::index_copy.out' is not currently implemented for the MPS + # device. + self.key_cache[layer_idx][:, :, cache_position] = key_states + self.value_cache[layer_idx][:, :, cache_position] = value_states + + return k_out, v_out + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states that were seen by the model.""" + + # TODO(gante): Remove this. + return self._seen_tokens + + def get_max_cache_shape(self) -> Optional[int]: + """Returns the maximum sequence length of the cached states.""" + + return self.max_cache_len + + def reset(self) -> None: + """Resets the cache values while preserving the objects.""" + + # For backwards compatibility. + # TODO(gante): Remove this. + self._seen_tokens = 0 + + # Zero out cache. + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address. + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + @property + def seen_tokens(self) -> int: + # For backwards compatibility. + # TODO(gante): Remove this. + return self._seen_tokens + + def _create_key_value_cache_tensors( + self, shape: Tuple[int, ...], device: torch.device + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Creates K/V cache tensors on a device. Pins memory for CPU tensors. Marks them as static + addresses for non-CPU tensors. + + Args: + shape (`Tuple[int, ...]`): Shape. + device (`torch.device`): Device. + + Returns: + Key and value cache tensors as a tuple. + """ + + is_cpu_device = device == torch.device("cpu") + + key_cache = torch.zeros(shape, dtype=self._dtype, device=device, pin_memory=is_cpu_device) + value_cache = torch.zeros(shape, dtype=self._dtype, device=device, pin_memory=is_cpu_device) + + # Note: `mark_static_address` is used to tag the cache as a fixed data pointer, + # preventing compiled graph breaks when updating the cache. + torch._dynamo.mark_static_address(key_cache) + torch._dynamo.mark_static_address(value_cache) + + return key_cache, value_cache + + def _prefetch_layer(self, layer_idx: int) -> None: + """Prefetch a layer to the device. Needs to be called in order of layer indices.""" + + # Don't fetch layers that do not exist. + if layer_idx >= len(self.key_cache): + return + + # Alternate between two on-device caches. + if self._prefetch_stream is not None: + with torch.cuda.stream(self._prefetch_stream): + self._prefetch_layer_in_context(layer_idx) + else: + self._prefetch_layer_in_context(layer_idx) + + def _prefetch_layer_in_context(self, layer_idx: int) -> None: + """Performs the actual copy of the layer to device cache.""" + + self._device_key_cache[layer_idx & 1].copy_(self.key_cache[layer_idx], non_blocking=True) + self._device_value_cache[layer_idx & 1].copy_(self.value_cache[layer_idx], non_blocking=True) diff --git a/lib/python3.12/site-packages/transformers/configuration_utils.py b/lib/python3.12/site-packages/transformers/configuration_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d5991cae8f9b7627b7925511046a09819f5027d2 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/configuration_utils.py @@ -0,0 +1,1201 @@ +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. 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. +"""Configuration base class and utilities.""" + +import copy +import json +import os +import warnings +from typing import Any, Optional, Union + +from packaging import version + +from . import __version__ +from .dynamic_module_utils import custom_object_save +from .modeling_gguf_pytorch_utils import load_gguf_checkpoint +from .utils import ( + CONFIG_NAME, + PushToHubMixin, + add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, + cached_file, + copy_func, + download_url, + extract_commit_hash, + is_remote_url, + is_torch_available, + logging, +) +from .utils.generic import is_timm_config_dict + + +logger = logging.get_logger(__name__) + + +class PretrainedConfig(PushToHubMixin): + # no-format + r""" + Base class for all configuration classes. Handles a few parameters common to all models' configurations as well as + methods for loading/downloading/saving configurations. + + + + A configuration file can be loaded and saved to disk. Loading the configuration file and using this file to + initialize a model does **not** load the model weights. It only affects the model's configuration. + + + + Class attributes (overridden by derived classes): + + - **model_type** (`str`) -- An identifier for the model type, serialized into the JSON file, and used to recreate + the correct object in [`~transformers.AutoConfig`]. + - **is_composition** (`bool`) -- Whether the config class is composed of multiple sub-configs. In this case the + config has to be initialized from two or more configs of type [`~transformers.PretrainedConfig`] like: + [`~transformers.EncoderDecoderConfig`] or [`~RagConfig`]. + - **keys_to_ignore_at_inference** (`List[str]`) -- A list of keys to ignore by default when looking at dictionary + outputs of the model during inference. + - **attribute_map** (`Dict[str, str]`) -- A dict that maps model specific attribute names to the standardized + naming of attributes. + - **base_model_tp_plan** (`Dict[str, Any]`) -- A dict that maps sub-modules FQNs of a base model to a tensor + parallel plan applied to the sub-module when `model.tensor_parallel` is called. + - **base_model_pp_plan** (`Dict[str, Tuple[List[str]]]`) -- A dict that maps child-modules of a base model to a + pipeline parallel plan that enables users to place the child-module on the appropriate device. + + Common attributes (present in all subclasses): + + - **vocab_size** (`int`) -- The number of tokens in the vocabulary, which is also the first dimension of the + embeddings matrix (this attribute may be missing for models that don't have a text modality like ViT). + - **hidden_size** (`int`) -- The hidden size of the model. + - **num_attention_heads** (`int`) -- The number of attention heads used in the multi-head attention layers of the + model. + - **num_hidden_layers** (`int`) -- The number of blocks in the model. + + + + Setting parameters for sequence generation in the model config is deprecated. For backward compatibility, loading + some of them will still be possible, but attempting to overwrite them will throw an exception -- you should set + them in a [~transformers.GenerationConfig]. Check the documentation of [~transformers.GenerationConfig] for more + information about the individual parameters. + + + + Arg: + name_or_path (`str`, *optional*, defaults to `""`): + Store the string that was passed to [`PreTrainedModel.from_pretrained`] or + [`TFPreTrainedModel.from_pretrained`] as `pretrained_model_name_or_path` if the configuration was created + with such a method. + output_hidden_states (`bool`, *optional*, defaults to `False`): + Whether or not the model should return all hidden-states. + output_attentions (`bool`, *optional*, defaults to `False`): + Whether or not the model should returns all attentions. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not the model should return a [`~transformers.utils.ModelOutput`] instead of a plain tuple. + is_encoder_decoder (`bool`, *optional*, defaults to `False`): + Whether the model is used as an encoder/decoder or not. + is_decoder (`bool`, *optional*, defaults to `False`): + Whether to only use the decoder in an encoder-decoder architecture, otherwise it has no effect on decoder-only or encoder-only architectures. + cross_attention_hidden_size** (`bool`, *optional*): + The hidden size of the cross-attention layer in case the model is used as a decoder in an encoder-decoder + setting and the cross-attention hidden dimension differs from `self.config.hidden_size`. + add_cross_attention (`bool`, *optional*, defaults to `False`): + Whether cross-attention layers should be added to the model. Note, this option is only relevant for models + that can be used as decoder models within the [`EncoderDecoderModel`] class, which consists of all models + in `AUTO_MODELS_FOR_CAUSAL_LM`. + tie_encoder_decoder (`bool`, *optional*, defaults to `False`): + Whether all encoder weights should be tied to their equivalent decoder weights. This requires the encoder + and decoder model to have the exact same parameter names. + prune_heads (`Dict[int, List[int]]`, *optional*, defaults to `{}`): + Pruned heads of the model. The keys are the selected layer indices and the associated values, the list of + heads to prune in said layer. + + For instance `{1: [0, 2], 2: [2, 3]}` will prune heads 0 and 2 on layer 1 and heads 2 and 3 on layer 2. + chunk_size_feed_forward (`int`, *optional*, defaults to `0`): + The chunk size of all feed forward layers in the residual attention blocks. A chunk size of `0` means that + the feed forward layer is not chunked. A chunk size of n means that the feed forward layer processes `n` < + sequence_length embeddings at a time. For more information on feed forward chunking, see [How does Feed + Forward Chunking work?](../glossary.html#feed-forward-chunking). + + > Parameters for fine-tuning tasks + + architectures (`List[str]`, *optional*): + Model architectures that can be used with the model pretrained weights. + finetuning_task (`str`, *optional*): + Name of the task used to fine-tune the model. This can be used when converting from an original (TensorFlow + or PyTorch) checkpoint. + id2label (`Dict[int, str]`, *optional*): + A map from index (for instance prediction index, or target index) to label. + label2id (`Dict[str, int]`, *optional*): A map from label to index for the model. + num_labels (`int`, *optional*): + Number of labels to use in the last layer added to the model, typically for a classification task. + task_specific_params (`Dict[str, Any]`, *optional*): + Additional keyword arguments to store for the current task. + problem_type (`str`, *optional*): + Problem type for `XxxForSequenceClassification` models. Can be one of `"regression"`, + `"single_label_classification"` or `"multi_label_classification"`. + + > Parameters linked to the tokenizer + + tokenizer_class (`str`, *optional*): + The name of the associated tokenizer class to use (if none is set, will use the tokenizer associated to the + model by default). + prefix (`str`, *optional*): + A specific prompt that should be added at the beginning of each text before calling the model. + bos_token_id (`int`, *optional*): The id of the _beginning-of-stream_ token. + pad_token_id (`int`, *optional*): The id of the _padding_ token. + eos_token_id (`int`, *optional*): The id of the _end-of-stream_ token. + decoder_start_token_id (`int`, *optional*): + If an encoder-decoder model starts decoding with a different token than _bos_, the id of that token. + sep_token_id (`int`, *optional*): The id of the _separation_ token. + + > PyTorch specific parameters + + torchscript (`bool`, *optional*, defaults to `False`): + Whether or not the model should be used with Torchscript. + tie_word_embeddings (`bool`, *optional*, defaults to `True`): + Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the + model has a output word embedding layer. + torch_dtype (`str`, *optional*): + The `dtype` of the weights. This attribute can be used to initialize the model to a non-default `dtype` + (which is normally `float32`) and thus allow for optimal storage allocation. For example, if the saved + model is `float16`, ideally we want to load it back using the minimal amount of memory needed to load + `float16` weights. Since the config object is stored in plain text, this attribute contains just the + floating type string without the `torch.` prefix. For example, for `torch.float16` ``torch_dtype` is the + `"float16"` string. + + This attribute is currently not being used during model loading time, but this may change in the future + versions. But we can already start preparing for the future by saving the dtype with save_pretrained. + + > TensorFlow specific parameters + + use_bfloat16 (`bool`, *optional*, defaults to `False`): + Whether or not the model should use BFloat16 scalars (only used by some TensorFlow models). + tf_legacy_loss (`bool`, *optional*, defaults to `False`): + Whether the model should use legacy TensorFlow losses. Legacy losses have variable output shapes and may + not be XLA-compatible. This option is here for backward compatibility and will be removed in Transformers + v5. + loss_type (`str`, *optional*): + The type of loss that the model should use. It should be in `LOSS_MAPPING`'s keys, otherwise the loss will + be automatically inferred from the model architecture. + """ + + model_type: str = "" + base_config_key: str = "" + sub_configs: dict[str, "PretrainedConfig"] = {} + is_composition: bool = False + attribute_map: dict[str, str] = {} + base_model_tp_plan: Optional[dict[str, Any]] = None + base_model_pp_plan: Optional[dict[str, tuple[list[str]]]] = None + _auto_class: Optional[str] = None + + def __setattr__(self, key, value): + if key in super().__getattribute__("attribute_map"): + key = super().__getattribute__("attribute_map")[key] + super().__setattr__(key, value) + + def __getattribute__(self, key): + if key != "attribute_map" and key in super().__getattribute__("attribute_map"): + key = super().__getattribute__("attribute_map")[key] + return super().__getattribute__(key) + + def __init__(self, **kwargs): + # Attributes with defaults + self.return_dict = kwargs.pop("return_dict", True) + self.output_hidden_states = kwargs.pop("output_hidden_states", False) + self.output_attentions = kwargs.pop("output_attentions", False) + self.torchscript = kwargs.pop("torchscript", False) # Only used by PyTorch models + self.torch_dtype = kwargs.pop("torch_dtype", None) # Only used by PyTorch models + self.use_bfloat16 = kwargs.pop("use_bfloat16", False) + self.tf_legacy_loss = kwargs.pop("tf_legacy_loss", False) # Only used by TensorFlow models + self.pruned_heads = kwargs.pop("pruned_heads", {}) + self.tie_word_embeddings = kwargs.pop( + "tie_word_embeddings", True + ) # Whether input and output word embeddings should be tied for all MLM, LM and Seq2Seq models. + self.chunk_size_feed_forward = kwargs.pop("chunk_size_feed_forward", 0) + + # Is decoder is used in encoder-decoder models to differentiate encoder from decoder + self.is_encoder_decoder = kwargs.pop("is_encoder_decoder", False) + self.is_decoder = kwargs.pop("is_decoder", False) + self.cross_attention_hidden_size = kwargs.pop("cross_attention_hidden_size", None) + self.add_cross_attention = kwargs.pop("add_cross_attention", False) + self.tie_encoder_decoder = kwargs.pop("tie_encoder_decoder", False) + + # Retrocompatibility: Parameters for sequence generation. While we will keep the ability to load these + # parameters, saving them will be deprecated. In a distant future, we won't need to load them. + for parameter_name, default_value in self._get_global_generation_defaults().items(): + setattr(self, parameter_name, kwargs.pop(parameter_name, default_value)) + + # Fine-tuning task arguments + self.architectures = kwargs.pop("architectures", None) + self.finetuning_task = kwargs.pop("finetuning_task", None) + self.id2label = kwargs.pop("id2label", None) + self.label2id = kwargs.pop("label2id", None) + if self.label2id is not None and not isinstance(self.label2id, dict): + raise ValueError("Argument label2id should be a dictionary.") + if self.id2label is not None: + if not isinstance(self.id2label, dict): + raise ValueError("Argument id2label should be a dictionary.") + num_labels = kwargs.pop("num_labels", None) + if num_labels is not None and len(self.id2label) != num_labels: + logger.warning( + f"You passed along `num_labels={num_labels}` with an incompatible id to label map: " + f"{self.id2label}. The number of labels will be overwritten to {self.num_labels}." + ) + self.id2label = {int(key): value for key, value in self.id2label.items()} + # Keys are always strings in JSON so convert ids to int here. + else: + self.num_labels = kwargs.pop("num_labels", 2) + + if self.torch_dtype is not None and isinstance(self.torch_dtype, str): + # we will start using self.torch_dtype in v5, but to be consistent with + # from_pretrained's torch_dtype arg convert it to an actual torch.dtype object + if is_torch_available(): + import torch + + self.torch_dtype = getattr(torch, self.torch_dtype) + + # Tokenizer arguments TODO: eventually tokenizer and models should share the same config + self.tokenizer_class = kwargs.pop("tokenizer_class", None) + self.prefix = kwargs.pop("prefix", None) + self.bos_token_id = kwargs.pop("bos_token_id", None) + self.pad_token_id = kwargs.pop("pad_token_id", None) + self.eos_token_id = kwargs.pop("eos_token_id", None) + self.sep_token_id = kwargs.pop("sep_token_id", None) + + self.decoder_start_token_id = kwargs.pop("decoder_start_token_id", None) + + # task specific arguments + self.task_specific_params = kwargs.pop("task_specific_params", None) + + # regression / multi-label classification + self.problem_type = kwargs.pop("problem_type", None) + allowed_problem_types = ("regression", "single_label_classification", "multi_label_classification") + if self.problem_type is not None and self.problem_type not in allowed_problem_types: + raise ValueError( + f"The config parameter `problem_type` was not understood: received {self.problem_type} " + "but only 'regression', 'single_label_classification' and 'multi_label_classification' are valid." + ) + + # TPU arguments + if kwargs.pop("xla_device", None) is not None: + logger.warning( + "The `xla_device` argument has been deprecated in v4.4.0 of Transformers. It is ignored and you can " + "safely remove it from your `config.json` file." + ) + + # Name or path to the pretrained checkpoint + self._name_or_path = str(kwargs.pop("name_or_path", "")) + # Config hash + self._commit_hash = kwargs.pop("_commit_hash", None) + + # Attention implementation to use, if relevant. + self._attn_implementation_internal = kwargs.pop("attn_implementation", None) + self._attn_implementation_autoset = False + + # Drop the transformers version info + self.transformers_version = kwargs.pop("transformers_version", None) + + # Deal with gradient checkpointing + if kwargs.get("gradient_checkpointing", False): + warnings.warn( + "Passing `gradient_checkpointing` to a config initialization is deprecated and will be removed in v5 " + "Transformers. Using `model.gradient_checkpointing_enable()` instead, or if you are using the " + "`Trainer` API, pass `gradient_checkpointing=True` in your `TrainingArguments`." + ) + + # Additional attributes without default values + 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 + + @property + def name_or_path(self) -> str: + return getattr(self, "_name_or_path", None) + + @name_or_path.setter + def name_or_path(self, value): + self._name_or_path = str(value) # Make sure that name_or_path is a string (for JSON encoding) + + @property + def use_return_dict(self) -> bool: + """ + `bool`: Whether or not return [`~utils.ModelOutput`] instead of tuples. + """ + # If torchscript is set, force `return_dict=False` to avoid jit errors + return self.return_dict and not self.torchscript + + @property + def num_labels(self) -> int: + """ + `int`: The number of labels for classification models. + """ + return len(self.id2label) + + @num_labels.setter + def num_labels(self, num_labels: int): + if not hasattr(self, "id2label") or self.id2label is None or len(self.id2label) != num_labels: + self.id2label = {i: f"LABEL_{i}" for i in range(num_labels)} + self.label2id = dict(zip(self.id2label.values(), self.id2label.keys())) + + @property + def _attn_implementation(self): + # This property is made private for now (as it cannot be changed and a PreTrainedModel.use_attn_implementation method needs to be implemented.) + if hasattr(self, "_attn_implementation_internal"): + if self._attn_implementation_internal is None: + # `config.attn_implementation` should never be None, for backward compatibility. + return "eager" + else: + return self._attn_implementation_internal + else: + return "eager" + + @_attn_implementation.setter + def _attn_implementation(self, value): + self._attn_implementation_internal = value + + def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs): + """ + Save a configuration object to the directory `save_directory`, so that it can be re-loaded using the + [`~PretrainedConfig.from_pretrained`] class method. + + Args: + save_directory (`str` or `os.PathLike`): + Directory where the configuration JSON file will be saved (will be created if it does not exist). + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + self._set_token_in_kwargs(kwargs) + + if os.path.isfile(save_directory): + raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") + + non_default_generation_parameters = self._get_non_default_generation_parameters() + if len(non_default_generation_parameters) > 0: + # TODO (joao): this should be an exception if the user has modified the loaded config. See #33886 + warnings.warn( + "Some non-default generation parameters are set in the model config. These should go into either a) " + "`model.generation_config` (as opposed to `model.config`); OR b) a GenerationConfig file " + "(https://huggingface.co/docs/transformers/generation_strategies#save-a-custom-decoding-strategy-with-your-model)." + "This warning will become an exception in the future." + f"\nNon-default generation parameters: {str(non_default_generation_parameters)}", + UserWarning, + ) + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = self._create_repo(repo_id, **kwargs) + files_timestamps = self._get_files_timestamps(save_directory) + + # If we have a custom config, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + custom_object_save(self, save_directory, config=self) + + # If we save using the predefined names, we can load using `from_pretrained` + output_config_file = os.path.join(save_directory, CONFIG_NAME) + + self.to_json_file(output_config_file, use_diff=True) + logger.info(f"Configuration saved in {output_config_file}") + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=kwargs.get("token"), + ) + + @staticmethod + def _set_token_in_kwargs(kwargs, token=None): + """Temporary method to deal with `token` and `use_auth_token`. + + This method is to avoid apply the same changes in all model config classes that overwrite `from_pretrained`. + + Need to clean up `use_auth_token` in a follow PR. + """ + # Some model config classes like CLIP define their own `from_pretrained` without the new argument `token` yet. + if token is None: + token = kwargs.pop("token", None) + use_auth_token = kwargs.pop("use_auth_token", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if token is not None: + kwargs["token"] = token + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Union[str, os.PathLike], + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + **kwargs, + ) -> "PretrainedConfig": + r""" + Instantiate a [`PretrainedConfig`] (or a derived class) from a pretrained model configuration. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + This can be either: + + - a string, the *model id* of a pretrained model configuration hosted inside a model repo on + huggingface.co. + - a path to a *directory* containing a configuration file saved using the + [`~PretrainedConfig.save_pretrained`] method, e.g., `./my_model_directory/`. + - a path or url to a saved configuration JSON *file*, e.g., `./my_model_directory/configuration.json`. + 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. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the configuration files and override the cached versions if + they exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`. + + + + return_unused_kwargs (`bool`, *optional*, defaults to `False`): + If `False`, then this function returns just the final configuration object. + + If `True`, then this functions returns a `Tuple(config, unused_kwargs)` where *unused_kwargs* is a + dictionary consisting of the key/value pairs whose keys are not configuration attributes: i.e., the + part of `kwargs` which has not been used to update `config` and is otherwise ignored. + 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. + kwargs (`Dict[str, Any]`, *optional*): + The values in kwargs of any keys which are configuration attributes will be used to override the loaded + values. Behavior concerning key/value pairs whose keys are *not* configuration attributes is controlled + by the `return_unused_kwargs` keyword parameter. + + Returns: + [`PretrainedConfig`]: The configuration object instantiated from this pretrained model. + + Examples: + + ```python + # We can't instantiate directly the base class *PretrainedConfig* so let's show the examples on a + # derived class: BertConfig + config = BertConfig.from_pretrained( + "google-bert/bert-base-uncased" + ) # Download configuration from huggingface.co and cache. + config = BertConfig.from_pretrained( + "./test/saved_model/" + ) # E.g. config (or model) was saved using *save_pretrained('./test/saved_model/')* + config = BertConfig.from_pretrained("./test/saved_model/my_configuration.json") + config = BertConfig.from_pretrained("google-bert/bert-base-uncased", output_attentions=True, foo=False) + assert config.output_attentions == True + config, unused_kwargs = BertConfig.from_pretrained( + "google-bert/bert-base-uncased", output_attentions=True, foo=False, return_unused_kwargs=True + ) + assert config.output_attentions == True + assert unused_kwargs == {"foo": False} + ```""" + kwargs["cache_dir"] = cache_dir + kwargs["force_download"] = force_download + kwargs["local_files_only"] = local_files_only + kwargs["revision"] = revision + + cls._set_token_in_kwargs(kwargs, token) + + config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs) + if cls.base_config_key and cls.base_config_key in config_dict: + config_dict = config_dict[cls.base_config_key] + + if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type: + # sometimes the config has no `base_config_key` if the config is used in several composite models + # e.g. LlamaConfig. In that case we try to see if there is match in `model_type` before raising a warning + for k, v in config_dict.items(): + if isinstance(v, dict) and v.get("model_type") == cls.model_type: + config_dict = v + + # raise warning only if we still can't see a match in `model_type` + if config_dict["model_type"] != cls.model_type: + logger.warning( + f"You are using a model of type {config_dict['model_type']} to instantiate a model of type " + f"{cls.model_type}. This is not supported for all configurations of models and can yield errors." + ) + + return cls.from_dict(config_dict, **kwargs) + + @classmethod + def get_config_dict( + cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs + ) -> tuple[dict[str, Any], dict[str, Any]]: + """ + From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a + [`PretrainedConfig`] using `from_dict`. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`): + The identifier of the pre-trained checkpoint from which we want the dictionary of parameters. + + Returns: + `Tuple[Dict, Dict]`: The dictionary(ies) that will be used to instantiate the configuration object. + + """ + cls._set_token_in_kwargs(kwargs) + + original_kwargs = copy.deepcopy(kwargs) + # Get config dict associated with the base config file + config_dict, kwargs = cls._get_config_dict(pretrained_model_name_or_path, **kwargs) + if config_dict is None: + return {}, kwargs + if "_commit_hash" in config_dict: + original_kwargs["_commit_hash"] = config_dict["_commit_hash"] + + # That config file may point us toward another config file to use. + if "configuration_files" in config_dict: + configuration_file = get_configuration_file(config_dict["configuration_files"]) + config_dict, kwargs = cls._get_config_dict( + pretrained_model_name_or_path, _configuration_file=configuration_file, **original_kwargs + ) + + return config_dict, kwargs + + @classmethod + def _get_config_dict( + cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs + ) -> tuple[dict[str, Any], dict[str, Any]]: + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + resume_download = kwargs.pop("resume_download", None) + proxies = kwargs.pop("proxies", None) + token = kwargs.pop("token", None) + local_files_only = kwargs.pop("local_files_only", False) + revision = kwargs.pop("revision", None) + trust_remote_code = kwargs.pop("trust_remote_code", None) + subfolder = kwargs.pop("subfolder", "") + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + commit_hash = kwargs.pop("_commit_hash", None) + + gguf_file = kwargs.get("gguf_file", None) + + if trust_remote_code is True: + logger.warning( + "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is" + " ignored." + ) + + user_agent = {"file_type": "config", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + + is_local = os.path.isdir(pretrained_model_name_or_path) + if os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)): + # Special case when pretrained_model_name_or_path is a local file + resolved_config_file = pretrained_model_name_or_path + is_local = True + elif is_remote_url(pretrained_model_name_or_path): + configuration_file = pretrained_model_name_or_path if gguf_file is None else gguf_file + resolved_config_file = download_url(pretrained_model_name_or_path) + else: + configuration_file = kwargs.pop("_configuration_file", CONFIG_NAME) if gguf_file is None else gguf_file + + try: + # Load from local folder or from cache or download from model Hub and cache + resolved_config_file = cached_file( + pretrained_model_name_or_path, + configuration_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _commit_hash=commit_hash, + ) + if resolved_config_file is None: + return None, kwargs + commit_hash = extract_commit_hash(resolved_config_file, commit_hash) + except OSError: + # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted to + # the original exception. + raise + except Exception: + # For any other exception, we throw a generic error. + raise OSError( + f"Can't load the configuration of '{pretrained_model_name_or_path}'. If you were trying to load it" + " from 'https://huggingface.co/models', make sure you don't have a local directory with the same" + f" name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory" + f" containing a {configuration_file} file" + ) + + try: + if gguf_file: + config_dict = load_gguf_checkpoint(resolved_config_file, return_tensors=False)["config"] + else: + # Load config dict + config_dict = cls._dict_from_json_file(resolved_config_file) + + config_dict["_commit_hash"] = commit_hash + except (json.JSONDecodeError, UnicodeDecodeError): + raise OSError(f"It looks like the config file at '{resolved_config_file}' is not a valid JSON file.") + + if is_local: + logger.info(f"loading configuration file {resolved_config_file}") + else: + logger.info(f"loading configuration file {configuration_file} from cache at {resolved_config_file}") + + if "auto_map" in config_dict and not is_local: + config_dict["auto_map"] = add_model_info_to_auto_map( + config_dict["auto_map"], pretrained_model_name_or_path + ) + if "custom_pipelines" in config_dict and not is_local: + config_dict["custom_pipelines"] = add_model_info_to_custom_pipelines( + config_dict["custom_pipelines"], pretrained_model_name_or_path + ) + + # timm models are not saved with the model_type in the config file + if "model_type" not in config_dict and is_timm_config_dict(config_dict): + config_dict["model_type"] = "timm_wrapper" + + return config_dict, kwargs + + @classmethod + def from_dict(cls, config_dict: dict[str, Any], **kwargs) -> "PretrainedConfig": + """ + Instantiates a [`PretrainedConfig`] from a Python dictionary of parameters. + + Args: + config_dict (`Dict[str, Any]`): + Dictionary that will be used to instantiate the configuration object. Such a dictionary can be + retrieved from a pretrained checkpoint by leveraging the [`~PretrainedConfig.get_config_dict`] method. + kwargs (`Dict[str, Any]`): + Additional parameters from which to initialize the configuration object. + + Returns: + [`PretrainedConfig`]: The configuration object instantiated from those parameters. + """ + return_unused_kwargs = kwargs.pop("return_unused_kwargs", False) + # Those arguments may be passed along for our internal telemetry. + # We remove them so they don't appear in `return_unused_kwargs`. + kwargs.pop("_from_auto", None) + kwargs.pop("_from_pipeline", None) + # The commit hash might have been updated in the `config_dict`, we don't want the kwargs to erase that update. + if "_commit_hash" in kwargs and "_commit_hash" in config_dict: + kwargs["_commit_hash"] = config_dict["_commit_hash"] + + # We remove it from kwargs so that it does not appear in `return_unused_kwargs`. + config_dict["attn_implementation"] = kwargs.pop("attn_implementation", None) + + config = cls(**config_dict) + + if hasattr(config, "pruned_heads"): + config.pruned_heads = {int(key): value for key, value in config.pruned_heads.items()} + + # Update config with kwargs if needed + if "num_labels" in kwargs and "id2label" in kwargs: + num_labels = kwargs["num_labels"] + id2label = kwargs["id2label"] if kwargs["id2label"] is not None else [] + if len(id2label) != num_labels: + raise ValueError( + f"You passed along `num_labels={num_labels}` with an incompatible id to label map: " + f"{kwargs['id2label']}. Since those arguments are inconsistent with each other, you should remove " + "one of them." + ) + to_remove = [] + for key, value in kwargs.items(): + if hasattr(config, key): + current_attr = getattr(config, key) + # To authorize passing a custom subconfig as kwarg in models that have nested configs. + if isinstance(current_attr, PretrainedConfig) and isinstance(value, dict): + value = current_attr.__class__(**value) + setattr(config, key, value) + if key != "torch_dtype": + to_remove.append(key) + for key in to_remove: + kwargs.pop(key, None) + + logger.info(f"Model config {config}") + if return_unused_kwargs: + return config, kwargs + else: + return config + + @classmethod + def from_json_file(cls, json_file: Union[str, os.PathLike]) -> "PretrainedConfig": + """ + Instantiates a [`PretrainedConfig`] from the path to a JSON file of parameters. + + Args: + json_file (`str` or `os.PathLike`): + Path to the JSON file containing the parameters. + + Returns: + [`PretrainedConfig`]: The configuration object instantiated from that JSON file. + + """ + config_dict = cls._dict_from_json_file(json_file) + return cls(**config_dict) + + @classmethod + def _dict_from_json_file(cls, json_file: Union[str, os.PathLike]): + with open(json_file, encoding="utf-8") as reader: + text = reader.read() + return json.loads(text) + + def __eq__(self, other): + return isinstance(other, PretrainedConfig) and (self.__dict__ == other.__dict__) + + def __repr__(self): + return f"{self.__class__.__name__} {self.to_json_string()}" + + def __iter__(self): + yield from self.__dict__ + + def to_diff_dict(self) -> dict[str, Any]: + """ + Removes all attributes from the configuration that correspond to the default config attributes for + better readability, while always retaining the `config` attribute from the class. Serializes to a + Python dictionary. + + Returns: + Dict[str, Any]: Dictionary of all the attributes that make up this configuration instance. + """ + config_dict = self.to_dict() + + # Get the default config dict (from a fresh PreTrainedConfig instance) + default_config_dict = PretrainedConfig().to_dict() + + # Get class-specific config dict if not part of a composition + class_config_dict = self.__class__().to_dict() if not self.is_composition else {} + + serializable_config_dict = {} + + # Only serialize values that differ from the default config, + # except always keep the 'config' attribute. + for key, value in config_dict.items(): + if ( + isinstance(getattr(self, key, None), PretrainedConfig) + and key in class_config_dict + and isinstance(class_config_dict[key], dict) + or key in self.sub_configs + ): + # For nested configs we need to clean the diff recursively + diff = recursive_diff_dict(value, default_config_dict, config_obj=getattr(self, key, None)) + if "model_type" in value: + # Needs to be set even if it's not in the diff + diff["model_type"] = value["model_type"] + serializable_config_dict[key] = diff + elif ( + key not in default_config_dict + or key == "transformers_version" + or key == "vocab_file" + or value != default_config_dict[key] + or (key in default_config_dict and value != class_config_dict.get(key, value)) + ): + serializable_config_dict[key] = value + + if hasattr(self, "quantization_config"): + serializable_config_dict["quantization_config"] = ( + self.quantization_config.to_dict() + if not isinstance(self.quantization_config, dict) + else self.quantization_config + ) + # Pop the `_pre_quantization_dtype` as torch.dtypes are not serializable. + _ = serializable_config_dict.pop("_pre_quantization_dtype", None) + + self.dict_torch_dtype_to_str(serializable_config_dict) + + if "_attn_implementation_internal" in serializable_config_dict: + del serializable_config_dict["_attn_implementation_internal"] + # Do not serialize `base_model_tp_plan` for now + if "base_model_tp_plan" in serializable_config_dict: + del serializable_config_dict["base_model_tp_plan"] + # Do not serialize `base_model_pp_plan` for now + if "base_model_pp_plan" in serializable_config_dict: + del serializable_config_dict["base_model_pp_plan"] + + if "_name_or_path" in serializable_config_dict: + del serializable_config_dict["_name_or_path"] + + return serializable_config_dict + + def to_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. + + Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + output = copy.deepcopy(self.__dict__) + if hasattr(self.__class__, "model_type"): + output["model_type"] = self.__class__.model_type + if "_auto_class" in output: + del output["_auto_class"] + if "_commit_hash" in output: + del output["_commit_hash"] + if "_attn_implementation_internal" in output: + del output["_attn_implementation_internal"] + # Do not serialize `base_model_tp_plan` for now + if "base_model_tp_plan" in output: + del output["base_model_tp_plan"] + # Do not serialize `base_model_pp_plan` for now + if "base_model_pp_plan" in output: + del output["base_model_pp_plan"] + + # Transformers version when serializing the model + output["transformers_version"] = __version__ + + for key, value in output.items(): + # Deal with nested configs like CLIP + if isinstance(value, PretrainedConfig): + value = value.to_dict() + del value["transformers_version"] + + output[key] = value + + if hasattr(self, "quantization_config"): + output["quantization_config"] = ( + self.quantization_config.to_dict() + if not isinstance(self.quantization_config, dict) + else self.quantization_config + ) + + # pop the `_pre_quantization_dtype` as torch.dtypes are not serializable. + _ = output.pop("_pre_quantization_dtype", None) + + self.dict_torch_dtype_to_str(output) + + return output + + def to_json_string(self, use_diff: bool = True) -> str: + """ + Serializes this instance to a JSON string. + + Args: + use_diff (`bool`, *optional*, defaults to `True`): + If set to `True`, only the difference between the config instance and the default `PretrainedConfig()` + is serialized to JSON string. + + Returns: + `str`: String containing all the attributes that make up this configuration instance in JSON format. + """ + if use_diff is True: + config_dict = self.to_diff_dict() + else: + config_dict = self.to_dict() + return json.dumps(config_dict, indent=2, sort_keys=True) + "\n" + + def to_json_file(self, json_file_path: Union[str, os.PathLike], use_diff: bool = True): + """ + Save this instance to a JSON file. + + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file in which this configuration instance's parameters will be saved. + use_diff (`bool`, *optional*, defaults to `True`): + If set to `True`, only the difference between the config instance and the default `PretrainedConfig()` + is serialized to JSON file. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + writer.write(self.to_json_string(use_diff=use_diff)) + + def update(self, config_dict: dict[str, Any]): + """ + Updates attributes of this class with attributes from `config_dict`. + + Args: + config_dict (`Dict[str, Any]`): Dictionary of attributes that should be updated for this class. + """ + for key, value in config_dict.items(): + setattr(self, key, value) + + def update_from_string(self, update_str: str): + """ + Updates attributes of this class with attributes from `update_str`. + + The expected format is ints, floats and strings as is, and for booleans use `true` or `false`. For example: + "n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index" + + The keys to change have to already exist in the config object. + + Args: + update_str (`str`): String with attributes that should be updated for this class. + + """ + + d = dict(x.split("=") for x in update_str.split(",")) + for k, v in d.items(): + if not hasattr(self, k): + raise ValueError(f"key {k} isn't in the original config dict") + + old_v = getattr(self, k) + if isinstance(old_v, bool): + if v.lower() in ["true", "1", "y", "yes"]: + v = True + elif v.lower() in ["false", "0", "n", "no"]: + v = False + else: + raise ValueError(f"can't derive true or false from {v} (key {k})") + elif isinstance(old_v, int): + v = int(v) + elif isinstance(old_v, float): + v = float(v) + elif not isinstance(old_v, str): + raise TypeError( + f"You can only update int, float, bool or string values in the config, got {v} for key {k}" + ) + + setattr(self, k, v) + + def dict_torch_dtype_to_str(self, d: dict[str, Any]) -> None: + """ + Checks whether the passed dictionary and its nested dicts have a *torch_dtype* key and if it's not None, + converts torch.dtype to a string of just the type. For example, `torch.float32` get converted into *"float32"* + string, which can then be stored in the json format. + """ + if d.get("torch_dtype", None) is not None: + if isinstance(d["torch_dtype"], dict): + d["torch_dtype"] = {k: str(v).split(".")[-1] for k, v in d["torch_dtype"].items()} + elif not isinstance(d["torch_dtype"], str): + d["torch_dtype"] = str(d["torch_dtype"]).split(".")[1] + for value in d.values(): + if isinstance(value, dict): + self.dict_torch_dtype_to_str(value) + + @classmethod + def register_for_auto_class(cls, auto_class="AutoConfig"): + """ + Register this class with a given auto class. This should only be used for custom configurations as the ones in + the library are already mapped with `AutoConfig`. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"AutoConfig"`): + The auto class to register this new configuration with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + @staticmethod + def _get_global_generation_defaults() -> dict[str, Any]: + return { + "max_length": 20, + "min_length": 0, + "do_sample": False, + "early_stopping": False, + "num_beams": 1, + "num_beam_groups": 1, + "diversity_penalty": 0.0, + "temperature": 1.0, + "top_k": 50, + "top_p": 1.0, + "typical_p": 1.0, + "repetition_penalty": 1.0, + "length_penalty": 1.0, + "no_repeat_ngram_size": 0, + "encoder_no_repeat_ngram_size": 0, + "bad_words_ids": None, + "num_return_sequences": 1, + "output_scores": False, + "return_dict_in_generate": False, + "forced_bos_token_id": None, + "forced_eos_token_id": None, + "remove_invalid_values": False, + "exponential_decay_length_penalty": None, + "suppress_tokens": None, + "begin_suppress_tokens": None, + } + + def _get_non_default_generation_parameters(self) -> dict[str, Any]: + """ + Gets the non-default generation parameters on the PretrainedConfig instance + """ + non_default_generation_parameters = {} + decoder_attribute_name = None + + # Composite models don't have a default config, use their decoder config as a fallback for default values + # If no known pattern is matched, then `default_config = None` -> check against the global generation defaults + try: + default_config = self.__class__() + except ValueError: + decoder_config = self.get_text_config(decoder=True) + if decoder_config is not self: + default_config = decoder_config.__class__() + else: + default_config = None + + # If it is a composite model, we want to check the subconfig that will be used for generation + self_decoder_config = self if decoder_attribute_name is None else getattr(self, decoder_attribute_name) + + for parameter_name, default_global_value in self._get_global_generation_defaults().items(): + if hasattr(self_decoder_config, parameter_name): + is_default_in_config = is_default_generation_value = None + parameter_value = getattr(self_decoder_config, parameter_name) + # Three cases in which is okay for the model config to hold generation config parameters: + # 1. The parameter is set to `None`, effectively delegating its value to the generation config + if parameter_value is None: + continue + # 2. If we have a default config, then the instance should hold the same generation defaults + if default_config is not None: + is_default_in_config = parameter_value == getattr(default_config, parameter_name) + # 3. if we don't have a default config, then the instance should hold the global generation defaults + else: + is_default_generation_value = parameter_value == default_global_value + + is_non_default = (is_default_in_config is False) or ( + is_default_in_config is None and is_default_generation_value is False + ) + if is_non_default: + non_default_generation_parameters[parameter_name] = getattr(self_decoder_config, parameter_name) + + return non_default_generation_parameters + + def get_text_config(self, decoder=False) -> "PretrainedConfig": + """ + Returns the config that is meant to be used with text IO. On most models, it is the original config instance + itself. On specific composite models, it is under a set of valid names. + + Args: + decoder (`Optional[bool]`, *optional*, defaults to `False`): + If set to `True`, then only search for decoder config names. + """ + decoder_possible_text_config_names = ("decoder", "generator", "text_config") + encoder_possible_text_config_names = ("text_encoder",) + if decoder: + possible_text_config_names = decoder_possible_text_config_names + else: + possible_text_config_names = encoder_possible_text_config_names + decoder_possible_text_config_names + + valid_text_config_names = [] + for text_config_name in possible_text_config_names: + if hasattr(self, text_config_name): + text_config = getattr(self, text_config_name, None) + if text_config is not None: + valid_text_config_names += [text_config_name] + + if len(valid_text_config_names) > 1: + raise ValueError( + f"Multiple valid text configs were found in the model config: {valid_text_config_names}. In this " + "case, using `get_text_config()` would be ambiguous. Please specify the desied text config directly." + ) + elif len(valid_text_config_names) == 1: + config_to_return = getattr(self, valid_text_config_names[0]) + else: + config_to_return = self + return config_to_return + + +def get_configuration_file(configuration_files: list[str]) -> str: + """ + Get the configuration file to use for this version of transformers. + + Args: + configuration_files (`List[str]`): The list of available configuration files. + + Returns: + `str`: The configuration file to use. + """ + configuration_files_map = {} + for file_name in configuration_files: + if file_name.startswith("config.") and file_name.endswith(".json") and file_name != "config.json": + v = file_name.removeprefix("config.").removesuffix(".json") + configuration_files_map[v] = file_name + available_versions = sorted(configuration_files_map.keys()) + + # Defaults to FULL_CONFIGURATION_FILE and then try to look at some newer versions. + configuration_file = CONFIG_NAME + transformers_version = version.parse(__version__) + for v in available_versions: + if version.parse(v) <= transformers_version: + configuration_file = configuration_files_map[v] + else: + # No point going further since the versions are sorted. + break + + return configuration_file + + +def recursive_diff_dict(dict_a, dict_b, config_obj=None): + """ + Helper function to recursively take the diff between two nested dictionaries. The resulting diff only contains the + values from `dict_a` that are different from values in `dict_b`. + + dict_b : the default config dictionary. We want to remove values that are in this one + """ + diff = {} + default = config_obj.__class__().to_dict() if config_obj is not None else {} + for key, value in dict_a.items(): + obj_value = getattr(config_obj, str(key), None) + if isinstance(obj_value, PretrainedConfig) and key in dict_b and isinstance(dict_b[key], dict): + diff_value = recursive_diff_dict(value, dict_b[key], config_obj=obj_value) + diff[key] = diff_value + elif key not in dict_b or (value != default[key]): + diff[key] = value + return diff + + +PretrainedConfig.push_to_hub = copy_func(PretrainedConfig.push_to_hub) +if PretrainedConfig.push_to_hub.__doc__ is not None: + PretrainedConfig.push_to_hub.__doc__ = PretrainedConfig.push_to_hub.__doc__.format( + object="config", object_class="AutoConfig", object_files="configuration file" + ) diff --git a/lib/python3.12/site-packages/transformers/convert_graph_to_onnx.py b/lib/python3.12/site-packages/transformers/convert_graph_to_onnx.py new file mode 100644 index 0000000000000000000000000000000000000000..922ece8c0f45203d92ad77bf7a8a439ea759c62e --- /dev/null +++ b/lib/python3.12/site-packages/transformers/convert_graph_to_onnx.py @@ -0,0 +1,551 @@ +# Copyright 2020 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. + +import warnings +from argparse import ArgumentParser +from os import listdir, makedirs +from pathlib import Path +from typing import Optional + +from packaging.version import Version, parse + +from transformers.pipelines import Pipeline, pipeline +from transformers.tokenization_utils import BatchEncoding +from transformers.utils import ModelOutput, is_tf_available, is_torch_available + + +# This is the minimal required version to +# support some ONNX Runtime features +ORT_QUANTIZE_MINIMUM_VERSION = parse("1.4.0") + + +SUPPORTED_PIPELINES = [ + "feature-extraction", + "ner", + "sentiment-analysis", + "fill-mask", + "question-answering", + "text-generation", + "translation_en_to_fr", + "translation_en_to_de", + "translation_en_to_ro", +] + + +class OnnxConverterArgumentParser(ArgumentParser): + """ + Wraps all the script arguments supported to export transformers models to ONNX IR + """ + + def __init__(self): + super().__init__("ONNX Converter") + + self.add_argument( + "--pipeline", + type=str, + choices=SUPPORTED_PIPELINES, + default="feature-extraction", + ) + self.add_argument( + "--model", + type=str, + required=True, + help="Model's id or path (ex: google-bert/bert-base-cased)", + ) + self.add_argument("--tokenizer", type=str, help="Tokenizer's id or path (ex: google-bert/bert-base-cased)") + self.add_argument( + "--framework", + type=str, + choices=["pt", "tf"], + help="Framework for loading the model", + ) + self.add_argument("--opset", type=int, default=11, help="ONNX opset to use") + self.add_argument( + "--check-loading", + action="store_true", + help="Check ONNX is able to load the model", + ) + self.add_argument( + "--use-external-format", + action="store_true", + help="Allow exporting model >= than 2Gb", + ) + self.add_argument( + "--quantize", + action="store_true", + help="Quantize the neural network to be run with int8", + ) + self.add_argument("output") + + +def generate_identified_filename(filename: Path, identifier: str) -> Path: + """ + Append a string-identifier at the end (before the extension, if any) to the provided filepath + + Args: + filename: pathlib.Path The actual path object we would like to add an identifier suffix + identifier: The suffix to add + + Returns: String with concatenated identifier at the end of the filename + """ + return filename.parent.joinpath(filename.stem + identifier).with_suffix(filename.suffix) + + +def check_onnxruntime_requirements(minimum_version: Version): + """ + Check onnxruntime is installed and if the installed version match is recent enough + + Raises: + ImportError: If onnxruntime is not installed or too old version is found + """ + try: + import onnxruntime + + # Parse the version of the installed onnxruntime + ort_version = parse(onnxruntime.__version__) + + # We require 1.4.0 minimum + if ort_version < ORT_QUANTIZE_MINIMUM_VERSION: + raise ImportError( + f"We found an older version of onnxruntime ({onnxruntime.__version__}) " + f"but we require onnxruntime to be >= {minimum_version} to enable all the conversions options.\n" + "Please update onnxruntime by running `pip install --upgrade onnxruntime`" + ) + + except ImportError: + raise ImportError( + "onnxruntime doesn't seem to be currently installed. " + "Please install the onnxruntime by running `pip install onnxruntime`" + " and relaunch the conversion." + ) + + +def ensure_valid_input(model, tokens, input_names): + """ + Ensure inputs are presented in the correct order, without any Non + + Args: + model: The model used to forward the input data + tokens: BatchEncoding holding the input data + input_names: The name of the inputs + + Returns: Tuple + + """ + print("Ensuring inputs are in correct order") + + model_args_name = model.forward.__code__.co_varnames + model_args, ordered_input_names = [], [] + for arg_name in model_args_name[1:]: # start at index 1 to skip "self" argument + if arg_name in input_names: + ordered_input_names.append(arg_name) + model_args.append(tokens[arg_name]) + else: + print(f"{arg_name} is not present in the generated input list.") + break + + print(f"Generated inputs order: {ordered_input_names}") + return ordered_input_names, tuple(model_args) + + +def infer_shapes(nlp: Pipeline, framework: str) -> tuple[list[str], list[str], dict, BatchEncoding]: + """ + Attempt to infer the static vs dynamic axes for each input and output tensors for a specific model + + Args: + nlp: The pipeline object holding the model to be exported + framework: The framework identifier to dispatch to the correct inference scheme (pt/tf) + + Returns: + + - List of the inferred input variable names + - List of the inferred output variable names + - Dictionary with input/output variables names as key and shape tensor as value + - a BatchEncoding reference which was used to infer all the above information + """ + + def build_shape_dict(name: str, tensor, is_input: bool, seq_len: int): + if isinstance(tensor, (tuple, list)): + return [build_shape_dict(name, t, is_input, seq_len) for t in tensor] + + else: + # Let's assume batch is the first axis with only 1 element (~~ might not be always true ...) + axes = {[axis for axis, numel in enumerate(tensor.shape) if numel == 1][0]: "batch"} + if is_input: + if len(tensor.shape) == 2: + axes[1] = "sequence" + else: + raise ValueError(f"Unable to infer tensor axes ({len(tensor.shape)})") + else: + seq_axes = [dim for dim, shape in enumerate(tensor.shape) if shape == seq_len] + axes.update(dict.fromkeys(seq_axes, "sequence")) + + print(f"Found {'input' if is_input else 'output'} {name} with shape: {axes}") + return axes + + tokens = nlp.tokenizer("This is a sample output", return_tensors=framework) + seq_len = tokens.input_ids.shape[-1] + outputs = nlp.model(**tokens) if framework == "pt" else nlp.model(tokens) + if isinstance(outputs, ModelOutput): + outputs = outputs.to_tuple() + if not isinstance(outputs, (list, tuple)): + outputs = (outputs,) + + # Generate input names & axes + input_vars = list(tokens.keys()) + input_dynamic_axes = {k: build_shape_dict(k, v, True, seq_len) for k, v in tokens.items()} + + # flatten potentially grouped outputs (past for gpt2, attentions) + outputs_flat = [] + for output in outputs: + if isinstance(output, (tuple, list)): + outputs_flat.extend(output) + else: + outputs_flat.append(output) + + # Generate output names & axes + output_names = [f"output_{i}" for i in range(len(outputs_flat))] + output_dynamic_axes = {k: build_shape_dict(k, v, False, seq_len) for k, v in zip(output_names, outputs_flat)} + + # Create the aggregated axes representation + dynamic_axes = dict(input_dynamic_axes, **output_dynamic_axes) + return input_vars, output_names, dynamic_axes, tokens + + +def load_graph_from_args( + pipeline_name: str, framework: str, model: str, tokenizer: Optional[str] = None, **models_kwargs +) -> Pipeline: + """ + Convert the set of arguments provided through the CLI to an actual pipeline reference (tokenizer + model + + Args: + pipeline_name: The kind of pipeline to use (ner, question-answering, etc.) + framework: The actual model to convert the pipeline from ("pt" or "tf") + model: The model name which will be loaded by the pipeline + tokenizer: The tokenizer name which will be loaded by the pipeline, default to the model's value + + Returns: Pipeline object + + """ + # If no tokenizer provided + if tokenizer is None: + tokenizer = model + + # Check the wanted framework is available + if framework == "pt" and not is_torch_available(): + raise Exception("Cannot convert because PyTorch is not installed. Please install torch first.") + if framework == "tf" and not is_tf_available(): + raise Exception("Cannot convert because TF is not installed. Please install tensorflow first.") + + print(f"Loading pipeline (model: {model}, tokenizer: {tokenizer})") + + # Allocate tokenizer and model + return pipeline(pipeline_name, model=model, tokenizer=tokenizer, framework=framework, model_kwargs=models_kwargs) + + +def convert_pytorch(nlp: Pipeline, opset: int, output: Path, use_external_format: bool): + """ + Export a PyTorch backed pipeline to ONNX Intermediate Representation (IR + + Args: + nlp: The pipeline to be exported + opset: The actual version of the ONNX operator set to use + output: Path where will be stored the generated ONNX model + use_external_format: Split the model definition from its parameters to allow model bigger than 2GB + + Returns: + + """ + if not is_torch_available(): + raise Exception("Cannot convert because PyTorch is not installed. Please install torch first.") + + import torch + from torch.onnx import export + + print(f"Using framework PyTorch: {torch.__version__}") + + with torch.no_grad(): + input_names, output_names, dynamic_axes, tokens = infer_shapes(nlp, "pt") + ordered_input_names, model_args = ensure_valid_input(nlp.model, tokens, input_names) + + export( + nlp.model, + model_args, + f=output.as_posix(), + input_names=ordered_input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + do_constant_folding=True, + opset_version=opset, + ) + + +def convert_tensorflow(nlp: Pipeline, opset: int, output: Path): + """ + Export a TensorFlow backed pipeline to ONNX Intermediate Representation (IR) + + Args: + nlp: The pipeline to be exported + opset: The actual version of the ONNX operator set to use + output: Path where will be stored the generated ONNX model + + Notes: TensorFlow cannot export model bigger than 2GB due to internal constraint from TensorFlow + + """ + if not is_tf_available(): + raise Exception("Cannot convert because TF is not installed. Please install tensorflow first.") + + print("/!\\ Please note TensorFlow doesn't support exporting model > 2Gb /!\\") + + try: + import tensorflow as tf + import tf2onnx + from tf2onnx import __version__ as t2ov + + print(f"Using framework TensorFlow: {tf.version.VERSION}, tf2onnx: {t2ov}") + + # Build + input_names, output_names, dynamic_axes, tokens = infer_shapes(nlp, "tf") + + # Forward + nlp.model.predict(tokens.data) + input_signature = [tf.TensorSpec.from_tensor(tensor, name=key) for key, tensor in tokens.items()] + model_proto, _ = tf2onnx.convert.from_keras( + nlp.model, input_signature, opset=opset, output_path=output.as_posix() + ) + + except ImportError as e: + raise Exception( + f"Cannot import {e.name} required to convert TF model to ONNX. Please install {e.name} first. {e}" + ) + + +def convert( + framework: str, + model: str, + output: Path, + opset: int, + tokenizer: Optional[str] = None, + use_external_format: bool = False, + pipeline_name: str = "feature-extraction", + **model_kwargs, +): + """ + Convert the pipeline object to the ONNX Intermediate Representation (IR) format + + Args: + framework: The framework the pipeline is backed by ("pt" or "tf") + model: The name of the model to load for the pipeline + output: The path where the ONNX graph will be stored + opset: The actual version of the ONNX operator set to use + tokenizer: The name of the model to load for the pipeline, default to the model's name if not provided + use_external_format: + Split the model definition from its parameters to allow model bigger than 2GB (PyTorch only) + pipeline_name: The kind of pipeline to instantiate (ner, question-answering, etc.) + model_kwargs: Keyword arguments to be forwarded to the model constructor + + Returns: + + """ + warnings.warn( + "The `transformers.convert_graph_to_onnx` package is deprecated and will be removed in version 5 of" + " Transformers", + FutureWarning, + ) + print(f"ONNX opset version set to: {opset}") + + # Load the pipeline + nlp = load_graph_from_args(pipeline_name, framework, model, tokenizer, **model_kwargs) + + if not output.parent.exists(): + print(f"Creating folder {output.parent}") + makedirs(output.parent.as_posix()) + elif len(listdir(output.parent.as_posix())) > 0: + raise Exception(f"Folder {output.parent.as_posix()} is not empty, aborting conversion") + + # Export the graph + if framework == "pt": + convert_pytorch(nlp, opset, output, use_external_format) + else: + convert_tensorflow(nlp, opset, output) + + +def optimize(onnx_model_path: Path) -> Path: + """ + Load the model at the specified path and let onnxruntime look at transformations on the graph to enable all the + optimizations possible + + Args: + onnx_model_path: filepath where the model binary description is stored + + Returns: Path where the optimized model binary description has been saved + + """ + from onnxruntime import InferenceSession, SessionOptions + + # Generate model name with suffix "optimized" + opt_model_path = generate_identified_filename(onnx_model_path, "-optimized") + sess_option = SessionOptions() + sess_option.optimized_model_filepath = opt_model_path.as_posix() + _ = InferenceSession(onnx_model_path.as_posix(), sess_option) + + print(f"Optimized model has been written at {opt_model_path}: \N{HEAVY CHECK MARK}") + print("/!\\ Optimized model contains hardware specific operators which might not be portable. /!\\") + + return opt_model_path + + +def quantize(onnx_model_path: Path) -> Path: + """ + Quantize the weights of the model from float32 to in8 to allow very efficient inference on modern CPU + + Args: + onnx_model_path: Path to location the exported ONNX model is stored + + Returns: The Path generated for the quantized + """ + import onnx + import onnxruntime + from onnx.onnx_pb import ModelProto + from onnxruntime.quantization import QuantizationMode + from onnxruntime.quantization.onnx_quantizer import ONNXQuantizer + from onnxruntime.quantization.registry import IntegerOpsRegistry + + # Load the ONNX model + onnx_model = onnx.load(onnx_model_path.as_posix()) + + if parse(onnx.__version__) < parse("1.5.0"): + print( + "Models larger than 2GB will fail to quantize due to protobuf constraint.\n" + "Please upgrade to onnxruntime >= 1.5.0." + ) + + # Copy it + copy_model = ModelProto() + copy_model.CopyFrom(onnx_model) + + # Construct quantizer + # onnxruntime renamed input_qType to activation_qType in v1.13.1, so we + # check the onnxruntime version to ensure backward compatibility. + # See also: https://github.com/microsoft/onnxruntime/pull/12873 + if parse(onnxruntime.__version__) < parse("1.13.1"): + quantizer = ONNXQuantizer( + model=copy_model, + per_channel=False, + reduce_range=False, + mode=QuantizationMode.IntegerOps, + static=False, + weight_qType=True, + input_qType=False, + tensors_range=None, + nodes_to_quantize=None, + nodes_to_exclude=None, + op_types_to_quantize=list(IntegerOpsRegistry), + ) + else: + quantizer = ONNXQuantizer( + model=copy_model, + per_channel=False, + reduce_range=False, + mode=QuantizationMode.IntegerOps, + static=False, + weight_qType=True, + activation_qType=False, + tensors_range=None, + nodes_to_quantize=None, + nodes_to_exclude=None, + op_types_to_quantize=list(IntegerOpsRegistry), + ) + + # Quantize and export + quantizer.quantize_model() + + # Append "-quantized" at the end of the model's name + quantized_model_path = generate_identified_filename(onnx_model_path, "-quantized") + + # Save model + print(f"Quantized model has been written at {quantized_model_path}: \N{HEAVY CHECK MARK}") + onnx.save_model(quantizer.model.model, quantized_model_path.as_posix()) + + return quantized_model_path + + +def verify(path: Path): + from onnxruntime import InferenceSession, SessionOptions + from onnxruntime.capi.onnxruntime_pybind11_state import RuntimeException + + print(f"Checking ONNX model loading from: {path} ...") + try: + onnx_options = SessionOptions() + _ = InferenceSession(path.as_posix(), onnx_options, providers=["CPUExecutionProvider"]) + print(f"Model {path} correctly loaded: \N{HEAVY CHECK MARK}") + except RuntimeException as re: + print(f"Error while loading the model {re}: \N{HEAVY BALLOT X}") + + +if __name__ == "__main__": + parser = OnnxConverterArgumentParser() + args = parser.parse_args() + + # Make sure output is absolute path + args.output = Path(args.output).absolute() + + try: + print("\n====== Converting model to ONNX ======") + # Convert + convert( + args.framework, + args.model, + args.output, + args.opset, + args.tokenizer, + args.use_external_format, + args.pipeline, + ) + + if args.quantize: + # Ensure requirements for quantization on onnxruntime is met + check_onnxruntime_requirements(ORT_QUANTIZE_MINIMUM_VERSION) + + # onnxruntime optimizations doesn't provide the same level of performances on TensorFlow than PyTorch + if args.framework == "tf": + print( + "\t Using TensorFlow might not provide the same optimization level compared to PyTorch.\n" + "\t For TensorFlow users you can try optimizing the model directly through onnxruntime_tools.\n" + "\t For more information, please refer to the onnxruntime documentation:\n" + "\t\thttps://github.com/microsoft/onnxruntime/tree/master/onnxruntime/python/tools/transformers\n" + ) + + print("\n====== Optimizing ONNX model ======") + + # Quantization works best when using the optimized version of the model + args.optimized_output = optimize(args.output) + + # Do the quantization on the right graph + args.quantized_output = quantize(args.optimized_output) + + # And verify + if args.check_loading: + print("\n====== Check exported ONNX model(s) ======") + verify(args.output) + + if hasattr(args, "optimized_output"): + verify(args.optimized_output) + + if hasattr(args, "quantized_output"): + verify(args.quantized_output) + + except Exception as e: + print(f"Error while converting the model: {e}") + exit(1) diff --git a/lib/python3.12/site-packages/transformers/convert_pytorch_checkpoint_to_tf2.py b/lib/python3.12/site-packages/transformers/convert_pytorch_checkpoint_to_tf2.py new file mode 100644 index 0000000000000000000000000000000000000000..fad6463e982667433c1235c470199d2a74031e84 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/convert_pytorch_checkpoint_to_tf2.py @@ -0,0 +1,440 @@ +# Copyright 2018 The HuggingFace Inc. team. +# +# 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. +"""Convert pytorch checkpoints to TensorFlow""" + +import argparse +import os + +from . import ( + AlbertConfig, + BartConfig, + BertConfig, + CamembertConfig, + CTRLConfig, + DistilBertConfig, + DPRConfig, + ElectraConfig, + FlaubertConfig, + GPT2Config, + LayoutLMConfig, + LxmertConfig, + OpenAIGPTConfig, + RobertaConfig, + T5Config, + TFAlbertForPreTraining, + TFBartForConditionalGeneration, + TFBartForSequenceClassification, + TFBertForPreTraining, + TFBertForQuestionAnswering, + TFBertForSequenceClassification, + TFCamembertForMaskedLM, + TFCTRLLMHeadModel, + TFDistilBertForMaskedLM, + TFDistilBertForQuestionAnswering, + TFDPRContextEncoder, + TFDPRQuestionEncoder, + TFDPRReader, + TFElectraForPreTraining, + TFFlaubertWithLMHeadModel, + TFGPT2LMHeadModel, + TFLayoutLMForMaskedLM, + TFLxmertForPreTraining, + TFLxmertVisualFeatureEncoder, + TFOpenAIGPTLMHeadModel, + TFRobertaForCausalLM, + TFRobertaForMaskedLM, + TFRobertaForSequenceClassification, + TFT5ForConditionalGeneration, + TFTransfoXLLMHeadModel, + TFWav2Vec2Model, + TFXLMRobertaForMaskedLM, + TFXLMWithLMHeadModel, + TFXLNetLMHeadModel, + TransfoXLConfig, + Wav2Vec2Config, + Wav2Vec2Model, + XLMConfig, + XLMRobertaConfig, + XLNetConfig, + is_torch_available, + load_pytorch_checkpoint_in_tf2_model, +) +from .utils import CONFIG_NAME, WEIGHTS_NAME, cached_file, logging + + +if is_torch_available(): + import numpy as np + import torch + + from . import ( + AlbertForPreTraining, + BartForConditionalGeneration, + BertForPreTraining, + BertForQuestionAnswering, + BertForSequenceClassification, + CamembertForMaskedLM, + CTRLLMHeadModel, + DistilBertForMaskedLM, + DistilBertForQuestionAnswering, + DPRContextEncoder, + DPRQuestionEncoder, + DPRReader, + ElectraForPreTraining, + FlaubertWithLMHeadModel, + GPT2LMHeadModel, + LayoutLMForMaskedLM, + LxmertForPreTraining, + LxmertVisualFeatureEncoder, + OpenAIGPTLMHeadModel, + RobertaForMaskedLM, + RobertaForSequenceClassification, + T5ForConditionalGeneration, + TransfoXLLMHeadModel, + XLMRobertaForMaskedLM, + XLMWithLMHeadModel, + XLNetLMHeadModel, + ) + + +logging.set_verbosity_info() + +MODEL_CLASSES = { + "bart": ( + BartConfig, + TFBartForConditionalGeneration, + TFBartForSequenceClassification, + BartForConditionalGeneration, + ), + "bert": ( + BertConfig, + TFBertForPreTraining, + BertForPreTraining, + ), + "google-bert/bert-large-uncased-whole-word-masking-finetuned-squad": ( + BertConfig, + TFBertForQuestionAnswering, + BertForQuestionAnswering, + ), + "google-bert/bert-large-cased-whole-word-masking-finetuned-squad": ( + BertConfig, + TFBertForQuestionAnswering, + BertForQuestionAnswering, + ), + "google-bert/bert-base-cased-finetuned-mrpc": ( + BertConfig, + TFBertForSequenceClassification, + BertForSequenceClassification, + ), + "dpr": ( + DPRConfig, + TFDPRQuestionEncoder, + TFDPRContextEncoder, + TFDPRReader, + DPRQuestionEncoder, + DPRContextEncoder, + DPRReader, + ), + "openai-community/gpt2": ( + GPT2Config, + TFGPT2LMHeadModel, + GPT2LMHeadModel, + ), + "xlnet": ( + XLNetConfig, + TFXLNetLMHeadModel, + XLNetLMHeadModel, + ), + "xlm": ( + XLMConfig, + TFXLMWithLMHeadModel, + XLMWithLMHeadModel, + ), + "xlm-roberta": ( + XLMRobertaConfig, + TFXLMRobertaForMaskedLM, + XLMRobertaForMaskedLM, + ), + "transfo-xl": ( + TransfoXLConfig, + TFTransfoXLLMHeadModel, + TransfoXLLMHeadModel, + ), + "openai-community/openai-gpt": ( + OpenAIGPTConfig, + TFOpenAIGPTLMHeadModel, + OpenAIGPTLMHeadModel, + ), + "roberta": ( + RobertaConfig, + TFRobertaForCausalLM, + TFRobertaForMaskedLM, + RobertaForMaskedLM, + ), + "layoutlm": ( + LayoutLMConfig, + TFLayoutLMForMaskedLM, + LayoutLMForMaskedLM, + ), + "FacebookAI/roberta-large-mnli": ( + RobertaConfig, + TFRobertaForSequenceClassification, + RobertaForSequenceClassification, + ), + "camembert": ( + CamembertConfig, + TFCamembertForMaskedLM, + CamembertForMaskedLM, + ), + "flaubert": ( + FlaubertConfig, + TFFlaubertWithLMHeadModel, + FlaubertWithLMHeadModel, + ), + "distilbert": ( + DistilBertConfig, + TFDistilBertForMaskedLM, + DistilBertForMaskedLM, + ), + "distilbert-base-distilled-squad": ( + DistilBertConfig, + TFDistilBertForQuestionAnswering, + DistilBertForQuestionAnswering, + ), + "lxmert": ( + LxmertConfig, + TFLxmertForPreTraining, + LxmertForPreTraining, + ), + "lxmert-visual-feature-encoder": ( + LxmertConfig, + TFLxmertVisualFeatureEncoder, + LxmertVisualFeatureEncoder, + ), + "Salesforce/ctrl": ( + CTRLConfig, + TFCTRLLMHeadModel, + CTRLLMHeadModel, + ), + "albert": ( + AlbertConfig, + TFAlbertForPreTraining, + AlbertForPreTraining, + ), + "t5": ( + T5Config, + TFT5ForConditionalGeneration, + T5ForConditionalGeneration, + ), + "electra": ( + ElectraConfig, + TFElectraForPreTraining, + ElectraForPreTraining, + ), + "wav2vec2": ( + Wav2Vec2Config, + TFWav2Vec2Model, + Wav2Vec2Model, + ), +} + + +def convert_pt_checkpoint_to_tf( + model_type, pytorch_checkpoint_path, config_file, tf_dump_path, compare_with_pt_model=False, use_cached_models=True +): + if model_type not in MODEL_CLASSES: + raise ValueError(f"Unrecognized model type, should be one of {list(MODEL_CLASSES.keys())}.") + + config_class, model_class, pt_model_class, aws_config_map = MODEL_CLASSES[model_type] + + # Initialise TF model + if config_file in aws_config_map: + config_file = cached_file(config_file, CONFIG_NAME, force_download=not use_cached_models) + config = config_class.from_json_file(config_file) + config.output_hidden_states = True + config.output_attentions = True + print(f"Building TensorFlow model from configuration: {config}") + tf_model = model_class(config) + + # Load weights from tf checkpoint + if pytorch_checkpoint_path in aws_config_map.keys(): + pytorch_checkpoint_path = cached_file( + pytorch_checkpoint_path, WEIGHTS_NAME, force_download=not use_cached_models + ) + # Load PyTorch checkpoint in tf2 model: + tf_model = load_pytorch_checkpoint_in_tf2_model(tf_model, pytorch_checkpoint_path) + + if compare_with_pt_model: + tfo = tf_model(tf_model.dummy_inputs, training=False) # build the network + + state_dict = torch.load(pytorch_checkpoint_path, map_location="cpu", weights_only=True) + pt_model = pt_model_class.from_pretrained( + pretrained_model_name_or_path=None, config=config, state_dict=state_dict + ) + + with torch.no_grad(): + pto = pt_model(**pt_model.dummy_inputs) + + np_pt = pto[0].numpy() + np_tf = tfo[0].numpy() + diff = np.amax(np.abs(np_pt - np_tf)) + print(f"Max absolute difference between models outputs {diff}") + assert diff <= 2e-2, f"Error, model absolute difference is >2e-2: {diff}" + + # Save pytorch-model + print(f"Save TensorFlow model to {tf_dump_path}") + tf_model.save_weights(tf_dump_path, save_format="h5") + + +def convert_all_pt_checkpoints_to_tf( + args_model_type, + tf_dump_path, + model_shortcut_names_or_path=None, + config_shortcut_names_or_path=None, + compare_with_pt_model=False, + use_cached_models=False, + remove_cached_files=False, + only_convert_finetuned_models=False, +): + if args_model_type is None: + model_types = list(MODEL_CLASSES.keys()) + else: + model_types = [args_model_type] + + for j, model_type in enumerate(model_types, start=1): + print("=" * 100) + print(f" Converting model type {j}/{len(model_types)}: {model_type}") + print("=" * 100) + if model_type not in MODEL_CLASSES: + raise ValueError(f"Unrecognized model type {model_type}, should be one of {list(MODEL_CLASSES.keys())}.") + + config_class, model_class, pt_model_class, aws_model_maps, aws_config_map = MODEL_CLASSES[model_type] + + if model_shortcut_names_or_path is None: + model_shortcut_names_or_path = list(aws_model_maps.keys()) + if config_shortcut_names_or_path is None: + config_shortcut_names_or_path = model_shortcut_names_or_path + + for i, (model_shortcut_name, config_shortcut_name) in enumerate( + zip(model_shortcut_names_or_path, config_shortcut_names_or_path), start=1 + ): + print("-" * 100) + if "-squad" in model_shortcut_name or "-mrpc" in model_shortcut_name or "-mnli" in model_shortcut_name: + if not only_convert_finetuned_models: + print(f" Skipping finetuned checkpoint {model_shortcut_name}") + continue + model_type = model_shortcut_name + elif only_convert_finetuned_models: + print(f" Skipping not finetuned checkpoint {model_shortcut_name}") + continue + print( + f" Converting checkpoint {i}/{len(aws_config_map)}: {model_shortcut_name} - model_type {model_type}" + ) + print("-" * 100) + + if config_shortcut_name in aws_config_map: + config_file = cached_file(config_shortcut_name, CONFIG_NAME, force_download=not use_cached_models) + else: + config_file = config_shortcut_name + + if model_shortcut_name in aws_model_maps: + model_file = cached_file(model_shortcut_name, WEIGHTS_NAME, force_download=not use_cached_models) + else: + model_file = model_shortcut_name + + if os.path.isfile(model_shortcut_name): + model_shortcut_name = "converted_model" + + convert_pt_checkpoint_to_tf( + model_type=model_type, + pytorch_checkpoint_path=model_file, + config_file=config_file, + tf_dump_path=os.path.join(tf_dump_path, model_shortcut_name + "-tf_model.h5"), + compare_with_pt_model=compare_with_pt_model, + ) + if remove_cached_files: + os.remove(config_file) + os.remove(model_file) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--tf_dump_path", default=None, type=str, required=True, help="Path to the output Tensorflow dump file." + ) + parser.add_argument( + "--model_type", + default=None, + type=str, + help=( + f"Model type selected in the list of {list(MODEL_CLASSES.keys())}. If not given, will download and " + "convert all the models from AWS." + ), + ) + parser.add_argument( + "--pytorch_checkpoint_path", + default=None, + type=str, + help=( + "Path to the PyTorch checkpoint path or shortcut name to download from AWS. " + "If not given, will download and convert all the checkpoints from AWS." + ), + ) + parser.add_argument( + "--config_file", + default=None, + type=str, + help=( + "The config json file corresponding to the pre-trained model. \n" + "This specifies the model architecture. If not given and " + "--pytorch_checkpoint_path is not given or is a shortcut name " + "use the configuration associated to the shortcut name on the AWS" + ), + ) + parser.add_argument( + "--compare_with_pt_model", action="store_true", help="Compare Tensorflow and PyTorch model predictions." + ) + parser.add_argument( + "--use_cached_models", + action="store_true", + help="Use cached models if possible instead of updating to latest checkpoint versions.", + ) + parser.add_argument( + "--remove_cached_files", + action="store_true", + help="Remove pytorch models after conversion (save memory when converting in batches).", + ) + parser.add_argument("--only_convert_finetuned_models", action="store_true", help="Only convert finetuned models.") + args = parser.parse_args() + + # if args.pytorch_checkpoint_path is not None: + # convert_pt_checkpoint_to_tf(args.model_type.lower(), + # args.pytorch_checkpoint_path, + # args.config_file if args.config_file is not None else args.pytorch_checkpoint_path, + # args.tf_dump_path, + # compare_with_pt_model=args.compare_with_pt_model, + # use_cached_models=args.use_cached_models) + # else: + convert_all_pt_checkpoints_to_tf( + args.model_type.lower() if args.model_type is not None else None, + args.tf_dump_path, + model_shortcut_names_or_path=[args.pytorch_checkpoint_path] + if args.pytorch_checkpoint_path is not None + else None, + config_shortcut_names_or_path=[args.config_file] if args.config_file is not None else None, + compare_with_pt_model=args.compare_with_pt_model, + use_cached_models=args.use_cached_models, + remove_cached_files=args.remove_cached_files, + only_convert_finetuned_models=args.only_convert_finetuned_models, + ) diff --git a/lib/python3.12/site-packages/transformers/convert_slow_tokenizer.py b/lib/python3.12/site-packages/transformers/convert_slow_tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..c8cc1cdbe97ba150f62230eea53b67ab80dd1edc --- /dev/null +++ b/lib/python3.12/site-packages/transformers/convert_slow_tokenizer.py @@ -0,0 +1,1736 @@ +# Copyright 2018 The HuggingFace Inc. team. +# +# 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. +""" +Utilities to convert slow tokenizers in their fast tokenizers counterparts. + +All the conversions are grouped here to gather SentencePiece dependencies outside of the fast tokenizers files and +allow to make our dependency on SentencePiece optional. +""" + +import warnings + +from packaging import version +from tokenizers import AddedToken, Regex, Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import BPE, Unigram, WordPiece + +from .utils import is_protobuf_available, is_sentencepiece_available, logging, requires_backends +from .utils.import_utils import PROTOBUF_IMPORT_ERROR + + +logger = logging.get_logger(__name__) + + +def import_protobuf(error_message=""): + if is_sentencepiece_available(): + from sentencepiece import sentencepiece_model_pb2 + + return sentencepiece_model_pb2 + if is_protobuf_available(): + import google.protobuf + + if version.parse(google.protobuf.__version__) < version.parse("4.0.0"): + from transformers.utils import sentencepiece_model_pb2 + else: + from transformers.utils import sentencepiece_model_pb2_new as sentencepiece_model_pb2 + return sentencepiece_model_pb2 + else: + raise ImportError(PROTOBUF_IMPORT_ERROR.format(error_message)) + + +def _get_prepend_scheme(add_prefix_space: bool, original_tokenizer) -> str: + if add_prefix_space: + prepend_scheme = "always" + if not getattr(original_tokenizer, "legacy", True): + prepend_scheme = "first" + else: + prepend_scheme = "never" + return prepend_scheme + + +def generate_merges(vocab, vocab_scores): + reverse = vocab_scores is not None + vocab_scores = dict(vocab_scores) if reverse else vocab + + merges = [] + for merge, piece_score in vocab_scores.items(): + local = [] + for index in range(1, len(merge)): + piece_l, piece_r = merge[:index], merge[index:] + if piece_l in vocab and piece_r in vocab: + local.append((piece_l, piece_r, piece_score)) + local = sorted(local, key=lambda x: (vocab[x[0]], vocab[x[1]])) + merges.extend(local) + + merges = sorted(merges, key=lambda val: (val[2], len(val[0]), len(val[1])), reverse=reverse) + merges = [(val[0], val[1]) for val in merges] + return merges + + +class SentencePieceExtractor: + """ + Extractor implementation for SentencePiece trained models. https://github.com/google/sentencepiece + """ + + def __init__(self, model: str): + requires_backends(self, "sentencepiece") + from sentencepiece import SentencePieceProcessor + + self.sp = SentencePieceProcessor() + self.sp.Load(model) + + def extract(self, vocab_scores=None) -> tuple[dict[str, int], list[tuple]]: + """ + By default will return vocab and merges with respect to their order, by sending `vocab_scores` we're going to + order the merges with respect to the piece scores instead. + """ + sp = self.sp + vocab = {sp.id_to_piece(index): index for index in range(sp.GetPieceSize())} + + merges = generate_merges(vocab, vocab_scores) + + return vocab, merges + + +class GemmaSentencePieceExtractor(SentencePieceExtractor): + def extract(self, vocab_scores=None) -> tuple[dict[str, int], list[tuple]]: + """ + By default will return vocab and merges with respect to their order, by sending `vocab_scores` we're going to + order the merges with respect to the piece scores instead. + """ + sp = self.sp + vocab = {sp.id_to_piece(index): index for index in range(sp.GetPieceSize())} + + # If "\t" is missing in the vocab, we have to do this to support merges + # "<0x09>" is the bytefallback for `\t` + if "\t" not in vocab: + vocab["\t"] = vocab.get("<0x09>") + merges = generate_merges(vocab, vocab_scores) + return vocab, merges + + +def check_number_comma(piece: str) -> bool: + return len(piece) < 2 or piece[-1] != "," or not piece[-2].isdigit() + + +class Converter: + def __init__(self, original_tokenizer): + self.original_tokenizer = original_tokenizer + + def converted(self) -> Tokenizer: + raise NotImplementedError() + + +class BertConverter(Converter): + def converted(self) -> Tokenizer: + vocab = self.original_tokenizer.vocab + tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token))) + + tokenize_chinese_chars = False + strip_accents = False + do_lower_case = False + if hasattr(self.original_tokenizer, "basic_tokenizer"): + tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars + strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents + do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case + + tokenizer.normalizer = normalizers.BertNormalizer( + clean_text=True, + handle_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + lowercase=do_lower_case, + ) + tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + + cls = str(self.original_tokenizer.cls_token) + sep = str(self.original_tokenizer.sep_token) + cls_token_id = self.original_tokenizer.cls_token_id + sep_token_id = self.original_tokenizer.sep_token_id + + tokenizer.post_processor = processors.TemplateProcessing( + single=f"{cls}:0 $A:0 {sep}:0", + pair=f"{cls}:0 $A:0 {sep}:0 $B:1 {sep}:1", + special_tokens=[ + (cls, cls_token_id), + (sep, sep_token_id), + ], + ) + tokenizer.decoder = decoders.WordPiece(prefix="##") + + return tokenizer + + +class SplinterConverter(Converter): + def converted(self) -> Tokenizer: + vocab = self.original_tokenizer.vocab + tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token))) + + tokenize_chinese_chars = False + strip_accents = False + do_lower_case = False + if hasattr(self.original_tokenizer, "basic_tokenizer"): + tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars + strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents + do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case + + tokenizer.normalizer = normalizers.BertNormalizer( + clean_text=True, + handle_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + lowercase=do_lower_case, + ) + tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + + cls = str(self.original_tokenizer.cls_token) + sep = str(self.original_tokenizer.sep_token) + question = str(self.original_tokenizer.question_token) + dot = "." + cls_token_id = self.original_tokenizer.cls_token_id + sep_token_id = self.original_tokenizer.sep_token_id + question_token_id = self.original_tokenizer.question_token_id + dot_token_id = self.original_tokenizer.convert_tokens_to_ids(".") + + if self.original_tokenizer.padding_side == "right": + pair = f"{cls}:0 $A:0 {question} {dot} {sep}:0 $B:1 {sep}:1" + else: + pair = f"{cls}:0 $A:0 {sep}:0 $B:1 {question} {dot} {sep}:1" + + tokenizer.post_processor = processors.TemplateProcessing( + single=f"{cls}:0 $A:0 {sep}:0", + pair=pair, + special_tokens=[ + (cls, cls_token_id), + (sep, sep_token_id), + (question, question_token_id), + (dot, dot_token_id), + ], + ) + tokenizer.decoder = decoders.WordPiece(prefix="##") + + return tokenizer + + +class FunnelConverter(Converter): + def converted(self) -> Tokenizer: + vocab = self.original_tokenizer.vocab + tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token))) + + tokenize_chinese_chars = False + strip_accents = False + do_lower_case = False + if hasattr(self.original_tokenizer, "basic_tokenizer"): + tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars + strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents + do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case + + tokenizer.normalizer = normalizers.BertNormalizer( + clean_text=True, + handle_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + lowercase=do_lower_case, + ) + tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + + cls = str(self.original_tokenizer.cls_token) + sep = str(self.original_tokenizer.sep_token) + cls_token_id = self.original_tokenizer.cls_token_id + sep_token_id = self.original_tokenizer.sep_token_id + + tokenizer.post_processor = processors.TemplateProcessing( + single=f"{cls}:2 $A:0 {sep}:0", # token_type_id is 2 for Funnel transformer + pair=f"{cls}:2 $A:0 {sep}:0 $B:1 {sep}:1", + special_tokens=[ + (cls, cls_token_id), + (sep, sep_token_id), + ], + ) + tokenizer.decoder = decoders.WordPiece(prefix="##") + + return tokenizer + + +class MPNetConverter(Converter): + def converted(self) -> Tokenizer: + vocab = self.original_tokenizer.vocab + tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token))) + + tokenize_chinese_chars = False + strip_accents = False + do_lower_case = False + if hasattr(self.original_tokenizer, "basic_tokenizer"): + tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars + strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents + do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case + + tokenizer.normalizer = normalizers.BertNormalizer( + clean_text=True, + handle_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + lowercase=do_lower_case, + ) + tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + + cls = str(self.original_tokenizer.cls_token) + sep = str(self.original_tokenizer.sep_token) + cls_token_id = self.original_tokenizer.cls_token_id + sep_token_id = self.original_tokenizer.sep_token_id + + tokenizer.post_processor = processors.TemplateProcessing( + single=f"{cls}:0 $A:0 {sep}:0", + pair=f"{cls}:0 $A:0 {sep}:0 {sep}:0 $B:1 {sep}:1", # MPNet uses two [SEP] tokens + special_tokens=[ + (cls, cls_token_id), + (sep, sep_token_id), + ], + ) + tokenizer.decoder = decoders.WordPiece(prefix="##") + + return tokenizer + + +class OpenAIGPTConverter(Converter): + def converted(self) -> Tokenizer: + vocab = self.original_tokenizer.encoder + merges = list(self.original_tokenizer.bpe_ranks.keys()) + unk_token = self.original_tokenizer.unk_token + + tokenizer = Tokenizer( + BPE( + vocab=vocab, + merges=merges, + dropout=None, + unk_token=str(unk_token), + end_of_word_suffix="", + fuse_unk=False, + ) + ) + + if tokenizer.token_to_id(str(unk_token)) is not None: + tokenizer.add_special_tokens([str(unk_token)]) + + tokenizer.normalizer = normalizers.BertNormalizer(lowercase=True) + tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + tokenizer.decoder = decoders.BPEDecoder(suffix="") + + return tokenizer + + +class GPT2Converter(Converter): + def converted(self, vocab: dict[str, int] = None, merges: list[tuple[str, str]] = None) -> Tokenizer: + if not vocab: + vocab = self.original_tokenizer.encoder + if not merges: + merges = list(self.original_tokenizer.bpe_ranks) + + tokenizer = Tokenizer( + BPE( + vocab=vocab, + merges=merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + ) + ) + + add_prefix_space = getattr(self.original_tokenizer, "add_prefix_space", False) + tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=add_prefix_space) + tokenizer.decoder = decoders.ByteLevel() + if getattr(self.original_tokenizer, "add_bos_token", False): + bos = self.original_tokenizer.bos_token + bos_token_id = self.original_tokenizer.bos_token_id + tokenizer.post_processor = processors.TemplateProcessing( + single=f"{bos}:0 $A:0", + pair=f"{bos}:0 $A:0 $B:1", + special_tokens=[ + (bos, bos_token_id), + ], + ) + else: + # XXX trim_offsets=False actually means this post_processor doesn't + # really do anything. + tokenizer.post_processor = processors.ByteLevel(trim_offsets=False) + return tokenizer + + +class HerbertConverter(Converter): + def converted(self) -> Tokenizer: + tokenizer_info_str = "#version:" + token_suffix = "" + + vocab = self.original_tokenizer.encoder + merges = list(self.original_tokenizer.bpe_ranks.keys()) + if tokenizer_info_str in merges[0][0]: + merges = merges[1:] + + tokenizer = Tokenizer( + BPE( + vocab, + merges, + dropout=None, + unk_token=self.original_tokenizer.unk_token, + end_of_word_suffix=token_suffix, + ) + ) + + tokenizer.normalizer = normalizers.BertNormalizer(lowercase=False, strip_accents=False) + tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + tokenizer.decoder = decoders.BPEDecoder(suffix=token_suffix) + tokenizer.post_processor = processors.BertProcessing( + sep=(self.original_tokenizer.sep_token, self.original_tokenizer.sep_token_id), + cls=(self.original_tokenizer.cls_token, self.original_tokenizer.cls_token_id), + ) + + return tokenizer + + +class Qwen2Converter(Converter): + def converted(self, vocab: dict[str, int] = None, merges: list[tuple[str, str]] = None) -> Tokenizer: + if not vocab: + vocab = self.original_tokenizer.encoder + if not merges: + merges = list(self.original_tokenizer.bpe_ranks.keys()) + + tokenizer = Tokenizer( + BPE( + vocab=vocab, + merges=merges, + dropout=None, + unk_token=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + byte_fallback=False, + ) + ) + + tokenizer.normalizer = normalizers.NFC() + + tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.Split( + Regex( + r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""" + ), + behavior="isolated", + invert=False, + ), + pre_tokenizers.ByteLevel( + add_prefix_space=getattr(self.original_tokenizer, "add_prefix_space", False), + use_regex=False, + ), + ] + ) + + tokenizer.decoder = decoders.ByteLevel() + tokenizer.post_processor = processors.ByteLevel(trim_offsets=False) + + return tokenizer + + +class RobertaConverter(Converter): + def converted(self) -> Tokenizer: + ot = self.original_tokenizer + vocab = ot.encoder + merges = list(ot.bpe_ranks.keys()) + + tokenizer = Tokenizer( + BPE( + vocab=vocab, + merges=merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + ) + ) + + tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=ot.add_prefix_space) + tokenizer.decoder = decoders.ByteLevel() + tokenizer.post_processor = processors.RobertaProcessing( + sep=(ot.sep_token, ot.sep_token_id), + cls=(ot.cls_token, ot.cls_token_id), + add_prefix_space=ot.add_prefix_space, + trim_offsets=True, # True by default on Roberta (historical) + ) + + return tokenizer + + +class RoFormerConverter(Converter): + def converted(self) -> Tokenizer: + from .models.roformer.tokenization_utils import JiebaPreTokenizer + + vocab = self.original_tokenizer.vocab + tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token))) + + strip_accents = False + do_lower_case = False + if hasattr(self.original_tokenizer, "basic_tokenizer"): + strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents + do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case + + tokenizer.normalizer = normalizers.BertNormalizer( + clean_text=True, + handle_chinese_chars=False, + strip_accents=strip_accents, + lowercase=do_lower_case, + ) + tokenizer.pre_tokenizer = pre_tokenizers.PreTokenizer.custom(JiebaPreTokenizer(vocab)) + + cls = str(self.original_tokenizer.cls_token) + sep = str(self.original_tokenizer.sep_token) + cls_token_id = self.original_tokenizer.cls_token_id + sep_token_id = self.original_tokenizer.sep_token_id + + tokenizer.post_processor = processors.TemplateProcessing( + single=f"{cls}:0 $A:0 {sep}:0", + pair=f"{cls}:0 $A:0 {sep}:0 $B:1 {sep}:1", + special_tokens=[ + (cls, cls_token_id), + (sep, sep_token_id), + ], + ) + tokenizer.decoder = decoders.WordPiece(prefix="##") + + return tokenizer + + +class DebertaConverter(Converter): + def converted(self) -> Tokenizer: + ot = self.original_tokenizer + vocab = ot.encoder + merges = list(ot.bpe_ranks.keys()) + + tokenizer = Tokenizer( + BPE( + vocab=vocab, + merges=merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + ) + ) + + tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=ot.add_prefix_space) + tokenizer.decoder = decoders.ByteLevel() + tokenizer.post_processor = processors.TemplateProcessing( + single="[CLS]:0 $A:0 [SEP]:0", + pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1", + special_tokens=[ + ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")), + ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")), + ], + ) + + return tokenizer + + +class SpmConverter(Converter): + handle_byte_fallback = False + SpmExtractor = SentencePieceExtractor + special_tokens = {} + + def __init__(self, *args): + requires_backends(self, "protobuf") + + super().__init__(*args) + + # from .utils import sentencepiece_model_pb2 as model_pb2 + model_pb2 = import_protobuf() + + m = model_pb2.ModelProto() + with open(self.original_tokenizer.vocab_file, "rb") as f: + m.ParseFromString(f.read()) + self.proto = m + + if self.proto.trainer_spec.byte_fallback and not self.handle_byte_fallback: + warnings.warn( + "The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option" + " which is not implemented in the fast tokenizers. In practice this means that the fast version of the" + " tokenizer can produce unknown tokens whereas the sentencepiece version would have converted these " + "unknown tokens into a sequence of byte tokens matching the original piece of text." + ) + + def vocab(self, proto): + return [(piece.piece, piece.score) for piece in proto.pieces] + + def unk_id(self, proto): + return proto.trainer_spec.unk_id + + def tokenizer(self, proto): + model_type = proto.trainer_spec.model_type + vocab_scores = self.vocab(proto) + + if model_type == 1: + tokenizer = Tokenizer( + Unigram( + vocab_scores, + unk_id=self.unk_id(proto), + byte_fallback=self.handle_byte_fallback, + ) + ) + + elif model_type == 2: + _, merges = self.SpmExtractor(self.original_tokenizer.vocab_file).extract(vocab_scores) + bpe_vocab = {word: i for i, (word, score) in enumerate(vocab_scores)} + tokenizer = Tokenizer( + BPE( + bpe_vocab, + merges, + unk_token=proto.trainer_spec.unk_piece, + fuse_unk=True, + byte_fallback=self.handle_byte_fallback, + dropout=None, + ) + ) + + else: + raise Exception( + "You're trying to run a `Unigram` model but you're file was trained with a different algorithm" + ) + + # control tokens are special + # user defined symbols are not + # both user and control tokens are AddedTokens + # Add user defined symbols (type == 4) from sentencepiece (https://github.com/google/sentencepiece/blob/6225e08edb2577757163b3f5dbba4c0b670ef445/src/sentencepiece_model.proto#L299C29-L299C33) + spm_added_tokens = [ + (id, p.piece, p.type == 3 or p.piece in self.special_tokens) + for id, p in enumerate(proto.pieces) + if p.type in [3, 4] + ] + tokenizer.add_tokens( + [ + AddedToken(token, normalized=False, special=special) + for id, token, special in sorted(spm_added_tokens, key=lambda x: x[0]) + ] + ) + + return tokenizer + + def normalizer(self, proto): + precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap + _normalizers = [ + normalizers.Strip(left=False, right=True), # stripping is important + normalizers.Replace(Regex(" {2,}"), "▁"), + ] + if not precompiled_charsmap: + return normalizers.Sequence(_normalizers) + else: + return normalizers.Sequence([normalizers.Precompiled(precompiled_charsmap)] + _normalizers) + + def pre_tokenizer(self, replacement, add_prefix_space): + prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer) + return pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme) + + def post_processor(self): + return None + + def decoder(self, replacement, add_prefix_space): + prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer) + return decoders.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme) + + def converted(self) -> Tokenizer: + tokenizer = self.tokenizer(self.proto) + + # Tokenizer assemble + normalizer = self.normalizer(self.proto) + if normalizer is not None: + tokenizer.normalizer = normalizer + + replacement = "▁" + add_prefix_space = True + if hasattr(self.original_tokenizer, "add_prefix_space"): + add_prefix_space = self.original_tokenizer.add_prefix_space + + pre_tokenizer = self.pre_tokenizer(replacement, add_prefix_space) + if pre_tokenizer is not None: + tokenizer.pre_tokenizer = pre_tokenizer + + tokenizer.decoder = self.decoder(replacement, add_prefix_space) + post_processor = self.post_processor() + if post_processor: + tokenizer.post_processor = post_processor + + return tokenizer + + +class AlbertConverter(SpmConverter): + def vocab(self, proto): + return [ + (piece.piece, piece.score) if check_number_comma(piece.piece) else (piece.piece, piece.score - 100) + for piece in proto.pieces + ] + + def normalizer(self, proto): + list_normalizers = [ + normalizers.Replace("``", '"'), + normalizers.Replace("''", '"'), + ] + if not self.original_tokenizer.keep_accents: + list_normalizers.append(normalizers.NFKD()) + list_normalizers.append(normalizers.StripAccents()) + if self.original_tokenizer.do_lower_case: + list_normalizers.append(normalizers.Lowercase()) + + precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap + + if precompiled_charsmap: + list_normalizers.append(normalizers.Precompiled(precompiled_charsmap)) + + list_normalizers.append(normalizers.Replace(Regex(" {2,}"), " ")) + return normalizers.Sequence(list_normalizers) + + def post_processor(self): + return processors.TemplateProcessing( + single="[CLS]:0 $A:0 [SEP]:0", + pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1", + special_tokens=[ + ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")), + ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")), + ], + ) + + +class BarthezConverter(SpmConverter): + def unk_id(self, proto): + unk_id = 3 + return unk_id + + def post_processor(self): + return processors.TemplateProcessing( + single=" $A ", + pair=" $A $B ", + special_tokens=[ + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class CamembertConverter(SpmConverter): + def vocab(self, proto): + vocab = [ + ("NOTUSED", 0.0), + ("", 0.0), + ("NOTUSED", 0.0), + ("", 0.0), + ("NOTUSED", -100), + ] + # We down-grade the original SentencePiece by -100 to avoid using it and use our added token instead + vocab += [(piece.piece, piece.score) for piece in proto.pieces[1:]] + vocab += [("", 0.0)] + return vocab + + def unk_id(self, proto): + # See vocab unk position + return 3 + + def post_processor(self): + return processors.TemplateProcessing( + single=" $A ", + pair=" $A $B ", + special_tokens=[ + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class DebertaV2Converter(SpmConverter): + def pre_tokenizer(self, replacement, add_prefix_space): + list_pretokenizers = [] + if self.original_tokenizer.split_by_punct: + list_pretokenizers.append(pre_tokenizers.Punctuation(behavior="isolated")) + prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer) + list_pretokenizers.append(pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme)) + return pre_tokenizers.Sequence(list_pretokenizers) + + def normalizer(self, proto): + list_normalizers = [] + if self.original_tokenizer.do_lower_case: + list_normalizers.append(normalizers.Lowercase()) + list_normalizers.append(normalizers.Strip()) + + precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap + if precompiled_charsmap: + list_normalizers.append(normalizers.Precompiled(precompiled_charsmap)) + list_normalizers.append(normalizers.Replace(Regex(" {2,}"), " ")) + + return normalizers.Sequence(list_normalizers) + + def post_processor(self): + return processors.TemplateProcessing( + single="[CLS]:0 $A:0 [SEP]:0", + pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1", + special_tokens=[ + ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")), + ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")), + ], + ) + + +class MBartConverter(SpmConverter): + def vocab(self, proto): + vocab = [ + ("", 0.0), + ("", 0.0), + ("", 0.0), + ("", 0.0), + ] + vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]] + vocab += [ + ("ar_AR", 0.0), + ("cs_CZ", 0.0), + ("de_DE", 0.0), + ("en_XX", 0.0), + ("es_XX", 0.0), + ("et_EE", 0.0), + ("fi_FI", 0.0), + ("fr_XX", 0.0), + ("gu_IN", 0.0), + ("hi_IN", 0.0), + ("it_IT", 0.0), + ("ja_XX", 0.0), + ("kk_KZ", 0.0), + ("ko_KR", 0.0), + ("lt_LT", 0.0), + ("lv_LV", 0.0), + ("my_MM", 0.0), + ("ne_NP", 0.0), + ("nl_XX", 0.0), + ("ro_RO", 0.0), + ("ru_RU", 0.0), + ("si_LK", 0.0), + ("tr_TR", 0.0), + ("vi_VN", 0.0), + ("zh_CN", 0.0), + ] + vocab += [("", 0.0)] + return vocab + + def unk_id(self, proto): + return 3 + + def post_processor(self): + return processors.TemplateProcessing( + single="$A en_XX", + pair="$A $B en_XX", + special_tokens=[ + ("en_XX", self.original_tokenizer.convert_tokens_to_ids("en_XX")), + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class MBart50Converter(SpmConverter): + def vocab(self, proto): + vocab = [ + ("", 0.0), + ("", 0.0), + ("", 0.0), + ("", 0.0), + ] + vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]] + vocab += [("ar_AR", 0.0), ("cs_CZ", 0.0), ("de_DE", 0.0), ("en_XX", 0.0), ("es_XX", 0.0), ("et_EE", 0.0), ("fi_FI", 0.0), ("fr_XX", 0.0), ("gu_IN", 0.0), ("hi_IN", 0.0), ("it_IT", 0.0), ("ja_XX", 0.0), ("kk_KZ", 0.0), ("ko_KR", 0.0), ("lt_LT", 0.0), ("lv_LV", 0.0), ("my_MM", 0.0), ("ne_NP", 0.0), ("nl_XX", 0.0), ("ro_RO", 0.0), ("ru_RU", 0.0), ("si_LK", 0.0), ("tr_TR", 0.0), ("vi_VN", 0.0), ("zh_CN", 0.0), ("af_ZA", 0.0), ("az_AZ", 0.0), ("bn_IN", 0.0), ("fa_IR", 0.0), ("he_IL", 0.0), ("hr_HR", 0.0), ("id_ID", 0.0), ("ka_GE", 0.0), ("km_KH", 0.0), ("mk_MK", 0.0), ("ml_IN", 0.0), ("mn_MN", 0.0), ("mr_IN", 0.0), ("pl_PL", 0.0), ("ps_AF", 0.0), ("pt_XX", 0.0), ("sv_SE", 0.0), ("sw_KE", 0.0), ("ta_IN", 0.0), ("te_IN", 0.0), ("th_TH", 0.0), ("tl_XX", 0.0), ("uk_UA", 0.0), ("ur_PK", 0.0), ("xh_ZA", 0.0), ("gl_ES", 0.0), ("sl_SI", 0.0)] # fmt: skip + vocab += [("", 0.0)] + return vocab + + def unk_id(self, proto): + return 3 + + def post_processor(self): + return processors.TemplateProcessing( + single="en_XX $A ", + pair="en_XX $A $B ", + special_tokens=[ + ("en_XX", self.original_tokenizer.convert_tokens_to_ids("en_XX")), + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class NllbConverter(SpmConverter): + def vocab(self, proto): + vocab = [ + ("", 0.0), + ("", 0.0), + ("", 0.0), + ("", 0.0), + ] + vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]] + return vocab + + def unk_id(self, proto): + return 3 + + def post_processor(self): + return processors.TemplateProcessing( + single="eng_Latn $A ", + pair="eng_Latn $A $B ", + special_tokens=[ + ("eng_Latn", self.original_tokenizer.convert_tokens_to_ids("eng_Latn")), + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class SeamlessM4TConverter(SpmConverter): + def vocab(self, proto): + vocab = [ + ("", 0.0), + ("", 0.0), + ("", 0.0), + ("", 0.0), + ] + vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]] + return vocab + + def unk_id(self, proto): + return self.original_tokenizer.unk_token_id + + def post_processor(self): + return processors.TemplateProcessing( + single="__eng__ $A ", + pair="__eng__ $A $B ", + special_tokens=[ + ("__eng__", self.original_tokenizer.convert_tokens_to_ids("__eng__")), + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class XLMRobertaConverter(SpmConverter): + def vocab(self, proto): + vocab = [ + ("", 0.0), + ("", 0.0), + ("", 0.0), + ("", 0.0), + ] + vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]] + vocab += [("", 0.0)] + return vocab + + def unk_id(self, proto): + unk_id = 3 + return unk_id + + def post_processor(self): + return processors.TemplateProcessing( + single=" $A ", + pair=" $A $B ", + special_tokens=[ + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class XLNetConverter(SpmConverter): + def vocab(self, proto): + return [ + (piece.piece, piece.score) if check_number_comma(piece.piece) else (piece.piece, piece.score - 100) + for piece in proto.pieces + ] + + def normalizer(self, proto): + list_normalizers = [ + normalizers.Replace("``", '"'), + normalizers.Replace("''", '"'), + ] + if not self.original_tokenizer.keep_accents: + list_normalizers.append(normalizers.NFKD()) + list_normalizers.append(normalizers.StripAccents()) + if self.original_tokenizer.do_lower_case: + list_normalizers.append(normalizers.Lowercase()) + + precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap + + if precompiled_charsmap: + list_normalizers.append(normalizers.Precompiled(precompiled_charsmap)) + + list_normalizers.append(normalizers.Replace(Regex(" {2,}"), " ")) + return normalizers.Sequence(list_normalizers) + + def post_processor(self): + return processors.TemplateProcessing( + single="$A:0 :0 :2", + pair="$A:0 :0 $B:1 :1 :2", + special_tokens=[ + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class ReformerConverter(SpmConverter): + pass + + +class RemBertConverter(SpmConverter): + # Inspired from AlbertConverter + def normalizer(self, proto): + list_normalizers = [ + normalizers.Replace("``", '"'), + normalizers.Replace("''", '"'), + normalizers.Replace(Regex(" {2,}"), " "), + ] + if not self.original_tokenizer.keep_accents: + list_normalizers.append(normalizers.NFKD()) + list_normalizers.append(normalizers.StripAccents()) + if self.original_tokenizer.do_lower_case: + list_normalizers.append(normalizers.Lowercase()) + + precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap + + if precompiled_charsmap: + list_normalizers.append(normalizers.Precompiled(precompiled_charsmap)) + + return normalizers.Sequence(list_normalizers) + + def post_processor(self): + return processors.TemplateProcessing( + single="[CLS]:0 $A:0 [SEP]:0", + pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1", + special_tokens=[ + ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")), + ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")), + ], + ) + + +class BertGenerationConverter(SpmConverter): + pass + + +class PegasusConverter(SpmConverter): + def vocab(self, proto): + vocab = [ + (self.original_tokenizer.pad_token, 0.0), + (self.original_tokenizer.eos_token, 0.0), + ] + + if self.original_tokenizer.mask_token_sent is not None: + vocab += [(self.original_tokenizer.mask_token_sent, 0.0)] + + if ( + self.original_tokenizer.mask_token is not None + and self.original_tokenizer.mask_token_id < self.original_tokenizer.offset + ): + vocab += [(self.original_tokenizer.mask_token, 0.0)] + + vocab += [(f"", -100.0) for i in range(2, self.original_tokenizer.offset)] + vocab += [(piece.piece, piece.score) for piece in proto.pieces[2:]] + return vocab + + def unk_id(self, proto): + return proto.trainer_spec.unk_id + self.original_tokenizer.offset + + def pre_tokenizer(self, replacement, add_prefix_space): + prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer) + return pre_tokenizers.Sequence( + [ + pre_tokenizers.WhitespaceSplit(), + pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme), + ] + ) + + def post_processor(self): + eos = self.original_tokenizer.eos_token + special_tokens = [ + (eos, self.original_tokenizer.eos_token_id), + ] + return processors.TemplateProcessing(single=["$A", eos], pair=["$A", "$B", eos], special_tokens=special_tokens) + + +class T5Converter(SpmConverter): + def vocab(self, proto): + num_extra_ids = self.original_tokenizer._extra_ids + vocab = [(piece.piece, piece.score) for piece in proto.pieces] + vocab += [(f"", 0.0) for i in range(num_extra_ids - 1, -1, -1)] + return vocab + + def post_processor(self): + return processors.TemplateProcessing( + single=["$A", ""], + pair=["$A", "", "$B", ""], + special_tokens=[ + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class UdopConverter(SpmConverter): + def post_processor(self): + return processors.TemplateProcessing( + single=["$A", ""], + pair=["$A", "", "$B", ""], + special_tokens=[ + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class WhisperConverter(Converter): + def converted(self) -> Tokenizer: + vocab = self.original_tokenizer.encoder + merges = list(self.original_tokenizer.bpe_ranks.keys()) + + tokenizer = Tokenizer( + BPE( + vocab=vocab, + merges=merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + ) + ) + + tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=self.original_tokenizer.add_prefix_space) + tokenizer.decoder = decoders.ByteLevel() + + prefix_token_ids = self.original_tokenizer.prefix_tokens + prefixes = self.original_tokenizer.convert_ids_to_tokens(prefix_token_ids) + eos = self.original_tokenizer.eos_token + eos_token_id = self.original_tokenizer.eos_token_id + prefix_template = " ".join([f"{token}:0" for token in prefixes]) + tokenizer.post_processor = processors.TemplateProcessing( + single=f"{prefix_template} $A:0 {eos}:0", + pair=f"{prefix_template} $A:0 $B:1 {eos}:1", + special_tokens=[ + (eos, eos_token_id), + *zip(prefixes, prefix_token_ids), + ], + ) + + return tokenizer + + +class BigBirdConverter(SpmConverter): + def post_processor(self): + return processors.TemplateProcessing( + single="[CLS]:0 $A:0 [SEP]:0", + pair="[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1", + special_tokens=[ + ("[CLS]", self.original_tokenizer.convert_tokens_to_ids("[CLS]")), + ("[SEP]", self.original_tokenizer.convert_tokens_to_ids("[SEP]")), + ], + ) + + +class CLIPConverter(Converter): + def converted(self) -> Tokenizer: + vocab = self.original_tokenizer.encoder + merges = list(self.original_tokenizer.bpe_ranks.keys()) + unk_token = self.original_tokenizer.unk_token + + tokenizer = Tokenizer( + BPE( + vocab=vocab, + merges=merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + unk_token=str(unk_token), + ) + ) + + tokenizer.normalizer = normalizers.Sequence( + [normalizers.NFC(), normalizers.Replace(Regex(r"\s+"), " "), normalizers.Lowercase()] + ) + tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.Split( + Regex(r"""'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+"""), + behavior="removed", + invert=True, + ), + pre_tokenizers.ByteLevel(add_prefix_space=False), + ] + ) + tokenizer.decoder = decoders.ByteLevel() + + # Hack to have a ByteLevel and TemplaceProcessor + tokenizer.post_processor = processors.RobertaProcessing( + sep=(self.original_tokenizer.eos_token, self.original_tokenizer.eos_token_id), + cls=(self.original_tokenizer.bos_token, self.original_tokenizer.bos_token_id), + add_prefix_space=False, + trim_offsets=False, + ) + return tokenizer + + +class LayoutLMv2Converter(Converter): + def converted(self) -> Tokenizer: + vocab = self.original_tokenizer.vocab + tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token))) + + tokenize_chinese_chars = False + strip_accents = False + do_lower_case = True + if hasattr(self.original_tokenizer, "basic_tokenizer"): + tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars + strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents + do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case + + tokenizer.normalizer = normalizers.BertNormalizer( + clean_text=True, + handle_chinese_chars=tokenize_chinese_chars, + strip_accents=strip_accents, + lowercase=do_lower_case, + ) + tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + + cls = str(self.original_tokenizer.cls_token) + sep = str(self.original_tokenizer.sep_token) + cls_token_id = self.original_tokenizer.cls_token_id + sep_token_id = self.original_tokenizer.sep_token_id + + tokenizer.post_processor = processors.TemplateProcessing( + single=f"{cls}:0 $A:0 {sep}:0", + pair=f"{cls}:0 $A:0 {sep}:0 $B:1 {sep}:1", + special_tokens=[ + (cls, cls_token_id), + (sep, sep_token_id), + ], + ) + tokenizer.decoder = decoders.WordPiece(prefix="##") + + return tokenizer + + +class BlenderbotConverter(Converter): + def converted(self) -> Tokenizer: + ot = self.original_tokenizer + vocab = ot.encoder + merges = list(ot.bpe_ranks.keys()) + + tokenizer = Tokenizer( + BPE( + vocab=vocab, + merges=merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + ) + ) + + tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=ot.add_prefix_space) + tokenizer.decoder = decoders.ByteLevel() + tokenizer.post_processor = processors.TemplateProcessing( + single=f"$A:0 {ot.eos_token}:0", + special_tokens=[ + (ot.eos_token, ot.eos_token_id), + ], + ) + + return tokenizer + + +class XGLMConverter(SpmConverter): + def vocab(self, proto): + vocab = [ + ("", 0.0), + ("", 0.0), + ("", 0.0), + ("", 0.0), + ] + vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]] + vocab += [("", 0.0), ("", 0.0), ("", 0.0), ("", 0.0), ("", 0.0), ("", 0.0), ("", 0.0)] # fmt: skip + return vocab + + def unk_id(self, proto): + unk_id = 3 + return unk_id + + def post_processor(self): + return processors.TemplateProcessing( + single=" $A", + pair=" $A $B", + special_tokens=[ + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ("", self.original_tokenizer.convert_tokens_to_ids("")), + ], + ) + + +class GemmaConverter(SpmConverter): + handle_byte_fallback = True + SpmExtractor = GemmaSentencePieceExtractor + # start and end of turn tokens must be marked as special + special_tokens = {"", ""} + + """" + split_by_unicode_script: true + split_by_number: true + split_by_whitespace: true + treat_whitespace_as_suffix: false + allow_whitespace_only_pieces: true + split_digits: true + byte_fallback: true + """ + + def normalizer(self, proto): + return normalizers.Replace(" ", "▁") + + def vocab(self, proto): + vocab = [ + (self.original_tokenizer.pad_token, 0.0), + (self.original_tokenizer.eos_token, 0.0), + (self.original_tokenizer.bos_token, 0.0), + ] + vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]] + + # Older gemma tokenizers had a missing tab token, so we fix that here + if not any(x[0] == "\t" for x in vocab): + override_index = next((i for i, x in enumerate(vocab) if x[0] == "<0x09>"), None) + if override_index is not None: + vocab[override_index] = ("\t", 0.0) + + return vocab + + def pre_tokenizer(self, replacement, add_prefix_space): + return pre_tokenizers.Split(" ", "merged_with_previous") + + def unk_id(self, proto): + unk_id = 3 + return unk_id + + def decoder(self, replacement, add_prefix_space): + return decoders.Sequence( + [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + ] + ) + + +class LlamaConverter(SpmConverter): + handle_byte_fallback = True + + def vocab(self, proto): + vocab = [ + (self.original_tokenizer.convert_ids_to_tokens(0), 0.0), + (self.original_tokenizer.convert_ids_to_tokens(1), 0.0), + (self.original_tokenizer.convert_ids_to_tokens(2), 0.0), + ] + vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]] + return vocab + + def unk_id(self, proto): + unk_id = 0 + return unk_id + + def decoder(self, replacement, add_prefix_space): + sequence = [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + ] + if add_prefix_space: + sequence += [decoders.Strip(content=" ", left=1)] + return decoders.Sequence(sequence) + + def normalizer(self, proto): + if getattr(self.original_tokenizer, "legacy", True): + sequence = [] + if getattr(self.original_tokenizer, "add_prefix_space", True): + sequence += [normalizers.Prepend(prepend="▁")] + sequence += [normalizers.Replace(pattern=" ", content="▁")] + return normalizers.Sequence(sequence) + return None # non-legacy, no normalizer + + def pre_tokenizer(self, replacement, add_prefix_space): + if not getattr(self.original_tokenizer, "legacy", True): # non-legacy, we need a replace + prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer) + return pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme, split=False) + return None + + def post_processor(self): + # the processor is defined in the LlamaTokenizerFast class. + return None + + +class MarkupLMConverter(Converter): + def converted(self) -> Tokenizer: + ot = self.original_tokenizer + vocab = ot.encoder + merges = list(ot.bpe_ranks.keys()) + + tokenizer = Tokenizer( + BPE( + vocab=vocab, + merges=merges, + dropout=None, + continuing_subword_prefix="", + end_of_word_suffix="", + fuse_unk=False, + unk_token=self.original_tokenizer.unk_token, + ) + ) + + tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=ot.add_prefix_space) + tokenizer.decoder = decoders.ByteLevel() + + cls = str(self.original_tokenizer.cls_token) + sep = str(self.original_tokenizer.sep_token) + cls_token_id = self.original_tokenizer.cls_token_id + sep_token_id = self.original_tokenizer.sep_token_id + + tokenizer.post_processor = processors.TemplateProcessing( + single=f"{cls} $A {sep}", + pair=f"{cls} $A {sep} $B {sep}", + special_tokens=[ + (cls, cls_token_id), + (sep, sep_token_id), + ], + ) + + return tokenizer + + +class MoshiConverter(SpmConverter): + handle_byte_fallback = True + + def __init__(self, vocab_file, model_max_length=None, **kwargs): + requires_backends(self, "protobuf") + + Converter.__init__(self, vocab_file) + + # from .utils import sentencepiece_model_pb2 as model_pb2 + model_pb2 = import_protobuf() + + m = model_pb2.ModelProto() + with open(vocab_file, "rb") as f: + m.ParseFromString(f.read()) + self.proto = m + + def normalizer(self, proto): + precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap + _normalizers = [ + normalizers.Replace(" ", "▁"), + ] + if not precompiled_charsmap: + return normalizers.Sequence(_normalizers) + else: + return normalizers.Sequence([normalizers.Precompiled(precompiled_charsmap)] + _normalizers) + + def decoder(self, replacement, add_prefix_space): + sequence = [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + ] + if add_prefix_space: + sequence += [decoders.Strip(content=" ", left=1)] + return decoders.Sequence(sequence) + + def pre_tokenizer(self, replacement, add_prefix_space): + prepend_scheme = "first" + return pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme, split=False) + + +class HeliumConverter(SpmConverter): + handle_byte_fallback = True + + def __init__(self, vocab_file=None, *args): + requires_backends(self, "protobuf") + + Converter.__init__(self, vocab_file) + + model_pb2 = import_protobuf() + + m = model_pb2.ModelProto() + with open(vocab_file, "rb") as f: + m.ParseFromString(f.read()) + self.proto = m + + def tokenizer(self, proto): + vocab_scores = self.vocab(proto) + tokenizer = Tokenizer( + Unigram( + vocab_scores, + unk_id=self.unk_id(proto), + byte_fallback=self.handle_byte_fallback, + ) + ) + # control tokens are special + # user defined symbols are not + # both user and control tokens are AddedTokens + # Add user defined symbols (type == 4) from sentencepiece (https://github.com/google/sentencepiece/blob/6225e08edb2577757163b3f5dbba4c0b670ef445/src/sentencepiece_model.proto#L299C29-L299C33) + spm_added_tokens = [ + (id, p.piece, p.type == 3 or p.piece in self.special_tokens) + for id, p in enumerate(proto.pieces) + if p.type in [3, 4] + ] + tokenizer.add_tokens( + [ + AddedToken(token, normalized=False, special=special, single_word=True) + for id, token, special in sorted(spm_added_tokens, key=lambda x: x[0]) + ] + ) + tokenizer.add_tokens([AddedToken("\n", normalized=False, special=False)]) + tokenizer.enable_padding(pad_token="", pad_id=3) + return tokenizer + + def vocab(self, proto): + vocab = [] + for piece in proto.pieces: + if piece.piece == "<0x0A>": + vocab += [("\n", piece.score)] + else: + vocab += [(piece.piece, piece.score)] + return vocab + + def unk_id(self, proto): + unk_id = 0 + return unk_id + + def decoder(self, replacement, add_prefix_space): + sequence = [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + ] + sequence += [decoders.Strip(content=" ", left=1)] + return decoders.Sequence(sequence) + + def normalizer(self, proto): + return normalizers.Sequence([normalizers.Prepend(" "), normalizers.Replace(r" ", "▁")]) + + def pre_tokenizer(self, replacement, add_prefix_space): + return pre_tokenizers.Sequence([pre_tokenizers.Split("\n", "contiguous")]) + + def post_processor(self): + return processors.TemplateProcessing( + single=[ + "", + "$A", + ], + pair=[ + "", + "$A", + "", + "$B", + ], + special_tokens=[ + ("", 1), + ], + ) + + +# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode +def bytes_to_unicode(): + """ + Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control + characters the bpe code barfs on. + + The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab + if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for + decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup + tables between utf-8 bytes and unicode strings. + """ + bs = ( + list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1)) + ) + cs = bs[:] + n = 0 + for b in range(2**8): + if b not in bs: + bs.append(b) + cs.append(2**8 + n) + n += 1 + cs = [chr(n) for n in cs] + return dict(zip(bs, cs)) + + +class TikTokenConverter: + """ + A general tiktoken converter. + """ + + def __init__( + self, + vocab_file=None, + pattern=r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""", + add_prefix_space=False, + additional_special_tokens=None, + *args, + **kwargs, + ): + super().__init__(*args) + self.vocab_file = vocab_file + self.pattern = pattern + self.add_prefix_space = add_prefix_space + self.additional_special_tokens = ( + additional_special_tokens.keys() if type(additional_special_tokens) is dict else additional_special_tokens + ) + + def extract_vocab_merges_from_model(self, tiktoken_url: str): + try: + from tiktoken.load import load_tiktoken_bpe + except Exception: + raise ValueError( + "`tiktoken` is required to read a `tiktoken` file. Install it with `pip install tiktoken`." + ) + + bpe_ranks = load_tiktoken_bpe(tiktoken_url) + byte_encoder = bytes_to_unicode() + + def token_bytes_to_string(b): + return "".join([byte_encoder[ord(char)] for char in b.decode("latin-1")]) + + merges = [] + vocab = {} + for token, rank in bpe_ranks.items(): + vocab[token_bytes_to_string(token)] = rank + if len(token) == 1: + continue + local = [] + for index in range(1, len(token)): + piece_l, piece_r = token[:index], token[index:] + if piece_l in bpe_ranks and piece_r in bpe_ranks and (piece_l + piece_r) in bpe_ranks: + local.append((piece_l, piece_r, rank)) + local = sorted(local, key=lambda x: (bpe_ranks[x[0]], bpe_ranks[x[1]]), reverse=False) + merges.extend(local) + merges = sorted(merges, key=lambda val: val[2], reverse=False) + merges = [(token_bytes_to_string(val[0]), token_bytes_to_string(val[1])) for val in merges] + return vocab, merges + + def tokenizer(self): + vocab_scores, merges = self.extract_vocab_merges_from_model(self.vocab_file) + tokenizer = Tokenizer(BPE(vocab_scores, merges, fuse_unk=False)) + if hasattr(tokenizer.model, "ignore_merges"): + tokenizer.model.ignore_merges = True + return tokenizer + + def converted(self) -> Tokenizer: + tokenizer = self.tokenizer() + tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.Split(Regex(self.pattern), behavior="isolated", invert=False), + pre_tokenizers.ByteLevel(add_prefix_space=self.add_prefix_space, use_regex=False), + ] + ) + tokenizer.decoder = decoders.ByteLevel() + + tokenizer.add_special_tokens( + [AddedToken(token, normalized=False, special=True) for token in self.additional_special_tokens] + ) + + tokenizer.post_processor = processors.ByteLevel(trim_offsets=False) + + return tokenizer + + +SLOW_TO_FAST_CONVERTERS = { + "AlbertTokenizer": AlbertConverter, + "BartTokenizer": RobertaConverter, + "BarthezTokenizer": BarthezConverter, + "BertTokenizer": BertConverter, + "BigBirdTokenizer": BigBirdConverter, + "BlenderbotTokenizer": BlenderbotConverter, + "CamembertTokenizer": CamembertConverter, + "CLIPTokenizer": CLIPConverter, + "CodeGenTokenizer": GPT2Converter, + "ConvBertTokenizer": BertConverter, + "DebertaTokenizer": DebertaConverter, + "DebertaV2Tokenizer": DebertaV2Converter, + "DistilBertTokenizer": BertConverter, + "DPRReaderTokenizer": BertConverter, + "DPRQuestionEncoderTokenizer": BertConverter, + "DPRContextEncoderTokenizer": BertConverter, + "ElectraTokenizer": BertConverter, + "FNetTokenizer": AlbertConverter, + "FunnelTokenizer": FunnelConverter, + "GPT2Tokenizer": GPT2Converter, + "HerbertTokenizer": HerbertConverter, + "LayoutLMTokenizer": BertConverter, + "LayoutLMv2Tokenizer": BertConverter, + "LayoutLMv3Tokenizer": RobertaConverter, + "LayoutXLMTokenizer": XLMRobertaConverter, + "LongformerTokenizer": RobertaConverter, + "LEDTokenizer": RobertaConverter, + "LxmertTokenizer": BertConverter, + "MarkupLMTokenizer": MarkupLMConverter, + "MBartTokenizer": MBartConverter, + "MBart50Tokenizer": MBart50Converter, + "MPNetTokenizer": MPNetConverter, + "MobileBertTokenizer": BertConverter, + "MvpTokenizer": RobertaConverter, + "NllbTokenizer": NllbConverter, + "OpenAIGPTTokenizer": OpenAIGPTConverter, + "PegasusTokenizer": PegasusConverter, + "Qwen2Tokenizer": Qwen2Converter, + "RealmTokenizer": BertConverter, + "ReformerTokenizer": ReformerConverter, + "RemBertTokenizer": RemBertConverter, + "RetriBertTokenizer": BertConverter, + "RobertaTokenizer": RobertaConverter, + "RoFormerTokenizer": RoFormerConverter, + "SeamlessM4TTokenizer": SeamlessM4TConverter, + "SqueezeBertTokenizer": BertConverter, + "T5Tokenizer": T5Converter, + "UdopTokenizer": UdopConverter, + "WhisperTokenizer": WhisperConverter, + "XLMRobertaTokenizer": XLMRobertaConverter, + "XLNetTokenizer": XLNetConverter, + "SplinterTokenizer": SplinterConverter, + "XGLMTokenizer": XGLMConverter, + "LlamaTokenizer": LlamaConverter, + "CodeLlamaTokenizer": LlamaConverter, + "GemmaTokenizer": GemmaConverter, + "Phi3Tokenizer": LlamaConverter, +} + + +def convert_slow_tokenizer(transformer_tokenizer, from_tiktoken=False) -> Tokenizer: + """ + Utilities to convert a slow tokenizer instance in a fast tokenizer instance. + + Args: + transformer_tokenizer ([`~tokenization_utils_base.PreTrainedTokenizer`]): + Instance of a slow tokenizer to convert in the backend tokenizer for + [`~tokenization_utils_base.PreTrainedTokenizerFast`]. + from_tiktoken (bool, optional): Whether to use the `tiktoken` library to convert the tokenizer instead of sentencepiece. + Defaults to False. + + Return: + A instance of [`~tokenizers.Tokenizer`] to be used as the backend tokenizer of a + [`~tokenization_utils_base.PreTrainedTokenizerFast`] + """ + + tokenizer_class_name = transformer_tokenizer.__class__.__name__ + if tokenizer_class_name in SLOW_TO_FAST_CONVERTERS and not from_tiktoken: + converter_class = SLOW_TO_FAST_CONVERTERS[tokenizer_class_name] + return converter_class(transformer_tokenizer).converted() + + else: + try: + logger.info("Converting from Tiktoken") + return TikTokenConverter( + vocab_file=transformer_tokenizer.vocab_file, + additional_special_tokens=transformer_tokenizer.additional_special_tokens, + ).converted() + except Exception: + raise ValueError( + f"Converting from SentencePiece and Tiktoken failed, if a converter for SentencePiece is available, provide a model path " + f"with a SentencePiece tokenizer.model file." + f"Currently available slow->fast converters: {list(SLOW_TO_FAST_CONVERTERS.keys())}" + ) diff --git a/lib/python3.12/site-packages/transformers/convert_slow_tokenizers_checkpoints_to_fast.py b/lib/python3.12/site-packages/transformers/convert_slow_tokenizers_checkpoints_to_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..855ab6381ee1de5dd2d13a279381160625fd6b59 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/convert_slow_tokenizers_checkpoints_to_fast.py @@ -0,0 +1,129 @@ +# Copyright 2018 The HuggingFace Inc. team. +# +# 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. +"""Convert slow tokenizers checkpoints in fast (serialization format of the `tokenizers` library)""" + +import argparse +import os + +import transformers + +from .convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS +from .utils import logging + + +logging.set_verbosity_info() + +logger = logging.get_logger(__name__) + + +TOKENIZER_CLASSES = { + # Phi3 uses Llama tokenizer + name: getattr(transformers, "LlamaTokenizerFast" if name == "Phi3Tokenizer" else name + "Fast") + for name in SLOW_TO_FAST_CONVERTERS +} + + +def convert_slow_checkpoint_to_fast(tokenizer_name, checkpoint_name, dump_path, force_download): + if tokenizer_name is not None and tokenizer_name not in TOKENIZER_CLASSES: + raise ValueError(f"Unrecognized tokenizer name, should be one of {list(TOKENIZER_CLASSES.keys())}.") + + if tokenizer_name is None: + tokenizer_names = TOKENIZER_CLASSES + else: + tokenizer_names = {tokenizer_name: getattr(transformers, tokenizer_name + "Fast")} + + logger.info(f"Loading tokenizer classes: {tokenizer_names}") + + for tokenizer_name in tokenizer_names: + tokenizer_class = TOKENIZER_CLASSES[tokenizer_name] + + add_prefix = True + if checkpoint_name is None: + checkpoint_names = list(tokenizer_class.max_model_input_sizes.keys()) + else: + checkpoint_names = [checkpoint_name] + + logger.info(f"For tokenizer {tokenizer_class.__class__.__name__} loading checkpoints: {checkpoint_names}") + + for checkpoint in checkpoint_names: + logger.info(f"Loading {tokenizer_class.__class__.__name__} {checkpoint}") + + # Load tokenizer + tokenizer = tokenizer_class.from_pretrained(checkpoint, force_download=force_download) + + # Save fast tokenizer + logger.info(f"Save fast tokenizer to {dump_path} with prefix {checkpoint} add_prefix {add_prefix}") + + # For organization names we create sub-directories + if "/" in checkpoint: + checkpoint_directory, checkpoint_prefix_name = checkpoint.split("/") + dump_path_full = os.path.join(dump_path, checkpoint_directory) + elif add_prefix: + checkpoint_prefix_name = checkpoint + dump_path_full = dump_path + else: + checkpoint_prefix_name = None + dump_path_full = dump_path + + logger.info(f"=> {dump_path_full} with prefix {checkpoint_prefix_name}, add_prefix {add_prefix}") + + if checkpoint in list(tokenizer.pretrained_vocab_files_map.values())[0]: + file_path = list(tokenizer.pretrained_vocab_files_map.values())[0][checkpoint] + next_char = file_path.split(checkpoint)[-1][0] + if next_char == "/": + dump_path_full = os.path.join(dump_path_full, checkpoint_prefix_name) + checkpoint_prefix_name = None + + logger.info(f"=> {dump_path_full} with prefix {checkpoint_prefix_name}, add_prefix {add_prefix}") + + file_names = tokenizer.save_pretrained( + dump_path_full, legacy_format=False, filename_prefix=checkpoint_prefix_name + ) + logger.info(f"=> File names {file_names}") + + for file_name in file_names: + if not file_name.endswith("tokenizer.json"): + os.remove(file_name) + logger.info(f"=> removing {file_name}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--dump_path", default=None, type=str, required=True, help="Path to output generated fast tokenizer files." + ) + parser.add_argument( + "--tokenizer_name", + default=None, + type=str, + help=( + f"Optional tokenizer type selected in the list of {list(TOKENIZER_CLASSES.keys())}. If not given, will " + "download and convert all the checkpoints from AWS." + ), + ) + parser.add_argument( + "--checkpoint_name", + default=None, + type=str, + help="Optional checkpoint name. If not given, will download and convert the canonical checkpoints from AWS.", + ) + parser.add_argument( + "--force_download", + action="store_true", + help="Re-download checkpoints.", + ) + args = parser.parse_args() + + convert_slow_checkpoint_to_fast(args.tokenizer_name, args.checkpoint_name, args.dump_path, args.force_download) diff --git a/lib/python3.12/site-packages/transformers/convert_tf_hub_seq_to_seq_bert_to_pytorch.py b/lib/python3.12/site-packages/transformers/convert_tf_hub_seq_to_seq_bert_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..e2c825a45b605c271c3912bf5b4f3ff4b5c1a32c --- /dev/null +++ b/lib/python3.12/site-packages/transformers/convert_tf_hub_seq_to_seq_bert_to_pytorch.py @@ -0,0 +1,86 @@ +# Copyright 2020 The HuggingFace Inc. team. +# +# 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. +"""Convert Seq2Seq TF Hub checkpoint.""" + +import argparse + +from . import ( + BertConfig, + BertGenerationConfig, + BertGenerationDecoder, + BertGenerationEncoder, + load_tf_weights_in_bert_generation, + logging, +) + + +logging.set_verbosity_info() + + +def convert_tf_checkpoint_to_pytorch(tf_hub_path, pytorch_dump_path, is_encoder_named_decoder, vocab_size, is_encoder): + # Initialise PyTorch model + bert_config = BertConfig.from_pretrained( + "google-bert/bert-large-cased", + vocab_size=vocab_size, + max_position_embeddings=512, + is_decoder=True, + add_cross_attention=True, + ) + bert_config_dict = bert_config.to_dict() + del bert_config_dict["type_vocab_size"] + config = BertGenerationConfig(**bert_config_dict) + if is_encoder: + model = BertGenerationEncoder(config) + else: + model = BertGenerationDecoder(config) + print(f"Building PyTorch model from configuration: {config}") + + # Load weights from tf checkpoint + load_tf_weights_in_bert_generation( + model, + tf_hub_path, + model_class="bert", + is_encoder_named_decoder=is_encoder_named_decoder, + is_encoder=is_encoder, + ) + + # Save pytorch-model + print(f"Save PyTorch model and config to {pytorch_dump_path}") + model.save_pretrained(pytorch_dump_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--tf_hub_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path." + ) + parser.add_argument( + "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model." + ) + parser.add_argument( + "--is_encoder_named_decoder", + action="store_true", + help="If decoder has to be renamed to encoder in PyTorch model.", + ) + parser.add_argument("--is_encoder", action="store_true", help="If model is an encoder.") + parser.add_argument("--vocab_size", default=50358, type=int, help="Vocab size of model") + args = parser.parse_args() + convert_tf_checkpoint_to_pytorch( + args.tf_hub_path, + args.pytorch_dump_path, + args.is_encoder_named_decoder, + args.vocab_size, + is_encoder=args.is_encoder, + ) diff --git a/lib/python3.12/site-packages/transformers/debug_utils.py b/lib/python3.12/site-packages/transformers/debug_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..91425d9fed7a3e4626d670b615565e6707b0f870 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/debug_utils.py @@ -0,0 +1,346 @@ +# Copyright 2020 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. + +import collections + +from .utils import ExplicitEnum, is_torch_available, logging + + +if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) + + +class DebugUnderflowOverflow: + """ + This debug class helps detect and understand where the model starts getting very large or very small, and more + importantly `nan` or `inf` weight and activation elements. + + There are 2 working modes: + + 1. Underflow/overflow detection (default) + 2. Specific batch absolute min/max tracing without detection + + Mode 1: Underflow/overflow detection + + To activate the underflow/overflow detection, initialize the object with the model : + + ```python + debug_overflow = DebugUnderflowOverflow(model) + ``` + + then run the training as normal and if `nan` or `inf` gets detected in at least one of the weight, input or output + elements this module will throw an exception and will print `max_frames_to_save` frames that lead to this event, + each frame reporting + + 1. the fully qualified module name plus the class name whose `forward` was run + 2. the absolute min and max value of all elements for each module weights, and the inputs and output + + For example, here is the header and the last few frames in detection report for `google/mt5-small` run in fp16 + mixed precision : + + ``` + Detected inf/nan during batch_number=0 + Last 21 forward frames: + abs min abs max metadata + [...] + encoder.block.2.layer.1.DenseReluDense.wi_0 Linear + 2.17e-07 4.50e+00 weight + 1.79e-06 4.65e+00 input[0] + 2.68e-06 3.70e+01 output + encoder.block.2.layer.1.DenseReluDense.wi_1 Linear + 8.08e-07 2.66e+01 weight + 1.79e-06 4.65e+00 input[0] + 1.27e-04 2.37e+02 output + encoder.block.2.layer.1.DenseReluDense.wo Linear + 1.01e-06 6.44e+00 weight + 0.00e+00 9.74e+03 input[0] + 3.18e-04 6.27e+04 output + encoder.block.2.layer.1.DenseReluDense T5DenseGatedGeluDense + 1.79e-06 4.65e+00 input[0] + 3.18e-04 6.27e+04 output + encoder.block.2.layer.1.dropout Dropout + 3.18e-04 6.27e+04 input[0] + 0.00e+00 inf output + ``` + + You can see here, that `T5DenseGatedGeluDense.forward` resulted in output activations, whose absolute max value was + around 62.7K, which is very close to fp16's top limit of 64K. In the next frame we have `Dropout` which + renormalizes the weights, after it zeroed some of the elements, which pushes the absolute max value to more than + 64K, and we get an overlow. + + As you can see it's the previous frames that we need to look into when the numbers start going into very large for + fp16 numbers. + + The tracking is done in a forward hook, which gets invoked immediately after `forward` has completed. + + By default the last 21 frames are printed. You can change the default to adjust for your needs. For example : + + ```python + debug_overflow = DebugUnderflowOverflow(model, max_frames_to_save=100) + ``` + + To validate that you have set up this debugging feature correctly, and you intend to use it in a training that + may take hours to complete, first run it with normal tracing enabled for one of a few batches as explained in + the next section. + + + Mode 2. Specific batch absolute min/max tracing without detection + + The second work mode is per-batch tracing with the underflow/overflow detection feature turned off. + + Let's say you want to watch the absolute min and max values for all the ingredients of each `forward` call of a + given batch, and only do that for batches 1 and 3. Then you instantiate this class as : + + ```python + debug_overflow = DebugUnderflowOverflow(model, trace_batch_nums=[1, 3]) + ``` + + And now full batches 1 and 3 will be traced using the same format as explained above. Batches are 0-indexed. + + This is helpful if you know that the program starts misbehaving after a certain batch number, so you can + fast-forward right to that area. + + + Early stopping: + + You can also specify the batch number after which to stop the training, with : + + ```python + debug_overflow = DebugUnderflowOverflow(model, trace_batch_nums=[1, 3], abort_after_batch_num=3) + ``` + + This feature is mainly useful in the tracing mode, but you can use it for any mode. + + + **Performance**: + + As this module measures absolute `min`/``max` of each weight of the model on every forward it'll slow the training + down. Therefore remember to turn it off once the debugging needs have been met. + + Args: + model (`nn.Module`): + The model to debug. + max_frames_to_save (`int`, *optional*, defaults to 21): + How many frames back to record + trace_batch_nums(`List[int]`, *optional*, defaults to `[]`): + Which batch numbers to trace (turns detection off) + abort_after_batch_num (`int``, *optional*): + Whether to abort after a certain batch number has finished + """ + + def __init__(self, model, max_frames_to_save=21, trace_batch_nums=[], abort_after_batch_num=None): + self.model = model + self.trace_batch_nums = trace_batch_nums + self.abort_after_batch_num = abort_after_batch_num + + # keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence + self.frames = collections.deque([], max_frames_to_save) + self.frame = [] + self.batch_number = 0 + self.total_calls = 0 + self.detected_overflow = False + self.prefix = " " + + self.analyse_model() + + self.register_forward_hook() + + def save_frame(self, frame=None): + if frame is not None: + self.expand_frame(frame) + self.frames.append("\n".join(self.frame)) + self.frame = [] # start a new frame + + def expand_frame(self, line): + self.frame.append(line) + + def trace_frames(self): + print("\n".join(self.frames)) + self.frames = [] + + def reset_saved_frames(self): + self.frames = [] + + def dump_saved_frames(self): + print(f"\nDetected inf/nan during batch_number={self.batch_number}") + print(f"Last {len(self.frames)} forward frames:") + print(f"{'abs min':8} {'abs max':8} metadata") + print("\n".join(self.frames)) + print("\n\n") + self.frames = [] + + def analyse_model(self): + # extract the fully qualified module names, to be able to report at run time. e.g.: + # encoder.block.2.layer.0.SelfAttention.o + # + # for shared weights only the first shared module name will be registered + self.module_names = {m: name for name, m in self.model.named_modules()} + # self.longest_module_name = max(len(v) for v in self.module_names.values()) + + def analyse_variable(self, var, ctx): + if torch.is_tensor(var): + self.expand_frame(get_abs_min_max(var, ctx)) + if detect_overflow(var, ctx): + self.detected_overflow = True + elif var is None: + self.expand_frame(f"{'None':>17} {ctx}") + else: + self.expand_frame(f"{'not a tensor':>17} {ctx}") + + def batch_start_frame(self): + self.expand_frame(f"\n\n{self.prefix} *** Starting batch number={self.batch_number} ***") + self.expand_frame(f"{'abs min':8} {'abs max':8} metadata") + + def batch_end_frame(self): + self.expand_frame(f"{self.prefix} *** Finished batch number={self.batch_number - 1} ***\n\n") + + def create_frame(self, module, input, output): + self.expand_frame(f"{self.prefix} {self.module_names[module]} {module.__class__.__name__}") + + # params + for name, p in module.named_parameters(recurse=False): + self.analyse_variable(p, name) + + # inputs + if isinstance(input, tuple): + for i, x in enumerate(input): + self.analyse_variable(x, f"input[{i}]") + else: + self.analyse_variable(input, "input") + + # outputs + if isinstance(output, tuple): + for i, x in enumerate(output): + # possibly a tuple of tuples + if isinstance(x, tuple): + for j, y in enumerate(x): + self.analyse_variable(y, f"output[{i}][{j}]") + else: + self.analyse_variable(x, f"output[{i}]") + else: + self.analyse_variable(output, "output") + + self.save_frame() + + def register_forward_hook(self): + self.model.apply(self._register_forward_hook) + + def _register_forward_hook(self, module): + module.register_forward_hook(self.forward_hook) + + def forward_hook(self, module, input, output): + # - input is a tuple of packed inputs (could be non-Tensors) + # - output could be a Tensor or a tuple of Tensors and non-Tensors + + last_frame_of_batch = False + + trace_mode = True if self.batch_number in self.trace_batch_nums else False + if trace_mode: + self.reset_saved_frames() + + if self.total_calls == 0: + self.batch_start_frame() + self.total_calls += 1 + + # count batch numbers - the very first forward hook of the batch will be called when the + # batch completes - i.e. it gets called very last - we know this batch has finished + if module == self.model: + self.batch_number += 1 + last_frame_of_batch = True + + self.create_frame(module, input, output) + + # if last_frame_of_batch: + # self.batch_end_frame() + + if trace_mode: + self.trace_frames() + + if last_frame_of_batch: + self.batch_start_frame() + + if self.detected_overflow and not trace_mode: + self.dump_saved_frames() + + # now we can abort, as it's pointless to continue running + raise ValueError( + "DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. " + "Please scroll up above this traceback to see the activation values prior to this event." + ) + + # abort after certain batch if requested to do so + if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num: + raise ValueError( + f"DebugUnderflowOverflow: aborting after {self.batch_number} batches due to" + f" `abort_after_batch_num={self.abort_after_batch_num}` arg" + ) + + +def get_abs_min_max(var, ctx): + abs_var = var.abs() + return f"{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}" + + +def detect_overflow(var, ctx): + """ + Report whether the tensor contains any `nan` or `inf` entries. + + This is useful for detecting overflows/underflows and best to call right after the function that did some math that + modified the tensor in question. + + This function contains a few other helper features that you can enable and tweak directly if you want to track + various other things. + + Args: + var: the tensor variable to check + ctx: the message to print as a context + + Return: + `True` if `inf` or `nan` was detected, `False` otherwise + """ + detected = False + if torch.isnan(var).any().item(): + detected = True + print(f"{ctx} has nans") + if torch.isinf(var).any().item(): + detected = True + print(f"{ctx} has infs") + + # if needed to monitor large elements can enable the following + if 0: # and detected: + n100 = var[torch.ge(var.abs(), 100)] + if n100.numel() > 0: + print(f"{ctx}: n100={n100.numel()}") + n1000 = var[torch.ge(var.abs(), 1000)] + if n1000.numel() > 0: + print(f"{ctx}: n1000={n1000.numel()}") + n10000 = var[torch.ge(var.abs(), 10000)] + if n10000.numel() > 0: + print(f"{ctx}: n10000={n10000.numel()}") + + if 0: + print(f"min={var.min():9.2e} max={var.max():9.2e}") + + if 0: + print(f"min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})") + + return detected + + +class DebugOption(ExplicitEnum): + UNDERFLOW_OVERFLOW = "underflow_overflow" + TPU_METRICS_DEBUG = "tpu_metrics_debug" diff --git a/lib/python3.12/site-packages/transformers/dependency_versions_check.py b/lib/python3.12/site-packages/transformers/dependency_versions_check.py new file mode 100644 index 0000000000000000000000000000000000000000..82d07850847ec357f36ff51088ddec36aceff093 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/dependency_versions_check.py @@ -0,0 +1,63 @@ +# Copyright 2020 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. + +from .dependency_versions_table import deps +from .utils.versions import require_version, require_version_core + + +# define which module versions we always want to check at run time +# (usually the ones defined in `install_requires` in setup.py) +# +# order specific notes: +# - tqdm must be checked before tokenizers + +pkgs_to_check_at_runtime = [ + "python", + "tqdm", + "regex", + "requests", + "packaging", + "filelock", + "numpy", + "tokenizers", + "huggingface-hub", + "safetensors", + "accelerate", + "pyyaml", +] + +for pkg in pkgs_to_check_at_runtime: + if pkg in deps: + if pkg == "tokenizers": + # must be loaded here, or else tqdm check may fail + from .utils import is_tokenizers_available + + if not is_tokenizers_available(): + continue # not required, check version only if installed + elif pkg == "accelerate": + # must be loaded here, or else tqdm check may fail + from .utils import is_accelerate_available + + # Maybe switch to is_torch_available in the future here so that Accelerate is hard dep of + # Transformers with PyTorch + if not is_accelerate_available(): + continue # not required, check version only if installed + + require_version_core(deps[pkg]) + else: + raise ValueError(f"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py") + + +def dep_version_check(pkg, hint=None): + require_version(deps[pkg], hint) diff --git a/lib/python3.12/site-packages/transformers/dependency_versions_table.py b/lib/python3.12/site-packages/transformers/dependency_versions_table.py new file mode 100644 index 0000000000000000000000000000000000000000..5c90fd0e595036d8d7f3d35c9e0ddf967396dee7 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/dependency_versions_table.py @@ -0,0 +1,106 @@ +# THIS FILE HAS BEEN AUTOGENERATED. To update: +# 1. modify the `_deps` dict in setup.py +# 2. run `make deps_table_update`` +deps = { + "Pillow": "Pillow>=10.0.1,<=15.0", + "accelerate": "accelerate>=0.26.0", + "av": "av", + "beautifulsoup4": "beautifulsoup4", + "blobfile": "blobfile", + "codecarbon": "codecarbon>=2.8.1", + "cookiecutter": "cookiecutter==1.7.3", + "dataclasses": "dataclasses", + "datasets": "datasets!=2.5.0", + "deepspeed": "deepspeed>=0.9.3", + "diffusers": "diffusers", + "dill": "dill<0.3.5", + "evaluate": "evaluate>=0.2.0", + "faiss-cpu": "faiss-cpu", + "fastapi": "fastapi", + "filelock": "filelock", + "flax": "flax>=0.4.1,<=0.7.0", + "fsspec": "fsspec<2023.10.0", + "ftfy": "ftfy", + "fugashi": "fugashi>=1.0", + "GitPython": "GitPython<3.1.19", + "hf-doc-builder": "hf-doc-builder>=0.3.0", + "hf_xet": "hf_xet", + "huggingface-hub": "huggingface-hub>=0.30.0,<1.0", + "importlib_metadata": "importlib_metadata", + "ipadic": "ipadic>=1.0.0,<2.0", + "isort": "isort>=5.5.4", + "jax": "jax>=0.4.1,<=0.4.13", + "jaxlib": "jaxlib>=0.4.1,<=0.4.13", + "jieba": "jieba", + "jinja2": "jinja2>=3.1.0", + "keras": "keras>2.9,<2.16", + "keras-nlp": "keras-nlp>=0.3.1,<0.14.0", + "kernels": "kernels>=0.3.2,<0.4", + "librosa": "librosa", + "natten": "natten>=0.14.6,<0.15.0", + "nltk": "nltk<=3.8.1", + "num2words": "num2words", + "numpy": "numpy>=1.17", + "onnxconverter-common": "onnxconverter-common", + "onnxruntime-tools": "onnxruntime-tools>=1.4.2", + "onnxruntime": "onnxruntime>=1.4.0", + "opencv-python": "opencv-python", + "optimum-benchmark": "optimum-benchmark>=0.3.0", + "optuna": "optuna", + "optax": "optax>=0.0.8,<=0.1.4", + "packaging": "packaging>=20.0", + "parameterized": "parameterized", + "phonemizer": "phonemizer", + "protobuf": "protobuf", + "psutil": "psutil", + "pyyaml": "pyyaml>=5.1", + "pydantic": "pydantic", + "pytest": "pytest>=7.2.0,<8.0.0", + "pytest-asyncio": "pytest-asyncio", + "pytest-rerunfailures": "pytest-rerunfailures", + "pytest-timeout": "pytest-timeout", + "pytest-xdist": "pytest-xdist", + "pytest-order": "pytest-order", + "python": "python>=3.9.0", + "ray[tune]": "ray[tune]>=2.7.0", + "regex": "regex!=2019.12.17", + "requests": "requests", + "rhoknp": "rhoknp>=1.1.0,<1.3.1", + "rjieba": "rjieba", + "rouge-score": "rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1", + "ruff": "ruff==0.11.2", + "sacrebleu": "sacrebleu>=1.4.12,<2.0.0", + "sacremoses": "sacremoses", + "safetensors": "safetensors>=0.4.3", + "sagemaker": "sagemaker>=2.31.0", + "schedulefree": "schedulefree>=1.2.6", + "scikit-learn": "scikit-learn", + "scipy": "scipy<1.13.0", + "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92", + "sigopt": "sigopt", + "starlette": "starlette", + "sudachipy": "sudachipy>=0.6.6", + "sudachidict_core": "sudachidict_core>=20220729", + "tensorboard": "tensorboard", + "tensorflow-cpu": "tensorflow-cpu>2.9,<2.16", + "tensorflow": "tensorflow>2.9,<2.16", + "tensorflow-text": "tensorflow-text<2.16", + "tensorflow-probability": "tensorflow-probability<0.24", + "tf2onnx": "tf2onnx", + "timeout-decorator": "timeout-decorator", + "tiktoken": "tiktoken", + "timm": "timm<=1.0.11", + "tokenizers": "tokenizers>=0.21,<0.22", + "torch": "torch>=2.0", + "torchaudio": "torchaudio", + "torchvision": "torchvision", + "pyctcdecode": "pyctcdecode>=0.4.0", + "tqdm": "tqdm>=4.27", + "unidic": "unidic>=1.0.2", + "unidic_lite": "unidic_lite>=1.0.7", + "urllib3": "urllib3<2.0.0", + "uvicorn": "uvicorn", + "pytest-rich": "pytest-rich", + "libcst": "libcst", + "rich": "rich", +} diff --git a/lib/python3.12/site-packages/transformers/dynamic_module_utils.py b/lib/python3.12/site-packages/transformers/dynamic_module_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9a72ab52e8fc7cad38ecaf6c719a325b3625093b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/dynamic_module_utils.py @@ -0,0 +1,702 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# 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. +"""Utilities to dynamically load objects from the Hub.""" + +import ast +import filecmp +import hashlib +import importlib +import importlib.util +import os +import re +import shutil +import signal +import sys +import threading +import warnings +from pathlib import Path +from types import ModuleType +from typing import Any, Optional, Union + +from huggingface_hub import try_to_load_from_cache + +from .utils import ( + HF_MODULES_CACHE, + TRANSFORMERS_DYNAMIC_MODULE_NAME, + cached_file, + extract_commit_hash, + is_offline_mode, + logging, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name +_HF_REMOTE_CODE_LOCK = threading.Lock() + + +def init_hf_modules(): + """ + Creates the cache directory for modules with an init, and adds it to the Python path. + """ + # This function has already been executed if HF_MODULES_CACHE already is in the Python path. + if HF_MODULES_CACHE in sys.path: + return + + sys.path.append(HF_MODULES_CACHE) + os.makedirs(HF_MODULES_CACHE, exist_ok=True) + init_path = Path(HF_MODULES_CACHE) / "__init__.py" + if not init_path.exists(): + init_path.touch() + importlib.invalidate_caches() + + +def create_dynamic_module(name: Union[str, os.PathLike]) -> None: + """ + Creates a dynamic module in the cache directory for modules. + + Args: + name (`str` or `os.PathLike`): + The name of the dynamic module to create. + """ + init_hf_modules() + dynamic_module_path = (Path(HF_MODULES_CACHE) / name).resolve() + # If the parent module does not exist yet, recursively create it. + if not dynamic_module_path.parent.exists(): + create_dynamic_module(dynamic_module_path.parent) + os.makedirs(dynamic_module_path, exist_ok=True) + init_path = dynamic_module_path / "__init__.py" + if not init_path.exists(): + init_path.touch() + # It is extremely important to invalidate the cache when we change stuff in those modules, or users end up + # with errors about module that do not exist. Same for all other `invalidate_caches` in this file. + importlib.invalidate_caches() + + +def get_relative_imports(module_file: Union[str, os.PathLike]) -> list[str]: + """ + Get the list of modules that are relatively imported in a module file. + + Args: + module_file (`str` or `os.PathLike`): The module file to inspect. + + Returns: + `list[str]`: The list of relative imports in the module. + """ + with open(module_file, encoding="utf-8") as f: + content = f.read() + + # Imports of the form `import .xxx` + relative_imports = re.findall(r"^\s*import\s+\.(\S+)\s*$", content, flags=re.MULTILINE) + # Imports of the form `from .xxx import yyy` + relative_imports += re.findall(r"^\s*from\s+\.(\S+)\s+import", content, flags=re.MULTILINE) + # Unique-ify + return list(set(relative_imports)) + + +def get_relative_import_files(module_file: Union[str, os.PathLike]) -> list[str]: + """ + Get the list of all files that are needed for a given module. Note that this function recurses through the relative + imports (if a imports b and b imports c, it will return module files for b and c). + + Args: + module_file (`str` or `os.PathLike`): The module file to inspect. + + Returns: + `list[str]`: The list of all relative imports a given module needs (recursively), which will give us the list + of module files a given module needs. + """ + no_change = False + files_to_check = [module_file] + all_relative_imports = [] + + # Let's recurse through all relative imports + while not no_change: + new_imports = [] + for f in files_to_check: + new_imports.extend(get_relative_imports(f)) + + module_path = Path(module_file).parent + new_import_files = [str(module_path / m) for m in new_imports] + new_import_files = [f for f in new_import_files if f not in all_relative_imports] + files_to_check = [f"{f}.py" for f in new_import_files] + + no_change = len(new_import_files) == 0 + all_relative_imports.extend(files_to_check) + + return all_relative_imports + + +def get_imports(filename: Union[str, os.PathLike]) -> list[str]: + """ + Extracts all the libraries (not relative imports this time) that are imported in a file. + + Args: + filename (`str` or `os.PathLike`): The module file to inspect. + + Returns: + `list[str]`: The list of all packages required to use the input module. + """ + with open(filename, encoding="utf-8") as f: + content = f.read() + imported_modules = set() + + def recursive_look_for_imports(node): + if isinstance(node, ast.Try): + return # Don't recurse into Try blocks and ignore imports in them + elif isinstance(node, ast.If): + test = node.test + for condition_node in ast.walk(test): + if isinstance(condition_node, ast.Call) and getattr(condition_node.func, "id", "").startswith( + "is_flash_attn" + ): + # Don't recurse into "if flash_attn_available()" blocks and ignore imports in them + return + elif isinstance(node, ast.Import): + # Handle 'import x' statements + for alias in node.names: + top_module = alias.name.split(".")[0] + if top_module: + imported_modules.add(top_module) + elif isinstance(node, ast.ImportFrom): + # Handle 'from x import y' statements, ignoring relative imports + if node.level == 0 and node.module: + top_module = node.module.split(".")[0] + if top_module: + imported_modules.add(top_module) + + # Recursively visit all children + for child in ast.iter_child_nodes(node): + recursive_look_for_imports(child) + + tree = ast.parse(content) + recursive_look_for_imports(tree) + + return sorted(imported_modules) + + +def check_imports(filename: Union[str, os.PathLike]) -> list[str]: + """ + Check if the current Python environment contains all the libraries that are imported in a file. Will raise if a + library is missing. + + Args: + filename (`str` or `os.PathLike`): The module file to check. + + Returns: + `list[str]`: The list of relative imports in the file. + """ + imports = get_imports(filename) + missing_packages = [] + for imp in imports: + try: + importlib.import_module(imp) + except ImportError as exception: + logger.warning(f"Encountered exception while importing {imp}: {exception}") + # Some packages can fail with an ImportError because of a dependency issue. + # This check avoids hiding such errors. + # See https://github.com/huggingface/transformers/issues/33604 + if "No module named" in str(exception): + missing_packages.append(imp) + else: + raise + + if len(missing_packages) > 0: + raise ImportError( + "This modeling file requires the following packages that were not found in your environment: " + f"{', '.join(missing_packages)}. Run `pip install {' '.join(missing_packages)}`" + ) + + return get_relative_imports(filename) + + +def get_class_in_module( + class_name: str, + module_path: Union[str, os.PathLike], + *, + force_reload: bool = False, +) -> type: + """ + Import a module on the cache directory for modules and extract a class from it. + + Args: + class_name (`str`): The name of the class to import. + module_path (`str` or `os.PathLike`): The path to the module to import. + force_reload (`bool`, *optional*, defaults to `False`): + Whether to reload the dynamic module from file if it already exists in `sys.modules`. + Otherwise, the module is only reloaded if the file has changed. + + Returns: + `typing.Type`: The class looked for. + """ + name = os.path.normpath(module_path) + if name.endswith(".py"): + name = name[:-3] + name = name.replace(os.path.sep, ".") + module_file: Path = Path(HF_MODULES_CACHE) / module_path + with _HF_REMOTE_CODE_LOCK: + if force_reload: + sys.modules.pop(name, None) + importlib.invalidate_caches() + cached_module: Optional[ModuleType] = sys.modules.get(name) + module_spec = importlib.util.spec_from_file_location(name, location=module_file) + + # Hash the module file and all its relative imports to check if we need to reload it + module_files: list[Path] = [module_file] + sorted(map(Path, get_relative_import_files(module_file))) + module_hash: str = hashlib.sha256(b"".join(bytes(f) + f.read_bytes() for f in module_files)).hexdigest() + + module: ModuleType + if cached_module is None: + module = importlib.util.module_from_spec(module_spec) + # insert it into sys.modules before any loading begins + sys.modules[name] = module + else: + module = cached_module + # reload in both cases, unless the module is already imported and the hash hits + if getattr(module, "__transformers_module_hash__", "") != module_hash: + module_spec.loader.exec_module(module) + module.__transformers_module_hash__ = module_hash + return getattr(module, class_name) + + +def get_cached_module_file( + pretrained_model_name_or_path: Union[str, os.PathLike], + module_file: str, + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + resume_download: Optional[bool] = None, + proxies: Optional[dict[str, str]] = None, + token: Optional[Union[bool, str]] = None, + revision: Optional[str] = None, + local_files_only: bool = False, + repo_type: Optional[str] = None, + _commit_hash: Optional[str] = None, + **deprecated_kwargs, +) -> str: + """ + Prepares Downloads a module from a local folder or a distant repo and returns its path inside the cached + Transformers module. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + This can be either: + + - a string, the *model id* of a pretrained model configuration hosted inside a model repo on + huggingface.co. + - a path to a *directory* containing a configuration file saved using the + [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`. + + module_file (`str`): + The name of the module file containing the class to look for. + 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. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the configuration files and override the cached versions if they + exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, will only try to load the tokenizer configuration from local files. + repo_type (`str`, *optional*): + Specify the repo type (useful when downloading from a space for instance). + + + + Passing `token=True` is required when you want to use a private model. + + + + Returns: + `str`: The path to the module inside the cache. + """ + use_auth_token = deprecated_kwargs.pop("use_auth_token", None) + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.") + token = use_auth_token + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + # Download and cache module_file from the repo `pretrained_model_name_or_path` of grab it if it's a local file. + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + is_local = os.path.isdir(pretrained_model_name_or_path) + if is_local: + submodule = os.path.basename(pretrained_model_name_or_path) + else: + submodule = pretrained_model_name_or_path.replace("/", os.path.sep) + cached_module = try_to_load_from_cache( + pretrained_model_name_or_path, module_file, cache_dir=cache_dir, revision=_commit_hash, repo_type=repo_type + ) + + new_files = [] + try: + # Load from URL or cache if already cached + resolved_module_file = cached_file( + pretrained_model_name_or_path, + module_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + revision=revision, + repo_type=repo_type, + _commit_hash=_commit_hash, + ) + if not is_local and cached_module != resolved_module_file: + new_files.append(module_file) + + except OSError: + logger.error(f"Could not locate the {module_file} inside {pretrained_model_name_or_path}.") + raise + + # Check we have all the requirements in our environment + modules_needed = check_imports(resolved_module_file) + + # Now we move the module inside our cached dynamic modules. + full_submodule = TRANSFORMERS_DYNAMIC_MODULE_NAME + os.path.sep + submodule + create_dynamic_module(full_submodule) + submodule_path = Path(HF_MODULES_CACHE) / full_submodule + if submodule == os.path.basename(pretrained_model_name_or_path): + # We copy local files to avoid putting too many folders in sys.path. This copy is done when the file is new or + # has changed since last copy. + if not (submodule_path / module_file).exists() or not filecmp.cmp( + resolved_module_file, str(submodule_path / module_file) + ): + shutil.copy(resolved_module_file, submodule_path / module_file) + importlib.invalidate_caches() + for module_needed in modules_needed: + module_needed = f"{module_needed}.py" + module_needed_file = os.path.join(pretrained_model_name_or_path, module_needed) + if not (submodule_path / module_needed).exists() or not filecmp.cmp( + module_needed_file, str(submodule_path / module_needed) + ): + shutil.copy(module_needed_file, submodule_path / module_needed) + importlib.invalidate_caches() + else: + # Get the commit hash + commit_hash = extract_commit_hash(resolved_module_file, _commit_hash) + + # The module file will end up being placed in a subfolder with the git hash of the repo. This way we get the + # benefit of versioning. + submodule_path = submodule_path / commit_hash + full_submodule = full_submodule + os.path.sep + commit_hash + create_dynamic_module(full_submodule) + + if not (submodule_path / module_file).exists(): + shutil.copy(resolved_module_file, submodule_path / module_file) + importlib.invalidate_caches() + # Make sure we also have every file with relative + for module_needed in modules_needed: + if not (submodule_path / f"{module_needed}.py").exists(): + get_cached_module_file( + pretrained_model_name_or_path, + f"{module_needed}.py", + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + token=token, + revision=revision, + local_files_only=local_files_only, + _commit_hash=commit_hash, + ) + new_files.append(f"{module_needed}.py") + + if len(new_files) > 0 and revision is None: + new_files = "\n".join([f"- {f}" for f in new_files]) + repo_type_str = "" if repo_type is None else f"{repo_type}s/" + url = f"https://huggingface.co/{repo_type_str}{pretrained_model_name_or_path}" + logger.warning( + f"A new version of the following files was downloaded from {url}:\n{new_files}" + "\n. Make sure to double-check they do not contain any added malicious code. To avoid downloading new " + "versions of the code file, you can pin a revision." + ) + + return os.path.join(full_submodule, module_file) + + +def get_class_from_dynamic_module( + class_reference: str, + pretrained_model_name_or_path: Union[str, os.PathLike], + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + resume_download: Optional[bool] = None, + proxies: Optional[dict[str, str]] = None, + token: Optional[Union[bool, str]] = None, + revision: Optional[str] = None, + local_files_only: bool = False, + repo_type: Optional[str] = None, + code_revision: Optional[str] = None, + **kwargs, +) -> type: + """ + Extracts a class from a module file, present in the local folder or repository of a model. + + + + Calling this function will execute the code in the module file found locally or downloaded from the Hub. It should + therefore only be called on trusted repos. + + + + + + Args: + class_reference (`str`): + The full name of the class to load, including its module and optionally its repo. + pretrained_model_name_or_path (`str` or `os.PathLike`): + This can be either: + + - a string, the *model id* of a pretrained model configuration hosted inside a model repo on + huggingface.co. + - a path to a *directory* containing a configuration file saved using the + [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`. + + This is used when `class_reference` does not specify another repo. + module_file (`str`): + The name of the module file containing the class to look for. + class_name (`str`): + The name of the class to import in the module. + 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. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the configuration files and override the cached versions if they + exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, will only try to load the tokenizer configuration from local files. + repo_type (`str`, *optional*): + Specify the repo type (useful when downloading from a space for instance). + code_revision (`str`, *optional*, defaults to `"main"`): + The specific revision to use for the code on the Hub, if the code leaves in a different repository than the + rest of the model. It can be a branch name, a tag name, or a commit id, since we use a git-based system for + storing models and other artifacts on huggingface.co, so `revision` can be any identifier allowed by git. + + + + Passing `token=True` is required when you want to use a private model. + + + + Returns: + `typing.Type`: The class, dynamically imported from the module. + + Examples: + + ```python + # Download module `modeling.py` from huggingface.co and cache then extract the class `MyBertModel` from this + # module. + cls = get_class_from_dynamic_module("modeling.MyBertModel", "sgugger/my-bert-model") + + # Download module `modeling.py` from a given repo and cache then extract the class `MyBertModel` from this + # module. + cls = get_class_from_dynamic_module("sgugger/my-bert-model--modeling.MyBertModel", "sgugger/another-bert-model") + ```""" + use_auth_token = kwargs.pop("use_auth_token", None) + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.") + token = use_auth_token + + # Catch the name of the repo if it's specified in `class_reference` + if "--" in class_reference: + repo_id, class_reference = class_reference.split("--") + else: + repo_id = pretrained_model_name_or_path + module_file, class_name = class_reference.split(".") + + if code_revision is None and pretrained_model_name_or_path == repo_id: + code_revision = revision + # And lastly we get the class inside our newly created module + final_module = get_cached_module_file( + repo_id, + module_file + ".py", + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + token=token, + revision=code_revision, + local_files_only=local_files_only, + repo_type=repo_type, + ) + return get_class_in_module(class_name, final_module, force_reload=force_download) + + +def custom_object_save(obj: Any, folder: Union[str, os.PathLike], config: Optional[dict] = None) -> list[str]: + """ + Save the modeling files corresponding to a custom model/configuration/tokenizer etc. in a given folder. Optionally + adds the proper fields in a config. + + Args: + obj (`Any`): The object for which to save the module files. + folder (`str` or `os.PathLike`): The folder where to save. + config (`PretrainedConfig` or dictionary, `optional`): + A config in which to register the auto_map corresponding to this custom object. + + Returns: + `List[str]`: The list of files saved. + """ + if obj.__module__ == "__main__": + logger.warning( + f"We can't save the code defining {obj} in {folder} as it's been defined in __main__. You should put " + "this code in a separate module so we can include it in the saved folder and make it easier to share via " + "the Hub." + ) + return + + def _set_auto_map_in_config(_config): + module_name = obj.__class__.__module__ + last_module = module_name.split(".")[-1] + full_name = f"{last_module}.{obj.__class__.__name__}" + # Special handling for tokenizers + if "Tokenizer" in full_name: + slow_tokenizer_class = None + fast_tokenizer_class = None + if obj.__class__.__name__.endswith("Fast"): + # Fast tokenizer: we have the fast tokenizer class and we may have the slow one has an attribute. + fast_tokenizer_class = f"{last_module}.{obj.__class__.__name__}" + if getattr(obj, "slow_tokenizer_class", None) is not None: + slow_tokenizer = getattr(obj, "slow_tokenizer_class") + slow_tok_module_name = slow_tokenizer.__module__ + last_slow_tok_module = slow_tok_module_name.split(".")[-1] + slow_tokenizer_class = f"{last_slow_tok_module}.{slow_tokenizer.__name__}" + else: + # Slow tokenizer: no way to have the fast class + slow_tokenizer_class = f"{last_module}.{obj.__class__.__name__}" + + full_name = (slow_tokenizer_class, fast_tokenizer_class) + + if isinstance(_config, dict): + auto_map = _config.get("auto_map", {}) + auto_map[obj._auto_class] = full_name + _config["auto_map"] = auto_map + elif getattr(_config, "auto_map", None) is not None: + _config.auto_map[obj._auto_class] = full_name + else: + _config.auto_map = {obj._auto_class: full_name} + + # Add object class to the config auto_map + if isinstance(config, (list, tuple)): + for cfg in config: + _set_auto_map_in_config(cfg) + elif config is not None: + _set_auto_map_in_config(config) + + result = [] + # Copy module file to the output folder. + object_file = sys.modules[obj.__module__].__file__ + dest_file = Path(folder) / (Path(object_file).name) + shutil.copy(object_file, dest_file) + result.append(dest_file) + + # Gather all relative imports recursively and make sure they are copied as well. + for needed_file in get_relative_import_files(object_file): + dest_file = Path(folder) / (Path(needed_file).name) + shutil.copy(needed_file, dest_file) + result.append(dest_file) + + return result + + +def _raise_timeout_error(signum, frame): + raise ValueError( + "Loading this model requires you to execute custom code contained in the model repository on your local " + "machine. Please set the option `trust_remote_code=True` to permit loading of this model." + ) + + +TIME_OUT_REMOTE_CODE = 15 + + +def resolve_trust_remote_code(trust_remote_code, model_name, has_local_code, has_remote_code): + if trust_remote_code is None: + if has_local_code: + trust_remote_code = False + elif has_remote_code and TIME_OUT_REMOTE_CODE > 0: + prev_sig_handler = None + try: + prev_sig_handler = signal.signal(signal.SIGALRM, _raise_timeout_error) + signal.alarm(TIME_OUT_REMOTE_CODE) + while trust_remote_code is None: + answer = input( + f"The repository for {model_name} contains custom code which must be executed to correctly " + f"load the model. You can inspect the repository content at https://hf.co/{model_name}.\n" + f"You can avoid this prompt in future by passing the argument `trust_remote_code=True`.\n\n" + f"Do you wish to run the custom code? [y/N] " + ) + if answer.lower() in ["yes", "y", "1"]: + trust_remote_code = True + elif answer.lower() in ["no", "n", "0", ""]: + trust_remote_code = False + signal.alarm(0) + except Exception: + # OS which does not support signal.SIGALRM + raise ValueError( + f"The repository for {model_name} contains custom code which must be executed to correctly " + f"load the model. You can inspect the repository content at https://hf.co/{model_name}.\n" + f"Please pass the argument `trust_remote_code=True` to allow custom code to be run." + ) + finally: + if prev_sig_handler is not None: + signal.signal(signal.SIGALRM, prev_sig_handler) + signal.alarm(0) + elif has_remote_code: + # For the CI which puts the timeout at 0 + _raise_timeout_error(None, None) + + if has_remote_code and not has_local_code and not trust_remote_code: + raise ValueError( + f"Loading {model_name} requires you to execute the configuration file in that" + " repo on your local machine. Make sure you have read the code there to avoid malicious use, then" + " set the option `trust_remote_code=True` to remove this error." + ) + + return trust_remote_code diff --git a/lib/python3.12/site-packages/transformers/feature_extraction_sequence_utils.py b/lib/python3.12/site-packages/transformers/feature_extraction_sequence_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c9a26bac9b3dcd9bb14d855f34494a38df3f7f71 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/feature_extraction_sequence_utils.py @@ -0,0 +1,371 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# 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. +""" +Sequence feature extraction class for common feature extractors to preprocess sequences. +""" + +from typing import Optional, Union + +import numpy as np + +from .feature_extraction_utils import BatchFeature, FeatureExtractionMixin +from .utils import PaddingStrategy, TensorType, is_tf_tensor, is_torch_tensor, logging, to_numpy + + +logger = logging.get_logger(__name__) + + +class SequenceFeatureExtractor(FeatureExtractionMixin): + """ + This is a general feature extraction class for speech recognition. + + Args: + feature_size (`int`): + The feature dimension of the extracted features. + sampling_rate (`int`): + The sampling rate at which the audio files should be digitalized expressed in hertz (Hz). + padding_value (`float`): + The value that is used to fill the padding values / vectors. + """ + + def __init__(self, feature_size: int, sampling_rate: int, padding_value: float, **kwargs): + self.feature_size = feature_size + self.sampling_rate = sampling_rate + self.padding_value = padding_value + + self.padding_side = kwargs.pop("padding_side", "right") + self.return_attention_mask = kwargs.pop("return_attention_mask", True) + + super().__init__(**kwargs) + + def pad( + self, + processed_features: Union[ + BatchFeature, + list[BatchFeature], + dict[str, BatchFeature], + dict[str, list[BatchFeature]], + list[dict[str, BatchFeature]], + ], + padding: Union[bool, str, PaddingStrategy] = True, + max_length: Optional[int] = None, + truncation: bool = False, + pad_to_multiple_of: Optional[int] = None, + return_attention_mask: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + ) -> BatchFeature: + """ + Pad input values / input vectors or a batch of input values / input vectors up to predefined length or to the + max sequence length in the batch. + + Padding side (left/right) padding values are defined at the feature extractor level (with `self.padding_side`, + `self.padding_value`) + + + + If the `processed_features` passed are dictionary of numpy arrays, PyTorch tensors or TensorFlow tensors, the + result will use the same type unless you provide a different tensor type with `return_tensors`. In the case of + PyTorch tensors, you will lose the specific device of your tensors however. + + + + Args: + processed_features ([`BatchFeature`], list of [`BatchFeature`], `Dict[str, List[float]]`, `Dict[str, List[List[float]]` or `List[Dict[str, List[float]]]`): + Processed inputs. Can represent one input ([`BatchFeature`] or `Dict[str, List[float]]`) or a batch of + input values / vectors (list of [`BatchFeature`], *Dict[str, List[List[float]]]* or *List[Dict[str, + List[float]]]*) so you can use this method during preprocessing as well as in a PyTorch Dataloader + collate function. + + Instead of `List[float]` you can have tensors (numpy arrays, PyTorch tensors or TensorFlow tensors), + see the note above for the return type. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + max_length (`int`, *optional*): + Maximum length of the returned list and optionally padding length (see above). + truncation (`bool`): + Activates truncation to cut input sequences longer than `max_length` to `max_length`. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific feature_extractor's default. + + [What are attention masks?](../glossary#attention-mask) + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + """ + # If we have a list of dicts, let's convert it in a dict of lists + # We do this to allow using this method as a collate_fn function in PyTorch Dataloader + if isinstance(processed_features, (list, tuple)) and isinstance(processed_features[0], (dict, BatchFeature)): + processed_features = { + key: [example[key] for example in processed_features] for key in processed_features[0].keys() + } + + # The model's main input name, usually `input_values`, has be passed for padding + if self.model_input_names[0] not in processed_features: + raise ValueError( + "You should supply an instance of `transformers.BatchFeature` or list of `transformers.BatchFeature`" + f" to this method that includes {self.model_input_names[0]}, but you provided" + f" {list(processed_features.keys())}" + ) + + required_input = processed_features[self.model_input_names[0]] + return_attention_mask = ( + return_attention_mask if return_attention_mask is not None else self.return_attention_mask + ) + + if len(required_input) == 0: + if return_attention_mask: + processed_features["attention_mask"] = [] + return processed_features + + # If we have PyTorch/TF tensors or lists as inputs, we cast them as Numpy arrays + # and rebuild them afterwards if no return_tensors is specified + # Note that we lose the specific device the tensor may be on for PyTorch + + first_element = required_input[0] + if isinstance(first_element, (list, tuple)): + # first_element might be an empty list/tuple in some edge cases so we grab the first non empty element. + index = 0 + while len(required_input[index]) == 0: + index += 1 + if index < len(required_input): + first_element = required_input[index][0] + + if return_tensors is None: + if is_tf_tensor(first_element): + return_tensors = "tf" + elif is_torch_tensor(first_element): + return_tensors = "pt" + elif isinstance(first_element, (int, float, list, tuple, np.ndarray)): + return_tensors = "np" + else: + raise ValueError( + f"type of {first_element} unknown: {type(first_element)}. " + "Should be one of a python, numpy, pytorch or tensorflow object." + ) + + for key, value in processed_features.items(): + if isinstance(value[0], (int, float)): + processed_features[key] = to_numpy(value) + else: + processed_features[key] = [to_numpy(v) for v in value] + + # Convert padding_strategy in PaddingStrategy + padding_strategy = self._get_padding_strategies(padding=padding, max_length=max_length) + + required_input = processed_features[self.model_input_names[0]] + + batch_size = len(required_input) + if not all(len(v) == batch_size for v in processed_features.values()): + raise ValueError("Some items in the output dictionary have a different batch size than others.") + + truncated_inputs = [] + for i in range(batch_size): + inputs = {k: v[i] for k, v in processed_features.items()} + # truncation + inputs_slice = self._truncate( + inputs, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + truncation=truncation, + ) + truncated_inputs.append(inputs_slice) + + if padding_strategy == PaddingStrategy.LONGEST: + # make sure that `max_length` cannot be longer than the longest truncated length + max_length = max(len(input_slice[self.model_input_names[0]]) for input_slice in truncated_inputs) + padding_strategy = PaddingStrategy.MAX_LENGTH + + batch_outputs = {} + for i in range(batch_size): + # padding + outputs = self._pad( + truncated_inputs[i], + max_length=max_length, + padding_strategy=padding_strategy, + pad_to_multiple_of=pad_to_multiple_of, + return_attention_mask=return_attention_mask, + ) + + for key, value in outputs.items(): + if key not in batch_outputs: + batch_outputs[key] = [] + if value.dtype is np.dtype(np.float64): + value = value.astype(np.float32) + batch_outputs[key].append(value) + + return BatchFeature(batch_outputs, tensor_type=return_tensors) + + def _pad( + self, + processed_features: Union[dict[str, np.ndarray], BatchFeature], + max_length: Optional[int] = None, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + pad_to_multiple_of: Optional[int] = None, + return_attention_mask: Optional[bool] = None, + ) -> dict: + """ + Pad inputs (on left/right and up to predefined length or max length in the batch) + + Args: + processed_features (`Union[Dict[str, np.ndarray], BatchFeature]`): + Dictionary of input values (`np.ndarray[float]`) / input vectors (`List[np.ndarray[float]]`) or batch + of inputs values (`List[np.ndarray[int]]`) / input vectors (`List[np.ndarray[int]]`) + max_length (`int`, *optional*): + Maximum length of the returned list and optionally padding length (see below) + padding_strategy (`PaddingStrategy`, *optional*, default to `PaddingStrategy.DO_NOT_PAD`): + PaddingStrategy to use for padding. + + - PaddingStrategy.LONGEST Pad to the longest sequence in the batch + - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) + - PaddingStrategy.DO_NOT_PAD: Do not pad + The feature_extractor padding sides are defined in self.padding_side: + + - 'left': pads on the left of the sequences + - 'right': pads on the right of the sequences + pad_to_multiple_of (`int`, *optional*): + Integer if set will pad the sequence to a multiple of the provided value. This is especially useful to + enable the use of Tensor Core on NVIDIA hardware with compute capability `>= 7.5` (Volta), or on TPUs + which benefit from having sequence lengths be a multiple of 128. + return_attention_mask (`bool`, *optional*): + Set to False to avoid returning attention mask (default: set to model specifics) + """ + required_input = processed_features[self.model_input_names[0]] + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = len(required_input) + + if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0): + max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of + + needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) < max_length + + if return_attention_mask and "attention_mask" not in processed_features: + processed_features["attention_mask"] = np.ones(len(required_input), dtype=np.int32) + + if needs_to_be_padded: + difference = max_length - len(required_input) + if self.padding_side == "right": + if return_attention_mask: + processed_features["attention_mask"] = np.pad( + processed_features["attention_mask"], (0, difference) + ) + padding_shape = ((0, difference), (0, 0)) if self.feature_size > 1 else (0, difference) + processed_features[self.model_input_names[0]] = np.pad( + required_input, padding_shape, "constant", constant_values=self.padding_value + ) + elif self.padding_side == "left": + if return_attention_mask: + processed_features["attention_mask"] = np.pad( + processed_features["attention_mask"], (difference, 0) + ) + padding_shape = ((difference, 0), (0, 0)) if self.feature_size > 1 else (difference, 0) + processed_features[self.model_input_names[0]] = np.pad( + required_input, padding_shape, "constant", constant_values=self.padding_value + ) + else: + raise ValueError("Invalid padding strategy:" + str(self.padding_side)) + + return processed_features + + def _truncate( + self, + processed_features: Union[dict[str, np.ndarray], BatchFeature], + max_length: Optional[int] = None, + pad_to_multiple_of: Optional[int] = None, + truncation: Optional[bool] = None, + ): + """ + Truncate inputs to predefined length or max length in the batch + + Args: + processed_features(`Union[Dict[str, np.ndarray], BatchFeature]`): + Dictionary of input values (`np.ndarray[float]`) / input vectors (`List[np.ndarray[float]]`) or batch + of inputs values (`List[np.ndarray[int]]`) / input vectors (`List[np.ndarray[int]]`) + max_length (`int`, *optional*): + maximum length of the returned list and optionally padding length (see below) + pad_to_multiple_of (`int`, *optional*) : + Integer if set will pad the sequence to a multiple of the provided value. This is especially useful to + enable the use of Tensor Core on NVIDIA hardware with compute capability `>= 7.5` (Volta), or on TPUs + which benefit from having sequence lengths be a multiple of 128. + truncation (`bool`, *optional*): + Activates truncation to cut input sequences longer than `max_length` to `max_length`. + """ + if not truncation: + return processed_features + elif truncation and max_length is None: + raise ValueError("When setting ``truncation=True``, make sure that ``max_length`` is defined.") + + required_input = processed_features[self.model_input_names[0]] + + # find `max_length` that fits `pad_to_multiple_of` + if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0): + max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of + + needs_to_be_truncated = len(required_input) > max_length + + if needs_to_be_truncated: + processed_features[self.model_input_names[0]] = processed_features[self.model_input_names[0]][:max_length] + if "attention_mask" in processed_features: + processed_features["attention_mask"] = processed_features["attention_mask"][:max_length] + + return processed_features + + def _get_padding_strategies(self, padding=False, max_length=None): + """ + Find the correct padding strategy + """ + + # Get padding strategy + if padding is not False: + if padding is True: + padding_strategy = PaddingStrategy.LONGEST # Default to pad to the longest sequence in the batch + elif not isinstance(padding, PaddingStrategy): + padding_strategy = PaddingStrategy(padding) + elif isinstance(padding, PaddingStrategy): + padding_strategy = padding + else: + padding_strategy = PaddingStrategy.DO_NOT_PAD + + # Set max length if needed + if max_length is None: + if padding_strategy == PaddingStrategy.MAX_LENGTH: + raise ValueError( + f"When setting ``padding={PaddingStrategy.MAX_LENGTH}``, make sure that max_length is defined" + ) + + # Test if we have a padding value + if padding_strategy != PaddingStrategy.DO_NOT_PAD and (self.padding_value is None): + raise ValueError( + "Asking to pad but the feature_extractor does not have a padding value. Please select a value to use" + " as `padding_value`. For example: `feature_extractor.padding_value = 0.0`." + ) + + return padding_strategy diff --git a/lib/python3.12/site-packages/transformers/feature_extraction_utils.py b/lib/python3.12/site-packages/transformers/feature_extraction_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ca2a3b5fde31d81554c76e26a24ebb4b806ed052 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/feature_extraction_utils.py @@ -0,0 +1,701 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# 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. +""" +Feature extraction saving/loading class for common feature extractors. +""" + +import copy +import json +import os +import warnings +from collections import UserDict +from typing import TYPE_CHECKING, Any, Optional, Union + +import numpy as np + +from .dynamic_module_utils import custom_object_save +from .utils import ( + FEATURE_EXTRACTOR_NAME, + PushToHubMixin, + TensorType, + add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, + cached_file, + copy_func, + download_url, + is_flax_available, + is_jax_tensor, + is_numpy_array, + is_offline_mode, + is_remote_url, + is_tf_available, + is_torch_available, + is_torch_device, + is_torch_dtype, + logging, + requires_backends, +) + + +if TYPE_CHECKING: + if is_torch_available(): + import torch # noqa + + +logger = logging.get_logger(__name__) + +PreTrainedFeatureExtractor = Union["SequenceFeatureExtractor"] # noqa: F821 + + +class BatchFeature(UserDict): + r""" + Holds the output of the [`~SequenceFeatureExtractor.pad`] and feature extractor specific `__call__` methods. + + This class is derived from a python dictionary and can be used as a dictionary. + + Args: + data (`dict`, *optional*): + Dictionary of lists/arrays/tensors returned by the __call__/pad methods ('input_values', 'attention_mask', + etc.). + tensor_type (`Union[None, str, TensorType]`, *optional*): + You can give a tensor_type here to convert the lists of integers in PyTorch/TensorFlow/Numpy Tensors at + initialization. + """ + + def __init__(self, data: Optional[dict[str, Any]] = None, tensor_type: Union[None, str, TensorType] = None): + super().__init__(data) + self.convert_to_tensors(tensor_type=tensor_type) + + def __getitem__(self, item: str) -> Union[Any]: + """ + If the key is a string, returns the value of the dict associated to `key` ('input_values', 'attention_mask', + etc.). + """ + if isinstance(item, str): + return self.data[item] + else: + raise KeyError("Indexing with integers is not available when using Python based feature extractors") + + def __getattr__(self, item: str): + try: + return self.data[item] + except KeyError: + raise AttributeError + + def __getstate__(self): + return {"data": self.data} + + def __setstate__(self, state): + if "data" in state: + self.data = state["data"] + + # Copied from transformers.tokenization_utils_base.BatchEncoding.keys + def keys(self): + return self.data.keys() + + # Copied from transformers.tokenization_utils_base.BatchEncoding.values + def values(self): + return self.data.values() + + # Copied from transformers.tokenization_utils_base.BatchEncoding.items + def items(self): + return self.data.items() + + def _get_is_as_tensor_fns(self, tensor_type: Optional[Union[str, TensorType]] = None): + if tensor_type is None: + return None, None + + # Convert to TensorType + if not isinstance(tensor_type, TensorType): + tensor_type = TensorType(tensor_type) + + # Get a function reference for the correct framework + if tensor_type == TensorType.TENSORFLOW: + if not is_tf_available(): + raise ImportError( + "Unable to convert output to TensorFlow tensors format, TensorFlow is not installed." + ) + import tensorflow as tf + + as_tensor = tf.constant + is_tensor = tf.is_tensor + elif tensor_type == TensorType.PYTORCH: + if not is_torch_available(): + raise ImportError("Unable to convert output to PyTorch tensors format, PyTorch is not installed.") + import torch # noqa + + def as_tensor(value): + if isinstance(value, (list, tuple)) and len(value) > 0: + if isinstance(value[0], np.ndarray): + value = np.array(value) + elif ( + isinstance(value[0], (list, tuple)) + and len(value[0]) > 0 + and isinstance(value[0][0], np.ndarray) + ): + value = np.array(value) + if isinstance(value, np.ndarray): + return torch.from_numpy(value) + else: + return torch.tensor(value) + + is_tensor = torch.is_tensor + elif tensor_type == TensorType.JAX: + if not is_flax_available(): + raise ImportError("Unable to convert output to JAX tensors format, JAX is not installed.") + import jax.numpy as jnp # noqa: F811 + + as_tensor = jnp.array + is_tensor = is_jax_tensor + else: + + def as_tensor(value, dtype=None): + if isinstance(value, (list, tuple)) and isinstance(value[0], (list, tuple, np.ndarray)): + value_lens = [len(val) for val in value] + if len(set(value_lens)) > 1 and dtype is None: + # we have a ragged list so handle explicitly + value = as_tensor([np.asarray(val) for val in value], dtype=object) + return np.asarray(value, dtype=dtype) + + is_tensor = is_numpy_array + return is_tensor, as_tensor + + def convert_to_tensors(self, tensor_type: Optional[Union[str, TensorType]] = None): + """ + Convert the inner content to tensors. + + Args: + tensor_type (`str` or [`~utils.TensorType`], *optional*): + The type of tensors to use. If `str`, should be one of the values of the enum [`~utils.TensorType`]. If + `None`, no modification is done. + """ + if tensor_type is None: + return self + + is_tensor, as_tensor = self._get_is_as_tensor_fns(tensor_type) + + # Do the tensor conversion in batch + for key, value in self.items(): + try: + if not is_tensor(value): + tensor = as_tensor(value) + + self[key] = tensor + except: # noqa E722 + if key == "overflowing_values": + raise ValueError("Unable to create tensor returning overflowing values of different lengths. ") + raise ValueError( + "Unable to create tensor, you should probably activate padding " + "with 'padding=True' to have batched tensors with the same length." + ) + + return self + + def to(self, *args, **kwargs) -> "BatchFeature": + """ + Send all values to device by calling `v.to(*args, **kwargs)` (PyTorch only). This should support casting in + different `dtypes` and sending the `BatchFeature` to a different `device`. + + Args: + args (`Tuple`): + Will be passed to the `to(...)` function of the tensors. + kwargs (`Dict`, *optional*): + Will be passed to the `to(...)` function of the tensors. + To enable asynchronous data transfer, set the `non_blocking` flag in `kwargs` (defaults to `False`). + + Returns: + [`BatchFeature`]: The same instance after modification. + """ + requires_backends(self, ["torch"]) + import torch # noqa + + new_data = {} + device = kwargs.get("device") + non_blocking = kwargs.get("non_blocking", False) + # Check if the args are a device or a dtype + if device is None and len(args) > 0: + # device should be always the first argument + arg = args[0] + if is_torch_dtype(arg): + # The first argument is a dtype + pass + elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int): + device = arg + else: + # it's something else + raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.") + # We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor` + for k, v in self.items(): + # check if v is a floating point + if isinstance(v, torch.Tensor) and torch.is_floating_point(v): + # cast and send to device + new_data[k] = v.to(*args, **kwargs) + elif isinstance(v, torch.Tensor) and device is not None: + new_data[k] = v.to(device=device, non_blocking=non_blocking) + else: + new_data[k] = v + self.data = new_data + return self + + +class FeatureExtractionMixin(PushToHubMixin): + """ + This is a feature extraction mixin used to provide saving/loading functionality for sequential and image feature + extractors. + """ + + _auto_class = None + + def __init__(self, **kwargs): + """Set elements of `kwargs` as attributes.""" + # Pop "processor_class" as it should be saved as private attribute + self._processor_class = kwargs.pop("processor_class", None) + # Additional attributes without default values + 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 _set_processor_class(self, processor_class: str): + """Sets processor class as an attribute.""" + self._processor_class = processor_class + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Union[str, os.PathLike], + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + **kwargs, + ): + r""" + Instantiate a type of [`~feature_extraction_utils.FeatureExtractionMixin`] from a feature extractor, *e.g.* a + derived class of [`SequenceFeatureExtractor`]. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + This can be either: + + - a string, the *model id* of a pretrained feature_extractor hosted inside a model repo on + huggingface.co. + - a path to a *directory* containing a feature extractor file saved using the + [`~feature_extraction_utils.FeatureExtractionMixin.save_pretrained`] method, e.g., + `./my_model_directory/`. + - a path or url to a saved feature extractor JSON *file*, e.g., + `./my_model_directory/preprocessor_config.json`. + cache_dir (`str` or `os.PathLike`, *optional*): + Path to a directory in which a downloaded pretrained model feature extractor should be cached if the + standard cache should not be used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the feature extractor files and override the cached versions + if they exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`. + + + + return_unused_kwargs (`bool`, *optional*, defaults to `False`): + If `False`, then this function returns just the final feature extractor object. If `True`, then this + functions returns a `Tuple(feature_extractor, unused_kwargs)` where *unused_kwargs* is a dictionary + consisting of the key/value pairs whose keys are not feature extractor attributes: i.e., the part of + `kwargs` which has not been used to update `feature_extractor` and is otherwise ignored. + kwargs (`Dict[str, Any]`, *optional*): + The values in kwargs of any keys which are feature extractor attributes will be used to override the + loaded values. Behavior concerning key/value pairs whose keys are *not* feature extractor attributes is + controlled by the `return_unused_kwargs` keyword parameter. + + Returns: + A feature extractor of type [`~feature_extraction_utils.FeatureExtractionMixin`]. + + Examples: + + ```python + # We can't instantiate directly the base class *FeatureExtractionMixin* nor *SequenceFeatureExtractor* so let's show the examples on a + # derived class: *Wav2Vec2FeatureExtractor* + feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained( + "facebook/wav2vec2-base-960h" + ) # Download feature_extraction_config from huggingface.co and cache. + feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained( + "./test/saved_model/" + ) # E.g. feature_extractor (or model) was saved using *save_pretrained('./test/saved_model/')* + feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("./test/saved_model/preprocessor_config.json") + feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained( + "facebook/wav2vec2-base-960h", return_attention_mask=False, foo=False + ) + assert feature_extractor.return_attention_mask is False + feature_extractor, unused_kwargs = Wav2Vec2FeatureExtractor.from_pretrained( + "facebook/wav2vec2-base-960h", return_attention_mask=False, foo=False, return_unused_kwargs=True + ) + assert feature_extractor.return_attention_mask is False + assert unused_kwargs == {"foo": False} + ```""" + kwargs["cache_dir"] = cache_dir + kwargs["force_download"] = force_download + kwargs["local_files_only"] = local_files_only + kwargs["revision"] = revision + + use_auth_token = kwargs.pop("use_auth_token", None) + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if token is not None: + kwargs["token"] = token + + feature_extractor_dict, kwargs = cls.get_feature_extractor_dict(pretrained_model_name_or_path, **kwargs) + + return cls.from_dict(feature_extractor_dict, **kwargs) + + def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs): + """ + Save a feature_extractor object to the directory `save_directory`, so that it can be re-loaded using the + [`~feature_extraction_utils.FeatureExtractionMixin.from_pretrained`] class method. + + Args: + save_directory (`str` or `os.PathLike`): + Directory where the feature extractor JSON file will be saved (will be created if it does not exist). + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + use_auth_token = kwargs.pop("use_auth_token", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if kwargs.get("token", None) is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + kwargs["token"] = use_auth_token + + if os.path.isfile(save_directory): + raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = self._create_repo(repo_id, **kwargs) + files_timestamps = self._get_files_timestamps(save_directory) + + # If we have a custom config, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + custom_object_save(self, save_directory, config=self) + + # If we save using the predefined names, we can load using `from_pretrained` + output_feature_extractor_file = os.path.join(save_directory, FEATURE_EXTRACTOR_NAME) + + self.to_json_file(output_feature_extractor_file) + logger.info(f"Feature extractor saved in {output_feature_extractor_file}") + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=kwargs.get("token"), + ) + + return [output_feature_extractor_file] + + @classmethod + def get_feature_extractor_dict( + cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs + ) -> tuple[dict[str, Any], dict[str, Any]]: + """ + From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a + feature extractor of type [`~feature_extraction_utils.FeatureExtractionMixin`] using `from_dict`. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`): + The identifier of the pre-trained checkpoint from which we want the dictionary of parameters. + + Returns: + `Tuple[Dict, Dict]`: The dictionary(ies) that will be used to instantiate the feature extractor object. + """ + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + resume_download = kwargs.pop("resume_download", None) + proxies = kwargs.pop("proxies", None) + subfolder = kwargs.pop("subfolder", None) + token = kwargs.pop("token", None) + use_auth_token = kwargs.pop("use_auth_token", None) + local_files_only = kwargs.pop("local_files_only", False) + revision = kwargs.pop("revision", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + + user_agent = {"file_type": "feature extractor", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + is_local = os.path.isdir(pretrained_model_name_or_path) + if os.path.isdir(pretrained_model_name_or_path): + feature_extractor_file = os.path.join(pretrained_model_name_or_path, FEATURE_EXTRACTOR_NAME) + if os.path.isfile(pretrained_model_name_or_path): + resolved_feature_extractor_file = pretrained_model_name_or_path + is_local = True + elif is_remote_url(pretrained_model_name_or_path): + feature_extractor_file = pretrained_model_name_or_path + resolved_feature_extractor_file = download_url(pretrained_model_name_or_path) + else: + feature_extractor_file = FEATURE_EXTRACTOR_NAME + try: + # Load from local folder or from cache or download from model Hub and cache + resolved_feature_extractor_file = cached_file( + pretrained_model_name_or_path, + feature_extractor_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + subfolder=subfolder, + token=token, + user_agent=user_agent, + revision=revision, + ) + except OSError: + # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted to + # the original exception. + raise + except Exception: + # For any other exception, we throw a generic error. + raise OSError( + f"Can't load feature extractor for '{pretrained_model_name_or_path}'. If you were trying to load" + " it from 'https://huggingface.co/models', make sure you don't have a local directory with the" + f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a" + f" directory containing a {FEATURE_EXTRACTOR_NAME} file" + ) + + try: + # Load feature_extractor dict + with open(resolved_feature_extractor_file, encoding="utf-8") as reader: + text = reader.read() + feature_extractor_dict = json.loads(text) + + except json.JSONDecodeError: + raise OSError( + f"It looks like the config file at '{resolved_feature_extractor_file}' is not a valid JSON file." + ) + + if is_local: + logger.info(f"loading configuration file {resolved_feature_extractor_file}") + else: + logger.info( + f"loading configuration file {feature_extractor_file} from cache at {resolved_feature_extractor_file}" + ) + + if not is_local: + if "auto_map" in feature_extractor_dict: + feature_extractor_dict["auto_map"] = add_model_info_to_auto_map( + feature_extractor_dict["auto_map"], pretrained_model_name_or_path + ) + if "custom_pipelines" in feature_extractor_dict: + feature_extractor_dict["custom_pipelines"] = add_model_info_to_custom_pipelines( + feature_extractor_dict["custom_pipelines"], pretrained_model_name_or_path + ) + + return feature_extractor_dict, kwargs + + @classmethod + def from_dict(cls, feature_extractor_dict: dict[str, Any], **kwargs) -> PreTrainedFeatureExtractor: + """ + Instantiates a type of [`~feature_extraction_utils.FeatureExtractionMixin`] from a Python dictionary of + parameters. + + Args: + feature_extractor_dict (`Dict[str, Any]`): + Dictionary that will be used to instantiate the feature extractor object. Such a dictionary can be + retrieved from a pretrained checkpoint by leveraging the + [`~feature_extraction_utils.FeatureExtractionMixin.to_dict`] method. + kwargs (`Dict[str, Any]`): + Additional parameters from which to initialize the feature extractor object. + + Returns: + [`~feature_extraction_utils.FeatureExtractionMixin`]: The feature extractor object instantiated from those + parameters. + """ + return_unused_kwargs = kwargs.pop("return_unused_kwargs", False) + + # Update feature_extractor with kwargs if needed + to_remove = [] + for key, value in kwargs.items(): + if key in feature_extractor_dict: + feature_extractor_dict[key] = value + to_remove.append(key) + for key in to_remove: + kwargs.pop(key, None) + + feature_extractor = cls(**feature_extractor_dict) + + logger.info(f"Feature extractor {feature_extractor}") + if return_unused_kwargs: + return feature_extractor, kwargs + else: + return feature_extractor + + def to_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + output = copy.deepcopy(self.__dict__) + output["feature_extractor_type"] = self.__class__.__name__ + if "mel_filters" in output: + del output["mel_filters"] + if "window" in output: + del output["window"] + return output + + @classmethod + def from_json_file(cls, json_file: Union[str, os.PathLike]) -> PreTrainedFeatureExtractor: + """ + Instantiates a feature extractor of type [`~feature_extraction_utils.FeatureExtractionMixin`] from the path to + a JSON file of parameters. + + Args: + json_file (`str` or `os.PathLike`): + Path to the JSON file containing the parameters. + + Returns: + A feature extractor of type [`~feature_extraction_utils.FeatureExtractionMixin`]: The feature_extractor + object instantiated from that JSON file. + """ + with open(json_file, encoding="utf-8") as reader: + text = reader.read() + feature_extractor_dict = json.loads(text) + return cls(**feature_extractor_dict) + + def to_json_string(self) -> str: + """ + Serializes this instance to a JSON string. + + Returns: + `str`: String containing all the attributes that make up this feature_extractor instance in JSON format. + """ + dictionary = self.to_dict() + + for key, value in dictionary.items(): + if isinstance(value, np.ndarray): + dictionary[key] = value.tolist() + + # make sure private name "_processor_class" is correctly + # saved as "processor_class" + _processor_class = dictionary.pop("_processor_class", None) + if _processor_class is not None: + dictionary["processor_class"] = _processor_class + + return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" + + def to_json_file(self, json_file_path: Union[str, os.PathLike]): + """ + Save this instance to a JSON file. + + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file in which this feature_extractor instance's parameters will be saved. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + writer.write(self.to_json_string()) + + def __repr__(self): + return f"{self.__class__.__name__} {self.to_json_string()}" + + @classmethod + def register_for_auto_class(cls, auto_class="AutoFeatureExtractor"): + """ + Register this class with a given auto class. This should only be used for custom feature extractors as the ones + in the library are already mapped with `AutoFeatureExtractor`. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"AutoFeatureExtractor"`): + The auto class to register this new feature extractor with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + +FeatureExtractionMixin.push_to_hub = copy_func(FeatureExtractionMixin.push_to_hub) +if FeatureExtractionMixin.push_to_hub.__doc__ is not None: + FeatureExtractionMixin.push_to_hub.__doc__ = FeatureExtractionMixin.push_to_hub.__doc__.format( + object="feature extractor", object_class="AutoFeatureExtractor", object_files="feature extractor file" + ) diff --git a/lib/python3.12/site-packages/transformers/file_utils.py b/lib/python3.12/site-packages/transformers/file_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ac6b36d2dbcf7fef3621a3b35f083ae89a39b1b4 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/file_utils.py @@ -0,0 +1,131 @@ +# Copyright 2020 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. +""" +File utilities: utilities related to download and cache models + +This module should not be update anymore and is only left for backward compatibility. +""" + +from huggingface_hub import get_full_repo_name # for backward compatibility +from huggingface_hub.constants import HF_HUB_DISABLE_TELEMETRY as DISABLE_TELEMETRY # for backward compatibility + +from . import __version__ + +# Backward compatibility imports, to make sure all those objects can be found in file_utils +from .utils import ( + CLOUDFRONT_DISTRIB_PREFIX, + CONFIG_NAME, + DUMMY_INPUTS, + DUMMY_MASK, + ENV_VARS_TRUE_AND_AUTO_VALUES, + ENV_VARS_TRUE_VALUES, + FEATURE_EXTRACTOR_NAME, + FLAX_WEIGHTS_NAME, + HF_MODULES_CACHE, + HUGGINGFACE_CO_PREFIX, + HUGGINGFACE_CO_RESOLVE_ENDPOINT, + MODEL_CARD_NAME, + MULTIPLE_CHOICE_DUMMY_INPUTS, + PYTORCH_PRETRAINED_BERT_CACHE, + PYTORCH_TRANSFORMERS_CACHE, + S3_BUCKET_PREFIX, + SENTENCEPIECE_UNDERLINE, + SPIECE_UNDERLINE, + TF2_WEIGHTS_NAME, + TF_WEIGHTS_NAME, + TORCH_FX_REQUIRED_VERSION, + TRANSFORMERS_CACHE, + TRANSFORMERS_DYNAMIC_MODULE_NAME, + USE_JAX, + USE_TF, + USE_TORCH, + WEIGHTS_INDEX_NAME, + WEIGHTS_NAME, + ContextManagers, + DummyObject, + EntryNotFoundError, + ExplicitEnum, + ModelOutput, + PaddingStrategy, + PushToHubMixin, + RepositoryNotFoundError, + RevisionNotFoundError, + TensorType, + _LazyModule, + add_code_sample_docstrings, + add_end_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + cached_property, + copy_func, + default_cache_path, + define_sagemaker_information, + get_torch_version, + has_file, + http_user_agent, + is_apex_available, + is_bs4_available, + is_coloredlogs_available, + is_datasets_available, + is_detectron2_available, + is_faiss_available, + is_flax_available, + is_ftfy_available, + is_g2p_en_available, + is_in_notebook, + is_ipex_available, + is_librosa_available, + is_offline_mode, + is_onnx_available, + is_pandas_available, + is_phonemizer_available, + is_protobuf_available, + is_psutil_available, + is_py3nvml_available, + is_pyctcdecode_available, + is_pytesseract_available, + is_pytorch_quantization_available, + is_rjieba_available, + is_sagemaker_dp_enabled, + is_sagemaker_mp_enabled, + is_scipy_available, + is_sentencepiece_available, + is_seqio_available, + is_sklearn_available, + is_soundfile_available, + is_spacy_available, + is_speech_available, + is_tensor, + is_tensorflow_probability_available, + is_tf2onnx_available, + is_tf_available, + is_timm_available, + is_tokenizers_available, + is_torch_available, + is_torch_bf16_available, + is_torch_cuda_available, + is_torch_fx_available, + is_torch_fx_proxy, + is_torch_mps_available, + is_torch_tf32_available, + is_torch_xla_available, + is_torchaudio_available, + is_training_run_on_sagemaker, + is_vision_available, + replace_return_docstrings, + requires_backends, + to_numpy, + to_py_obj, + torch_only_method, +) diff --git a/lib/python3.12/site-packages/transformers/generation/__init__.py b/lib/python3.12/site-packages/transformers/generation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cf1fa3661e0c89b805b2a690df75e0f924cd27cd --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/__init__.py @@ -0,0 +1,348 @@ +# Copyright 2022 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. + +from typing import TYPE_CHECKING + +from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available + + +_import_structure = { + "configuration_utils": [ + "BaseWatermarkingConfig", + "CompileConfig", + "GenerationConfig", + "GenerationMode", + "SynthIDTextWatermarkingConfig", + "WatermarkingConfig", + ], + "streamers": ["AsyncTextIteratorStreamer", "BaseStreamer", "TextIteratorStreamer", "TextStreamer"], +} + +try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["beam_constraints"] = [ + "Constraint", + "ConstraintListState", + "DisjunctiveConstraint", + "PhrasalConstraint", + ] + _import_structure["beam_search"] = [ + "BeamHypotheses", + "BeamScorer", + "BeamSearchScorer", + "ConstrainedBeamSearchScorer", + ] + _import_structure["candidate_generator"] = [ + "AssistedCandidateGenerator", + "CandidateGenerator", + "EarlyExitCandidateGenerator", + "PromptLookupCandidateGenerator", + ] + _import_structure["logits_process"] = [ + "AlternatingCodebooksLogitsProcessor", + "ClassifierFreeGuidanceLogitsProcessor", + "EncoderNoRepeatNGramLogitsProcessor", + "EncoderRepetitionPenaltyLogitsProcessor", + "EpsilonLogitsWarper", + "EtaLogitsWarper", + "ExponentialDecayLengthPenalty", + "ForcedBOSTokenLogitsProcessor", + "ForcedEOSTokenLogitsProcessor", + "HammingDiversityLogitsProcessor", + "InfNanRemoveLogitsProcessor", + "LogitNormalization", + "LogitsProcessor", + "LogitsProcessorList", + "MinLengthLogitsProcessor", + "MinNewTokensLengthLogitsProcessor", + "MinPLogitsWarper", + "NoBadWordsLogitsProcessor", + "NoRepeatNGramLogitsProcessor", + "PrefixConstrainedLogitsProcessor", + "RepetitionPenaltyLogitsProcessor", + "SequenceBiasLogitsProcessor", + "SuppressTokensLogitsProcessor", + "SuppressTokensAtBeginLogitsProcessor", + "SynthIDTextWatermarkLogitsProcessor", + "TemperatureLogitsWarper", + "TopKLogitsWarper", + "TopPLogitsWarper", + "TypicalLogitsWarper", + "UnbatchedClassifierFreeGuidanceLogitsProcessor", + "WhisperTimeStampLogitsProcessor", + "WatermarkLogitsProcessor", + ] + _import_structure["stopping_criteria"] = [ + "MaxLengthCriteria", + "MaxTimeCriteria", + "ConfidenceCriteria", + "EosTokenCriteria", + "StoppingCriteria", + "StoppingCriteriaList", + "validate_stopping_criteria", + "StopStringCriteria", + ] + _import_structure["utils"] = [ + "GenerationMixin", + "GreedySearchEncoderDecoderOutput", + "GreedySearchDecoderOnlyOutput", + "SampleEncoderDecoderOutput", + "SampleDecoderOnlyOutput", + "BeamSearchEncoderDecoderOutput", + "BeamSearchDecoderOnlyOutput", + "BeamSampleEncoderDecoderOutput", + "BeamSampleDecoderOnlyOutput", + "ContrastiveSearchEncoderDecoderOutput", + "ContrastiveSearchDecoderOnlyOutput", + "GenerateBeamDecoderOnlyOutput", + "GenerateBeamEncoderDecoderOutput", + "GenerateDecoderOnlyOutput", + "GenerateEncoderDecoderOutput", + ] + _import_structure["watermarking"] = [ + "WatermarkDetector", + "WatermarkDetectorOutput", + "BayesianDetectorModel", + "BayesianDetectorConfig", + "SynthIDTextWatermarkDetector", + ] + +try: + if not is_tf_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["tf_logits_process"] = [ + "TFForcedBOSTokenLogitsProcessor", + "TFForcedEOSTokenLogitsProcessor", + "TFForceTokensLogitsProcessor", + "TFLogitsProcessor", + "TFLogitsProcessorList", + "TFLogitsWarper", + "TFMinLengthLogitsProcessor", + "TFNoBadWordsLogitsProcessor", + "TFNoRepeatNGramLogitsProcessor", + "TFRepetitionPenaltyLogitsProcessor", + "TFSuppressTokensAtBeginLogitsProcessor", + "TFSuppressTokensLogitsProcessor", + "TFTemperatureLogitsWarper", + "TFTopKLogitsWarper", + "TFTopPLogitsWarper", + ] + _import_structure["tf_utils"] = [ + "TFGenerationMixin", + "TFGreedySearchDecoderOnlyOutput", + "TFGreedySearchEncoderDecoderOutput", + "TFSampleEncoderDecoderOutput", + "TFSampleDecoderOnlyOutput", + "TFBeamSearchEncoderDecoderOutput", + "TFBeamSearchDecoderOnlyOutput", + "TFBeamSampleEncoderDecoderOutput", + "TFBeamSampleDecoderOnlyOutput", + "TFContrastiveSearchEncoderDecoderOutput", + "TFContrastiveSearchDecoderOnlyOutput", + ] + +try: + if not is_flax_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["flax_logits_process"] = [ + "FlaxForcedBOSTokenLogitsProcessor", + "FlaxForcedEOSTokenLogitsProcessor", + "FlaxForceTokensLogitsProcessor", + "FlaxLogitsProcessor", + "FlaxLogitsProcessorList", + "FlaxLogitsWarper", + "FlaxMinLengthLogitsProcessor", + "FlaxSuppressTokensAtBeginLogitsProcessor", + "FlaxSuppressTokensLogitsProcessor", + "FlaxTemperatureLogitsWarper", + "FlaxTopKLogitsWarper", + "FlaxTopPLogitsWarper", + "FlaxWhisperTimeStampLogitsProcessor", + "FlaxNoRepeatNGramLogitsProcessor", + ] + _import_structure["flax_utils"] = [ + "FlaxGenerationMixin", + "FlaxGreedySearchOutput", + "FlaxSampleOutput", + "FlaxBeamSearchOutput", + ] + +if TYPE_CHECKING: + from .configuration_utils import ( + BaseWatermarkingConfig, + CompileConfig, + GenerationConfig, + GenerationMode, + SynthIDTextWatermarkingConfig, + WatermarkingConfig, + ) + from .streamers import AsyncTextIteratorStreamer, BaseStreamer, TextIteratorStreamer, TextStreamer + + try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint + from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer + from .candidate_generator import ( + AssistedCandidateGenerator, + CandidateGenerator, + EarlyExitCandidateGenerator, + PromptLookupCandidateGenerator, + ) + from .logits_process import ( + AlternatingCodebooksLogitsProcessor, + ClassifierFreeGuidanceLogitsProcessor, + EncoderNoRepeatNGramLogitsProcessor, + EncoderRepetitionPenaltyLogitsProcessor, + EpsilonLogitsWarper, + EtaLogitsWarper, + ExponentialDecayLengthPenalty, + ForcedBOSTokenLogitsProcessor, + ForcedEOSTokenLogitsProcessor, + HammingDiversityLogitsProcessor, + InfNanRemoveLogitsProcessor, + LogitNormalization, + LogitsProcessor, + LogitsProcessorList, + MinLengthLogitsProcessor, + MinNewTokensLengthLogitsProcessor, + MinPLogitsWarper, + NoBadWordsLogitsProcessor, + NoRepeatNGramLogitsProcessor, + PrefixConstrainedLogitsProcessor, + RepetitionPenaltyLogitsProcessor, + SequenceBiasLogitsProcessor, + SuppressTokensAtBeginLogitsProcessor, + SuppressTokensLogitsProcessor, + SynthIDTextWatermarkLogitsProcessor, + TemperatureLogitsWarper, + TopKLogitsWarper, + TopPLogitsWarper, + TypicalLogitsWarper, + UnbatchedClassifierFreeGuidanceLogitsProcessor, + WatermarkLogitsProcessor, + WhisperTimeStampLogitsProcessor, + ) + from .stopping_criteria import ( + ConfidenceCriteria, + EosTokenCriteria, + MaxLengthCriteria, + MaxTimeCriteria, + StoppingCriteria, + StoppingCriteriaList, + StopStringCriteria, + validate_stopping_criteria, + ) + from .utils import ( + BeamSampleDecoderOnlyOutput, + BeamSampleEncoderDecoderOutput, + BeamSearchDecoderOnlyOutput, + BeamSearchEncoderDecoderOutput, + ContrastiveSearchDecoderOnlyOutput, + ContrastiveSearchEncoderDecoderOutput, + GenerateBeamDecoderOnlyOutput, + GenerateBeamEncoderDecoderOutput, + GenerateDecoderOnlyOutput, + GenerateEncoderDecoderOutput, + GenerationMixin, + GreedySearchDecoderOnlyOutput, + GreedySearchEncoderDecoderOutput, + SampleDecoderOnlyOutput, + SampleEncoderDecoderOutput, + ) + from .watermarking import ( + BayesianDetectorConfig, + BayesianDetectorModel, + SynthIDTextWatermarkDetector, + WatermarkDetector, + WatermarkDetectorOutput, + ) + + try: + if not is_tf_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .tf_logits_process import ( + TFForcedBOSTokenLogitsProcessor, + TFForcedEOSTokenLogitsProcessor, + TFForceTokensLogitsProcessor, + TFLogitsProcessor, + TFLogitsProcessorList, + TFLogitsWarper, + TFMinLengthLogitsProcessor, + TFNoBadWordsLogitsProcessor, + TFNoRepeatNGramLogitsProcessor, + TFRepetitionPenaltyLogitsProcessor, + TFSuppressTokensAtBeginLogitsProcessor, + TFSuppressTokensLogitsProcessor, + TFTemperatureLogitsWarper, + TFTopKLogitsWarper, + TFTopPLogitsWarper, + ) + from .tf_utils import ( + TFBeamSampleDecoderOnlyOutput, + TFBeamSampleEncoderDecoderOutput, + TFBeamSearchDecoderOnlyOutput, + TFBeamSearchEncoderDecoderOutput, + TFContrastiveSearchDecoderOnlyOutput, + TFContrastiveSearchEncoderDecoderOutput, + TFGenerationMixin, + TFGreedySearchDecoderOnlyOutput, + TFGreedySearchEncoderDecoderOutput, + TFSampleDecoderOnlyOutput, + TFSampleEncoderDecoderOutput, + ) + + try: + if not is_flax_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .flax_logits_process import ( + FlaxForcedBOSTokenLogitsProcessor, + FlaxForcedEOSTokenLogitsProcessor, + FlaxForceTokensLogitsProcessor, + FlaxLogitsProcessor, + FlaxLogitsProcessorList, + FlaxLogitsWarper, + FlaxMinLengthLogitsProcessor, + FlaxNoRepeatNGramLogitsProcessor, + FlaxSuppressTokensAtBeginLogitsProcessor, + FlaxSuppressTokensLogitsProcessor, + FlaxTemperatureLogitsWarper, + FlaxTopKLogitsWarper, + FlaxTopPLogitsWarper, + FlaxWhisperTimeStampLogitsProcessor, + ) + from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c57d68c804b719a00a604ba0ce7ef63855f2c2a Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/beam_constraints.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/beam_constraints.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8fcc5e31d5b30c536a371e41e5a1de7587720fc0 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/beam_constraints.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/beam_search.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/beam_search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25dad44f9b657b0d9615195312bcb57138029347 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/beam_search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/candidate_generator.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/candidate_generator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63456277dd015c574db14cb7cb3209ab1bdd3b5f Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/candidate_generator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/configuration_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/configuration_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b97c896dc15d3fca9bfbcf791238d8c57dca618d Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/configuration_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/flax_logits_process.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/flax_logits_process.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..48c522f033b0b0deae69d6b69478bcb378699ae7 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/flax_logits_process.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/flax_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/flax_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..66358b687e73c875e0cc8f11bfdb3d00c7e0c41b Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/flax_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/stopping_criteria.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/stopping_criteria.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..81873a7309d6ea1ba4f9a8574725fb8560c46e7d Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/stopping_criteria.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/streamers.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/streamers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd2623212733a9b348ca7cfd71a060e56e12b000 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/streamers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/tf_logits_process.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/tf_logits_process.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f94bee0786e451ca7743b0732857ffd53c7a5f43 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/tf_logits_process.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/__pycache__/watermarking.cpython-312.pyc b/lib/python3.12/site-packages/transformers/generation/__pycache__/watermarking.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..04eacf5876e87437269f10d69788015ff51132d1 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/generation/__pycache__/watermarking.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/generation/beam_constraints.py b/lib/python3.12/site-packages/transformers/generation/beam_constraints.py new file mode 100644 index 0000000000000000000000000000000000000000..daf64209b796772ac5c8da0dc139d944c2e82430 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/beam_constraints.py @@ -0,0 +1,524 @@ +from abc import ABC, abstractmethod +from typing import List, Optional + + +class Constraint(ABC): + r"""Abstract base class for all constraints that can be applied during generation. + It must define how the constraint can be satisfied. + + All classes that inherit Constraint must follow the requirement that + + ```py + completed = False + while not completed: + _, completed = constraint.update(constraint.advance()) + ``` + + will always terminate (halt). + """ + + def __init__(self): + # test for the above condition + self.test() + + def test(self): + """ + Tests whether this constraint has been properly defined. + """ + counter = 0 + completed = False + while not completed: + if counter == 1: + self.reset() + advance = self.advance() + if not self.does_advance(advance): + raise Exception( + "Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true." + ) + + stepped, completed, reset = self.update(advance) + counter += 1 + + if counter > 10000: + raise Exception("update() does not fulfill the constraint.") + + if self.remaining() != 0: + raise Exception("Custom Constraint is not defined correctly.") + + @abstractmethod + def advance(self): + """ + When called, returns the token(s) that would take this constraint one step closer to being fulfilled. + + Return: + token_ids (Union[int, List[int], None]): + - A single token ID (int) that advances the constraint, or + - A list of token IDs that could advance the constraint + - None if the constraint is completed or cannot be advanced + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + @abstractmethod + def does_advance(self, token_id: int): + """ + Reads in a token and returns whether it creates progress. + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + @abstractmethod + def update(self, token_id: int): + """ + Reads in a token and returns booleans that indicate the progress made by it. This function will update the + state of this object unlikes `does_advance(self, token_id: int)`. + + This isn't to test whether a certain token will advance the progress; it's to update its state as if it has + been generated. This becomes important if token_id != desired token (refer to else statement in + PhrasalConstraint) + + Args: + token_id(`int`): + The id of a newly generated token in the beam search. + Return: + stepped(`bool`): + Whether this constraint has become one step closer to being fulfuilled. + completed(`bool`): + Whether this constraint has been completely fulfilled by this token being generated. + reset (`bool`): + Whether this constraint has reset its progress by this token being generated. + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + @abstractmethod + def reset(self): + """ + Resets the state of this constraint to its initialization. We would call this in cases where the fulfillment of + a constraint is abrupted by an unwanted token. + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + @abstractmethod + def remaining(self): + """ + Returns the number of remaining steps of `advance()` in order to complete this constraint. + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + @abstractmethod + def copy(self, stateful=False): + """ + Creates a new instance of this constraint. + + Args: + stateful(`bool`): Whether to not only copy the constraint for new instance, but also its state. + + Return: + constraint(`Constraint`): The same constraint as the one being called from. + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + +class PhrasalConstraint(Constraint): + r""" + [`Constraint`] enforcing that an ordered sequence of tokens is included in the output. + + Args: + token_ids (`List[int]`): + The id of the token that must be generated by the output. + """ + + def __init__(self, token_ids: List[int]): + super(Constraint, self).__init__() + + if not isinstance(token_ids, list) or len(token_ids) == 0: + raise ValueError(f"`token_ids` has to be a non-empty list, but is {token_ids}.") + if any((not isinstance(token_id, int) or token_id < 0) for token_id in token_ids): + raise ValueError(f"Each list in `token_ids` has to be a list of positive integers, but is {token_ids}.") + + self.token_ids = token_ids + + self.seqlen = len(self.token_ids) + self.fulfilled_idx = -1 # the index of the currently fulfilled step + self.completed = False + + def advance(self): + if self.completed: + return None + return self.token_ids[self.fulfilled_idx + 1] + + def does_advance(self, token_id: int): + if not isinstance(token_id, int): + raise TypeError(f"`token_id` has to be an `int`, but is {token_id} of type {type(token_id)}") + + if self.completed: + return False + + return token_id == self.token_ids[self.fulfilled_idx + 1] + + def update(self, token_id: int): + if not isinstance(token_id, int): + raise TypeError(f"`token_id` has to be an `int`, but is {token_id} of type {type(token_id)}") + + stepped = False + completed = False + reset = False + + if self.does_advance(token_id): + self.fulfilled_idx += 1 + stepped = True + if self.fulfilled_idx == (self.seqlen - 1): + completed = True + self.completed = completed + else: + # failed to make progress. + reset = True + self.reset() + return stepped, completed, reset + + def reset(self): + self.completed = False + self.fulfilled_idx = 0 + + def remaining(self): + return self.seqlen - (self.fulfilled_idx + 1) + + def copy(self, stateful=False): + new_constraint = PhrasalConstraint(self.token_ids) + + if stateful: + new_constraint.seq_len = self.seqlen + new_constraint.fulfilled_idx = self.fulfilled_idx + new_constraint.completed = self.completed + + return new_constraint + + +class DisjunctiveTrie: + def __init__(self, nested_token_ids: List[List[int]], no_subsets=True): + r""" + A helper class that builds a trie with the words represented in `nested_token_ids`. + """ + self.max_height = max([len(one) for one in nested_token_ids]) + + root = {} + for token_ids in nested_token_ids: + level = root + for tidx, token_id in enumerate(token_ids): + if token_id not in level: + level[token_id] = {} + + level = level[token_id] + + if no_subsets and self.has_subsets(root, nested_token_ids): + raise ValueError( + "Each list in `nested_token_ids` can't be a complete subset of another list, but is" + f" {nested_token_ids}." + ) + + self.trie = root + + def next_tokens(self, current_seq): + """ + The next possible tokens that will progress the trie, given the current sequence of tokens in `current_seq`. + """ + start = self.trie + + for current_token in current_seq: + start = start[current_token] + + next_tokens = list(start.keys()) + + return next_tokens + + def reached_leaf(self, current_seq): + next_tokens = self.next_tokens(current_seq) + + return len(next_tokens) == 0 + + def count_leaves(self, root): + next_nodes = list(root.values()) + if len(next_nodes) == 0: + return 1 + else: + return sum([self.count_leaves(nn) for nn in next_nodes]) + + def has_subsets(self, trie, nested_token_ids): + """ + Returns whether # of leaves == # of words. Otherwise some word is a subset of another. + """ + leaf_count = self.count_leaves(trie) + return len(nested_token_ids) != leaf_count + + +class DisjunctiveConstraint(Constraint): + r""" + A special [`Constraint`] that is fulfilled by fulfilling just one of several constraints. + + Args: + nested_token_ids (`List[List[int]]`): + A list of words, where each word is a list of ids. This constraint is fulfilled by generating just one from + the list of words. + """ + + def __init__(self, nested_token_ids: List[List[int]]): + super(Constraint, self).__init__() + + if not isinstance(nested_token_ids, list) or len(nested_token_ids) == 0: + raise ValueError(f"`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}.") + if any(not isinstance(token_ids, list) for token_ids in nested_token_ids): + raise ValueError(f"`nested_token_ids` has to be a list of lists, but is {nested_token_ids}.") + if any( + any((not isinstance(token_id, int) or token_id < 0) for token_id in token_ids) + for token_ids in nested_token_ids + ): + raise ValueError( + f"Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}." + ) + + self.trie = DisjunctiveTrie(nested_token_ids) + self.token_ids = nested_token_ids + + self.seqlen = self.trie.max_height + self.current_seq = [] + self.completed = False + + def advance(self): + token_list = self.trie.next_tokens(self.current_seq) + + if len(token_list) == 0: + return None + else: + return token_list + + def does_advance(self, token_id: int): + if not isinstance(token_id, int): + raise TypeError(f"`token_id` is supposed to be type `int`, but is {token_id} of type {type(token_id)}") + + next_tokens = self.trie.next_tokens(self.current_seq) + + return token_id in next_tokens + + def update(self, token_id: int): + if not isinstance(token_id, int): + raise TypeError(f"`token_id` is supposed to be type `int`, but is {token_id} of type {type(token_id)}") + + stepped = False + completed = False + reset = False + + if self.does_advance(token_id): + self.current_seq.append(token_id) + stepped = True + else: + reset = True + self.reset() + + completed = self.trie.reached_leaf(self.current_seq) + self.completed = completed + + return stepped, completed, reset + + def reset(self): + self.completed = False + self.current_seq = [] + + def remaining(self): + if self.completed: + # since this can be completed without reaching max height + return 0 + else: + return self.seqlen - len(self.current_seq) + + def copy(self, stateful=False): + new_constraint = DisjunctiveConstraint(self.token_ids) + + if stateful: + new_constraint.seq_len = self.seqlen + new_constraint.current_seq = self.current_seq + new_constraint.completed = self.completed + + return new_constraint + + +class ConstraintListState: + r""" + A class for beam scorers to track its progress through a list of constraints. + + Args: + constraints (`List[Constraint]`): + A list of [`Constraint`] objects that must be fulfilled by the beam scorer. + """ + + def __init__(self, constraints: List[Constraint]): + self.constraints = constraints + + # max # of steps required to fulfill a given constraint + self.max_seqlen = max([c.seqlen for c in constraints]) + self.n_constraints = len(constraints) + self.completed = False + + self.init_state() + + def init_state(self): + self.complete_constraints = [] + self.inprogress_constraint = None + self.pending_constraints = [constraint.copy(stateful=False) for constraint in self.constraints] + + def get_bank(self): + add = 0 + if self.inprogress_constraint: + # extra points for having a constraint mid-fulfilled + add += self.max_seqlen - self.inprogress_constraint.remaining() + + return (len(self.complete_constraints) * self.max_seqlen) + add + + def advance(self): + """The list of tokens to generate such that we can make progress. + By "list" we don't mean the list of token that will fully fulfill a constraint. + + Given constraints `c_i = {t_ij | j == # of tokens}`, If we're not in the middle of progressing through a + specific constraint `c_i`, we return: + + `[t_k1 for k in indices of unfulfilled constraints]` + + If we are in the middle of a constraint, then we return: + `[t_ij]`, where `i` is the index of the inprogress constraint, `j` is the next step for the constraint. + + Though we don't care which constraint is fulfilled first, if we are in the progress of fulfilling a constraint, + that's the only one we'll return. + """ + token_list = [] + if self.inprogress_constraint is None: + for constraint in self.pending_constraints: # "pending" == "unfulfilled yet" + advance = constraint.advance() + if isinstance(advance, int): + token_list.append(advance) + elif isinstance(advance, list): + token_list.extend(advance) + else: + advance = self.inprogress_constraint.advance() + if isinstance(advance, int): + token_list.append(advance) + elif isinstance(advance, list): + token_list.extend(advance) + + if len(token_list) == 0: + return None + else: + return token_list + + def reset(self, token_ids: Optional[List[int]]): + """ + token_ids: the tokens generated thus far to reset the state of the progress through constraints. + """ + self.init_state() + + if token_ids is not None: + for token in token_ids: + # completes or steps **one** constraint + complete, stepped = self.add(token) + + # the entire list of constraints are fulfilled + if self.completed: + break + + def add(self, token_id: int): + if not isinstance(token_id, int): + raise TypeError(f"`token_id` should be an `int`, but is `{token_id}`.") + + complete, stepped = False, False + + if self.completed: + complete = True + stepped = False + return complete, stepped + + if self.inprogress_constraint is not None: + # In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current + # job, simply update the state + + stepped, complete, reset = self.inprogress_constraint.update(token_id) + if reset: + # 1. If the next token breaks the progress, then we must restart. + # e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books". + + # But that doesn't mean we self.init_state(), since we only reset the state for this particular + # constraint, not the full list of constraints. + + self.pending_constraints.append(self.inprogress_constraint.copy(stateful=False)) + self.inprogress_constraint = None + + if complete: + # 2. If the next token completes the constraint, move it to completed list, set + # inprogress to None. If there are no pending constraints either, then this full list of constraints + # is complete. + + self.complete_constraints.append(self.inprogress_constraint) + self.inprogress_constraint = None + + if len(self.pending_constraints) == 0: + # we're done! + self.completed = True + + else: + # Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list + # of constraints? + + for cidx, pending_constraint in enumerate(self.pending_constraints): + if pending_constraint.does_advance(token_id): + stepped, complete, reset = pending_constraint.update(token_id) + + if not stepped: + raise Exception( + "`constraint.update(token_id)` is not yielding incremental progress, " + "even though `constraint.does_advance(token_id)` is true." + ) + + if complete: + self.complete_constraints.append(pending_constraint) + self.inprogress_constraint = None + + if not complete and stepped: + self.inprogress_constraint = pending_constraint + + if complete or stepped: + # If we made any progress at all, then it's at least not a "pending constraint". + + self.pending_constraints = ( + self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :] + ) + + if len(self.pending_constraints) == 0 and self.inprogress_constraint is None: + # If there's no longer any pending after this and no inprogress either, then we must be + # complete. + + self.completed = True + + break # prevent accidentally stepping through multiple constraints with just one token. + + return complete, stepped + + def copy(self, stateful=True): + new_state = ConstraintListState(self.constraints) # we actually never though self.constraints objects + # throughout this process. So it's at initialization state. + + if stateful: + new_state.complete_constraints = [ + constraint.copy(stateful=True) for constraint in self.complete_constraints + ] + if self.inprogress_constraint is not None: + new_state.inprogress_constraint = self.inprogress_constraint.copy(stateful=True) + new_state.pending_constraints = [constraint.copy() for constraint in self.pending_constraints] + + return new_state diff --git a/lib/python3.12/site-packages/transformers/generation/beam_search.py b/lib/python3.12/site-packages/transformers/generation/beam_search.py new file mode 100644 index 0000000000000000000000000000000000000000..3938deb4826025834967ec7782b47242c5d77899 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/beam_search.py @@ -0,0 +1,1013 @@ +# coding=utf-8 +# Copyright 2020 The HuggingFace Inc. team +# +# 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 abc import ABC, abstractmethod +from collections import UserDict +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import torch + +from ..utils import add_start_docstrings +from .beam_constraints import Constraint, ConstraintListState + + +PROCESS_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size * num_beams, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using any class inheriting from [`PreTrainedTokenizer`]. See + [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + next_scores (`torch.FloatTensor` of shape `(batch_size, 2 * num_beams)`): + Current scores of the top `2 * num_beams` non-finished beam hypotheses. + next_tokens (`torch.LongTensor` of shape `(batch_size, 2 * num_beams)`): + `input_ids` of the tokens corresponding to the top `2 * num_beams` non-finished beam hypotheses. + next_indices (`torch.LongTensor` of shape `(batch_size, 2 * num_beams)`): + Beam indices indicating to which beam hypothesis the `next_tokens` correspond. + pad_token_id (`int`, *optional*): + The id of the *padding* token. + eos_token_id (`Union[int, List[int]]`, *optional*): + The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens. + beam_indices (`torch.LongTensor`, *optional*): + Beam indices indicating to which beam hypothesis each token correspond. + group_index (`int`, *optional*): + The index of the group of beams. Used with [`~PreTrainedModel.group_beam_search`]. + + Return: + `UserDict`: A dictionary composed of the fields as defined above: + + - **next_beam_scores** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Updated scores of all + non-finished beams. + - **next_beam_tokens** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Next tokens to be added + to the non-finished beam_hypotheses. + - **next_beam_indices** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Beam indices + indicating to which beam the next tokens shall be added. + +""" + +FINALIZE_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size * num_beams, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using any class inheriting from [`PreTrainedTokenizer`]. See + [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + final_beam_scores (`torch.FloatTensor` of shape `(batch_size * num_beams)`): + The final scores of all non-finished beams. + final_beam_tokens (`torch.FloatTensor` of shape `(batch_size * num_beams)`): + The last tokens to be added to the non-finished beam_hypotheses. + final_beam_indices (`torch.FloatTensor` of shape `(batch_size * num_beams)`): + The beam indices indicating to which beam the `final_beam_tokens` shall be added. + pad_token_id (`int`, *optional*): + The id of the *padding* token. + eos_token_id (`Union[int, List[int]]`, *optional*): + The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens. + + Return: + `torch.LongTensor` of shape `(batch_size * num_return_sequences, sequence_length)`: The generated sequences. + The second dimension (sequence_length) is either equal to `max_length` or shorter if all batches finished early + due to the `eos_token_id`. + +""" + + +class BeamScorer(ABC): + """ + Abstract base class for all beam scorers that are used for [`~PreTrainedModel.beam_search`] and + [`~PreTrainedModel.beam_sample`]. + """ + + @abstractmethod + @add_start_docstrings(PROCESS_INPUTS_DOCSTRING) + def process( + self, + input_ids: torch.LongTensor, + next_scores: torch.FloatTensor, + next_tokens: torch.LongTensor, + next_indices: torch.LongTensor, + **kwargs, + ) -> Tuple[torch.Tensor]: + raise NotImplementedError("This is an abstract method.") + + @abstractmethod + @add_start_docstrings(FINALIZE_INPUTS_DOCSTRING) + def finalize( + self, + input_ids: torch.LongTensor, + next_scores: torch.FloatTensor, + next_tokens: torch.LongTensor, + next_indices: torch.LongTensor, + max_length: int, + **kwargs, + ) -> torch.LongTensor: + raise NotImplementedError("This is an abstract method.") + + +class BeamSearchScorer(BeamScorer): + r""" + [`BeamScorer`] implementing standard beam search decoding. + + Adapted in part from [Facebook's XLM beam search + code](https://github.com/facebookresearch/XLM/blob/9e6f6814d17be4fe5b15f2e6c43eb2b2d76daeb4/src/model/transformer.py#L529). + + Reference for the diverse beam search algorithm and implementation [Ashwin Kalyan's DBS + implementation](https://github.com/ashwinkalyan/dbs/blob/master/dbs/beam_utils.lua) + + Args: + batch_size (`int`): + Batch Size of `input_ids` for which standard beam search decoding is run in parallel. + num_beams (`int`): + Number of beams for beam search. + device (`torch.device`): + Defines the device type (*e.g.*, `"cpu"` or `"cuda"`) on which this instance of `BeamSearchScorer` will be + allocated. + length_penalty (`float`, *optional*, defaults to 1.0): + Exponential penalty to the length that is used with beam-based generation. It is applied as an exponent to + the sequence length, which in turn is used to divide the score of the sequence. Since the score is the log + likelihood of the sequence (i.e. negative), `length_penalty` > 0.0 promotes longer sequences, while + `length_penalty` < 0.0 encourages shorter sequences. + do_early_stopping (`bool` or `str`, *optional*, defaults to `False`): + Controls the stopping condition for beam-based methods, like beam-search. It accepts the following values: + `True`, where the generation stops as soon as there are `num_beams` complete candidates; `False`, where an + heuristic is applied and the generation stops when is it very unlikely to find better candidates; + `"never"`, where the beam search procedure only stops when there cannot be better candidates (canonical + beam search algorithm). + num_beam_hyps_to_keep (`int`, *optional*, defaults to 1): + The number of beam hypotheses that shall be returned upon calling + [`~transformers.BeamSearchScorer.finalize`]. + num_beam_groups (`int`, *optional*, defaults to 1): + Number of groups to divide `num_beams` into in order to ensure diversity among different groups of beams. + See [this paper](https://arxiv.org/pdf/1610.02424.pdf) for more details. + max_length (`int`, *optional*): + The maximum length of the sequence to be generated. + """ + + def __init__( + self, + batch_size: int, + num_beams: int, + device: torch.device, + length_penalty: Optional[float] = 1.0, + do_early_stopping: Optional[Union[bool, str]] = False, + num_beam_hyps_to_keep: Optional[int] = 1, + num_beam_groups: Optional[int] = 1, + max_length: Optional[int] = None, + ): + self.num_beams = num_beams + self.device = device + self.length_penalty = length_penalty + self.do_early_stopping = do_early_stopping + self.num_beam_hyps_to_keep = num_beam_hyps_to_keep + self.num_beam_groups = num_beam_groups + self.group_size = self.num_beams // self.num_beam_groups + + self._is_init = False + # self._beam_hyps[i*self.num_beam_groups+j] is the beam_hyps of the j-th group in the i-th mini-batch. + # If group_beam_search is not used, the list consists of `batch_size` beam_hyps. + self._beam_hyps = [ + BeamHypotheses( + num_beams=self.group_size, + length_penalty=self.length_penalty, + early_stopping=self.do_early_stopping, + max_length=max_length, + ) + for _ in range(batch_size * self.num_beam_groups) + ] + # self._done[i*self.num_beam_groups+j] indicates whether the generation of the beam_hyps of the j-th group + # in the i-th mini-batch is complete. + self._done = torch.tensor( + [False for _ in range(batch_size * self.num_beam_groups)], dtype=torch.bool, device=self.device + ) + + if not isinstance(num_beams, int) or num_beams <= 1: + raise ValueError( + f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1," + " one should make use of `greedy_search` instead." + ) + + if not isinstance(num_beam_groups, int) or (num_beam_groups > num_beams) or (num_beams % num_beam_groups != 0): + raise ValueError( + "`num_beam_groups` has to be an integer smaller or equal than `num_beams` and `num_beams` has to be" + f" divisible by `num_beam_groups`, but is {num_beam_groups} with `num_beams` being {num_beams}." + ) + + @property + def is_done(self) -> bool: + return self._done.all() + + def process( + self, + input_ids: torch.LongTensor, + next_scores: torch.FloatTensor, + next_tokens: torch.LongTensor, + next_indices: torch.LongTensor, + pad_token_id: Optional[Union[int, torch.Tensor]] = None, + eos_token_id: Optional[Union[int, List[int], torch.Tensor]] = None, + beam_indices: Optional[torch.LongTensor] = None, + group_index: Optional[int] = 0, + decoder_prompt_len: Optional[int] = 0, + ) -> Dict[str, torch.Tensor]: + # add up to the length which the next_scores is calculated on (including decoder prompt) + cur_len = input_ids.shape[-1] + 1 + batch_size = len(self._beam_hyps) // self.num_beam_groups + + if not (batch_size == (input_ids.shape[0] // self.group_size)): + if self.num_beam_groups > 1: + raise ValueError( + f"A group beam size of {input_ids.shape[0]} is used as the input, but a group beam " + f"size of {self.group_size} is expected by the beam scorer." + ) + else: + raise ValueError( + f"A beam size of {input_ids.shape[0]} is used as the input, but a beam size of " + f"{self.group_size} is expected by the beam scorer." + ) + + device = input_ids.device + next_beam_scores = torch.zeros((batch_size, self.group_size), dtype=next_scores.dtype, device=device) + next_beam_tokens = torch.zeros((batch_size, self.group_size), dtype=next_tokens.dtype, device=device) + next_beam_indices = torch.zeros((batch_size, self.group_size), dtype=next_indices.dtype, device=device) + + if eos_token_id is not None and not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id) + + for batch_idx in range(batch_size): + batch_group_idx = batch_idx * self.num_beam_groups + group_index + if self._done[batch_group_idx]: + if self.num_beams < len(self._beam_hyps[batch_group_idx]): + raise ValueError(f"Batch can only be done if at least {self.num_beams} beams have been generated") + if eos_token_id is None or pad_token_id is None: + raise ValueError("Generated beams >= num_beams -> eos_token_id and pad_token have to be defined") + # pad the batch + next_beam_scores[batch_idx, :] = 0 + next_beam_tokens[batch_idx, :] = pad_token_id + next_beam_indices[batch_idx, :] = 0 + continue + + # next tokens for this sentence + beam_idx = 0 + for beam_token_rank, (next_token, next_score, next_index) in enumerate( + zip(next_tokens[batch_idx], next_scores[batch_idx], next_indices[batch_idx]) + ): + batch_beam_idx = batch_idx * self.group_size + next_index + # add to generated hypotheses if end of sentence + if (eos_token_id is not None) and (next_token.item() in eos_token_id): + # if beam_token does not belong to top num_beams tokens, it should not be added + is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size + if is_beam_token_worse_than_top_num_beams: + continue + if beam_indices is not None: + beam_index = beam_indices[batch_beam_idx] + beam_index = beam_index + (batch_beam_idx,) + else: + beam_index = None + + self._beam_hyps[batch_group_idx].add( + input_ids[batch_beam_idx].clone(), + next_score.item(), + beam_indices=beam_index, + generated_len=cur_len - decoder_prompt_len, + ) + else: + # add next predicted token since it is not eos_token + next_beam_scores[batch_idx, beam_idx] = next_score + next_beam_tokens[batch_idx, beam_idx] = next_token + next_beam_indices[batch_idx, beam_idx] = batch_beam_idx + beam_idx += 1 + + # once the beam for next step is full, don't add more tokens to it. + if beam_idx == self.group_size: + break + + if beam_idx < self.group_size: + raise ValueError( + f"At most {self.group_size} tokens in {next_tokens[batch_idx]} can be equal to `eos_token_id:" + f" {eos_token_id}`. Make sure {next_tokens[batch_idx]} are corrected." + ) + + # Check if we are done so that we can save a pad step if all(done) + self._done[batch_group_idx] = self._done[batch_group_idx] or self._beam_hyps[batch_group_idx].is_done( + next_scores[batch_idx].max().item(), cur_len, decoder_prompt_len + ) + + return UserDict( + { + "next_beam_scores": next_beam_scores.view(-1), + "next_beam_tokens": next_beam_tokens.view(-1), + "next_beam_indices": next_beam_indices.view(-1), + } + ) + + def finalize( + self, + input_ids: torch.LongTensor, + final_beam_scores: torch.FloatTensor, + final_beam_tokens: torch.LongTensor, + final_beam_indices: torch.LongTensor, + max_length: int, + pad_token_id: Optional[Union[int, torch.Tensor]] = None, + eos_token_id: Optional[Union[int, List[int], torch.Tensor]] = None, + beam_indices: Optional[torch.LongTensor] = None, + decoder_prompt_len: Optional[int] = 0, + ) -> Tuple[torch.LongTensor]: + batch_size = len(self._beam_hyps) // self.num_beam_groups + + if eos_token_id is not None and not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id) + + # finalize all open beam hypotheses and add to generated hypotheses + for batch_group_idx, beam_hyp in enumerate(self._beam_hyps): + if self._done[batch_group_idx]: + continue + + # all open beam hypotheses are added to the beam hypothesis + # beam hypothesis class automatically keeps the best beams + for index_per_group in range(self.group_size): + batch_beam_idx = batch_group_idx * self.group_size + index_per_group + final_score = final_beam_scores[batch_beam_idx].item() + final_tokens = input_ids[batch_beam_idx] + beam_index = beam_indices[batch_beam_idx] if beam_indices is not None else None + generated_len = final_tokens.shape[-1] - decoder_prompt_len + beam_hyp.add(final_tokens, final_score, beam_indices=beam_index, generated_len=generated_len) + + # select the best hypotheses + sent_lengths = input_ids.new(batch_size * self.num_beam_hyps_to_keep) + best = [] + best_indices = [] + best_scores = torch.zeros(batch_size * self.num_beam_hyps_to_keep, device=self.device, dtype=torch.float32) + + # retrieve best hypotheses + for i in range(batch_size): + beam_hyps_in_batch = self._beam_hyps[i * self.num_beam_groups : (i + 1) * self.num_beam_groups] + candidate_beams = [beam for beam_hyp in beam_hyps_in_batch for beam in beam_hyp.beams] + sorted_hyps = sorted(candidate_beams, key=lambda x: x[0]) + for j in range(self.num_beam_hyps_to_keep): + best_hyp_tuple = sorted_hyps.pop() + best_score = best_hyp_tuple[0] + best_hyp = best_hyp_tuple[1] + best_index = best_hyp_tuple[2] + sent_lengths[self.num_beam_hyps_to_keep * i + j] = len(best_hyp) + + # append hyp to lists + best.append(best_hyp) + + # append indices to list + best_indices.append(best_index) + + best_scores[i * self.num_beam_hyps_to_keep + j] = best_score + + # prepare for adding eos + sent_lengths_max = sent_lengths.max().item() + 1 + sent_max_len = min(sent_lengths_max, max_length) if max_length is not None else sent_lengths_max + decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len) + + if len(best_indices) > 0 and best_indices[0] is not None: + indices: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len) + else: + indices = None + + # shorter batches are padded if needed + if sent_lengths.min().item() != sent_lengths.max().item(): + if pad_token_id is None: + raise ValueError("`pad_token_id` has to be defined") + decoded.fill_(pad_token_id) + + if indices is not None: + indices.fill_(-1) + + # fill with hypotheses and eos_token_id if the latter fits in + for i, (hypo, best_idx) in enumerate(zip(best, best_indices)): + decoded[i, : sent_lengths[i]] = hypo + + if indices is not None: + indices[i, : len(best_idx)] = torch.tensor(best_idx) + + if sent_lengths[i] < sent_max_len: + # inserting only the first eos_token_id + decoded[i, sent_lengths[i]] = eos_token_id[0] + + return UserDict( + { + "sequences": decoded, + "sequence_scores": best_scores, + "beam_indices": indices, + } + ) + + +class ConstrainedBeamSearchScorer(BeamScorer): + r""" + [`BeamScorer`] implementing constrained beam search decoding. + + + Args: + batch_size (`int`): + Batch Size of `input_ids` for which standard beam search decoding is run in parallel. + num_beams (`int`): + Number of beams for beam search. + constraints (`List[Constraint]`): + A list of positive constraints represented as `Constraint` objects that must be fulfilled in the generation + output. For more information, the documentation of [`Constraint`] should be read. + device (`torch.device`): + Defines the device type (*e.g.*, `"cpu"` or `"cuda"`) on which this instance of `BeamSearchScorer` will be + allocated. + length_penalty (`float`, *optional*, defaults to 1.0): + Exponential penalty to the length that is used with beam-based generation. It is applied as an exponent to + the sequence length, which in turn is used to divide the score of the sequence. Since the score is the log + likelihood of the sequence (i.e. negative), `length_penalty` > 0.0 promotes longer sequences, while + `length_penalty` < 0.0 encourages shorter sequences. + do_early_stopping (`bool` or `str`, *optional*, defaults to `False`): + Controls the stopping condition for beam-based methods, like beam-search. It accepts the following values: + `True`, where the generation stops as soon as there are `num_beams` complete candidates; `False`, where an + heuristic is applied and the generation stops when is it very unlikely to find better candidates; + `"never"`, where the beam search procedure only stops when there cannot be better candidates (canonical + beam search algorithm). + num_beam_hyps_to_keep (`int`, *optional*, defaults to 1): + The number of beam hypotheses that shall be returned upon calling + [`~transformers.BeamSearchScorer.finalize`]. + num_beam_groups (`int`, *optional*, defaults to 1): + Number of groups to divide `num_beams` into in order to ensure diversity among different groups of beams. + See [this paper](https://arxiv.org/pdf/1610.02424.pdf) for more details. + max_length (`int`, *optional*): + The maximum length of the sequence to be generated. + """ + + def __init__( + self, + batch_size: int, + num_beams: int, + constraints: List[Constraint], + device: torch.device, + length_penalty: Optional[float] = 1.0, + do_early_stopping: Optional[Union[bool, str]] = False, + num_beam_hyps_to_keep: Optional[int] = 1, + num_beam_groups: Optional[int] = 1, + max_length: Optional[int] = None, + ): + self.num_beams = num_beams + self.device = device + self.length_penalty = length_penalty + self.do_early_stopping = do_early_stopping + self.num_beam_hyps_to_keep = num_beam_hyps_to_keep + self.num_beam_groups = num_beam_groups + self.group_size = self.num_beams // self.num_beam_groups + self.constraints = constraints + + self._is_init = False + self._beam_hyps = [ + BeamHypotheses( + num_beams=self.num_beams, + length_penalty=self.length_penalty, + early_stopping=self.do_early_stopping, + max_length=max_length, + ) + for _ in range(batch_size) + ] + self._done = torch.tensor([False for _ in range(batch_size)], dtype=torch.bool, device=self.device) + + if not isinstance(num_beams, int) or num_beams <= 1: + raise ValueError( + f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1," + " one should make use of `greedy_search` instead." + ) + + if not isinstance(num_beam_groups, int) or (num_beam_groups > num_beams) or (num_beams % num_beam_groups != 0): + raise ValueError( + "`num_beam_groups` has to be an integer smaller or equal than `num_beams` and `num_beams` has to be" + f" divisible by `num_beam_groups`, but is {num_beam_groups} with `num_beams` being {num_beams}." + ) + + @property + def is_done(self) -> bool: + return self._done.all() + + def make_constraint_states(self, n): + return [ConstraintListState([constraint.copy() for constraint in self.constraints]) for _ in range(n)] + + def check_completes_constraints(self, sequence): + new_state = self.make_constraint_states(1)[0] + new_state.reset(sequence) + return new_state.completed + + def process( + self, + input_ids: torch.LongTensor, + next_scores: torch.FloatTensor, + next_tokens: torch.LongTensor, + next_indices: torch.LongTensor, + scores_for_all_vocab: torch.FloatTensor, + pad_token_id: Optional[Union[int, torch.Tensor]] = None, + eos_token_id: Optional[Union[int, List[int], torch.Tensor]] = None, + beam_indices: Optional[torch.LongTensor] = None, + decoder_prompt_len: Optional[int] = 0, + ) -> Tuple[torch.Tensor]: + r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size * num_beams, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using any class inheriting from [`PreTrainedTokenizer`]. See + [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + next_scores (`torch.FloatTensor` of shape `(batch_size, 2 * num_beams)`): + Current scores of the top `2 * num_beams` non-finished beam hypotheses. + next_tokens (`torch.LongTensor` of shape `(batch_size, 2 * num_beams)`): + `input_ids` of the tokens corresponding to the top `2 * num_beams` non-finished beam hypotheses. + next_indices (`torch.LongTensor` of shape `(batch_size, 2 * num_beams)`): + Beam indices indicating to which beam hypothesis the `next_tokens` correspond. + scores_for_all_vocab (`torch.FloatTensor` of shape `(batch_size * num_beams, sequence_length)`): + The scores of all tokens in the vocabulary for each of the beam hypotheses. + pad_token_id (`int`, *optional*): + The id of the *padding* token. + eos_token_id (`Union[int, List[int]]`, *optional*): + The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens. + beam_indices (`torch.LongTensor`, *optional*): + Beam indices indicating to which beam hypothesis each token correspond. + decoder_prompt_len (`int`, *optional*): + The length of prompt that is included in the input to decoder. + Return: + `UserDict`: A dictionary composed of the fields as defined above: + + - **next_beam_scores** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Updated scores of + all + non-finished beams. + + - **next_beam_tokens** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Next tokens to be + added + to the non-finished beam_hypotheses. + - **next_beam_indices** (`torch.FloatTensor` of shape `(batch_size * num_beams)`) -- Beam indices + indicating to which beam the next tokens shall be added. + """ + + # add up to the length which the next_scores is calculated on (including decoder prompt) + cur_len = input_ids.shape[-1] + 1 + batch_size = len(self._beam_hyps) + if not (batch_size == (input_ids.shape[0] // self.group_size)): + if self.num_beam_groups > 1: + raise ValueError( + f"A group beam size of {input_ids.shape[0]} is used as the input, but a group beam " + f"size of {self.group_size} is expected by the beam scorer." + ) + else: + raise ValueError( + f"A beam size of {input_ids.shape[0]} is used as the input, but a beam size of " + f"{self.group_size} is expected by the beam scorer." + ) + + device = input_ids.device + + next_beam_scores = torch.zeros((batch_size, self.group_size), dtype=next_scores.dtype, device=device) + next_beam_tokens = torch.zeros((batch_size, self.group_size), dtype=next_tokens.dtype, device=device) + next_beam_indices = torch.zeros((batch_size, self.group_size), dtype=next_indices.dtype, device=device) + + if eos_token_id is not None and not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id) + + for batch_idx, beam_hyp in enumerate(self._beam_hyps): + if self._done[batch_idx]: + if self.num_beams < len(beam_hyp): + raise ValueError(f"Batch can only be done if at least {self.num_beams} beams have been generated") + if eos_token_id is None or pad_token_id is None: + raise ValueError("Generated beams >= num_beams -> eos_token_id and pad_token have to be defined") + # pad the batch + next_beam_scores[batch_idx, :] = 0 + next_beam_tokens[batch_idx, :] = pad_token_id + next_beam_indices[batch_idx, :] = 0 + continue + + # next tokens for this sentence. + beam_idx = 0 + for beam_token_rank, (next_token, next_score, next_index) in enumerate( + zip(next_tokens[batch_idx], next_scores[batch_idx], next_indices[batch_idx]) + ): + batch_beam_idx = batch_idx * self.group_size + next_index + # add to generated hypotheses if end of sentence + if (eos_token_id is not None) and (next_token.item() in eos_token_id): + # if beam_token does not belong to top num_beams tokens, it should not be added + is_beam_token_worse_than_top_num_beams = beam_token_rank >= self.group_size + if is_beam_token_worse_than_top_num_beams: + continue + + completes_constraint = self.check_completes_constraints(input_ids[batch_beam_idx].tolist()) + if completes_constraint: + if beam_indices is not None: + beam_index = beam_indices[batch_beam_idx] + beam_index = beam_index + (batch_beam_idx,) + else: + beam_index = None + + beam_hyp.add( + input_ids[batch_beam_idx].clone(), + next_score.item(), + beam_indices=beam_index, + generated_len=cur_len - decoder_prompt_len, + ) + else: + # add next predicted token since it is not eos_token + next_beam_scores[batch_idx, beam_idx] = next_score + next_beam_tokens[batch_idx, beam_idx] = next_token + next_beam_indices[batch_idx, beam_idx] = batch_beam_idx + beam_idx += 1 + + # once the beam for next step is full, don't add more tokens to it. + if beam_idx == self.group_size: + break + + new_scores, new_tokens, new_indices = self.step_sentence_constraint( + batch_idx, + input_ids, + scores_for_all_vocab, + next_beam_scores[batch_idx], + next_beam_tokens[batch_idx], + next_beam_indices[batch_idx], + ) + + next_beam_scores[batch_idx] = new_scores + next_beam_tokens[batch_idx] = new_tokens + next_beam_indices[batch_idx] = new_indices + + if beam_idx < self.group_size: + raise ValueError( + f"At most {self.group_size} tokens in {next_tokens[batch_idx]} can be equal to `eos_token_id:" + f" {eos_token_id}`. Make sure {next_tokens[batch_idx]} are corrected." + ) + + # Check if we are done so that we can save a pad step if all(done) + self._done[batch_idx] = self._done[batch_idx] or beam_hyp.is_done( + next_scores[batch_idx].max().item(), cur_len, decoder_prompt_len + ) + + return UserDict( + { + "next_beam_scores": next_beam_scores.view(-1), + "next_beam_tokens": next_beam_tokens.view(-1), + "next_beam_indices": next_beam_indices.view(-1), + } + ) + + def step_sentence_constraint( + self, + batch_idx: int, + input_ids: torch.LongTensor, + vocab_scores: torch.FloatTensor, + sent_beam_scores: torch.FloatTensor, + sent_beam_tokens: torch.LongTensor, + sent_beam_indices: torch.LongTensor, + push_progress: bool = False, + ): + # sent_beam_tokens are the next {num_beams} number of tokens that are under consideration for this beam + # (candidate next tokens) + + # 1. Adding "advance_tokens" + # using ConstraintStateList.advance(), we propose new tokens to be added into this "candidate list" that will + # advance us in fulfilling the constraints. + + # 2. Selecting best candidates such that we end up with highest probable candidates + # that fulfill our constraints. + + orig_len = sent_beam_indices.size(0) + device = sent_beam_indices.device + + # initialize states + topk_contraint_states = self.make_constraint_states(orig_len) + advance_constraint_states = self.make_constraint_states(orig_len) + + sidx, eidx = batch_idx * orig_len, (batch_idx + 1) * orig_len + this_batch_input_ids = input_ids[sidx:eidx] + this_batch_token_scores = vocab_scores[sidx:eidx] + full_hypotheses = torch.cat((input_ids[sent_beam_indices], sent_beam_tokens.unsqueeze(-1)), dim=-1) + + # need to make new hypothesis that advance the constraints + track_new = { + "new_seqs": full_hypotheses.tolist(), + "new_states": [], + "new_indices": [], + "new_tokens": [], + "new_scores": [], + } + for seq_idx, pre_seq in enumerate(this_batch_input_ids): + # pre_seq = ith sequence generated before this step. + + # input_ids -> (topk) generic beam search best model next tokens + # -> (advance) constraints forcing the next token + # either way, we need to sort them into "banks" later, so store a "ConstraintListState" for all types of + # hypotheses. + + topk_state = topk_contraint_states[seq_idx] + topk_state.reset(full_hypotheses[seq_idx].tolist()) + + advance_state = advance_constraint_states[seq_idx] + advance_state.reset(pre_seq.tolist()) + + if not advance_state.completed: + advance_tokens = torch.tensor(advance_state.advance(), dtype=torch.long, device=device) + for advance_token in advance_tokens: + # since adding each `advance_token` leads to a different hypothesis, create new state instance. + new_state = advance_state.copy(stateful=True) + new_state.add(advance_token.tolist()) + + advance_seq = torch.cat((pre_seq, advance_token.unsqueeze(0)), -1).tolist() + if advance_seq not in track_new["new_seqs"]: + # prevent duplicates, which are basically bound to happen in this process. + track_new["new_seqs"].append(advance_seq) + track_new["new_indices"].append(sidx + seq_idx) # idx -> global idx across all the batches + track_new["new_tokens"].append(advance_token) + track_new["new_scores"].append(this_batch_token_scores[seq_idx].take(advance_token)) + track_new["new_states"].append(new_state) + elif push_progress: + # Basically, `sent_beam_indices` often chooses very little among `input_ids` the generated sequences that + # actually fulfill our constraints. For example, let constraints == ["loves pies"] and + + # pre_seq_1 = "The child loves pies and" pre_seq_2 = "The child plays in the playground and" + + # Without this step, if `sent_beam_indices` is something like [1,1], then + # 1. `pre_seq_1` won't be added to the list of (topk) hypothesis since it's not in the indices and + # 2. it won't be added to the list of (advance) hypothesis since it's completed already. (this is + # the else part of `if constraints_completed[seq_idx]`) + # 3. it ends up simply getting removed from consideration. + + # #3 might be fine and actually desired, since it's likely that it's a low-probability output anyways, + # especially if it's not in the list of `sent_beam_indices`. But this often leads to lengthened beam + # search times, since completed sequences keep getting removed after all this effort for constrained + # generation. + + # Here, we basically take `pre_seq_1` and to "push" it into the considered list of hypotheses, by simply + # appending the next likely token in the vocabulary and adding it to the list of hypotheses. + + new_score, new_token = torch.max(this_batch_token_scores[seq_idx], 0) # some next probable token + advance_seq = torch.cat((pre_seq, new_token.unsqueeze(0)), -1) + + advance_state = advance_constraint_states[seq_idx] + + advance_seq = advance_seq.tolist() + + advance_state.reset(advance_seq) + if advance_seq not in track_new["new_seqs"]: + # but still don't want to have duplicates + track_new["new_seqs"].append(advance_seq) + track_new["new_indices"].append(seq_idx) + track_new["new_tokens"].append(new_token) + track_new["new_scores"].append(new_score) + track_new["new_states"].append(advance_state) + + if len(track_new["new_indices"]) > 0: + new_indices = torch.tensor(track_new["new_indices"], device=device) + new_tokens = torch.stack(track_new["new_tokens"]).to(device) + new_scores = torch.stack(track_new["new_scores"]).to(device) + + all_states = topk_contraint_states + track_new["new_states"] + all_tokens = torch.cat((sent_beam_tokens, new_tokens), -1) + all_scores = torch.cat((sent_beam_scores, new_scores), -1) + all_banks = torch.tensor([one.get_bank() for one in all_states], device=device) + + zipped = all_banks * 100 + all_scores + indices = zipped.sort(descending=True).indices + sorted_banks = all_banks[indices] + + # Then we end up with {sorted among bank C}, {sorted among bank C-1}, ..., {sorted among bank 0} + + counter = -1 + cur_bank = sorted_banks[0] + increments = [] + for bank in sorted_banks: + if bank == cur_bank: + counter += 1 + else: + counter = 0 + cur_bank = bank + increments.append(counter) + rearrangers = torch.tensor(np.argsort(increments, kind="mergesort")) + + indices = indices[rearrangers][:orig_len] + + sent_beam_scores = all_scores[indices] + sent_beam_tokens = all_tokens[indices] + sent_beam_indices = torch.cat((sent_beam_indices, new_indices))[indices] + + return sent_beam_scores, sent_beam_tokens, sent_beam_indices + + def finalize( + self, + input_ids: torch.LongTensor, + final_beam_scores: torch.FloatTensor, + final_beam_tokens: torch.LongTensor, + final_beam_indices: torch.LongTensor, + max_length: int, + pad_token_id: Optional[Union[int, torch.Tensor]] = None, + eos_token_id: Optional[Union[int, List[int], torch.Tensor]] = None, + beam_indices: Optional[torch.LongTensor] = None, + decoder_prompt_len: Optional[int] = 0, + ) -> Tuple[torch.LongTensor]: + batch_size = len(self._beam_hyps) + + if eos_token_id is not None and not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id) + + # finalize all open beam hypotheses and add to generated hypotheses + for batch_idx, beam_hyp in enumerate(self._beam_hyps): + if self._done[batch_idx]: + continue + + # all open beam hypotheses are added to the beam hypothesis + # beam hypothesis class automatically keeps the best beams + + ids_collect = [] + for beam_id in range(self.num_beams): + batch_beam_idx = batch_idx * self.num_beams + beam_id + final_score = final_beam_scores[batch_beam_idx].item() + final_tokens = input_ids[batch_beam_idx] + + completes_constraint = self.check_completes_constraints(final_tokens.tolist()) + if completes_constraint: + beam_index = beam_indices[batch_beam_idx] if beam_indices is not None else None + generated_len = final_tokens.shape[-1] - decoder_prompt_len + beam_hyp.add(final_tokens, final_score, beam_indices=beam_index, generated_len=generated_len) + ids_collect.append(beam_id) + + # due to overly complex constraints or other factors, sometimes we can't gaurantee a successful + # generation. In these cases we simply return the highest scoring outputs. + if len(ids_collect) < self.num_beam_hyps_to_keep: + for beam_id in range(self.num_beams): + if beam_id not in ids_collect: + batch_beam_idx = batch_idx * self.num_beams + beam_id + final_score = final_beam_scores[batch_beam_idx].item() + final_tokens = input_ids[batch_beam_idx] + generated_len = final_tokens.shape[-1] - decoder_prompt_len + beam_hyp.add(final_tokens, final_score, generated_len=generated_len) + if len(ids_collect) >= self.num_beam_hyps_to_keep: + break + + # select the best hypotheses + sent_lengths = input_ids.new(batch_size * self.num_beam_hyps_to_keep) + best = [] + best_indices = [] + best_scores = torch.zeros(batch_size * self.num_beam_hyps_to_keep, device=self.device, dtype=torch.float32) + + # retrieve best hypotheses + for i, beam_hyp in enumerate(self._beam_hyps): + sorted_hyps = sorted(beam_hyp.beams, key=lambda x: x[0]) + for j in range(self.num_beam_hyps_to_keep): + best_hyp_tuple = sorted_hyps.pop() + best_score = best_hyp_tuple[0] + best_hyp = best_hyp_tuple[1] + best_index = best_hyp_tuple[2] + sent_lengths[self.num_beam_hyps_to_keep * i + j] = len(best_hyp) + + # append to lists + best.append(best_hyp) + + # append indices to list + best_indices.append(best_index) + + best_scores[i * self.num_beam_hyps_to_keep + j] = best_score + + # prepare for adding eos + sent_lengths_max = sent_lengths.max().item() + 1 + + sent_max_len = min(sent_lengths_max, max_length) if max_length is not None else sent_lengths_max + decoded: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len) + + if len(best_indices) > 0 and best_indices[0] is not None: + indices: torch.LongTensor = input_ids.new(batch_size * self.num_beam_hyps_to_keep, sent_max_len) + else: + indices = None + + # shorter batches are padded if needed + if sent_lengths.min().item() != sent_lengths.max().item(): + if pad_token_id is None: + raise ValueError("`pad_token_id` has to be defined") + decoded.fill_(pad_token_id) + + if indices is not None: + indices.fill_(-1) + + # fill with hypotheses and eos_token_id if the latter fits in + for i, (hypo, best_idx) in enumerate(zip(best, best_indices)): + decoded[i, : sent_lengths[i]] = hypo + + if indices is not None: + indices[i, : len(best_idx)] = torch.tensor(best_idx) + + if sent_lengths[i] < sent_max_len: + # inserting only the first eos_token_id + decoded[i, sent_lengths[i]] = eos_token_id[0] + + return UserDict( + { + "sequences": decoded, + "sequence_scores": best_scores, + "beam_indices": indices, + } + ) + + +class BeamHypotheses: + def __init__(self, num_beams: int, length_penalty: float, early_stopping: bool, max_length: Optional[int] = None): + """ + Initialize n-best list of hypotheses. + """ + self.length_penalty = length_penalty + self.early_stopping = early_stopping + self.max_length = max_length + self.num_beams = num_beams + self.beams = [] + self.worst_score = 1e9 + + if not isinstance(self.early_stopping, bool) and self.max_length is None: + raise ValueError( + "When `do_early_stopping` is set to a string, `max_length` must be defined. Ensure it is passed to the" + " BeamScorer class instance at initialization time." + ) + + def __len__(self): + """ + Number of hypotheses in the list. + """ + return len(self.beams) + + def add( + self, + hyp: torch.LongTensor, + sum_logprobs: float, + beam_indices: Optional[torch.LongTensor] = None, + generated_len: Optional[int] = None, + ): + """ + Add a new hypothesis to the list. + """ + if generated_len is not None: + score = sum_logprobs / (generated_len**self.length_penalty) + # This 'else' case exists for retrocompatibility + else: + score = sum_logprobs / (hyp.shape[-1] ** self.length_penalty) + + if len(self) < self.num_beams or score > self.worst_score: + self.beams.append((score, hyp, beam_indices)) + if len(self) > self.num_beams: + sorted_next_scores = sorted([(s, idx) for idx, (s, _, _) in enumerate(self.beams)]) + del self.beams[sorted_next_scores[0][1]] + self.worst_score = sorted_next_scores[1][0] + else: + self.worst_score = min(score, self.worst_score) + + def is_done(self, best_sum_logprobs: float, cur_len: int, decoder_prompt_len: Optional[int] = 0) -> bool: + """ + If there are enough hypotheses and that none of the hypotheses being generated can become better than the worst + one in the heap, then we are done with this sentence. + """ + + if len(self) < self.num_beams: + return False + + # `True`: stop as soon as at least `num_beams` hypotheses are finished + if self.early_stopping is True: + return True + # `False`: heuristic -- compute best possible score from `cur_len`, even though it is not entirely accurate + # when `length_penalty` is positive. See the discussion below for more details. + # https://github.com/huggingface/transformers/pull/20901#issuecomment-1369845565 + elif self.early_stopping is False: + highest_attainable_score = best_sum_logprobs / (cur_len - decoder_prompt_len) ** self.length_penalty + ret = self.worst_score >= highest_attainable_score + return ret + # `"never"`: compute the best possible score, depending on the signal of `length_penalty` + else: + # `length_penalty` > 0.0 -> max denominator is obtaned from `max_length`, not from `cur_len` -> min + # abs(`highest_attainable_score`) is obtained -> `highest_attainable_score` is negative, hence we obtain + # its max this way + if self.length_penalty > 0.0: + if self.max_length <= decoder_prompt_len: + raise ValueError("max_length is not larger than decoder prompt length") + highest_attainable_score = ( + best_sum_logprobs / (self.max_length - decoder_prompt_len) ** self.length_penalty + ) + # the opposite logic applies here (max `highest_attainable_score` from `cur_len`) + else: + highest_attainable_score = best_sum_logprobs / (cur_len - decoder_prompt_len) ** self.length_penalty + ret = self.worst_score >= highest_attainable_score + return ret diff --git a/lib/python3.12/site-packages/transformers/generation/candidate_generator.py b/lib/python3.12/site-packages/transformers/generation/candidate_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..fe57f532e68750af1e38bb0387dabf594b3d3b82 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/candidate_generator.py @@ -0,0 +1,1160 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. +# +# 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. + +import copy +import weakref +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple + +import numpy as np +import torch + +from ..utils import is_sklearn_available + + +if is_sklearn_available(): + from sklearn.metrics import roc_curve + +from ..cache_utils import DynamicCache +from ..pytorch_utils import isin_mps_friendly +from .logits_process import LogitsProcessorList, MinLengthLogitsProcessor, SuppressTokensLogitsProcessor + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + from ..tokenization_utils_base import PreTrainedTokenizerBase + from .configuration_utils import GenerationConfig + + +class CandidateGenerator: + """Abstract base class for all candidate generators that can be applied during assisted generation.""" + + def get_candidates(self, input_ids: torch.LongTensor) -> Tuple[torch.LongTensor, Optional[torch.FloatTensor]]: + """ + Fetches the candidates to be tried for the current input. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + + Return: + `torch.LongTensor` of shape `(batch_size, candidate_length)` containing the candidate sequences to be + assessed by the model and, optionally, a `torch.FloatTensor` of shape `(batch_size, candidate_length, + vocabulary_size)` containing the logits associated to each candidate. + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can call `get_candidates`." + ) + + def update_candidate_strategy(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, num_matches: int): + """ + Updates the candidate generation strategy based on the outcomes. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + scores (`torch.FloatTensor` of shape `(batch_size, candidate_length, config.vocab_size)`): + Prediction scores of a language modeling head. These can be logits for each vocabulary when not using + beam search or log softmax for each vocabulary token when using beam search + num_matches (`int`): + The number of matches between the candidate sequences and the model predictions. + """ + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can call " + "`update_candidate_strategy`." + ) + + +class AssistedCandidateGenerator(CandidateGenerator): + """ + `CandidateGenerator` class to be used for assisted generation and speculative decoding. This class generates + candidates through the use of a smaller model. Read the following blog post for more information: + https://huggingface.co/blog/assisted-generation + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + assistant_model (`PreTrainedModel`): + The model to be used for generating candidates. This model should be smaller than the main model. + generation_config (`~generation.GenerationConfig`, *optional*): + The generation configuration to be used as base parametrization for the generation call. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + model_kwargs (`Dict`): + The keyword arguments that will be passed to the main model, and are used as base inputs for the assistant + model as well. + inputs_tensor (`torch.Tensor`, *optional*): + The model input tensor. In encoder-decoder models, this is the encoder input. + """ + + def __init__( + self, + input_ids: torch.LongTensor, + assistant_model: "PreTrainedModel", + generation_config: "GenerationConfig", + model_kwargs: Dict, + inputs_tensor: Optional[torch.Tensor] = None, + logits_processor: "LogitsProcessorList" = None, + ): + # Make sure all data at the same device as assistant model + device = assistant_model.device + input_ids = input_ids.to(device) + if inputs_tensor is not None: + inputs_tensor = inputs_tensor.to(device) + + # Prepare the assistant and the starting number of candidate tokens + self.assistant_model = assistant_model + self.num_assistant_tokens = assistant_model.generation_config.num_assistant_tokens + self.assistant_confidence_threshold = assistant_model.generation_config.assistant_confidence_threshold + + # Set eos in assistant same as in target model + self.assistant_model.generation_config.eos_token_id = generation_config.eos_token_id + + # Prepare the kwargs for the assistant model + assistant_kwargs = {} + for key, value in model_kwargs.items(): # deepcopy crashes if we attempt to copy encoder outputs with grads + if key not in ("encoder_outputs", "past_key_values"): + assistant_kwargs[key] = ( + value.detach().to(device) if isinstance(value, torch.Tensor) else copy.deepcopy(value) + ) + + # Remove potential default "logits_to_keep" key + if "logits_to_keep" in assistant_kwargs.keys() and not assistant_model._supports_logits_to_keep(): + del assistant_kwargs["logits_to_keep"] + + # If the assistant is an encoder-decoder model, assume the encoder is different on the assistant. + if assistant_model.config.is_encoder_decoder: + inputs_tensor, model_input_name, assistant_kwargs = assistant_model._prepare_model_inputs( + inputs_tensor, assistant_model.generation_config.bos_token_id, assistant_kwargs + ) + assistant_kwargs = assistant_model._prepare_encoder_decoder_kwargs_for_generation( + inputs_tensor, assistant_kwargs, model_input_name, assistant_model.generation_config + ) + elif "encoder_outputs" in model_kwargs: + assistant_kwargs["encoder_outputs"] = model_kwargs["encoder_outputs"] + self.assistant_kwargs = assistant_kwargs + + # Prepare assistant model's keys of inputs + if assistant_model.config.is_encoder_decoder: + # both are encoder-decoder + self.input_ids_key = "decoder_input_ids" + elif "encoder_outputs" in assistant_kwargs: + # special case for encoder-decoder with decoder-only assistant (like DistilWhisper) + self.input_ids_key = "input_ids" + self.assistant_kwargs["attention_mask"] = self.assistant_kwargs.get( + "decoder_attention_mask", + torch.ones((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.long), + ) + else: + # both are decoder-only + self.input_ids_key = "input_ids" + + # Prepare generation-related options. + self.logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() + self.generation_config = copy.deepcopy(generation_config) + + self.generation_config.return_dict_in_generate = True + self.generation_config.output_scores = True + self.generation_config.assistant_confidence_threshold = self.assistant_confidence_threshold + # this flag allow us set the confidence stopping criteria for assistant model generation. + self.generation_config.is_assistant = True + + # avoid unnecessary warnings that min_length is larger than max_new_tokens + # remove the `MinLengthLogitsProcessor` if exists (NOTE: no need to check for `MinNewTokensLogitsProcessor`) + self.main_model_min_length = self.generation_config.min_length + self.generation_config.min_length = 0 + self.generation_config.min_new_tokens = None + for processor in self.logits_processor: + if isinstance(processor, MinLengthLogitsProcessor): + raise ValueError( + "Passing `MinLengthLogitsProcessor` when using `assisted_generation is disabled. " + "Please pass in `min_length` into `.generate()` instead" + ) + + # We need to roll back the cache in assisted generation, only DynamicCache is supported + self.generation_config.cache_implementation = None + + if ( + is_sklearn_available() + and self.assistant_model.generation_config.assistant_confidence_threshold + and type(self) is AssistedCandidateGenerator + ): + self.probs = [] + self.matches = [] + + def get_candidates(self, input_ids: torch.LongTensor) -> Tuple[torch.LongTensor, Optional[torch.FloatTensor]]: + """ + Fetches the candidates to be tried for the current input. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + + Return: + `torch.LongTensor` of shape `(batch_size, candidate_length)` containing the candidate sequences to be + assessed by the model and a `torch.FloatTensor` of shape `(batch_size, candidate_length, + vocabulary_size)` containing the logits associated to each candidate. + """ + input_ids = input_ids.to(self.assistant_model.device) + # Calculate new tokens to generate + min_new_tokens, max_new_tokens = self._calculate_new_tokens(input_ids) + if max_new_tokens == 0: + return input_ids, None + # Update past key values and masks + self._update_past_and_masks(input_ids) + # Generate candidates + generation_args = self._prepare_generation_args(input_ids, min_new_tokens, max_new_tokens) + candidate_ids, candidate_logits = self._generate_candidates(generation_args) + return candidate_ids, candidate_logits + + def update_candidate_strategy(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, num_matches: int): + """ + Updates the candidate generation strategy based on the outcomes. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + scores (`torch.FloatTensor` of shape `(batch_size, candidate_length, config.vocab_size)`): + Prediction scores of a language modeling head. These can be logits for each vocabulary when not using + beam search or log softmax for each vocabulary token when using beam search + num_matches (`int`): + The number of matches between the candidate sequences and the model predictions. + """ + # Adjust the max number of assistant tokens to use in the next iteration. This is a simple heuristic, + # probably can be improved -- we want to balance the benefits of getting assistant tokens correct with the + # cost of forecasting incorrect assistant tokens. + if self.assistant_model.generation_config.num_assistant_tokens_schedule in { + "heuristic", + "heuristic_transient", + }: + # len(scores[0])-1 is the number of candidates according to the target tokenizer. + if num_matches == len(scores[0]) - 1: + self.num_assistant_tokens += 2.0 + else: + self.num_assistant_tokens = max(1.0, self.num_assistant_tokens - 1.0) + + # The assistant's confidence threshold is adjusted throughout the speculative iterations to reduce the number of unnecessary draft and target forward passes. The costs are estimated based on the ROC curve, which considers the probability of the draft token and its match with the target. A cost of 25% is assigned to false positives and 75% to false negatives. + # This adaptation is not compatible with UAG, as it relies on the number of matched tokens based on the draft vocabulary, which is unavailable in UAG. + if ( + is_sklearn_available() + and self.assistant_model.generation_config.assistant_confidence_threshold + and type(self) is AssistedCandidateGenerator + ): + # update self.matches + self.matches.extend([1] * num_matches) + if len(self.probs) > len(self.matches): + self.matches.append(0) + + # update self.probs + excess_length = len(self.probs) - len(self.matches) + if excess_length > 0: + del self.probs[-excess_length:] + + if ( + len(self.probs) > 5 and {0, 1}.issubset(self.matches) + ): # require at least 5 samples to calculate the ROC curve and at least one positive and one negative sample + fpr, tpr, thresholds = roc_curve(self.matches, self.probs) + fnr = 1 - tpr + + # Calculate the cost for each threshold + costs = fpr + 3 * fnr + + # Find the threshold that minimizes the cost + optimal_threshold_index = np.argmin(costs) + best_threshold = thresholds[optimal_threshold_index] + + self.assistant_model.generation_config.assistant_confidence_threshold = best_threshold + + def _calculate_new_tokens(self, input_ids: torch.LongTensor) -> Tuple[int, int]: + """Calculate the minimum and maximum number of new tokens to generate.""" + new_cur_len = input_ids.shape[-1] + max_new_tokens = min(int(self.num_assistant_tokens), self.generation_config.max_length - new_cur_len - 1) + min_new_tokens = max(min(max_new_tokens, self.main_model_min_length - new_cur_len), 0) + return min_new_tokens, max_new_tokens + + def _update_past_and_masks( + self, input_ids: torch.LongTensor, remove_from_pkv: int = 0, num_added_tokens: int = 1 + ) -> bool: + """Update past key values and attention masks for subsequent generation rounds.""" + has_past_key_values = self.assistant_kwargs.get("past_key_values", None) is not None + if has_past_key_values: + new_cache_size = input_ids.shape[-1] - 1 - remove_from_pkv + self.assistant_kwargs["past_key_values"] = _crop_past_key_values( + self.assistant_model, self.assistant_kwargs["past_key_values"], new_cache_size - num_added_tokens + ) + self.assistant_kwargs = _prepare_attention_mask( + self.assistant_kwargs, input_ids.shape[-1], self.assistant_model.config.is_encoder_decoder + ) + self.assistant_kwargs = _prepare_token_type_ids(self.assistant_kwargs, input_ids.shape[-1]) + + return has_past_key_values + + def _prepare_generation_args(self, input_ids: torch.LongTensor, min_new_tokens: int, max_new_tokens: int) -> Dict: + """Prepare arguments for the generation call.""" + return { + self.input_ids_key: input_ids, + "min_new_tokens": min_new_tokens, + "max_new_tokens": max_new_tokens, + "generation_config": self.generation_config, + "logits_processor": self.logits_processor, + } + + def _generate_candidates(self, generation_args: Dict) -> Tuple[torch.LongTensor, Optional[torch.FloatTensor]]: + """Generate candidate sequences using the assistant model.""" + assistant_output = self.assistant_model.generate(**generation_args, **self.assistant_kwargs) + self.assistant_kwargs["past_key_values"] = assistant_output.past_key_values + if ( + is_sklearn_available() + and self.assistant_model.generation_config.assistant_confidence_threshold + and type(self) is AssistedCandidateGenerator + ): + scores_tensor = torch.cat(assistant_output.scores, dim=0) + scores_softmax = torch.softmax(scores_tensor, dim=-1) + ids = assistant_output.sequences[-1, -len(assistant_output.scores) :] + p = scores_softmax[range(len(ids)), ids] + self.probs.extend(p.tolist()) + candidate_logits = torch.stack(assistant_output.scores, dim=1) + candidate_ids = assistant_output.sequences + return candidate_ids, candidate_logits + + +class AssistedCandidateGeneratorDifferentTokenizers(AssistedCandidateGenerator): + """ + `CandidateGenerator` class to be used for Universal Assisted Generation (UAD): assisted generation with different tokenizers + for the assistant and main models. This class generates candidates through the use of a smaller + model. + + The main model input tokens are re-encoded into assistant model tokens, then candidate tokens are generated in the assistant encoding, which are + in turn re-encoded into main model candidate tokens. Validation then proceeds as explained above. + The re-encoding steps involve decoding token ids into text and then encoding the text using a different tokenizer. + Since re-encoding the tokens may result in tokenization discrepancies, UAD finds the longest common subsequence between the source and target encodings, + to ensure the new tokens include the correct prompt suffix. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + assistant_model (`PreTrainedModel`): + The model to be used for generating candidates. This model should be smaller than the main model. + target_tokenizer (`PreTrainedTokenizerBase`): + The tokenizer used for the target model. + assistant_tokenizer (`PreTrainedTokenizerBase`): + The tokenizer used for the assistant model. + generation_config (`~generation.GenerationConfig`, *optional*): + The generation configuration to be used as base parametrization for the generation call. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + model_kwargs (`Dict`): + The keyword arguments that will be passed to the main model, and are used as base inputs for the assistant + model as well. + inputs_tensor (`torch.Tensor`, *optional*): + The model input tensor. In encoder-decoder models, this is the encoder input. + """ + + def __init__( + self, + input_ids: torch.LongTensor, + assistant_model: "PreTrainedModel", + target_tokenizer: "PreTrainedTokenizerBase", + assistant_tokenizer: "PreTrainedTokenizerBase", + generation_config: "GenerationConfig", + model_kwargs: Dict, + inputs_tensor: Optional[torch.Tensor] = None, + logits_processor: "LogitsProcessorList" = None, + ): + super().__init__(input_ids, assistant_model, generation_config, model_kwargs, inputs_tensor, logits_processor) + + self.target_tokenizer = target_tokenizer + self.assistant_tokenizer = assistant_tokenizer + self.prev_target_ids_len: Optional[int] = None + self.prev_assistant_ids = None + self.target_lookbehind = assistant_model.generation_config.target_lookbehind + self.assistant_lookbehind = assistant_model.generation_config.assistant_lookbehind + + @staticmethod + def _get_longest_diag_dict(input_matrix, nonzero_idx): + """ + Calculates the length of the longest diagonal sequence in a given matrix. + Args: + input_matrix (torch.Tensor): The input matrix. + nonzero_idx (torch.Tensor): The indices of the non-zero elements in the matrix. + Returns: + dict: A dictionary where the keys are the indices of the non-zero elements and the values are the lengths of the longest diagonal sequences starting from those indices. + """ + + visited = set() + diags = {} + for idx in nonzero_idx: + start_idx = torch.clone(idx) + tuple_start_idx = tuple(start_idx.tolist()) + + if tuple_start_idx in visited: + continue + + visited.add(tuple_start_idx) + cur_diag_len = 1 + start_idx += 1 + while start_idx[0] < input_matrix.shape[0] and start_idx[1] < input_matrix.shape[1]: + tuple_start_idx = tuple(start_idx.tolist()) + visited.add(tuple_start_idx) + + if input_matrix[start_idx[0], start_idx[1]] == 1: + cur_diag_len += 1 + start_idx += 1 + else: + break + + diags[idx] = cur_diag_len + return diags + + @staticmethod + def _get_longest_diag_index(input_matrix): + """ + Returns the start index and length of the longest diagonal in the given input. + Args: + input_matrix (numpy.ndarray): The input matrix. + Returns: + tuple: A tuple containing the start index and length of the longest diagonal. + """ + + diags = AssistedCandidateGeneratorDifferentTokenizers._get_longest_diag_dict( + input_matrix, input_matrix.nonzero() + ) + diags_values = list(diags.values()) + diags_keys = list(diags.keys()) + best_diag = np.argmax(diags_values) + diag_start_index = diags_keys[best_diag] + diag_start_length = diags_values[best_diag] + return diag_start_index, diag_start_length + + @staticmethod + def _get_tokens_diag(prompt, prompt_plus_new_tokens): + """ + Input: + prompt: 2D array of shape (batch_size, prompt_length), represents the original prompt tokens + prompt_plus_new_tokens: 2D array of shape (batch_size, prompt_length), represents the suffix of the original prompt, with additional new tokens. + Output: + discrepancy_length: int, represents the number of tokens that need to be replaced from prompt + new_tokens_only: 2D array of shape (batch_size, new_token_length), represents the new tokens that are not in prompt + discrepancy_only: 2D array of shape (batch_size, discrepancy_length), represents the new tokens that are in prompt but not in prompt_plus_new_tokens + """ + compare_mat = prompt_plus_new_tokens.T == prompt + if not torch.is_tensor(compare_mat): + compare_mat = torch.tensor(compare_mat) + + compare_mat_int = compare_mat.to(int) + + if not compare_mat_int.any().item(): + # empty intersection between prompt and prompt_plus_new_tokens + return None, None, None + + longest_location, longest_diag_length = AssistedCandidateGeneratorDifferentTokenizers._get_longest_diag_index( + compare_mat_int + ) + new_token_start_index = longest_location[0] + longest_diag_length + discrepancy_with_old = longest_location[1] + longest_diag_length + discrepancy_length = (prompt.shape[1] - discrepancy_with_old).item() + new_tokens_only = prompt_plus_new_tokens[:, new_token_start_index + discrepancy_length :] + discrepancy_only = prompt_plus_new_tokens[ + :, new_token_start_index : new_token_start_index + discrepancy_length + ] + return discrepancy_length, new_tokens_only, discrepancy_only + + def convert_source_tokens_to_target_tokens( + self, + input_ids, + source_tokenizer, + destination_tokenizer, + ): + """ + Convert token IDs from one tokenizer to another. + Args: + input_ids: The input token IDs. + source_tokenizer: The source tokenizer. + destination_tokenizer: The destination tokenizer. + Returns: + The converted token IDs. + """ + text = source_tokenizer.batch_decode(input_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True) + dest_ids = destination_tokenizer(text, add_special_tokens=True, return_tensors="pt")["input_ids"] + return dest_ids.to(input_ids.device) + + def get_candidates(self, input_ids: torch.LongTensor) -> Tuple[torch.LongTensor, Optional[torch.FloatTensor]]: + """ + Fetches the candidates to be tried for the current input. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + + Return: + `torch.LongTensor` of shape `(batch_size, candidate_length)` containing the candidate sequences to be + assessed by the model and a `torch.FloatTensor` of shape `(batch_size, candidate_length, + vocabulary_size)` containing the logits associated to each candidate. + """ + max_new_tokens = int(self.num_assistant_tokens) + if max_new_tokens == 0: + return input_ids, None + + input_ids = input_ids.to(self.assistant_model.device) + remove_from_pkv = 0 + + assistant_input_ids, remove_from_pkv = self._prepare_assistant_input_ids(input_ids) + self.prev_assistant_ids = assistant_input_ids + + min_new_tokens = max(min(max_new_tokens, self.main_model_min_length - assistant_input_ids.shape[-1]), 0) + + self._update_past_and_masks(assistant_input_ids, remove_from_pkv) + generation_args = self._prepare_generation_args(assistant_input_ids, min_new_tokens, max_new_tokens) + self.assistant_kwargs.pop("attention_mask", None) + + assistant_output = self.assistant_model.generate(**generation_args, **self.assistant_kwargs) + new_target_ids = self._process_assistant_outputs(input_ids, assistant_output.sequences, assistant_input_ids) + + # Update state + self.prev_target_ids_len = input_ids.shape[1] + self.assistant_kwargs["past_key_values"] = assistant_output.past_key_values + self.prev_assistant_ids = assistant_output.sequences + + if self.prev_target_ids_len >= new_target_ids.shape[1]: + return input_ids, None + + return new_target_ids, None + + def _prepare_assistant_input_ids(self, input_ids: torch.LongTensor) -> Tuple[torch.LongTensor, int]: + """Converts target input IDs to assistant input IDs, handling discrepancies.""" + convert_kwargs = { + "source_tokenizer": self.target_tokenizer, + "destination_tokenizer": self.assistant_tokenizer, + } + remove_from_pkv = 0 + + if self.prev_assistant_ids is not None and self.prev_target_ids_len > self.target_lookbehind: + # input_ids contains all target prompt input ids and some new target input ids + start_index_in_target_window = self.prev_target_ids_len - self.target_lookbehind + + new_assistant_ids = self.convert_source_tokens_to_target_tokens( + input_ids[:, start_index_in_target_window:], **convert_kwargs + ) + prompt_use_length = new_assistant_ids.shape[1] + prompt_use = self.prev_assistant_ids[:, -prompt_use_length:] + + discrepancy_length, new_tokens_only, discrepancy_only = self._get_tokens_diag( + prompt_use, new_assistant_ids + ) + assistant_input_ids = self.prev_assistant_ids + + if new_tokens_only is not None: + if discrepancy_length > 0 and discrepancy_only.shape[1] > 0: + if discrepancy_length == discrepancy_only.shape[1]: + assistant_input_ids[:, -discrepancy_length:] = discrepancy_only + + elif discrepancy_length > discrepancy_only.shape[1]: + discrepancy_length_diff = discrepancy_length - discrepancy_only.shape[1] + assistant_input_ids = assistant_input_ids[:, :-discrepancy_length_diff] + assistant_input_ids[:, -discrepancy_only.shape[1] :] = discrepancy_only + + remove_from_pkv = discrepancy_length + + if new_tokens_only.shape[1] > 0: + assistant_input_ids = torch.cat([assistant_input_ids, new_tokens_only], dim=-1) + else: + # edge case: in case of no intersection between prompt and new_assistant_ids + assistant_input_ids = torch.cat([assistant_input_ids, new_assistant_ids], dim=-1) + else: + assistant_input_ids = self.convert_source_tokens_to_target_tokens(input_ids, **convert_kwargs) + self.prev_target_ids_len = input_ids.shape[1] + + return assistant_input_ids, remove_from_pkv + + def _process_assistant_outputs( + self, input_ids: torch.LongTensor, assistant_sequences: torch.LongTensor, assistant_input_ids: torch.LongTensor + ) -> torch.LongTensor: + """Processes assistant outputs to obtain target input IDs.""" + num_prev_assistant = self.prev_assistant_ids.shape[1] + start_assistant_look_index = num_prev_assistant - self.assistant_lookbehind + + new_target_ids_from_window = self.convert_source_tokens_to_target_tokens( + assistant_sequences[:, start_assistant_look_index:], + source_tokenizer=self.assistant_tokenizer, + destination_tokenizer=self.target_tokenizer, + ) + target_prompt_use_length = new_target_ids_from_window.shape[1] + + target_prompt_use = input_ids[:, -target_prompt_use_length:] + + _, target_new_tokens_only, _ = self._get_tokens_diag(target_prompt_use, new_target_ids_from_window) + + new_target_ids = input_ids + + if target_new_tokens_only is not None: + if target_new_tokens_only.shape[1] > 0: + new_target_ids = torch.cat([new_target_ids, target_new_tokens_only], dim=-1) + else: + # edge case: in case of no intersection between prompt and new_target_ids + new_target_ids = torch.cat([new_target_ids, new_target_ids_from_window], dim=-1) + + if hasattr(self.generation_config, "max_length"): + new_target_ids = new_target_ids[:, : self.generation_config.max_length] + + return new_target_ids + + +class AssistantToTargetTranslator: + """ + Translates token ids and logits between assistant and target model vocabularies. This class is used to handle + vocabulary mismatches when using different tokenizers for the assistant and target models in speculative decoding, + as introduced in the paper "Lossless Speculative Decoding Algorithms for Heterogeneous Vocabularies" + (https://www.arxiv.org/abs/2502.05202). + It maintains mappings between the two vocabularies and handles token/logit conversion. + + Args: + target_tokenizer (`PreTrainedTokenizerBase`): + The tokenizer used by the target (main) model. + assistant_tokenizer (`PreTrainedTokenizerBase`): + The tokenizer used by the assistant model. + assistant_model_device (`str`, defaults to "cpu"): + The device where the assistant model is located. Used for placing tensors. + target_vocab_size (`int`, *optional*): + The size of the target model's vocabulary. If not provided, will be inferred from the target tokenizer. + """ + + FILTER_VALUE: float = -float("Inf") # The value used to filter out unmapped tokens in the logits. + SUPPRESS_TOKEN_ID: int = -1 # The ID used to mark suppressed tokens in the mapping. + + def __init__( + self, + target_tokenizer: "PreTrainedTokenizerBase", + assistant_tokenizer: "PreTrainedTokenizerBase", + target_vocab_size: int, # required since target_vocab_size can be different from the length of target_tokenizer.get_vocab() + assistant_model_device: str = "cpu", + ): + self._target_tokenizer: "PreTrainedTokenizerBase" = target_tokenizer + self._assistant_tokenizer: "PreTrainedTokenizerBase" = assistant_tokenizer + self._assistant_model_device: str = assistant_model_device + self.target_vocab_size: int = target_vocab_size + self._assistant_to_target_input_ids, self.target_to_assistant_input_ids = ( + self._get_assistant_to_target_input_ids() + ) + self._suppress_input_ids: list[int] = self._get_suppress_input_ids() + self.logits_processors: Optional[LogitsProcessorList] = None + if len(self._suppress_input_ids) > 0: + # len(self._suppress_input_ids) = 0 if the assistant vocab is a subset of the target vocab + self.logits_processors = LogitsProcessorList( + [SuppressTokensLogitsProcessor(self._get_suppress_input_ids(), self._assistant_model_device)] + ) + + def _get_assistant_to_target_input_ids(self): + target_vocab = self._target_tokenizer.get_vocab() + assistant_vocab = self._assistant_tokenizer.get_vocab() + + space_str = " " + target_space_ids = self._target_tokenizer(space_str, add_special_tokens=False)["input_ids"] + if len(target_space_ids) > 0: + target_space_sign = self._target_tokenizer.convert_ids_to_tokens(target_space_ids)[0][0] + + assistant_space_ids = self._assistant_tokenizer(space_str, add_special_tokens=False)["input_ids"] + if len(assistant_space_ids) > 0: + assistant_space_sign = self._assistant_tokenizer.convert_ids_to_tokens(assistant_space_ids)[0][0] + + if target_space_sign != assistant_space_sign: + # If the assistant tokenizer has a different space sign than the target tokenizer, + # we need to replace the assistant space sign with the target space sign in the assistant_vocab. + assistant_vocab = { + ( + tok.replace(assistant_space_sign, target_space_sign, 1) + if tok.startswith(assistant_space_sign) + else tok + ): idx + for tok, idx in assistant_vocab.items() + } + + max_assistant_index = max(assistant_vocab.values()) + assistant_to_target_input_ids = torch.full((max_assistant_index + 1,), self.SUPPRESS_TOKEN_ID, dtype=int) + target_to_assistant_input_ids: Dict[int, int] = {} + for tok, assistant_id in assistant_vocab.items(): + target_id = target_vocab.get(tok) + if target_id is not None: + assistant_to_target_input_ids[assistant_id] = target_id + target_to_assistant_input_ids[target_id] = assistant_id + return assistant_to_target_input_ids.to(self._assistant_model_device), target_to_assistant_input_ids + + def _get_suppress_input_ids(self) -> list[int]: + """ + Get the input ids that are in the assistant vocab but not in the target vocab. + """ + return torch.where(self._assistant_to_target_input_ids == self.SUPPRESS_TOKEN_ID)[0] + + def get_target_ids( + self, assistant_input_ids, target_input_ids, assistant_candidate_ids: torch.LongTensor + ) -> torch.LongTensor: + """ + Return the target candidate ids that correspond to the assistant candidate ids. + Note that we have already the target ids for the prompt and we only need to find the target ids for the new tokens. + Moreover, assistant ids of the original prompt does not necessarily appear in _assistant_to_target_input_ids. + """ + + num_new_tokens = len(assistant_candidate_ids[0]) - assistant_input_ids.shape[1] + if num_new_tokens == 0: + return target_input_ids + else: + transformed_slice = self._assistant_to_target_input_ids[assistant_candidate_ids[0, -num_new_tokens:]] + return torch.cat((target_input_ids, transformed_slice.unsqueeze(0)), dim=1) + + def get_target_logits(self, assistant_logits: torch.FloatTensor) -> torch.FloatTensor: + """ + Return the target logits that correspond to the assistant logits. + """ + + target_shape: tuple[int, ...] = (*assistant_logits.shape[:-1], self.target_vocab_size) + target_logits: torch.FloatTensor = torch.full( + target_shape, self.FILTER_VALUE, device=self._assistant_model_device + ) + # Mask for valid indices + assistant_indices_mask = self._assistant_to_target_input_ids != self.SUPPRESS_TOKEN_ID + # Exclude invalid indices + target_logits_supported_indices = self._assistant_to_target_input_ids[assistant_indices_mask] + valid_assistant_logits = assistant_logits[..., : self._assistant_to_target_input_ids.shape[0]] + + target_logits[..., target_logits_supported_indices] = valid_assistant_logits[..., assistant_indices_mask] + + return target_logits + + +class AssistantVocabTranslatorCache: + """ + Cache for `AssistantToTargetTranslator` instances. The instances are computed at + pre-processing time, and this cache allows us to avoid recomputing them. + """ + + _cache = weakref.WeakKeyDictionary() + + @classmethod + def get_translator( + cls, + target_tokenizer: "PreTrainedTokenizerBase", + assistant_tokenizer: "PreTrainedTokenizerBase", + target_vocab_size: int, + assistant_model_device: str = "cpu", + ) -> AssistantToTargetTranslator: + assistant_dict = cls._cache.get(target_tokenizer) + if assistant_dict is None: + assistant_dict = weakref.WeakKeyDictionary() + cls._cache[target_tokenizer] = assistant_dict + + mapping = assistant_dict.get(assistant_tokenizer) + if mapping is None: + mapping = AssistantToTargetTranslator( + target_tokenizer, assistant_tokenizer, target_vocab_size, assistant_model_device + ) + assistant_dict[assistant_tokenizer] = mapping + + return mapping + + @classmethod + def cleanup(cls): + """ + Clean up dead references in the cache. + This removes entries where either the target_tokenizer or assistant_tokenizer + has been garbage collected. + """ + # Remove entries from the outer cache where the target_tokenizer is no longer alive + dead_keys = [key for key in cls._cache if key is None] + for key in dead_keys: + del cls._cache[key] + + # For each assistant_dict, remove entries where assistant_tokenizer is no longer alive + for assistant_dict in cls._cache.values(): + dead_keys = [key for key in assistant_dict if key is None] + for key in dead_keys: + del assistant_dict[key] + + +class UniversalSpeculativeDecodingGenerator(AssistedCandidateGeneratorDifferentTokenizers): + """ + `CandidateGenerator` class to be used for Universal Speculative Decoding (USD): speculative decoding with different tokenizers + for the assistant and main models. This class generates candidates through the use of a smaller model. + """ + + def __init__( + self, + input_ids: torch.LongTensor, + assistant_model: "PreTrainedModel", + target_tokenizer: "PreTrainedTokenizerBase", + assistant_tokenizer: "PreTrainedTokenizerBase", + generation_config: "GenerationConfig", + model_kwargs: Dict, + atm_translator: AssistantToTargetTranslator, + inputs_tensor: Optional[torch.Tensor] = None, + logits_processor: "LogitsProcessorList" = None, + ): + # Initialize translator before parent class + self._atm_translator = atm_translator + super().__init__( + input_ids, + assistant_model, + target_tokenizer, + assistant_tokenizer, + generation_config, + model_kwargs, + inputs_tensor, + logits_processor, + ) + # Track sequence lengths and previous assistant IDs + self._target_seq_len_with_candidates: int = 0 + self._prev_assistant_ids: Optional[torch.LongTensor] = None + + def get_candidates(self, input_ids: torch.LongTensor) -> Tuple[torch.LongTensor, Optional[torch.FloatTensor]]: + """ + Simplified version of get_candidates that uses the translator cache for token conversion. + """ + target_input_ids = input_ids.to(self.assistant_model.device) + assistant_input_ids, num_added_tokens = self._prepare_assistant_input_ids(target_input_ids) + min_new_tokens, max_new_tokens = self._calculate_new_tokens(target_input_ids) + + if max_new_tokens == 0: + return input_ids, None + + self._update_past_and_masks(assistant_input_ids, num_added_tokens=num_added_tokens) + generation_args = self._prepare_generation_args(assistant_input_ids, min_new_tokens, max_new_tokens) + + # Ensure scores are returned + generation_args["generation_config"].output_scores = True + generation_args["generation_config"].return_dict_in_generate = True + + # Generate and process outputs using translator + if self._atm_translator.logits_processors is not None: + generation_args["logits_processor"] = self._atm_translator.logits_processors + self._prev_assistant_ids, assistant_candidate_logits = self._generate_candidates(generation_args) + + # Use translator to convert tokens and logits + target_candidate_ids = self._atm_translator.get_target_ids( + assistant_input_ids, target_input_ids, self._prev_assistant_ids + ) + self._target_seq_len_with_candidates = target_candidate_ids.shape[-1] + target_candidate_logits = self._atm_translator.get_target_logits(assistant_candidate_logits) + + return target_candidate_ids, target_candidate_logits + + def _update_past_and_masks(self, assistant_input_ids: torch.LongTensor, num_added_tokens: int = 1) -> bool: + if self._prev_assistant_ids is None: + # Prepare attention mask for the first generation. + # For subsequent generations, the attention mask is updated in super()_update_past_and_masks. + self.assistant_kwargs = _prepare_attention_mask( + self.assistant_kwargs, assistant_input_ids.shape[-1], self.assistant_model.config.is_encoder_decoder + ) + return super()._update_past_and_masks(assistant_input_ids, num_added_tokens=num_added_tokens) + + def _prepare_assistant_input_ids(self, target_input_ids: torch.LongTensor) -> torch.LongTensor: + """ + Simplified token conversion that only processes new tokens. + """ + # Calculate new tokens since last call + target_seq_len = target_input_ids.shape[-1] + if self._target_seq_len_with_candidates == 0: + new_token_count = target_seq_len + else: + new_token_count = 1 + target_new_ids = target_input_ids[:, -new_token_count:] + + # Convert the new tokens + assistant_new_ids = None + if self._target_seq_len_with_candidates > 0: + # we have only one new token and we can directly convert it + assistant_new_ids = self._atm_translator.target_to_assistant_input_ids.get(target_new_ids[0].item()) + if assistant_new_ids is None: + target_new_text = self.target_tokenizer.batch_decode( + target_new_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True + ) + assistant_new_ids = self.assistant_tokenizer( + target_new_text, add_special_tokens=False, return_tensors="pt" + )["input_ids"].to(self.assistant_model.device) + else: + assistant_new_ids = torch.tensor([[assistant_new_ids]], device=self.assistant_model.device) + + # Update or initialize assistant IDs + if self._prev_assistant_ids is None: + assistant_input_ids = assistant_new_ids + else: + tokens_to_remove = self._target_seq_len_with_candidates + 1 - target_seq_len + # If the number of new tokens is greater than zero, truncate the previous assistant IDs + if tokens_to_remove > 0: + self._prev_assistant_ids = self._prev_assistant_ids[:, :-tokens_to_remove] + assistant_input_ids = torch.cat([self._prev_assistant_ids, assistant_new_ids], dim=-1) + assistant_input_ids = assistant_input_ids.to(dtype=torch.long) + + return assistant_input_ids, len(assistant_new_ids[0]) + + +class PromptLookupCandidateGenerator(CandidateGenerator): + """ + `CandidateGenerator` class to be used for prompt lookup generation. This class generates candidates by looking up + likely continuations in the provided prompt (input_ids) itself. + Read the following blog post for more information: https://github.com/apoorvumang/prompt-lookup-decoding + + Args: + max_matching_ngram_size (`int`): + The maximum ngram size to be considered for matching in the prompt + num_output_tokens (`int`): + The number of tokens to be output as candidate tokens. + max_length (`int`): + The number of total maximum tokens that can be generated. For decoder-only models that includes the prompt length. + Defaults to 20, which is the max length used as default in generation config. + """ + + def __init__( + self, + eos_token_id: Optional[torch.Tensor] = None, + num_output_tokens: int = 10, + max_matching_ngram_size: Optional[int] = None, + max_length: int = 20, + ): + self.num_output_tokens = num_output_tokens + self.max_matching_ngram_size = max_matching_ngram_size if max_matching_ngram_size else 2 + self.max_length = max_length + self.eos_token_id = eos_token_id + + if self.max_matching_ngram_size <= 0 or self.num_output_tokens <= 0: + raise ValueError("Invalid max_matching_ngram_size or num_output_tokens") + + def get_candidates(self, input_ids: torch.LongTensor) -> Tuple[torch.LongTensor, Optional[torch.FloatTensor]]: + """ + Fetches the candidates to be tried for the current input. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + + Return: + `torch.LongTensor` of shape `(num_candidates, candidate_length)`: The candidate sequences to be tried. + """ + input_length = input_ids.size(1) + + # Don't generate more than `max_length - 1` candidates since the target model generates one extra token. + if self.max_length == input_length + 1: + return input_ids, None + + chosen_ids = None + match_found = False + for ngram_size in range(min(self.max_matching_ngram_size, input_length - 1), 0, -1): + # Create sliding windows of size ngram_size + windows = input_ids.unfold(dimension=1, size=ngram_size, step=1) + + # Convert ngram to a tensor for comparison + ngram_tensor = input_ids[0, -ngram_size:] + + # Find where the windows match the ngram + matches = (windows == ngram_tensor).all(dim=2) + + # Get the indices of matches + match_indices = matches.nonzero(as_tuple=True)[1] + + # Iterate through match indices to find a valid continuation + for idx in match_indices: + start_idx = idx + ngram_size + end_idx = start_idx + self.num_output_tokens + end_idx = min(end_idx, input_length, self.max_length) + + if start_idx < end_idx: + chosen_ids = input_ids[0, start_idx:end_idx] + match_found = True + + # remove remaining candidate ids if an "eos" token is found, otherwise the target model may + # accept eos and the rest as valid, thus not stopping generation after "eos" + # NOTE: below code is written based on the fact that assisted decoding supports only bs=1 + mask = isin_mps_friendly(chosen_ids, self.eos_token_id) + match_indices_eos = torch.nonzero(mask) + if match_indices_eos.numel() > 0: + first_eos_index = match_indices_eos[0].item() + chosen_ids = chosen_ids[:first_eos_index] + break + if match_found: + break + + if chosen_ids is None or len(chosen_ids) == 0: + # In case we didn't find a match return the input sequence unchanged, reverts back to autoregressive decoding + return input_ids, None + + # Now need extend input_ids with chosen_ids + chosen_ids = chosen_ids.unsqueeze(0) + candidate_input_ids = torch.cat((input_ids, chosen_ids), dim=1) + # assisted_generation expects logits as well, but we don't have those here, so returning None + return candidate_input_ids, None + + def update_candidate_strategy(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, num_matches: int): + """ + Updates the candidate generation strategy based on the outcomes. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + scores (`torch.FloatTensor` of shape `(batch_size, candidate_length, config.vocab_size)`): + Prediction scores of a language modeling head. These can be logits for each vocabulary when not using + beam search or log softmax for each vocabulary token when using beam search + num_matches (`int`): + The number of matches between the candidate sequences and the model predictions. + """ + # Currently does nothing + return + + +class EarlyExitCandidateGenerator(AssistedCandidateGenerator): + """ + `CandidateGenerator` class to be used for assisted generation and speculative decoding. This class generates + candidates through the use of **the model itself**, exiting early. Can only be used with models that support early + exit, e.g., `facebook/layerskip-llama3.2-1B`. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + assistant_model (`PreTrainedModel`): + The original model. This model must support early exit (i.e. is trained to compute logits in earlier + layers). + generation_config (`~generation.GenerationConfig`, *optional*): + The generation configuration to be used as base parametrization for the generation call. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + model_kwargs (`Dict`): + The keyword arguments that will be passed to the main model, and are used as base inputs for the assistant + model as well. + inputs_tensor (`torch.Tensor`, *optional*): + The model input tensor. In encoder-decoder models, this is the encoder input. + """ + + def __init__( + self, + input_ids: torch.LongTensor, + assistant_model: "PreTrainedModel", + generation_config: "GenerationConfig", + model_kwargs: Dict, + inputs_tensor: Optional[torch.Tensor] = None, + logits_processor: "LogitsProcessorList" = None, + ): + super().__init__( + input_ids=input_ids, + assistant_model=assistant_model, + generation_config=generation_config, + model_kwargs=model_kwargs, + inputs_tensor=inputs_tensor, + logits_processor=logits_processor, + ) + # We have to move early exit out of the generation config, otherwise the assistant will also call `generate` + # with early exit + self.assistant_early_exit = self.generation_config.assistant_early_exit + self.generation_config.assistant_early_exit = None + + def get_candidates(self, input_ids: torch.LongTensor) -> Tuple[torch.LongTensor, Optional[torch.FloatTensor]]: + # Temporarily sets the number of hidden layers to the early exit value + base_model = getattr(self.assistant_model, self.assistant_model.base_model_prefix) + original_num_hidden_layers = base_model.config.num_hidden_layers + base_model.config.num_hidden_layers = self.assistant_early_exit + candidate_ids, candidate_logits = super().get_candidates(input_ids) + base_model.config.num_hidden_layers = original_num_hidden_layers + return candidate_ids, candidate_logits + + +def _crop_past_key_values(model, past_key_values, max_length): + """Crops the past key values up to a certain maximum length.""" + new_past = [] + if model.config.is_encoder_decoder: + for idx in range(len(past_key_values)): + new_past.append( + ( + past_key_values[idx][0][:, :, :max_length, :], + past_key_values[idx][1][:, :, :max_length, :], + past_key_values[idx][2], + past_key_values[idx][3], + ) + ) + past_key_values = tuple(new_past) + # gptbigcode is special and stores kv in shape (batch_size, seq_len, dim), if it's a multi_query model + elif "gptbigcode" in model.__class__.__name__.lower() or ( + model.config.architectures is not None and "gptbigcode" in model.config.architectures[0].lower() + ): + if model.config.multi_query: + for idx in range(len(past_key_values)): + past_key_values[idx] = past_key_values[idx][:, :max_length, :] + else: + for idx in range(len(past_key_values)): + past_key_values[idx] = past_key_values[idx][:, :, :max_length, :] + elif isinstance(past_key_values, DynamicCache): + past_key_values.crop(max_length) + elif past_key_values is not None: + for idx in range(len(past_key_values)): + if past_key_values[idx] != ([], []): + new_past.append( + ( + past_key_values[idx][0][:, :, :max_length, :], + past_key_values[idx][1][:, :, :max_length, :], + ) + ) + else: + new_past.append((past_key_values[idx][0], past_key_values[idx][1])) + past_key_values = tuple(new_past) + return past_key_values + + +def _prepare_attention_mask(model_kwargs: Dict[str, Any], new_length: int, is_encoder_decoder: bool) -> Dict[str, Any]: + """Expands or crops the model's mask for decoding purposes, to the defined length""" + + mask_key = "decoder_attention_mask" if is_encoder_decoder else "attention_mask" + if mask_key not in model_kwargs: + return model_kwargs + + mask = model_kwargs[mask_key] + mask_length_diff = new_length - mask.shape[1] + + if mask_length_diff < 0: + model_kwargs[mask_key] = mask[:, :mask_length_diff] + elif mask_length_diff > 0: + model_kwargs[mask_key] = torch.cat([mask, mask.new_ones((mask.shape[0], mask_length_diff))], dim=-1) + + # Handle cross attention models + if "cross_attention_mask" in model_kwargs: + # Mllama case + cross_mask = model_kwargs["cross_attention_mask"] + if mask_length_diff < 0: + model_kwargs["cross_attention_mask"] = cross_mask[:, :mask_length_diff] + elif mask_length_diff > 0: + new_mask = cross_mask[:, -1:, :, :].repeat(1, mask_length_diff, 1, 1) + model_kwargs["cross_attention_mask"] = torch.cat([cross_mask, new_mask], dim=1) + elif "image_attention_mask" in model_kwargs: + # IDEFICS case + cross_mask = model_kwargs["image_attention_mask"] + if mask_length_diff < 0: + model_kwargs["image_attention_mask"] = cross_mask[:, :mask_length_diff] + elif mask_length_diff > 0: + new_mask = cross_mask[:, -1:, :].repeat(1, mask_length_diff, 1) + model_kwargs["image_attention_mask"] = torch.cat([cross_mask, new_mask], dim=1) + + return model_kwargs + + +def _prepare_token_type_ids(model_kwargs: Dict[str, Any], new_length: int) -> Dict[str, Any]: + """Expands or crops the model's token_type_ids for decoding purposes, to the defined length""" + if "token_type_ids" not in model_kwargs or model_kwargs["token_type_ids"] is None: + return model_kwargs + + token_type_ids = model_kwargs["token_type_ids"] + final_token_type = token_type_ids[:, -1].unsqueeze(-1) + type_length_diff = new_length - token_type_ids.shape[1] + + if type_length_diff < 0: + token_type_ids = token_type_ids[:, :type_length_diff] + elif type_length_diff > 0: + token_type_copies = final_token_type.repeat(1, type_length_diff) + model_kwargs["token_type_ids"] = torch.cat([model_kwargs["token_type_ids"], token_type_copies], dim=-1) + return model_kwargs diff --git a/lib/python3.12/site-packages/transformers/generation/configuration_utils.py b/lib/python3.12/site-packages/transformers/generation/configuration_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c743480d786b71a7dcf2e8f8896db69dba588851 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/configuration_utils.py @@ -0,0 +1,1632 @@ +# coding=utf-8 +# Copyright 2022 The HuggingFace Inc. team. +# +# 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. +"""Generation configuration class and utilities.""" + +import copy +import json +import os +import warnings +from abc import ABC, abstractmethod +from dataclasses import dataclass, is_dataclass +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union + +from .. import __version__ +from ..configuration_utils import PretrainedConfig +from ..utils import ( + GENERATION_CONFIG_NAME, + ExplicitEnum, + PushToHubMixin, + cached_file, + download_url, + extract_commit_hash, + is_remote_url, + is_torch_available, + logging, +) + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + + +logger = logging.get_logger(__name__) +METADATA_FIELDS = ("_from_model_config", "_commit_hash", "_original_object_hash", "transformers_version") +CACHE_CONFIG_MAPPING = {} +NEED_SETUP_CACHE_CLASSES_MAPPING = {} +QUANT_BACKEND_CLASSES_MAPPING = {} +ALL_CACHE_IMPLEMENTATIONS = [] + +if is_torch_available(): + from ..cache_utils import ( + HQQQuantizedCache, + HybridCache, + HybridChunkedCache, + MambaCache, + OffloadedStaticCache, + QuantizedCacheConfig, + QuantoQuantizedCache, + SlidingWindowCache, + StaticCache, + StaticCacheConfig, + ) + from .logits_process import SynthIDTextWatermarkLogitsProcessor, WatermarkLogitsProcessor + + CACHE_CONFIG_MAPPING["quantized"] = QuantizedCacheConfig + CACHE_CONFIG_MAPPING["static"] = StaticCacheConfig + NEED_SETUP_CACHE_CLASSES_MAPPING = { + "static": StaticCache, + "offloaded_static": OffloadedStaticCache, + "sliding_window": SlidingWindowCache, + "hybrid": HybridCache, + "hybrid_chunked": HybridChunkedCache, + "mamba": MambaCache, + } + QUANT_BACKEND_CLASSES_MAPPING = {"quanto": QuantoQuantizedCache, "HQQ": HQQQuantizedCache} + ALL_CACHE_IMPLEMENTATIONS = ( + list(NEED_SETUP_CACHE_CLASSES_MAPPING.keys()) + list(CACHE_CONFIG_MAPPING.keys()) + ["offloaded", "dynamic"] + ) + + +class GenerationMode(ExplicitEnum): + """ + Possible generation modes, downstream of the [`~generation.GenerationMixin.generate`] method. + """ + + # Non-beam methods + CONTRASTIVE_SEARCH = "contrastive_search" + GREEDY_SEARCH = "greedy_search" + SAMPLE = "sample" + ASSISTED_GENERATION = "assisted_generation" + DOLA_GENERATION = "dola_generation" + # Beam methods + BEAM_SEARCH = "beam_search" + BEAM_SAMPLE = "beam_sample" + CONSTRAINED_BEAM_SEARCH = "constrained_beam_search" + GROUP_BEAM_SEARCH = "group_beam_search" + + +class GenerationConfig(PushToHubMixin): + # no-format + """ + Class that holds a configuration for a generation task. A `generate` call supports the following generation methods + for text-decoder, text-to-text, speech-to-text, and vision-to-text models: + + - *greedy decoding* if `num_beams=1` and `do_sample=False` + - *contrastive search* if `penalty_alpha>0.` and `top_k>1` + - *multinomial sampling* if `num_beams=1` and `do_sample=True` + - *beam-search decoding* if `num_beams>1` and `do_sample=False` + - *beam-search multinomial sampling* if `num_beams>1` and `do_sample=True` + - *diverse beam-search decoding* if `num_beams>1` and `num_beam_groups>1` + - *constrained beam-search decoding* if `constraints!=None` or `force_words_ids!=None` + - *assisted decoding* if `assistant_model` or `prompt_lookup_num_tokens` is passed to `.generate()` + - *dola decoding* if `dola_layers` is passed to `.generate()` + + To learn more about decoding strategies refer to the [text generation strategies guide](../generation_strategies). + + + + A large number of these flags control the logits or the stopping criteria of the generation. Make sure you check + the [generate-related classes](https://huggingface.co/docs/transformers/internal/generation_utils) for a full + description of the possible manipulations, as well as examples of their usage. + + + + Arg: + > Parameters that control the length of the output + + max_length (`int`, *optional*, defaults to 20): + The maximum length the generated tokens can have. Corresponds to the length of the input prompt + + `max_new_tokens`. Its effect is overridden by `max_new_tokens`, if also set. + max_new_tokens (`int`, *optional*): + The maximum numbers of tokens to generate, ignoring the number of tokens in the prompt. + min_length (`int`, *optional*, defaults to 0): + The minimum length of the sequence to be generated. Corresponds to the length of the input prompt + + `min_new_tokens`. Its effect is overridden by `min_new_tokens`, if also set. + min_new_tokens (`int`, *optional*): + The minimum numbers of tokens to generate, ignoring the number of tokens in the prompt. + early_stopping (`bool` or `str`, *optional*, defaults to `False`): + Controls the stopping condition for beam-based methods, like beam-search. It accepts the following values: + `True`, where the generation stops as soon as there are `num_beams` complete candidates; `False`, where an + heuristic is applied and the generation stops when is it very unlikely to find better candidates; + `"never"`, where the beam search procedure only stops when there cannot be better candidates (canonical + beam search algorithm). + max_time (`float`, *optional*): + The maximum amount of time you allow the computation to run for in seconds. generation will still finish + the current pass after allocated time has been passed. + stop_strings (`str or List[str]`, *optional*): + A string or a list of strings that should terminate generation if the model outputs them. + + > Parameters that control the generation strategy used + + do_sample (`bool`, *optional*, defaults to `False`): + Whether or not to use sampling ; use greedy decoding otherwise. + num_beams (`int`, *optional*, defaults to 1): + Number of beams for beam search. 1 means no beam search. + num_beam_groups (`int`, *optional*, defaults to 1): + Number of groups to divide `num_beams` into in order to ensure diversity among different groups of beams. + [this paper](https://arxiv.org/pdf/1610.02424.pdf) for more details. + penalty_alpha (`float`, *optional*): + The values balance the model confidence and the degeneration penalty in contrastive search decoding. + dola_layers (`str` or `List[int]`, *optional*): + The layers to use for DoLa decoding. If `None`, DoLa decoding is not used. If a string, it must + be one of "low" or "high", which means using the lower part or higher part of the model layers, respectively. + "low" means the first half of the layers up to the first 20 layers, and "high" means the last half of the + layers up to the last 20 layers. + If a list of integers, it must contain the indices of the layers to use for candidate premature layers in DoLa. + The 0-th layer is the word embedding layer of the model. Set to `'low'` to improve long-answer reasoning tasks, + `'high'` to improve short-answer tasks. Check the [documentation](https://github.com/huggingface/transformers/blob/main/docs/source/en/generation_strategies.md) + or [the paper](https://arxiv.org/abs/2309.03883) for more details. + + > Parameters that control the cache + + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should use the past last key/values attentions (if applicable to the model) to + speed up decoding. + cache_implementation (`str`, *optional*, default to `None`): + Name of the cache class that will be instantiated in `generate`, for faster decoding. Possible values are: + + - `"dynamic"`: [`DynamicCache`] + - `"static"`: [`StaticCache`] + - `"offloaded_static"`: [`OffloadedStaticCache`] + - `"sliding_window"`: [`SlidingWindowCache`] + - `"hybrid"`: [`HybridCache`] + - `"mamba"`: [`MambaCache`] + - `"quantized"`: [`QuantizedCache`] + + If none is specified, we will use the default cache for the model (which is often [`DynamicCache`]). See + our [cache documentation](https://huggingface.co/docs/transformers/en/kv_cache) for further information. + cache_config (`CacheConfig` or `dict`, *optional*, default to `None`): + Arguments used in the key-value cache class can be passed in `cache_config`. Can be passed as a `Dict` and + it will be converted to its repsective `CacheConfig` internally. + Otherwise can be passed as a `CacheConfig` class matching the indicated `cache_implementation`. + return_legacy_cache (`bool`, *optional*, default to `True`): + Whether to return the legacy or new format of the cache when `DynamicCache` is used by default. + + > Parameters for manipulation of the model output logits + + temperature (`float`, *optional*, defaults to 1.0): + The value used to module the next token probabilities. This value is set in a model's `generation_config.json` file. If it isn't set, the default value is 1.0 + top_k (`int`, *optional*, defaults to 50): + The number of highest probability vocabulary tokens to keep for top-k-filtering. This value is set in a model's `generation_config.json` file. If it isn't set, the default value is 50. + top_p (`float`, *optional*, defaults to 1.0): + If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to + `top_p` or higher are kept for generation. This value is set in a model's `generation_config.json` file. If it isn't set, the default value is 1.0 + min_p (`float`, *optional*): + Minimum token probability, which will be scaled by the probability of the most likely token. It must be a + value between 0 and 1. Typical values are in the 0.01-0.2 range, comparably selective as setting `top_p` in + the 0.99-0.8 range (use the opposite of normal `top_p` values). + typical_p (`float`, *optional*, defaults to 1.0): + Local typicality measures how similar the conditional probability of predicting a target token next is to + the expected conditional probability of predicting a random token next, given the partial text already + generated. If set to float < 1, the smallest set of the most locally typical tokens with probabilities that + add up to `typical_p` or higher are kept for generation. See [this + paper](https://arxiv.org/pdf/2202.00666.pdf) for more details. + epsilon_cutoff (`float`, *optional*, defaults to 0.0): + If set to float strictly between 0 and 1, only tokens with a conditional probability greater than + `epsilon_cutoff` will be sampled. In the paper, suggested values range from 3e-4 to 9e-4, depending on the + size of the model. See [Truncation Sampling as Language Model + Desmoothing](https://arxiv.org/abs/2210.15191) for more details. + eta_cutoff (`float`, *optional*, defaults to 0.0): + Eta sampling is a hybrid of locally typical sampling and epsilon sampling. If set to float strictly between + 0 and 1, a token is only considered if it is greater than either `eta_cutoff` or `sqrt(eta_cutoff) * + exp(-entropy(softmax(next_token_logits)))`. The latter term is intuitively the expected next token + probability, scaled by `sqrt(eta_cutoff)`. In the paper, suggested values range from 3e-4 to 2e-3, + depending on the size of the model. See [Truncation Sampling as Language Model + Desmoothing](https://arxiv.org/abs/2210.15191) for more details. + diversity_penalty (`float`, *optional*, defaults to 0.0): + This value is subtracted from a beam's score if it generates a token same as any beam from other group at a + particular time. Note that `diversity_penalty` is only effective if `group beam search` is enabled. + repetition_penalty (`float`, *optional*, defaults to 1.0): + The parameter for repetition penalty. 1.0 means no penalty. See [this + paper](https://arxiv.org/pdf/1909.05858.pdf) for more details. + encoder_repetition_penalty (`float`, *optional*, defaults to 1.0): + The paramater for encoder_repetition_penalty. An exponential penalty on sequences that are not in the + original input. 1.0 means no penalty. + length_penalty (`float`, *optional*, defaults to 1.0): + Exponential penalty to the length that is used with beam-based generation. It is applied as an exponent to + the sequence length, which in turn is used to divide the score of the sequence. Since the score is the log + likelihood of the sequence (i.e. negative), `length_penalty` > 0.0 promotes longer sequences, while + `length_penalty` < 0.0 encourages shorter sequences. + no_repeat_ngram_size (`int`, *optional*, defaults to 0): + If set to int > 0, all ngrams of that size can only occur once. + bad_words_ids (`List[List[int]]`, *optional*): + List of list of token ids that are not allowed to be generated. Check + [`~generation.NoBadWordsLogitsProcessor`] for further documentation and examples. + force_words_ids (`List[List[int]]` or `List[List[List[int]]]`, *optional*): + List of token ids that must be generated. If given a `List[List[int]]`, this is treated as a simple list of + words that must be included, the opposite to `bad_words_ids`. If given `List[List[List[int]]]`, this + triggers a [disjunctive constraint](https://github.com/huggingface/transformers/issues/14081), where one + can allow different forms of each word. + renormalize_logits (`bool`, *optional*, defaults to `False`): + Whether to renormalize the logits after applying all the logits processors (including the custom + ones). It's highly recommended to set this flag to `True` as the search algorithms suppose the score logits + are normalized but some logit processors break the normalization. + constraints (`List[Constraint]`, *optional*): + Custom constraints that can be added to the generation to ensure that the output will contain the use of + certain tokens as defined by `Constraint` objects, in the most sensible way possible. + forced_bos_token_id (`int`, *optional*, defaults to `model.config.forced_bos_token_id`): + The id of the token to force as the first generated token after the `decoder_start_token_id`. Useful for + multilingual models like [mBART](../model_doc/mbart) where the first generated token needs to be the target + language token. + forced_eos_token_id (`int` or List[int]`, *optional*, defaults to `model.config.forced_eos_token_id`): + The id of the token to force as the last generated token when `max_length` is reached. Optionally, use a + list to set multiple *end-of-sequence* tokens. + remove_invalid_values (`bool`, *optional*, defaults to `model.config.remove_invalid_values`): + Whether to remove possible *nan* and *inf* outputs of the model to prevent the generation method to crash. + Note that using `remove_invalid_values` can slow down generation. + exponential_decay_length_penalty (`tuple(int, float)`, *optional*): + This Tuple adds an exponentially increasing length penalty, after a certain amount of tokens have been + generated. The tuple shall consist of: `(start_index, decay_factor)` where `start_index` indicates where + penalty starts and `decay_factor` represents the factor of exponential decay + suppress_tokens (`List[int]`, *optional*): + A list of tokens that will be suppressed at generation. The `SupressTokens` logit processor will set their + log probs to `-inf` so that they are not sampled. + begin_suppress_tokens (`List[int]`, *optional*): + A list of tokens that will be suppressed at the beginning of the generation. The `SupressBeginTokens` logit + processor will set their log probs to `-inf` so that they are not sampled. + forced_decoder_ids (`List[List[int]]`, *optional*): + A list of pairs of integers which indicates a mapping from generation indices to token indices that will be + forced before sampling. For example, `[[1, 123]]` means the second generated token will always be a token + of index 123. + sequence_bias (`Dict[Tuple[int], float]`, *optional*)): + Dictionary that maps a sequence of tokens to its bias term. Positive biases increase the odds of the + sequence being selected, while negative biases do the opposite. Check + [`~generation.SequenceBiasLogitsProcessor`] for further documentation and examples. + token_healing (`bool`, *optional*, defaults to `False`): + Heal tail tokens of prompts by replacing them with their appropriate extensions. + This enhances the quality of completions for prompts affected by greedy tokenization bias. + guidance_scale (`float`, *optional*): + The guidance scale for classifier free guidance (CFG). CFG is enabled by setting `guidance_scale > 1`. + Higher guidance scale encourages the model to generate samples that are more closely linked to the input + prompt, usually at the expense of poorer quality. + low_memory (`bool`, *optional*): + Switch to sequential beam search and sequential topk for contrastive search to reduce peak memory. + Used with beam search and contrastive search. + watermarking_config (`BaseWatermarkingConfig` or `dict`, *optional*): + Arguments used to watermark the model outputs by adding a small bias to randomly selected set of "green" + tokens. See the docs of [`SynthIDTextWatermarkingConfig`] and [`WatermarkingConfig`] for more + details. If passed as `Dict`, it will be converted to a `WatermarkingConfig` internally. + + > Parameters that define the output variables of generate + + num_return_sequences (`int`, *optional*, defaults to 1): + The number of independently computed returned sequences for each element in the batch. + output_attentions (`bool`, *optional*, defaults to `False`): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more details. + output_hidden_states (`bool`, *optional*, defaults to `False`): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more details. + output_scores (`bool`, *optional*, defaults to `False`): + Whether or not to return the prediction scores. See `scores` under returned tensors for more details. + output_logits (`bool`, *optional*): + Whether or not to return the unprocessed prediction logit scores. See `logits` under returned tensors for + more details. + return_dict_in_generate (`bool`, *optional*, defaults to `False`): + Whether or not to return a [`~utils.ModelOutput`], as opposed to returning exclusively the generated + sequence. This flag must be set to `True` to return the generation cache (when `use_cache` is `True`) + or optional outputs (see flags starting with `output_`) + + > Special tokens that can be used at generation time + + pad_token_id (`int`, *optional*): + The id of the *padding* token. + bos_token_id (`int`, *optional*): + The id of the *beginning-of-sequence* token. + eos_token_id (`Union[int, List[int]]`, *optional*): + The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens. + + > Generation parameters exclusive to encoder-decoder models + + encoder_no_repeat_ngram_size (`int`, *optional*, defaults to 0): + If set to int > 0, all ngrams of that size that occur in the `encoder_input_ids` cannot occur in the + `decoder_input_ids`. + decoder_start_token_id (`int` or `List[int]`, *optional*): + If an encoder-decoder model starts decoding with a different token than *bos*, the id of that token or a list of length + `batch_size`. Indicating a list enables different start ids for each element in the batch + (e.g. multilingual models with different target languages in one batch) + + > Generation parameters exclusive to assistant generation + is_assistant (`bool`, *optional*, defaults to `False`): + Whether the model is an assistant (draft) model. + num_assistant_tokens (`int`, *optional*, defaults to 20): + Defines the number of _speculative tokens_ that shall be generated by the assistant model before being + checked by the target model at each iteration. Higher values for `num_assistant_tokens` make the generation + more _speculative_ : If the assistant model is performant larger speed-ups can be reached, if the assistant + model requires lots of corrections, lower speed-ups are reached. + num_assistant_tokens_schedule (`str`, *optional*, defaults to `"constant"`): + Defines the schedule at which max assistant tokens shall be changed during inference. + - `"heuristic"`: When all speculative tokens are correct, increase `num_assistant_tokens` by 2 else + reduce by 1. `num_assistant_tokens` value is persistent over multiple generation calls with the same assistant model. + - `"heuristic_transient"`: Same as `"heuristic"` but `num_assistant_tokens` is reset to its initial value after each generation call. + - `"constant"`: `num_assistant_tokens` stays unchanged during generation + assistant_confidence_threshold (`float`, *optional*, defaults to 0.4): + The confidence threshold for the assistant model. If the assistant model's confidence in its prediction for the current token is lower + than this threshold, the assistant model stops the current token generation iteration, even if the number of _speculative tokens_ + (defined by `num_assistant_tokens`) is not yet reached. The assistant's confidence threshold is adjusted throughout the speculative iterations to reduce the number of unnecessary draft and target forward passes, biased towards avoiding false negatives. + `assistant_confidence_threshold` value is persistent over multiple generation calls with the same assistant model. + It is an unsupervised version of the dynamic speculation lookahead + from Dynamic Speculation Lookahead Accelerates Speculative Decoding of Large Language Models . + prompt_lookup_num_tokens (`int`, *optional*): + The number of tokens to be output as candidate tokens. + max_matching_ngram_size (`int`, *optional*): + The maximum ngram size to be considered for matching in the prompt. Default to 2 if not provided. + assistant_early_exit(`int`, *optional*): + If set to a positive integer, early exit of the model will be used as an assistant. Can only be used with + models that support early exit (i.e. models where logits from intermediate layers can be interpreted by the LM head). + assistant_lookbehind(`int`, *optional*, defaults to 10): + If set to a positive integer, the re-encodeing process will additionally consider the last `assistant_lookbehind` assistant tokens + to correctly align tokens. Can only be used with different tokenizers in speculative decoding. + See this [blog](https://huggingface.co/blog/universal_assisted_generation) for more details. + target_lookbehind(`int`, *optional*, defaults to 10): + If set to a positive integer, the re-encodeing process will additionally consider the last `target_lookbehind` target tokens + to correctly align tokens. Can only be used with different tokenizers in speculative decoding. + See this [blog](https://huggingface.co/blog/universal_assisted_generation) for more details. + + > Parameters related to performances and compilation + + compile_config (CompileConfig, *optional*): + If using a static cache, this controls how `generate` will `compile` the forward pass for performance + gains. + + disable_compile (`bool`, *optional*): Whether to disable the automatic compilation of the forward pass. Automatic compilation happens when specific criteria are met, including using a compileable cache. Please open an issue if you find the need to use this flag. + + > Wild card + + generation_kwargs: + Additional generation kwargs will be forwarded to the `generate` function of the model. Kwargs that are not + present in `generate`'s signature will be used in the model forward pass. + """ + + extra_output_flags = ("output_attentions", "output_hidden_states", "output_scores", "output_logits") + + def __init__(self, **kwargs): + # Parameters that control the length of the output + self.max_length = kwargs.pop("max_length", 20) + self.max_new_tokens = kwargs.pop("max_new_tokens", None) + self.min_length = kwargs.pop("min_length", 0) + self.min_new_tokens = kwargs.pop("min_new_tokens", None) + self.early_stopping = kwargs.pop("early_stopping", False) + self.max_time = kwargs.pop("max_time", None) + self.stop_strings = kwargs.pop("stop_strings", None) + + # Parameters that control the generation strategy used + self.do_sample = kwargs.pop("do_sample", False) + self.num_beams = kwargs.pop("num_beams", 1) + self.num_beam_groups = kwargs.pop("num_beam_groups", 1) + self.penalty_alpha = kwargs.pop("penalty_alpha", None) + self.dola_layers = kwargs.pop("dola_layers", None) + + # Parameters that control the cache + self.use_cache = kwargs.pop("use_cache", True) + self.cache_implementation = kwargs.pop("cache_implementation", None) + self.cache_config = kwargs.pop("cache_config", None) + if self.cache_implementation is not None and self.cache_implementation in CACHE_CONFIG_MAPPING: + cache_config_class = CACHE_CONFIG_MAPPING[self.cache_implementation] + if isinstance(self.cache_config, dict): + self.cache_config = cache_config_class.from_dict(self.cache_config) + self.return_legacy_cache = kwargs.pop("return_legacy_cache", None) + self.prefill_chunk_size = kwargs.pop("prefill_chunk_size", None) + + # Parameters for manipulation of the model output logits + self.temperature = kwargs.pop("temperature", 1.0) + self.top_k = kwargs.pop("top_k", 50) + self.top_p = kwargs.pop("top_p", 1.0) + self.min_p = kwargs.pop("min_p", None) + self.typical_p = kwargs.pop("typical_p", 1.0) + self.epsilon_cutoff = kwargs.pop("epsilon_cutoff", 0.0) + self.eta_cutoff = kwargs.pop("eta_cutoff", 0.0) + self.diversity_penalty = kwargs.pop("diversity_penalty", 0.0) + self.repetition_penalty = kwargs.pop("repetition_penalty", 1.0) + self.encoder_repetition_penalty = kwargs.pop("encoder_repetition_penalty", 1.0) + self.length_penalty = kwargs.pop("length_penalty", 1.0) + self.no_repeat_ngram_size = kwargs.pop("no_repeat_ngram_size", 0) + self.bad_words_ids = kwargs.pop("bad_words_ids", None) + self.force_words_ids = kwargs.pop("force_words_ids", None) + self.renormalize_logits = kwargs.pop("renormalize_logits", False) + self.constraints = kwargs.pop("constraints", None) + self.forced_bos_token_id = kwargs.pop("forced_bos_token_id", None) + self.forced_eos_token_id = kwargs.pop("forced_eos_token_id", None) + self.remove_invalid_values = kwargs.pop("remove_invalid_values", False) + self.exponential_decay_length_penalty = kwargs.pop("exponential_decay_length_penalty", None) + self.suppress_tokens = kwargs.pop("suppress_tokens", None) + self.begin_suppress_tokens = kwargs.pop("begin_suppress_tokens", None) + self.forced_decoder_ids = kwargs.pop("forced_decoder_ids", None) + self.sequence_bias = kwargs.pop("sequence_bias", None) + self.token_healing = kwargs.pop("token_healing", False) + self.guidance_scale = kwargs.pop("guidance_scale", None) + self.low_memory = kwargs.pop("low_memory", None) + watermarking_config = kwargs.pop("watermarking_config", None) + if watermarking_config is None: + self.watermarking_config = None + elif isinstance(watermarking_config, BaseWatermarkingConfig): + self.watermarking_config = watermarking_config + else: + self.watermarking_config = WatermarkingConfig.from_dict(watermarking_config) + + # Parameters that define the output variables of `generate` + self.num_return_sequences = kwargs.pop("num_return_sequences", 1) + self.output_attentions = kwargs.pop("output_attentions", False) + self.output_hidden_states = kwargs.pop("output_hidden_states", False) + self.output_scores = kwargs.pop("output_scores", False) + self.output_logits = kwargs.pop("output_logits", None) + self.return_dict_in_generate = kwargs.pop("return_dict_in_generate", False) + + # Special tokens that can be used at generation time + self.pad_token_id = kwargs.pop("pad_token_id", None) + self.bos_token_id = kwargs.pop("bos_token_id", None) + self.eos_token_id = kwargs.pop("eos_token_id", None) + + # Generation parameters exclusive to encoder-decoder models + self.encoder_no_repeat_ngram_size = kwargs.pop("encoder_no_repeat_ngram_size", 0) + self.decoder_start_token_id = kwargs.pop("decoder_start_token_id", None) + + # Assistant generation + self.is_assistant = False + self.num_assistant_tokens = kwargs.pop("num_assistant_tokens", 20) + self.num_assistant_tokens_schedule = kwargs.pop("num_assistant_tokens_schedule", "constant") + self.assistant_confidence_threshold = kwargs.pop("assistant_confidence_threshold", 0.4) + self.prompt_lookup_num_tokens = kwargs.pop("prompt_lookup_num_tokens", None) + self.max_matching_ngram_size = kwargs.pop("max_matching_ngram_size", None) + self.assistant_early_exit = kwargs.pop("assistant_early_exit", None) + ## assistant generation for different tokenizers, the windows size for assistant/target model + self.assistant_lookbehind = kwargs.pop("assistant_lookbehind", 10) + self.target_lookbehind = kwargs.pop("target_lookbehind", 10) + + # Performance + self.compile_config = kwargs.pop("compile_config", CompileConfig()) + self.disable_compile = kwargs.pop("disable_compile", False) + # Wild card + self.generation_kwargs = kwargs.pop("generation_kwargs", {}) + + # The remaining attributes do not parametrize `.generate()`, but are informative and/or used by the hub + # interface. + self._from_model_config = kwargs.pop("_from_model_config", False) + self._commit_hash = kwargs.pop("_commit_hash", None) + self.transformers_version = kwargs.pop("transformers_version", __version__) + + # Additional attributes without default values + if not self._from_model_config: + # we don't want to copy values from the model config if we're initializing a `GenerationConfig` from a + # model's default configuration file + 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 + + # Validate the values of the attributes + self.validate(is_init=True) + + def __hash__(self): + return hash(self.to_json_string(ignore_metadata=True)) + + def __eq__(self, other): + if not isinstance(other, GenerationConfig): + return False + + self_without_metadata = self.to_json_string(use_diff=False, ignore_metadata=True) + other_without_metadata = other.to_json_string(use_diff=False, ignore_metadata=True) + return self_without_metadata == other_without_metadata + + def __repr__(self): + return f"{self.__class__.__name__} {self.to_json_string(ignore_metadata=True)}" + + def get_generation_mode(self, assistant_model: Optional["PreTrainedModel"] = None) -> GenerationMode: + """ + Returns the generation mode triggered by the [`GenerationConfig`] instance. + + Arg: + assistant_model (`PreTrainedModel`, *optional*): + The assistant model to be used for assisted generation. If set, the generation mode will be + assisted generation. + + Returns: + `GenerationMode`: The generation mode triggered by the instance. + """ + # TODO joao: find out a way of not depending on external fields (e.g. `assistant_model`), then make this a + # property and part of the `__repr__` + if self.constraints is not None or self.force_words_ids is not None: + generation_mode = GenerationMode.CONSTRAINED_BEAM_SEARCH + elif self.num_beams == 1: + if self.do_sample is False: + if ( + self.top_k is not None + and self.top_k > 1 + and self.penalty_alpha is not None + and self.penalty_alpha > 0 + ): + generation_mode = GenerationMode.CONTRASTIVE_SEARCH + else: + generation_mode = GenerationMode.GREEDY_SEARCH + else: + generation_mode = GenerationMode.SAMPLE + else: + if self.num_beam_groups > 1: + generation_mode = GenerationMode.GROUP_BEAM_SEARCH + elif self.do_sample is True: + generation_mode = GenerationMode.BEAM_SAMPLE + else: + generation_mode = GenerationMode.BEAM_SEARCH + + # Assisted generation may extend some generation modes + if ( + assistant_model is not None + or self.prompt_lookup_num_tokens is not None + or self.assistant_early_exit is not None + ): + if generation_mode in ("greedy_search", "sample"): + generation_mode = GenerationMode.ASSISTED_GENERATION + else: + raise ValueError( + "You've set `assistant_model`, which triggers assisted generate. Currently, assisted generate " + "is only supported with Greedy Search and Sample." + ) + + # DoLa generation may extend some generation modes + if self.dola_layers is not None: + if generation_mode in ("greedy_search", "sample"): + generation_mode = GenerationMode.DOLA_GENERATION + else: + raise ValueError( + "You've set `dola_layers`, which triggers DoLa generate. Currently, DoLa generate " + "is only supported with Greedy Search and Sample." + ) + return generation_mode + + def validate(self, is_init=False): + """ + Validates the values of the attributes of the [`GenerationConfig`] instance. Raises exceptions in the presence + of parameterization that can be detected as incorrect from the configuration instance alone. + + Note that some parameters not validated here are best validated at generate runtime, as they may depend on + other inputs and/or the model, such as parameters related to the generation length. + + Arg: + is_init (`bool`, *optional*, defaults to `False`): + Whether the validation is performed during the initialization of the instance. + """ + + # Validation of individual attributes + if self.early_stopping not in {True, False, "never"}: + raise ValueError(f"`early_stopping` must be a boolean or 'never', but is {self.early_stopping}.") + if self.max_new_tokens is not None and self.max_new_tokens <= 0: + raise ValueError(f"`max_new_tokens` must be greater than 0, but is {self.max_new_tokens}.") + if self.pad_token_id is not None and self.pad_token_id < 0: + warnings.warn( + f"`pad_token_id` should be positive but got {self.pad_token_id}. This will cause errors when batch " + "generating, if there is padding. Please set `pad_token_id` explicitly as " + "`model.generation_config.pad_token_id=PAD_TOKEN_ID` to avoid errors in generation" + ) + + # Validation of attribute relations: + fix_location = "" + if is_init: + fix_location = ( + " This was detected when initializing the generation config instance, which means the corresponding " + "file may hold incorrect parameterization and should be fixed." + ) + + # 1. detect sampling-only parameterization when not in sampling mode + if self.do_sample is False: + greedy_wrong_parameter_msg = ( + "`do_sample` is set to `False`. However, `{flag_name}` is set to `{flag_value}` -- this flag is only " + "used in sample-based generation modes. You should set `do_sample=True` or unset `{flag_name}`." + + fix_location + ) + if self.temperature is not None and self.temperature != 1.0: + warnings.warn( + greedy_wrong_parameter_msg.format(flag_name="temperature", flag_value=self.temperature), + UserWarning, + ) + if self.top_p is not None and self.top_p != 1.0: + warnings.warn( + greedy_wrong_parameter_msg.format(flag_name="top_p", flag_value=self.top_p), + UserWarning, + ) + if self.min_p is not None: + warnings.warn( + greedy_wrong_parameter_msg.format(flag_name="min_p", flag_value=self.min_p), + UserWarning, + ) + if self.typical_p is not None and self.typical_p != 1.0: + warnings.warn( + greedy_wrong_parameter_msg.format(flag_name="typical_p", flag_value=self.typical_p), + UserWarning, + ) + if ( + self.top_k is not None and self.top_k != 50 and self.penalty_alpha is None + ): # contrastive search uses top_k + warnings.warn( + greedy_wrong_parameter_msg.format(flag_name="top_k", flag_value=self.top_k), + UserWarning, + ) + if self.epsilon_cutoff is not None and self.epsilon_cutoff != 0.0: + warnings.warn( + greedy_wrong_parameter_msg.format(flag_name="epsilon_cutoff", flag_value=self.epsilon_cutoff), + UserWarning, + ) + if self.eta_cutoff is not None and self.eta_cutoff != 0.0: + warnings.warn( + greedy_wrong_parameter_msg.format(flag_name="eta_cutoff", flag_value=self.eta_cutoff), + UserWarning, + ) + + # 2. detect beam-only parameterization when not in beam mode + if self.num_beams is None: + warnings.warn("`num_beams` is set to None - defaulting to 1.", UserWarning) + self.num_beams = 1 + + if self.num_beams == 1: + single_beam_wrong_parameter_msg = ( + "`num_beams` is set to 1. However, `{flag_name}` is set to `{flag_value}` -- this flag is only used " + "in beam-based generation modes. You should set `num_beams>1` or unset `{flag_name}`." + fix_location + ) + if self.early_stopping is not False: + warnings.warn( + single_beam_wrong_parameter_msg.format(flag_name="early_stopping", flag_value=self.early_stopping), + UserWarning, + ) + if self.num_beam_groups is not None and self.num_beam_groups != 1: + warnings.warn( + single_beam_wrong_parameter_msg.format( + flag_name="num_beam_groups", flag_value=self.num_beam_groups + ), + UserWarning, + ) + if self.diversity_penalty is not None and self.diversity_penalty != 0.0: + warnings.warn( + single_beam_wrong_parameter_msg.format( + flag_name="diversity_penalty", flag_value=self.diversity_penalty + ), + UserWarning, + ) + if self.length_penalty is not None and self.length_penalty != 1.0: + warnings.warn( + single_beam_wrong_parameter_msg.format(flag_name="length_penalty", flag_value=self.length_penalty), + UserWarning, + ) + if self.constraints is not None: + warnings.warn( + single_beam_wrong_parameter_msg.format(flag_name="constraints", flag_value=self.constraints), + UserWarning, + ) + + # 3. detect incorrect paramaterization specific to advanced beam modes + else: + # constrained beam search + if self.constraints is not None or self.force_words_ids is not None: + constrained_wrong_parameter_msg = ( + "one of `constraints`, `force_words_ids` is not `None`, triggering constrained beam search. However, " + "`{flag_name}` is set to `{flag_value}`, which is incompatible with this generation mode. Set " + "`constraints` and `force_words_ids` to `None` or unset `{flag_name}` to continue." + fix_location + ) + if self.do_sample is True: + raise ValueError( + constrained_wrong_parameter_msg.format(flag_name="do_sample", flag_value=self.do_sample) + ) + if self.num_beam_groups is not None and self.num_beam_groups != 1: + raise ValueError( + constrained_wrong_parameter_msg.format( + flag_name="num_beam_groups", flag_value=self.num_beam_groups + ) + ) + # group beam search + if self.diversity_penalty != 0.0 or self.num_beam_groups != 1: + group_error_prefix = ( + "`diversity_penalty` is not 0.0 or `num_beam_groups` is not 1, triggering group beam search. In " + "this generation mode, " + ) + if self.do_sample is True: + raise ValueError(group_error_prefix + "`do_sample` must be set to `False`") + if self.num_beams % self.num_beam_groups != 0: + raise ValueError(group_error_prefix + "`num_beams` should be divisible by `num_beam_groups`") + if self.diversity_penalty == 0.0: + raise ValueError( + group_error_prefix + + "`diversity_penalty` should be greater than `0.0`, otherwise your groups will be identical." + ) + # DoLa generation + if self.dola_layers is not None and (self.repetition_penalty is None or self.repetition_penalty < 1.2): + warnings.warn( + "`dola_layers` is set to trigger DoLa decoding, but `repetition_penalty` is set to a value of " + f"{self.repetition_penalty}, which could induce unwanted repetition. The recommended value for " + "DoLa decoding is `repetition_penalty>=1.2`.", + UserWarning, + ) + + # 4. check `num_return_sequences` + if self.num_return_sequences != 1: + if self.num_beams == 1: + if self.do_sample is False: + raise ValueError( + "Greedy methods without beam search do not support `num_return_sequences` different than 1 " + f"(got {self.num_return_sequences})." + ) + elif self.num_return_sequences > self.num_beams: + raise ValueError( + f"`num_return_sequences` ({self.num_return_sequences}) has to be smaller or equal to `num_beams` " + f"({self.num_beams})." + ) + + # 5. check cache-related arguments + if self.cache_implementation is not None and self.cache_implementation not in ALL_CACHE_IMPLEMENTATIONS: + raise ValueError( + f"Invalid `cache_implementation` ({self.cache_implementation}). Choose one of: " + f"{ALL_CACHE_IMPLEMENTATIONS}" + ) + if self.cache_config is not None: + cache_class = CACHE_CONFIG_MAPPING.get(self.cache_implementation) + if cache_class is None: + raise ValueError( + "You provided a `cache_config` but the cache implementation you are using " + f"({self.cache_implementation}) does not require any config. Make sure to use the " + "correct cache implementation matching your cache config." + ) + if not isinstance(self.cache_config, cache_class): + self.cache_config = cache_class.from_dict(self.cache_config) + self.cache_config.validate() + if self.use_cache is False: + # In this case, all cache-related arguments should be unset. However, since `use_cache=False` is often used + # passed to `generate` directly to hot-fix cache issues, let's raise a warning instead of an error + # (otherwise a user might need to overwrite several parameters). + no_cache_warning = ( + "You have set `use_cache` to `False`, but {cache_arg} is set to {cache_arg_value}. {cache_arg} will " + "have no effect." + ) + for arg_name in ("cache_implementation", "cache_config", "return_legacy_cache"): + if getattr(self, arg_name) is not None: + logger.warning_once( + no_cache_warning.format(cache_arg=arg_name, cache_arg_value=getattr(self, arg_name)) + ) + + # 6. check watermarking arguments + if self.watermarking_config is not None: + if not ( + isinstance(self.watermarking_config, WatermarkingConfig) + or isinstance(self.watermarking_config, SynthIDTextWatermarkingConfig) + ): + warnings.warn( + "`watermarking_config` as a dict is deprecated. Please construct `watermarking_config` object with " + "`WatermarkingConfig` or `SynthIDTextWatermarkingConfig` class.", + FutureWarning, + ) + self.watermarking_config = WatermarkingConfig.from_dict(self.watermarking_config) + self.watermarking_config.validate() + + # 7. performances arguments + if not isinstance(self.compile_config, CompileConfig): + raise ValueError( + f"You provided `compile_config` as an instance of {type(self.compile_config)}, but it must be an instance of `CompileConfig`." + ) + + # 8. other incorrect combinations + if self.return_dict_in_generate is not True: + for extra_output_flag in self.extra_output_flags: + if getattr(self, extra_output_flag) is True: + warnings.warn( + f"`return_dict_in_generate` is NOT set to `True`, but `{extra_output_flag}` is. When " + f"`return_dict_in_generate` is not `True`, `{extra_output_flag}` is ignored.", + UserWarning, + ) + + # 8. check common issue: passing `generate` arguments inside the generation config + generate_arguments = ( + "logits_processor", + "stopping_criteria", + "prefix_allowed_tokens_fn", + "synced_gpus", + "assistant_model", + "streamer", + "negative_prompt_ids", + "negative_prompt_attention_mask", + ) + for arg in generate_arguments: + if hasattr(self, arg): + raise ValueError( + f"Argument `{arg}` is not a valid argument of `GenerationConfig`. It should be passed to " + "`generate()` (or a pipeline) directly." + ) + + def save_pretrained( + self, + save_directory: Union[str, os.PathLike], + config_file_name: Optional[Union[str, os.PathLike]] = None, + push_to_hub: bool = False, + **kwargs, + ): + r""" + Save a generation configuration object to the directory `save_directory`, so that it can be re-loaded using the + [`~GenerationConfig.from_pretrained`] class method. + + Args: + save_directory (`str` or `os.PathLike`): + Directory where the configuration JSON file will be saved (will be created if it does not exist). + config_file_name (`str` or `os.PathLike`, *optional*, defaults to `"generation_config.json"`): + Name of the generation configuration JSON file to be saved in `save_directory`. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + + # At save time, validate the instance -- if any warning/exception is thrown, we refuse to save the instance. + # This strictness is enforced to prevent bad configurations from being saved and re-used. + try: + with warnings.catch_warnings(record=True) as caught_warnings: + self.validate() + if len(caught_warnings) > 0: + raise ValueError(str([w.message for w in caught_warnings])) + except ValueError as exc: + raise ValueError( + "The generation config instance is invalid -- `.validate()` throws warnings and/or exceptions. " + "Fix these issues to save the configuration.\n\nThrown during validation:\n" + str(exc) + ) + + use_auth_token = kwargs.pop("use_auth_token", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. " + "Please use `token` instead.", + FutureWarning, + ) + if kwargs.get("token", None) is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + kwargs["token"] = use_auth_token + + config_file_name = config_file_name if config_file_name is not None else GENERATION_CONFIG_NAME + + if os.path.isfile(save_directory): + raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = self._create_repo(repo_id, **kwargs) + files_timestamps = self._get_files_timestamps(save_directory) + + output_config_file = os.path.join(save_directory, config_file_name) + + self.to_json_file(output_config_file, use_diff=True) + logger.info(f"Configuration saved in {output_config_file}") + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=kwargs.get("token"), + ) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name: Union[str, os.PathLike], + config_file_name: Optional[Union[str, os.PathLike]] = None, + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + **kwargs, + ) -> "GenerationConfig": + r""" + Instantiate a [`GenerationConfig`] from a generation configuration file. + + Args: + pretrained_model_name (`str` or `os.PathLike`): + This can be either: + + - a string, the *model id* of a pretrained model configuration hosted inside a model repo on + huggingface.co. + - a path to a *directory* containing a configuration file saved using the + [`~GenerationConfig.save_pretrained`] method, e.g., `./my_model_directory/`. + config_file_name (`str` or `os.PathLike`, *optional*, defaults to `"generation_config.json"`): + Name of the generation configuration JSON file to be loaded from `pretrained_model_name`. + 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. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the configuration files and override the cached versions if + they exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`. + + + + return_unused_kwargs (`bool`, *optional*, defaults to `False`): + If `False`, then this function returns just the final configuration object. + + If `True`, then this functions returns a `Tuple(config, unused_kwargs)` where *unused_kwargs* is a + dictionary consisting of the key/value pairs whose keys are not configuration attributes: i.e., the + part of `kwargs` which has not been used to update `config` and is otherwise ignored. + 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. + kwargs (`Dict[str, Any]`, *optional*): + The values in kwargs of any keys which are configuration attributes will be used to override the loaded + values. Behavior concerning key/value pairs whose keys are *not* configuration attributes is controlled + by the `return_unused_kwargs` keyword parameter. + + Returns: + [`GenerationConfig`]: The configuration object instantiated from this pretrained model. + + Examples: + + ```python + >>> from transformers import GenerationConfig + + >>> # Download configuration from huggingface.co and cache. + >>> generation_config = GenerationConfig.from_pretrained("openai-community/gpt2") + + >>> # E.g. config was saved using *save_pretrained('./test/saved_model/')* + >>> generation_config.save_pretrained("./test/saved_model/") + >>> generation_config = GenerationConfig.from_pretrained("./test/saved_model/") + + >>> # You can also specify configuration names to your generation configuration file + >>> generation_config.save_pretrained("./test/saved_model/", config_file_name="my_configuration.json") + >>> generation_config = GenerationConfig.from_pretrained("./test/saved_model/", "my_configuration.json") + + >>> # If you'd like to try a minor variation to an existing configuration, you can also pass generation + >>> # arguments to `.from_pretrained()`. Be mindful that typos and unused arguments will be ignored + >>> generation_config, unused_kwargs = GenerationConfig.from_pretrained( + ... "openai-community/gpt2", top_k=1, foo=False, do_sample=True, return_unused_kwargs=True + ... ) + >>> generation_config.top_k + 1 + + >>> unused_kwargs + {'foo': False} + ```""" + config_file_name = config_file_name if config_file_name is not None else GENERATION_CONFIG_NAME + + resume_download = kwargs.pop("resume_download", None) + proxies = kwargs.pop("proxies", None) + use_auth_token = kwargs.pop("use_auth_token", None) + subfolder = kwargs.pop("subfolder", "") + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + commit_hash = kwargs.pop("_commit_hash", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + user_agent = {"file_type": "config", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + config_path = os.path.join(pretrained_model_name, config_file_name) + config_path = str(config_path) + + is_local = os.path.exists(config_path) + if os.path.isfile(os.path.join(subfolder, config_path)): + # Special case when config_path is a local file + resolved_config_file = config_path + is_local = True + elif is_remote_url(config_path): + configuration_file = config_path + resolved_config_file = download_url(config_path) + else: + configuration_file = config_file_name + try: + # Load from local folder or from cache or download from model Hub and cache + resolved_config_file = cached_file( + pretrained_model_name, + configuration_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _commit_hash=commit_hash, + ) + commit_hash = extract_commit_hash(resolved_config_file, commit_hash) + except EnvironmentError: + # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted to + # the original exception. + raise + except Exception: + # For any other exception, we throw a generic error. + raise EnvironmentError( + f"Can't load the configuration of '{pretrained_model_name}'. If you were trying to load it" + " from 'https://huggingface.co/models', make sure you don't have a local directory with the same" + f" name. Otherwise, make sure '{pretrained_model_name}' is the correct path to a directory" + f" containing a {configuration_file} file" + ) + + try: + # Load config dict + config_dict = cls._dict_from_json_file(resolved_config_file) + config_dict["_commit_hash"] = commit_hash + except (json.JSONDecodeError, UnicodeDecodeError): + raise EnvironmentError( + f"It looks like the config file at '{resolved_config_file}' is not a valid JSON file." + ) + + if is_local: + logger.info(f"loading configuration file {resolved_config_file}") + else: + logger.info(f"loading configuration file {configuration_file} from cache at {resolved_config_file}") + + if kwargs.get("return_unused_kwargs") is True: + config, unused_kwargs = cls.from_dict(config_dict, **kwargs) + config._original_object_hash = hash(config) # Hash to detect whether the instance was modified + return config, unused_kwargs + else: + config = cls.from_dict(config_dict, **kwargs) + config._original_object_hash = hash(config) # Hash to detect whether the instance was modified + return config + + @classmethod + def _dict_from_json_file(cls, json_file: Union[str, os.PathLike]): + with open(json_file, "r", encoding="utf-8") as reader: + text = reader.read() + return json.loads(text) + + @classmethod + def from_dict(cls, config_dict: Dict[str, Any], **kwargs) -> "GenerationConfig": + """ + Instantiates a [`GenerationConfig`] from a Python dictionary of parameters. + + Args: + config_dict (`Dict[str, Any]`): + Dictionary that will be used to instantiate the configuration object. + kwargs (`Dict[str, Any]`): + Additional parameters from which to initialize the configuration object. + + Returns: + [`GenerationConfig`]: The configuration object instantiated from those parameters. + """ + return_unused_kwargs = kwargs.pop("return_unused_kwargs", False) + # Those arguments may be passed along for our internal telemetry. + # We remove them so they don't appear in `return_unused_kwargs`. + kwargs.pop("_from_auto", None) + kwargs.pop("_from_pipeline", None) + # The commit hash might have been updated in the `config_dict`, we don't want the kwargs to erase that update. + if "_commit_hash" in kwargs and "_commit_hash" in config_dict: + kwargs["_commit_hash"] = config_dict["_commit_hash"] + + # The line below allows model-specific config to be loaded as well through kwargs, with safety checks. + # See https://github.com/huggingface/transformers/pull/21269 + config = cls(**{**config_dict, **kwargs}) + unused_kwargs = config.update(**kwargs) + + logger.info(f"Generate config {config}") + if return_unused_kwargs: + return config, unused_kwargs + else: + return config + + def dict_torch_dtype_to_str(self, d: Dict[str, Any]) -> None: + """ + Checks whether the passed dictionary and its nested dicts have a *torch_dtype* key and if it's not None, + converts torch.dtype to a string of just the type. For example, `torch.float32` get converted into *"float32"* + string, which can then be stored in the json format. + """ + if d.get("torch_dtype", None) is not None and not isinstance(d["torch_dtype"], str): + d["torch_dtype"] = str(d["torch_dtype"]).split(".")[1] + for value in d.values(): + if isinstance(value, dict): + self.dict_torch_dtype_to_str(value) + + def to_diff_dict(self) -> Dict[str, Any]: + """ + Removes all attributes from config which correspond to the default config attributes for better readability and + serializes to a Python dictionary. + + Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance, + """ + config_dict = self.to_dict() + + # get the default config dict + default_config_dict = GenerationConfig().to_dict() + + serializable_config_dict = {} + + # only serialize values that differ from the default config + for key, value in config_dict.items(): + if key not in default_config_dict or key == "transformers_version" or value != default_config_dict[key]: + serializable_config_dict[key] = value + + self.dict_torch_dtype_to_str(serializable_config_dict) + return serializable_config_dict + + def to_dict(self) -> Dict[str, Any]: + """ + Serializes this instance to a Python dictionary. + + Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + output = copy.deepcopy(self.__dict__) + + # Fields to ignore at serialization time + if "_commit_hash" in output: + del output["_commit_hash"] + if "_original_object_hash" in output: + del output["_original_object_hash"] + if "compile_config" in output: + del output["compile_config"] + + # Transformers version when serializing this file + output["transformers_version"] = __version__ + + self.dict_torch_dtype_to_str(output) + return output + + def to_json_string(self, use_diff: bool = True, ignore_metadata: bool = False) -> str: + """ + Serializes this instance to a JSON string. + + Args: + use_diff (`bool`, *optional*, defaults to `True`): + If set to `True`, only the difference between the config instance and the default `GenerationConfig()` + is serialized to JSON string. + ignore_metadata (`bool`, *optional*, defaults to `False`): + Whether to ignore the metadata fields present in the instance + + Returns: + `str`: String containing all the attributes that make up this configuration instance in JSON format. + """ + if use_diff is True: + config_dict = self.to_diff_dict() + else: + config_dict = self.to_dict() + + if ignore_metadata: + for metadata_field in METADATA_FIELDS: + config_dict.pop(metadata_field, None) + + def convert_keys_to_string(obj): + if isinstance(obj, dict): + return {str(key): convert_keys_to_string(value) for key, value in obj.items()} + elif isinstance(obj, list): + return [convert_keys_to_string(item) for item in obj] + else: + return obj + + def convert_dataclass_to_dict(obj): + if isinstance(obj, dict): + return {key: convert_dataclass_to_dict(value) for key, value in obj.items()} + elif is_dataclass(obj): + return obj.to_dict() + else: + return obj + + config_dict = convert_keys_to_string(config_dict) + config_dict = convert_dataclass_to_dict(config_dict) + + return json.dumps(config_dict, indent=2, sort_keys=True) + "\n" + + def to_json_file(self, json_file_path: Union[str, os.PathLike], use_diff: bool = True): + """ + Save this instance to a JSON file. + + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file in which this configuration instance's parameters will be saved. + use_diff (`bool`, *optional*, defaults to `True`): + If set to `True`, only the difference between the config instance and the default `GenerationConfig()` + is serialized to JSON file. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + writer.write(self.to_json_string(use_diff=use_diff)) + + @classmethod + def from_model_config(cls, model_config: PretrainedConfig) -> "GenerationConfig": + """ + Instantiates a [`GenerationConfig`] from a [`PretrainedConfig`]. This function is useful to convert legacy + [`PretrainedConfig`] objects, which may contain generation parameters, into a stand-alone [`GenerationConfig`]. + + Args: + model_config (`PretrainedConfig`): + The model config that will be used to instantiate the generation config. + + Returns: + [`GenerationConfig`]: The configuration object instantiated from those parameters. + """ + config_dict = model_config.to_dict() + config_dict.pop("_from_model_config", None) + + # Removes all `None` from the model config dict -- this lets the generation config defaults to take hold + config_dict = {key: value for key, value in config_dict.items() if value is not None} + + generation_config = cls.from_dict(config_dict, return_unused_kwargs=False, _from_model_config=True) + + # Special case: some models have generation attributes set in the decoder. Use them if still unset in the + # generation config (which in turn is defined from the outer attributes of model config). + decoder_config = model_config.get_text_config(decoder=True) + if decoder_config is not model_config: + default_generation_config = GenerationConfig() + decoder_config_dict = decoder_config.to_dict() + for attr in generation_config.to_dict().keys(): + is_unset = getattr(generation_config, attr) == getattr(default_generation_config, attr) + if attr in decoder_config_dict and is_unset: + setattr(generation_config, attr, decoder_config_dict[attr]) + + # If any `output_...` flag is set to `True`, we ensure `return_dict_in_generate` is set to `True`. + if generation_config.return_dict_in_generate is False: + if any( + getattr(generation_config, extra_output_flag, False) + for extra_output_flag in generation_config.extra_output_flags + ): + generation_config.return_dict_in_generate = True + + # Hash to detect whether the instance was modified + generation_config._original_object_hash = hash(generation_config) + return generation_config + + def update(self, **kwargs): + """ + Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes, + returning all the unused kwargs. + + Args: + kwargs (`Dict[str, Any]`): + Dictionary of attributes to tentatively update this class. + + Returns: + `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance. + """ + to_remove = [] + for key, value in kwargs.items(): + if hasattr(self, key): + setattr(self, key, value) + to_remove.append(key) + + # Confirm that the updated instance is still valid + self.validate() + + # Remove all the attributes that were updated, without modifying the input dict + unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove} + return unused_kwargs + + +@dataclass +class BaseWatermarkingConfig(ABC): + """Generic watermarking config""" + + @classmethod + def from_dict(cls, config_dict, **kwargs): + """ + Constructs a BaseWatermarkingConfig instance from a dictionary of parameters. + + Args: + config_dict (Dict[str, Any]): Dictionary containing configuration parameters. + **kwargs: Additional keyword arguments to override dictionary values. + + Returns: + BaseWatermarkingConfig: Instance of BaseWatermarkingConfig constructed from the dictionary. + """ + config = cls(**config_dict) + to_remove = [] + for key, value in kwargs.items(): + if hasattr(config, key): + setattr(config, key, value) + to_remove.append(key) + for key in to_remove: + kwargs.pop(key, None) + return config + + def to_json_file(self, json_file_path: Union[str, os.PathLike]): + """ + Save this instance to a JSON file. + + Args: + json_file_path (Union[str, os.PathLike]): Path to the JSON file in which this configuration instance's parameters will be saved. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + config_dict = self.to_dict() + json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n" + + writer.write(json_string) + + def to_dict(self) -> Dict[str, Any]: + """ + Serializes this instance to a Python dictionary. + + Returns: + Dict[str, Any]: Dictionary of all the attributes that make up this configuration instance. + """ + output = copy.deepcopy(self.__dict__) + return output + + def __iter__(self): + for attr, value in copy.deepcopy(self.__dict__).items(): + yield attr, value + + def __repr__(self): + return f"{self.__class__.__name__} {self.to_json_string()}" + + def to_json_string(self): + """ + Serializes this instance to a JSON formatted string. + + Returns: + str: JSON formatted string representing the configuration instance. + """ + return json.dumps(self.__dict__, indent=2) + "\n" + + def update(self, **kwargs): + """ + Update the configuration attributes with new values. + + Args: + **kwargs: Keyword arguments representing configuration attributes and their new values. + """ + for key, value in kwargs.items(): + if hasattr(self, key): + setattr(self, key, value) + + @abstractmethod + def validate(self): ... + + @abstractmethod + def construct_processor(self, vocab_size): ... + + +@dataclass +class WatermarkingConfig(BaseWatermarkingConfig): + """ + Class that holds arguments for watermark generation and should be passed into `GenerationConfig` during `generate`. + See [this paper](https://arxiv.org/abs/2306.04634) for more details on the arguments. + + Accepts the following keys: + - greenlist_ratio (`float`): + Used for watermarking. The ratio of "green" tokens used to the vocabulary size. Defaults to 0.25. + - bias (`float`): + Used with watermarking. The bias added to the selected "green" tokens' logits. Defaults to 2.0. + - hashing_key (`int`): + Hashing key used for watermarking. Defaults to 15485863 (the millionth prime). + - seeding_scheme (`str`): + Algorithm to use for watermarking. Accepts values: + - "lefthash" (default): "green" tokens selection depend on the last token (Algorithm 2 from the paper) + - "selfhash": "green" tokens selection depends on the current token itself (Algorithm 3 from the paper) + The downside of this scheme is that it considers all possible next tokens and can be slower than "lefthash". + - context_width(`int`): + The context length of previous tokens to use in seeding. Higher context length makes watermarking more robust. + """ + + def __init__( + self, + greenlist_ratio: Optional[float] = 0.25, + bias: Optional[float] = 2.0, + hashing_key: Optional[int] = 15485863, + seeding_scheme: Optional[str] = "lefthash", + context_width: Optional[int] = 1, + ): + self.greenlist_ratio = greenlist_ratio + self.bias = bias + self.hashing_key = hashing_key + self.seeding_scheme = seeding_scheme + self.context_width = context_width + + def validate(self): + watermark_missing_arg_msg = ( + "Some of the keys in `watermarking_config` are defined incorrectly. `{key}` should be {correct_value}` " + "but found {found_value}" + ) + if self.seeding_scheme not in ["selfhash", "lefthash"]: + raise ValueError( + watermark_missing_arg_msg.format( + key="seeding_scheme", + correct_value="[`selfhash`, `lefthash`]", + found_value=self.seeding_scheme, + ), + ) + if not 0.0 <= self.greenlist_ratio <= 1.0: + raise ValueError( + watermark_missing_arg_msg.format( + key="greenlist_ratio", + correct_value="in range between 0.0 and 1.0", + found_value=self.seeding_scheme, + ), + ) + if not self.context_width >= 1: + raise ValueError( + watermark_missing_arg_msg.format( + key="context_width", + correct_value="a positive integer", + found_value=self.context_width, + ), + ) + + def construct_processor(self, vocab_size: int, device) -> "WatermarkLogitsProcessor": + return WatermarkLogitsProcessor( + vocab_size=vocab_size, + device=device, + greenlist_ratio=self.greenlist_ratio, + bias=self.bias, + hashing_key=self.hashing_key, + seeding_scheme=self.seeding_scheme, + context_width=self.context_width, + ) + + +@dataclass +class SynthIDTextWatermarkingConfig(BaseWatermarkingConfig): + """ + Class that holds arguments for watermark generation and should be passed into `GenerationConfig` during `generate`. + See [this paper](https://www.nature.com/articles/s41586-024-08025-4) for more details on the arguments. + + Args: + ngram_len (`int`): + Ngram length. + keys (`List[int]`): + A sequence of watermarking keys, one for each depth. + context_history_size (`int`, *optional*, defaults to 1024): + Size of the tensor to keep track of seen contexts. + sampling_table_seed (`int`, *optional*, defaults to 0): + Random seed to generate the sampling table. + sampling_table_size (`int`, *optional*, defaults to 65536): + Size of the sampling table. + skip_first_ngram_calls (`bool`, *optional*, defaults to `False`): + Whether to skip first ngram calls. + debug_mode (`bool`, optional, *optional*, defaults to `False`): + Logits are modified to uniform one got before watermarking modification is applied. This is to test the + implementation. + + Examples: + ```python + >>> from transformers import AutoModelForCausalLM, AutoTokenizer, SynthIDTextWatermarkingConfig + + >>> tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b', padding_side="left") + >>> model = AutoModelForCausalLM.from_pretrained('google/gemma-2-2b') + + >>> # SynthID Text configuration + >>> watermarking_config = SynthIDTextWatermarkingConfig( + ... keys=[654, 400, 836, 123, 340, 443, 597, 160, 57], + ... ngram_len=5, + ... ) + + >>> # Generation with watermarking + >>> tokenized_prompts = tokenizer(["Once upon a time, "], return_tensors="pt", padding=True) + >>> output_sequences = model.generate( + ... **tokenized_prompts, watermarking_config=watermarking_config, do_sample=True, max_new_tokens=10 + ... ) + >>> watermarked_text = tokenizer.batch_decode(output_sequences, skip_special_tokens=True) + ``` + """ + + def __init__( + self, + ngram_len: int, + keys: List[int], + context_history_size: int = 1024, + sampling_table_seed: int = 0, + sampling_table_size: int = 2**16, + skip_first_ngram_calls: bool = False, + debug_mode: bool = False, + ): + self.ngram_len = ngram_len + self.keys = keys + self.sampling_table_size = sampling_table_size + self.sampling_table_seed = sampling_table_seed + self.context_history_size = context_history_size + self.skip_first_ngram_calls = skip_first_ngram_calls + self.debug_mode = debug_mode + + def validate(self): + watermark_missing_arg_msg = ( + "Some of the keys in `watermarking_config` are defined incorrectly. `{key}` should be {correct_value}` " + "but found {found_value}" + ) + if self.sampling_table_size > 2**24: + raise ValueError( + watermark_missing_arg_msg.format( + key="sampling_table_size", + correct_value="< 2**24", + found_value=self.sampling_table_size, + ), + ) + + def construct_processor(self, vocab_size: int, device) -> "WatermarkLogitsProcessor": + return SynthIDTextWatermarkLogitsProcessor( + ngram_len=self.ngram_len, + keys=self.keys, + sampling_table_size=self.sampling_table_size, + sampling_table_seed=self.sampling_table_seed, + context_history_size=self.context_history_size, + device=device, + skip_first_ngram_calls=self.skip_first_ngram_calls, + debug_mode=self.debug_mode, + ) + + +@dataclass +class CompileConfig: + """ + Class that holds arguments relative to `torch.compile` behavior, when using automatic compilation in `generate`. + See [`torch.compile`](https://pytorch.org/docs/stable/generated/torch.compile.html) for more details on the arguments. + + Args: + fullgraph (`bool`, *optional*, defaults to `True`): + If `True`, requires that the whole forward be capturable in a single graph. + dynamic (`bool` or `None`, *optional*): + Whether to try to use dynamic shape graphs. + backend (`str` or `Callable`, *optional*, defaults to `"inductor"`): + Backend to be used. + mode (`str`, *optional*, defaults to `"reduce-overhead"`): + Controls balance between performance and overhead. + options (`dict`, *optional*): + A dictionary of options to pass to the backend. + + Examples: + ```python + >>> from transformers import AutoModelForCausalLM, AutoTokenizer, CompileConfig + + >>> tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b') + >>> model = AutoModelForCausalLM.from_pretrained('google/gemma-2-2b').cuda() + + >>> # Automatic compile configuration, used with static cache + >>> compile_config = CompileConfig(dynamic=True) + + >>> # Generation with static cache and compile config + >>> input = tokenizer.encode("Hello there, how", return_tensors="pt").cuda() + >>> output = model.generate( + ... input, do_sample=False, max_new_tokens=300, cache_implementation="static", compile_config=compile_config + ... ) + >>> output_text = tokenizer.batch_decode(output, skip_special_tokens=True)[0] + ``` + """ + + fullgraph: bool = True + dynamic: Optional[bool] = None + backend: Union[str, Callable] = "inductor" + mode: str = "reduce-overhead" + options: Optional[dict] = None + # Used to flag our `generate` call to compile on e.g. CPU. Often not optimal, but useful for testing purposes. + _compile_all_devices = None + + def to_dict(self) -> Dict[str, Any]: + """Serializes this instance to a Python dictionary.""" + return copy.deepcopy({key: value for key, value in self.__dict__.items() if key != "_compile_all_devices"}) diff --git a/lib/python3.12/site-packages/transformers/generation/flax_logits_process.py b/lib/python3.12/site-packages/transformers/generation/flax_logits_process.py new file mode 100644 index 0000000000000000000000000000000000000000..d106c32defa4db3577d3a72158550d355aa5ece7 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/flax_logits_process.py @@ -0,0 +1,544 @@ +# coding=utf-8 +# Copyright 2021 The HuggingFace Inc. team +# +# 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. + +import inspect + +import jax +import jax.lax as lax +import jax.numpy as jnp +from jax.experimental import sparse + +from ..utils import add_start_docstrings +from ..utils.logging import get_logger + + +logger = get_logger(__name__) + + +LOGITS_PROCESSOR_INPUTS_DOCSTRING = r""" + Args: + input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`): + Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam + search or log softmax for each vocabulary token when using beam search + kwargs (`Dict[str, Any]`, *optional*): + Additional logits processor specific kwargs. + + Return: + `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores. + +""" + + +class FlaxLogitsProcessor: + """Abstract base class for all logit processors that can be applied during generation.""" + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray) -> jnp.ndarray: + """Flax method for processing logits.""" + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + +class FlaxLogitsWarper: + """Abstract base class for all logit warpers that can be applied during generation with multinomial sampling.""" + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray) -> jnp.ndarray: + """Flax method for warping logits.""" + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + +class FlaxLogitsProcessorList(list): + """ + This class can be used to create a list of [`FlaxLogitsProcessor`] or [`FlaxLogitsWarper`] to subsequently process + a `scores` input tensor. This class inherits from list and adds a specific *__call__* method to apply each + [`FlaxLogitsProcessor`] or [`FlaxLogitsWarper`] to the inputs. + """ + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int, **kwargs) -> jnp.ndarray: + for processor in self: + function_args = inspect.signature(processor.__call__).parameters + if len(function_args) > 3: + if not all(arg in kwargs for arg in list(function_args.keys())[2:]): + raise ValueError( + f"Make sure that all the required parameters: {list(function_args.keys())} for " + f"{processor.__class__} are passed to the logits processor." + ) + scores = processor(input_ids, scores, cur_len, **kwargs) + else: + scores = processor(input_ids, scores, cur_len) + return scores + + +class FlaxTemperatureLogitsWarper(FlaxLogitsWarper): + r""" + [`FlaxLogitsWarper`] for temperature (exponential scaling output probability distribution). + + Args: + temperature (`float`): + The value used to module the logits distribution. + """ + + def __init__(self, temperature: float): + if not isinstance(temperature, float) or not (temperature > 0): + raise ValueError(f"`temperature` has to be a strictly positive float, but is {temperature}") + + self.temperature = temperature + + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray: + scores = scores / self.temperature + return scores + + +class FlaxTopPLogitsWarper(FlaxLogitsWarper): + """ + [`FlaxLogitsWarper`] that performs top-p, i.e. restricting to top tokens summing to prob_cut_off <= prob_cut_off. + + Args: + top_p (`float`): + If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or + higher are kept for generation. + filter_value (`float`, *optional*, defaults to -inf): + All filtered values will be set to this float value. + min_tokens_to_keep (`int`, *optional*, defaults to 1): + Minimum number of tokens that cannot be filtered. + """ + + def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1): + if not isinstance(top_p, float) or (top_p < 0 or top_p > 1.0): + raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}") + if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1): + raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}") + + self.top_p = top_p + self.filter_value = filter_value + self.min_tokens_to_keep = min_tokens_to_keep + + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray: + topk_scores, topk_indices = lax.top_k(scores, scores.shape[-1]) + + mask_scores = jnp.full_like(scores, self.filter_value) + cumulative_probs = jax.nn.softmax(topk_scores, axis=-1).cumsum(axis=-1) + score_mask = cumulative_probs < self.top_p + + # include the token that is higher than top_p as well + score_mask = jnp.roll(score_mask, 1) + score_mask |= score_mask.at[:, 0].set(True) + + # min tokens to keep + score_mask = score_mask.at[:, : self.min_tokens_to_keep].set(True) + + topk_next_scores = jnp.where(score_mask, topk_scores, mask_scores) + next_scores = jax.lax.sort_key_val(topk_indices, topk_next_scores)[-1] + + return next_scores + + +class FlaxTopKLogitsWarper(FlaxLogitsWarper): + r""" + [`FlaxLogitsWarper`] that performs top-k, i.e. restricting to the k highest probability elements. + + Args: + top_k (`int`): + The number of highest probability vocabulary tokens to keep for top-k-filtering. + filter_value (`float`, *optional*, defaults to -inf): + All filtered values will be set to this float value. + min_tokens_to_keep (`int`, *optional*, defaults to 1): + Minimum number of tokens that cannot be filtered. + """ + + def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1): + if not isinstance(top_k, int) or top_k <= 0: + raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}") + + self.top_k = max(top_k, min_tokens_to_keep) + self.filter_value = filter_value + + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray: + batch_size, vocab_size = scores.shape + next_scores_flat = jnp.full(batch_size * vocab_size, self.filter_value) + + topk = min(self.top_k, scores.shape[-1]) # Safety check + topk_scores, topk_indices = lax.top_k(scores, topk) + shift = jnp.broadcast_to((jnp.arange(batch_size) * vocab_size)[:, None], (batch_size, topk)).flatten() + topk_scores_flat = topk_scores.flatten() + topk_indices_flat = topk_indices.flatten() + shift + + next_scores_flat = next_scores_flat.at[topk_indices_flat].set(topk_scores_flat) + next_scores = next_scores_flat.reshape(batch_size, vocab_size) + return next_scores + + +class FlaxForcedBOSTokenLogitsProcessor(FlaxLogitsProcessor): + r""" + [`FlaxLogitsProcessor`] that enforces the specified token as the first generated token. + + Args: + bos_token_id (`int`): + The id of the token to force as the first generated token. + """ + + def __init__(self, bos_token_id: int): + self.bos_token_id = bos_token_id + + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray: + new_scores = jnp.full(scores.shape, -float("inf")) + + apply_penalty = 1 - jnp.bool_(cur_len - 1) + + scores = jnp.where(apply_penalty, new_scores.at[:, self.bos_token_id].set(0), scores) + + return scores + + +class FlaxForcedEOSTokenLogitsProcessor(FlaxLogitsProcessor): + r""" + [`FlaxLogitsProcessor`] that enforces the specified token as the last generated token when `max_length` is reached. + + Args: + max_length (`int`): + The maximum length of the sequence to be generated. + eos_token_id (`int`): + The id of the token to force as the last generated token when `max_length` is reached. + """ + + def __init__(self, max_length: int, eos_token_id: int): + self.max_length = max_length + self.eos_token_id = eos_token_id + + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray: + new_scores = jnp.full(scores.shape, -float("inf")) + + apply_penalty = 1 - jnp.bool_(cur_len - self.max_length + 1) + + scores = jnp.where(apply_penalty, new_scores.at[:, self.eos_token_id].set(0), scores) + + return scores + + +class FlaxMinLengthLogitsProcessor(FlaxLogitsProcessor): + r""" + [`FlaxLogitsProcessor`] enforcing a min-length by setting EOS probability to 0. + + Args: + min_length (`int`): + The minimum length below which the score of `eos_token_id` is set to `-float("Inf")`. + eos_token_id (`int`): + The id of the *end-of-sequence* token. + """ + + def __init__(self, min_length: int, eos_token_id: int): + if not isinstance(min_length, int) or min_length < 0: + raise ValueError(f"`min_length` has to be a positive integer, but is {min_length}") + + if not isinstance(eos_token_id, int) or eos_token_id < 0: + raise ValueError(f"`eos_token_id` has to be a positive integer, but is {eos_token_id}") + + self.min_length = min_length + self.eos_token_id = eos_token_id + + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray: + # create boolean flag to decide if min length penalty should be applied + apply_penalty = 1 - jnp.clip(cur_len - self.min_length, 0, 1) + + scores = jnp.where(apply_penalty, scores.at[:, self.eos_token_id].set(-float("inf")), scores) + + return scores + + +class FlaxSuppressTokensAtBeginLogitsProcessor(FlaxLogitsProcessor): + r""" + [`FlaxLogitsProcessor`] supressing a list of tokens as soon as the `generate` function starts generating using + `begin_index` tokens. This should ensure that the tokens defined by `begin_suppress_tokens` are not sampled at the + beginning of the generation. + + Args: + begin_suppress_tokens (`List[int]`): + Tokens to not sample. + begin_index (`int`): + Index where the tokens are suppressed. + """ + + def __init__(self, begin_suppress_tokens, begin_index): + self.begin_suppress_tokens = list(begin_suppress_tokens) + self.begin_index = begin_index + + def __call__(self, input_ids, scores, cur_len: int): + apply_penalty = 1 - jnp.bool_(cur_len - self.begin_index) + + scores = jnp.where(apply_penalty, scores.at[:, self.begin_suppress_tokens].set(-float("inf")), scores) + + return scores + + +class FlaxSuppressTokensLogitsProcessor(FlaxLogitsProcessor): + r""" + [`FlaxLogitsProcessor`] suppressing a list of tokens at each decoding step. The processor will set their log probs + to be `-inf` so they are not sampled. + + Args: + suppress_tokens (`list`): + Tokens to not sample. + """ + + def __init__(self, suppress_tokens: list): + self.suppress_tokens = list(suppress_tokens) + + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray: + scores = scores.at[..., self.suppress_tokens].set(-float("inf")) + + return scores + + +class FlaxForceTokensLogitsProcessor(FlaxLogitsProcessor): + r""" + [`FlaxLogitsProcessor`] that takes a list of pairs of integers which indicates a mapping from generation indices to + token indices that will be forced before sampling. The processor will set their log probs to 0 and all other tokens + to `-inf` so that they are sampled at their corresponding index. + + Args: + force_token_map (`list`): + Map giving token ids and indices where they will be forced to be sampled. + """ + + def __init__(self, force_token_map): + force_token_map = dict(force_token_map) + # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the + # index of the array corresponds to the index of the token to be forced, for XLA compatibility. + # Indexes without forced tokens will have a negative value. + force_token_array = jnp.ones((max(force_token_map.keys()) + 1), dtype=jnp.int32) * -1 + for index, token in force_token_map.items(): + if token is not None: + force_token_array = force_token_array.at[index].set(token) + self.force_token_array = jnp.int32(force_token_array) + + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray: + def _force_token(generation_idx): + batch_size = scores.shape[0] + current_token = self.force_token_array[generation_idx] + + new_scores = jnp.ones_like(scores, dtype=scores.dtype) * -float("inf") + updates = jnp.zeros((batch_size, 1), dtype=scores.dtype) + new_scores = lax.dynamic_update_slice(new_scores, updates, (0, current_token)) + return new_scores + + scores = lax.cond( + cur_len >= self.force_token_array.shape[0], + # If the current length is geq than the length of force_token_array, the processor does nothing. + lambda: scores, + # Otherwise, it may force a certain token. + lambda: lax.cond( + self.force_token_array[cur_len] >= 0, + # Only valid (positive) tokens are forced + lambda: _force_token(cur_len), + # Otherwise, the processor does nothing. + lambda: scores, + ), + ) + return scores + + +class FlaxWhisperTimeStampLogitsProcessor(FlaxLogitsProcessor): + r""" + Whisper specific Processor. This processor can be used to force a list of tokens. The processor will set their log + probs to `inf` so that they are sampled at their corresponding index. + + Args: + generate_config (`GenerateConfig`): + The generate config used to generate the output. The following parameters are required: + eos_token_id (`int`, *optional*, defaults to 50257): + The id of the *end-of-sequence* token. + no_timestamps_token_id (`int`, *optional*, defaults to 50363): + The id of the `"<|notimestamps|>"` token. + max_initial_timestamp_index (`int`, *optional*, defaults to 1): + Used to set the maximum value of the initial timestamp. This is used to prevent the model from + predicting timestamps that are too far in the future. + """ + + def __init__(self, generate_config, model_config, decoder_input_length): + self.eos_token_id = generate_config.eos_token_id + self.no_timestamps_token_id = generate_config.no_timestamps_token_id + self.timestamp_begin = generate_config.no_timestamps_token_id + 1 + + self.begin_index = decoder_input_length + 1 + + if generate_config.is_multilingual: + # room for language token and task token + self.begin_index += 2 + if hasattr(generate_config, "max_initial_timestamp_index"): + self.max_initial_timestamp_index = generate_config.max_initial_timestamp_index + else: + self.max_initial_timestamp_index = model_config.vocab_size + if self.max_initial_timestamp_index is None: + self.max_initial_timestamp_index = model_config.vocab_size + + def __call__(self, input_ids, scores, cur_len): + # suppress <|notimestamps|> which is handled by without_timestamps + scores = scores.at[:, self.no_timestamps_token_id].set(-float("inf")) + + def handle_pairs(input_ids_k, scores_k): + last_was_timestamp = jnp.where((cur_len - self.begin_index) >= 1, True, False) + last_was_timestamp = jnp.where( + input_ids_k[cur_len - 1] >= self.timestamp_begin, + True and last_was_timestamp, + False, + ) + + penultimate_was_timestamp = jnp.where((cur_len - self.begin_index) < 2, True, False) + penultimate_was_timestamp = jnp.where( + input_ids_k[cur_len - 2] >= self.timestamp_begin, + True, + penultimate_was_timestamp, + ) + + return jnp.where( + last_was_timestamp, + jnp.where( + penultimate_was_timestamp > 0, + scores_k.at[self.timestamp_begin :].set(-float("inf")), + scores_k.at[: self.eos_token_id].set(-float("inf")), + ), + scores_k, + ) + + scores = jax.vmap(handle_pairs)(input_ids, scores) + + apply_max_initial_timestamp = jnp.where(cur_len == self.begin_index, True, False) + apply_max_initial_timestamp = jnp.where( + self.max_initial_timestamp_index is not None, + True and apply_max_initial_timestamp, + False, + ) + + last_allowed = self.timestamp_begin + self.max_initial_timestamp_index + + scores = jnp.where( + apply_max_initial_timestamp, + scores.at[:, last_allowed + 1 :].set(-float("inf")), + scores, + ) + + # if sum of probability over timestamps is above any other token, sample timestamp + logprobs = jax.nn.log_softmax(scores, axis=-1) + + def handle_cumulative_probs(logprobs_k, scores_k): + timestamp_logprob = jax.nn.logsumexp(logprobs_k[self.timestamp_begin :], axis=-1) + max_text_token_logprob = jnp.max(logprobs_k[: self.timestamp_begin]) + return jnp.where( + timestamp_logprob > max_text_token_logprob, + scores_k.at[: self.timestamp_begin].set(-float("inf")), + scores_k, + ) + + scores = jax.vmap(handle_cumulative_probs)(logprobs, scores) + + return scores + + +class FlaxNoRepeatNGramLogitsProcessor(FlaxLogitsProcessor): + r""" + [`FlaxLogitsProcessor`] that enforces no repetition of n-grams. See + [Fairseq](https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345). + + Args: + ngram_size (`int`): + All ngrams of size `ngram_size` can only occur once. + """ + + def __init__(self, ngram_size: int): + if not isinstance(ngram_size, int) or ngram_size <= 0: + raise ValueError(f"`ngram_size` has to be a strictly positive integer, but is {ngram_size}") + self.ngram_size = ngram_size + + def get_previous_ngrams(self, input_ids: jnp.ndarray, vocab_size: int, cur_len: int): + """ + get a matrix of size (batch_size,) + (vocab_size,)*n (for n-grams) that + represent the n-grams that occurred previously. + The BCOO representation allow to store only the few non-zero entries, instead of the full (huge) matrix + """ + batch_size, seq_len = input_ids.shape + # number of n-grams in the whole sequence + seq_ngrams = seq_len - (self.ngram_size - 1) + # number of n-grams in the currently generated sequence + cur_ngrams = cur_len - (self.ngram_size - 1) + + def body_fun(i, val): + b = i % batch_size + pos = i // batch_size + return val.at[i].set( + jnp.array( + [ + b, + ] + + [jnp.array(input_ids)[b, pos + j] for j in range(self.ngram_size)] + ) + ) + + shape = (batch_size * seq_ngrams, self.ngram_size + 1) + all_update_indices = jax.lax.fori_loop( + 0, batch_size * cur_ngrams, body_fun, jnp.zeros(shape, dtype=input_ids.dtype) + ) + + # ignore the n-grams not yet generated + data = (jnp.arange(batch_size * seq_ngrams) < batch_size * cur_ngrams).astype("float32") + + return sparse.BCOO((data, all_update_indices), shape=(batch_size,) + (vocab_size,) * self.ngram_size) + + def get_banned_tokens_mask(self, latest_tokens: jnp.ndarray, previous_ngrams) -> jnp.ndarray: + """ + Determines which tokens must be banned given latest tokens and the previously seen + ngrams. + """ + + @sparse.sparsify + @jax.vmap + def inner_fn(latest_tokens, previous_ngrams): + return previous_ngrams[tuple(latest_tokens)] + + return sparse.bcoo_todense(inner_fn(latest_tokens, previous_ngrams)) + + def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray: + def true_fn(): + _, vocab_size = scores.shape + # store the previously seen n-grams + previous_ngrams = self.get_previous_ngrams(input_ids, vocab_size, cur_len) + + # get the n-1 last tokens that prefix the n-gram being generated + latest_tokens = jnp.zeros((input_ids.shape[0], self.ngram_size - 1), dtype=input_ids.dtype) + latest_tokens = jax.lax.dynamic_update_slice( + latest_tokens, + jax.lax.dynamic_slice( + input_ids, (0, cur_len - (self.ngram_size - 1)), (input_ids.shape[0], (self.ngram_size - 1)) + ), + (0, 0), + ) + + # compute the banned tokens, ie all the tokens that when added to the latest tokens lead to a n-gram that was previously generated + banned_tokens_indices_mask = self.get_banned_tokens_mask(latest_tokens, previous_ngrams).astype("bool") + return jnp.where(banned_tokens_indices_mask, -float("inf"), scores) + + output = jax.lax.cond((cur_len >= self.ngram_size - 1), true_fn, lambda: scores) + return output diff --git a/lib/python3.12/site-packages/transformers/generation/flax_utils.py b/lib/python3.12/site-packages/transformers/generation/flax_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ddd718cbb8a63646d39cc12718ca606a4b77405f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/flax_utils.py @@ -0,0 +1,1029 @@ +# coding=utf-8 +# Copyright 2021 The Google AI Flax Team Authors, and The HuggingFace Inc. team. +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + + +import copy +import inspect +import warnings +from functools import partial +from typing import Any, Dict, Optional, Union + +import flax +import jax +import jax.numpy as jnp +import numpy as np +from jax import lax + +from ..models.auto import ( + FLAX_MODEL_FOR_CAUSAL_LM_MAPPING, + FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, + FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING, +) +from ..utils import ModelOutput, logging +from .configuration_utils import GenerationConfig +from .flax_logits_process import ( + FlaxForcedBOSTokenLogitsProcessor, + FlaxForcedEOSTokenLogitsProcessor, + FlaxForceTokensLogitsProcessor, + FlaxLogitsProcessorList, + FlaxMinLengthLogitsProcessor, + FlaxNoRepeatNGramLogitsProcessor, + FlaxSuppressTokensAtBeginLogitsProcessor, + FlaxSuppressTokensLogitsProcessor, + FlaxTemperatureLogitsWarper, + FlaxTopKLogitsWarper, + FlaxTopPLogitsWarper, +) + + +logger = logging.get_logger(__name__) + + +@flax.struct.dataclass +class FlaxGreedySearchOutput(ModelOutput): + """ + Flax Base class for outputs of decoder-only generation models using greedy search. + + + Args: + sequences (`jnp.ndarray` of shape `(batch_size, max_length)`): + The generated sequences. + """ + + sequences: Optional[jnp.ndarray] = None + + +@flax.struct.dataclass +class FlaxSampleOutput(ModelOutput): + """ + Flax Base class for outputs of decoder-only generation models using sampling. + + + Args: + sequences (`jnp.ndarray` of shape `(batch_size, max_length)`): + The generated sequences. + """ + + sequences: Optional[jnp.ndarray] = None + + +@flax.struct.dataclass +class FlaxBeamSearchOutput(ModelOutput): + """ + Flax Base class for outputs of decoder-only generation models using greedy search. + + + Args: + sequences (`jnp.ndarray` of shape `(batch_size, max_length)`): + The generated sequences. + scores (`jnp.ndarray` of shape `(batch_size,)`): + The scores (log probabilities) of the generated sequences. + """ + + sequences: Optional[jnp.ndarray] = None + scores: Optional[jnp.ndarray] = None + + +@flax.struct.dataclass +class GreedyState: + cur_len: jnp.ndarray + sequences: jnp.ndarray + running_token: jnp.ndarray + is_sent_finished: jnp.ndarray + model_kwargs: Dict[str, jnp.ndarray] + + +@flax.struct.dataclass +class SampleState: + cur_len: jnp.ndarray + sequences: jnp.ndarray + running_token: jnp.ndarray + is_sent_finished: jnp.ndarray + prng_key: jnp.ndarray + model_kwargs: Dict[str, jnp.ndarray] + + +@flax.struct.dataclass +class BeamSearchState: + cur_len: jnp.ndarray + running_sequences: jnp.ndarray + running_scores: jnp.ndarray + sequences: jnp.ndarray + scores: jnp.ndarray + is_sent_finished: jnp.ndarray + model_kwargs: Dict[str, jnp.ndarray] + + +class FlaxGenerationMixin: + """ + A class containing all functions for auto-regressive text generation, to be used as a mixin in + [`FlaxPreTrainedModel`]. + + The class exposes [`~generation.FlaxGenerationMixin.generate`], which can be used for: + - *greedy decoding* by calling [`~generation.FlaxGenerationMixin._greedy_search`] if `num_beams=1` and + `do_sample=False` + - *multinomial sampling* by calling [`~generation.FlaxGenerationMixin._sample`] if `num_beams=1` and + `do_sample=True` + - *beam-search decoding* by calling [`~generation.FlaxGenerationMixin._beam_search`] if `num_beams>1` and + `do_sample=False` + + You do not need to call any of the above methods directly. Pass custom parameter values to 'generate' instead. To + learn more about decoding strategies refer to the [text generation strategies guide](../generation_strategies). + """ + + def prepare_inputs_for_generation(self, *args, **kwargs): + raise NotImplementedError( + "A model class needs to define a `prepare_inputs_for_generation` method in order to use `generate`." + ) + + @staticmethod + def _run_loop_in_debug(cond_fn, body_fn, init_state): + """ + Run generation in untraced mode. This should only be used for debugging purposes. + """ + state = init_state + while cond_fn(state): + state = body_fn(state) + return state + + def _prepare_encoder_decoder_kwargs_for_generation(self, input_ids, params, model_kwargs): + encoder_kwargs = { + argument: value + for argument, value in model_kwargs.items() + if not (argument.startswith("decoder_") or argument.startswith("cross_attn")) + } + model_kwargs["encoder_outputs"] = self.encode(input_ids, params=params, return_dict=True, **encoder_kwargs) + return model_kwargs + + def _prepare_decoder_input_ids_for_generation( + self, + batch_size: int, + decoder_start_token_id: Optional[int] = None, + bos_token_id: Optional[int] = None, + model_kwargs: Optional[Dict[str, jnp.ndarray]] = None, + ) -> jnp.ndarray: + if model_kwargs is not None and "decoder_input_ids" in model_kwargs: + # Only use this arg if not None, otherwise just remove from model_kwargs + decoder_input_ids = model_kwargs.pop("decoder_input_ids") + if decoder_input_ids is not None: + return decoder_input_ids + decoder_start_token_id = self._get_decoder_start_token_id(decoder_start_token_id, bos_token_id) + return jnp.array(decoder_start_token_id, dtype="i4").reshape(1, -1).repeat(batch_size, axis=0) + + def _get_decoder_start_token_id( + self, decoder_start_token_id: Optional[int] = None, bos_token_id: Optional[int] = None + ) -> int: + # retrieve decoder_start_token_id for encoder-decoder models + # fall back to bos_token_id if necessary + decoder_start_token_id = ( + decoder_start_token_id + if decoder_start_token_id is not None + else self.generation_config.decoder_start_token_id + ) + bos_token_id = bos_token_id if bos_token_id is not None else self.generation_config.bos_token_id + if decoder_start_token_id is not None: + return decoder_start_token_id + elif ( + hasattr(self.config, "decoder") + and hasattr(self.config.decoder, "decoder_start_token_id") + and self.config.decoder.decoder_start_token_id is not None + ): + return self.config.decoder.decoder_start_token_id + elif bos_token_id is not None: + return bos_token_id + elif ( + hasattr(self.config, "decoder") + and hasattr(self.config.decoder, "bos_token_id") + and self.config.decoder.bos_token_id is not None + ): + return self.config.decoder.bos_token_id + raise ValueError( + "`decoder_start_token_id` or `bos_token_id` has to be defined for encoder-decoder generation." + ) + + @staticmethod + def _expand_to_num_beams(tensor, num_beams): + return jnp.broadcast_to(tensor[:, None], (tensor.shape[0], num_beams) + tensor.shape[1:]) + + def _adapt_logits_for_beam_search(self, logits): + """ + This function can be overwritten in the specific modeling_flax_.py classes to allow for custom beam + search behavior. Note that the only model that overwrites this method is [`~transformes.FlaxMarianMTModel`]. + """ + return logits + + def _validate_model_class(self): + """ + Confirms that the model class is compatible with generation. If not, raises an exception that points to the + right class to use. + """ + if not self.can_generate(): + generate_compatible_mappings = [ + FLAX_MODEL_FOR_CAUSAL_LM_MAPPING, + FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING, + FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, + ] + generate_compatible_classes = set() + for model_mapping in generate_compatible_mappings: + supported_models = model_mapping.get(type(self.config), default=None) + if supported_models is not None: + generate_compatible_classes.add(supported_models.__name__) + exception_message = ( + f"The current model class ({self.__class__.__name__}) is not compatible with `.generate()`, as " + "it doesn't have a language model head." + ) + if generate_compatible_classes: + exception_message += f" Please use one of the following classes instead: {generate_compatible_classes}" + raise TypeError(exception_message) + + def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]): + """Validates model kwargs for generation. Generate argument typos will also be caught here.""" + unused_model_args = [] + model_args = set(inspect.signature(self.prepare_inputs_for_generation).parameters) + # `kwargs`/`model_kwargs` is often used to handle optional forward pass inputs like `attention_mask`. If + # `prepare_inputs_for_generation` doesn't accept them, then a stricter check can be made ;) + if "kwargs" in model_args or "model_kwargs" in model_args: + model_args |= set(inspect.signature(self.__call__).parameters) + for key, value in model_kwargs.items(): + if value is not None and key not in model_args: + unused_model_args.append(key) + + if unused_model_args: + raise ValueError( + f"The following `model_kwargs` are not used by the model: {unused_model_args} (note: typos in the" + " generate arguments will also show up in this list)" + ) + + def generate( + self, + input_ids: jnp.ndarray, + generation_config: Optional[GenerationConfig] = None, + prng_key: Optional[jnp.ndarray] = None, + trace: bool = True, + params: Optional[Dict[str, jnp.ndarray]] = None, + logits_processor: Optional[FlaxLogitsProcessorList] = None, + **kwargs, + ): + r""" + Generates sequences of token ids for models with a language modeling head. + + Parameters: + input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + generation_config (`~generation.GenerationConfig`, *optional*): + The generation configuration to be used as base parametrization for the generation call. `**kwargs` + passed to generate matching the attributes of `generation_config` will override them. If + `generation_config` is not provided, the default will be used, which had the following loading + priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model + configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s + default values, whose documentation should be checked to parameterize generation. + trace (`bool`, *optional*, defaults to `True`): + Whether to trace generation. Setting `trace=False` should only be used for debugging and will lead to a + considerably slower runtime. + params (`Dict[str, jnp.ndarray]`, *optional*): + Optionally the model parameters can be passed. Can be useful for parallelized generation. + logits_processor (`FlaxLogitsProcessorList `, *optional*): + Custom logits processors that complement the default logits processors built from arguments and + generation config. If a logit processor is passed that is already created with the arguments or a + generation config an error is thrown. This feature is intended for advanced users. + kwargs (`Dict[str, Any]`, *optional*): + Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be + forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder + specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*. + + Return: + [`~utils.ModelOutput`]. + + """ + # Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call + self._validate_model_class() + + # priority: `generation_config` argument > `model.generation_config` (the default generation config) + if generation_config is None: + # legacy: users may modify the model configuration to control generation. To trigger this legacy behavior, + # two conditions must be met + # 1) the generation config must have been created from the model config (`_from_model_config` field); + # 2) the generation config must have seen no modification since its creation (the hash is the same). + if self.generation_config._from_model_config and self.generation_config._original_object_hash == hash( + self.generation_config + ): + new_generation_config = GenerationConfig.from_model_config(self.config) + if new_generation_config != self.generation_config: + warnings.warn( + "You have modified the pretrained model configuration to control generation. This is a" + " deprecated strategy to control generation and will be removed soon, in a future version." + " Please use and modify the model generation configuration (see" + " https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )" + ) + self.generation_config = new_generation_config + generation_config = self.generation_config + + generation_config = copy.deepcopy(generation_config) + model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs + self._validate_model_kwargs(model_kwargs.copy()) + + logits_processor = logits_processor if logits_processor is not None else FlaxLogitsProcessorList() + + # set init values + prng_key = prng_key if prng_key is not None else jax.random.PRNGKey(0) + + if generation_config.pad_token_id is None and generation_config.eos_token_id is not None: + if model_kwargs.get("attention_mask") is None: + logger.warning( + "The attention mask and the pad token id were not set. As a consequence, you may observe " + "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results." + ) + eos_token_id = generation_config.eos_token_id + if isinstance(eos_token_id, list): + eos_token_id = eos_token_id[0] + generation_config.pad_token_id = eos_token_id + + if generation_config.decoder_start_token_id is None and self.config.is_encoder_decoder: + raise ValueError("`decoder_start_token_id` has to be defined for encoder-decoder generation.") + + # decoder-only models should use left-padding for generation (can't be checked with `trace=True`) + if not self.config.is_encoder_decoder and not trace: + if ( + generation_config.pad_token_id is not None + and jnp.sum(input_ids[:, -1] == generation_config.pad_token_id) > 0 + ): + logger.warning( + "A decoder-only architecture is being used, but right-padding was detected! For correct " + "generation results, please set `padding_side='left'` when initializing the tokenizer." + ) + + batch_size = input_ids.shape[0] + + if self.config.is_encoder_decoder: + # add encoder_outputs to model_kwargs + if model_kwargs.get("encoder_outputs") is None: + model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(input_ids, params, model_kwargs) + # prepare decoder_input_ids for generation + input_ids = self._prepare_decoder_input_ids_for_generation( + batch_size, + decoder_start_token_id=generation_config.decoder_start_token_id, + bos_token_id=generation_config.bos_token_id, + model_kwargs=model_kwargs, + ) + + # Prepare `max_length` depending on other stopping criteria. + input_ids_seq_length = input_ids.shape[-1] + has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None + if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20: + # 20 is the default max_length of the generation config + warnings.warn( + f"Using the model-agnostic default `max_length` (={generation_config.max_length}) " + "to control the generation length. recommend setting `max_new_tokens` to control the maximum length of the generation.", + UserWarning, + ) + elif generation_config.max_new_tokens is not None: + if not has_default_max_length and generation_config.max_length is not None: + logger.warning( + f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(=" + f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. " + "Please refer to the documentation for more information. " + "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)" + ) + generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length + else: # by default let's always generate 20 new tokens + if generation_config.max_length == GenerationConfig().max_length: + generation_config.max_length = generation_config.max_length + input_ids_seq_length + max_position_embeddings = getattr(self.config, "max_position_embeddings", None) + if max_position_embeddings is not None: + generation_config.max_length = min(generation_config.max_length, max_position_embeddings) + + if generation_config.min_length is not None and generation_config.min_length > generation_config.max_length: + raise ValueError( + f"Unfeasable length constraints: the minimum length ({generation_config.min_length}) is larger than" + f" the maximum length ({generation_config.max_length})" + ) + if input_ids_seq_length >= generation_config.max_length: + input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids" + logger.warning( + f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to" + f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider" + " increasing`max_new_tokens`." + ) + + logits_processor = self._get_logits_processor( + generation_config=generation_config, + input_ids_seq_length=input_ids_seq_length, + logits_processor=logits_processor, + ) + + if not generation_config.do_sample and generation_config.num_beams == 1: + return self._greedy_search( + input_ids, + generation_config.max_length, + generation_config.pad_token_id, + generation_config.eos_token_id, + logits_processor=logits_processor, + trace=trace, + params=params, + model_kwargs=model_kwargs, + ) + elif generation_config.do_sample and generation_config.num_beams == 1: + logits_warper = self._get_logits_warper(generation_config=generation_config) + return self._sample( + input_ids, + generation_config.max_length, + generation_config.pad_token_id, + generation_config.eos_token_id, + prng_key, + logits_warper=logits_warper, + logits_processor=logits_processor, + trace=trace, + params=params, + model_kwargs=model_kwargs, + ) + elif not generation_config.do_sample and generation_config.num_beams > 1: + # broadcast input_ids & encoder_outputs + input_ids = self._expand_to_num_beams(input_ids, num_beams=generation_config.num_beams) + + if "encoder_outputs" in model_kwargs: + model_kwargs["encoder_outputs"]["last_hidden_state"] = self._expand_to_num_beams( + model_kwargs["encoder_outputs"]["last_hidden_state"], num_beams=generation_config.num_beams + ) + + for kwarg in ["attention_mask", "decoder_attention_mask"]: + if kwarg in model_kwargs: + model_kwargs[kwarg] = self._expand_to_num_beams( + model_kwargs[kwarg], num_beams=generation_config.num_beams + ) + + return self._beam_search( + input_ids, + generation_config.max_length, + generation_config.pad_token_id, + generation_config.eos_token_id, + length_penalty=generation_config.length_penalty, + early_stopping=generation_config.early_stopping, + logits_processor=logits_processor, + trace=trace, + params=params, + num_return_sequences=generation_config.num_return_sequences, + model_kwargs=model_kwargs, + ) + else: + raise NotImplementedError("`Beam sampling is currently not implemented.") + + def _get_logits_warper(self, generation_config: GenerationConfig) -> FlaxLogitsProcessorList: + """ + This class returns a [`FlaxLogitsProcessorList`] list object that contains all relevant [`FlaxLogitsWarper`] + instances used for multinomial sampling. + """ + warpers = FlaxLogitsProcessorList() + + if generation_config.temperature is not None and generation_config.temperature != 1.0: + warpers.append(FlaxTemperatureLogitsWarper(generation_config.temperature)) + if generation_config.top_k is not None and generation_config.top_k != 0: + warpers.append(FlaxTopKLogitsWarper(top_k=generation_config.top_k, min_tokens_to_keep=1)) + if generation_config.top_p is not None and generation_config.top_p < 1.0: + warpers.append(FlaxTopPLogitsWarper(top_p=generation_config.top_p, min_tokens_to_keep=1)) + + return warpers + + def _get_logits_processor( + self, + generation_config: GenerationConfig, + input_ids_seq_length: int, + logits_processor: Optional[FlaxLogitsProcessorList], + ) -> FlaxLogitsProcessorList: + """ + This class returns a [`FlaxLogitsProcessorList`] list object that contains all relevant [`FlaxLogitsProcessor`] + instances used to modify the scores of the language model head. + """ + processors = FlaxLogitsProcessorList() + + if ( + generation_config.min_length is not None + and generation_config.eos_token_id is not None + and generation_config.min_length > -1 + ): + processors.append( + FlaxMinLengthLogitsProcessor(generation_config.min_length, generation_config.eos_token_id) + ) + if generation_config.forced_bos_token_id is not None: + processors.append(FlaxForcedBOSTokenLogitsProcessor(generation_config.forced_bos_token_id)) + if generation_config.forced_eos_token_id is not None: + processors.append( + FlaxForcedEOSTokenLogitsProcessor(generation_config.max_length, generation_config.forced_eos_token_id) + ) + if generation_config.suppress_tokens is not None: + processors.append(FlaxSuppressTokensLogitsProcessor(generation_config.suppress_tokens)) + if generation_config.begin_suppress_tokens is not None: + begin_index = input_ids_seq_length + begin_index = ( + begin_index + if (input_ids_seq_length > 1 or generation_config.forced_bos_token_id is None) + else begin_index + 1 + ) + if generation_config.forced_decoder_ids is not None and len(generation_config.forced_decoder_ids) > 0: + # generation starts after the last token that is forced + begin_index += generation_config.forced_decoder_ids[-1][0] + processors.append( + FlaxSuppressTokensAtBeginLogitsProcessor(generation_config.begin_suppress_tokens, begin_index) + ) + if generation_config.forced_decoder_ids is not None: + forced_decoder_ids = [ + [input_ids_seq_length + i[0] - 1, i[1]] for i in generation_config.forced_decoder_ids + ] + processors.append(FlaxForceTokensLogitsProcessor(forced_decoder_ids)) + if generation_config.no_repeat_ngram_size is not None and generation_config.no_repeat_ngram_size > 0: + processors.append(FlaxNoRepeatNGramLogitsProcessor(generation_config.no_repeat_ngram_size)) + processors = self._merge_criteria_processor_list(processors, logits_processor) + + return processors + + def _merge_criteria_processor_list( + self, + default_list: FlaxLogitsProcessorList, + custom_list: FlaxLogitsProcessorList, + ) -> FlaxLogitsProcessorList: + if len(custom_list) == 0: + return default_list + for default in default_list: + for custom in custom_list: + if type(custom) is type(default): + object_type = "logits processor" + raise ValueError( + f"A custom {object_type} of type {type(custom)} with values {custom} has been passed to" + f" `generate`, but it has already been created with the values {default}. {default} has been" + " created by passing the corresponding arguments to generate or by the model's config default" + f" values. If you just want to change the default values of {object_type} consider passing" + f" them as arguments to `generate` instead of using a custom {object_type}." + ) + default_list.extend(custom_list) + return default_list + + def _greedy_search( + self, + input_ids: None, + max_length: Optional[int] = None, + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + logits_processor: Optional[FlaxLogitsProcessorList] = None, + trace: bool = True, + params: Optional[Dict[str, jnp.ndarray]] = None, + model_kwargs: Optional[Dict[str, jnp.ndarray]] = None, + ): + # init values + max_length = max_length if max_length is not None else self.generation_config.max_length + pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id + eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id + + batch_size, cur_len = input_ids.shape + + eos_token_id = jnp.array(eos_token_id, dtype=jnp.int32 if eos_token_id is not None else None) + pad_token_id = jnp.array(pad_token_id, dtype=jnp.int32) + cur_len = jnp.array(cur_len) + + # per batch-item holding current token in loop. + sequences = jnp.full((batch_size, max_length), pad_token_id, dtype=jnp.int32) + sequences = lax.dynamic_update_slice(sequences, input_ids, (0, 0)) + + # per batch-item state bit indicating if sentence has finished. + is_sent_finished = jnp.zeros((batch_size,), dtype=jnp.bool_) + + # For Seq2Seq generation, we only need to use the decoder instead of the whole model in generation loop + # and pass it the `encoder_outputs`, which are part of the `model_kwargs`. + model = self.decode if self.config.is_encoder_decoder else self + # initialize model specific kwargs + model_kwargs = self.prepare_inputs_for_generation(input_ids, max_length, **model_kwargs) + + # initialize state + state = GreedyState( + cur_len=cur_len, + sequences=sequences, + running_token=input_ids, + is_sent_finished=is_sent_finished, + model_kwargs=model_kwargs, + ) + + def greedy_search_cond_fn(state): + """state termination condition fn.""" + has_reached_max_length = state.cur_len == max_length + all_sequence_finished = jnp.all(state.is_sent_finished) + finish_generation = jnp.logical_or(has_reached_max_length, all_sequence_finished) + return ~finish_generation + + def greedy_search_body_fn(state): + """state update fn.""" + model_outputs = model(state.running_token, params=params, **state.model_kwargs) + logits = model_outputs.logits[:, -1] + + # apply min_length, ... + logits = logits_processor(state.sequences, logits, state.cur_len) + + next_token = jnp.argmax(logits, axis=-1) + + next_token = next_token * ~state.is_sent_finished + pad_token_id * state.is_sent_finished + next_is_sent_finished = state.is_sent_finished | (next_token == eos_token_id) + next_token = next_token[:, None] + + next_sequences = lax.dynamic_update_slice(state.sequences, next_token, (0, state.cur_len)) + next_model_kwargs = self.update_inputs_for_generation(model_outputs, state.model_kwargs) + return GreedyState( + cur_len=state.cur_len + 1, + sequences=next_sequences, + running_token=next_token, + is_sent_finished=next_is_sent_finished, + model_kwargs=next_model_kwargs, + ) + + # The very first prompt often has sequence length > 1, so run outside of `lax.while_loop` to comply with TPU + if input_ids.shape[1] > 1: + state = greedy_search_body_fn(state) + + if not trace: + state = self._run_loop_in_debug(greedy_search_cond_fn, greedy_search_body_fn, state) + else: + state = lax.while_loop(greedy_search_cond_fn, greedy_search_body_fn, state) + + return FlaxGreedySearchOutput(sequences=state.sequences) + + def _sample( + self, + input_ids: None, + max_length: Optional[int] = None, + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + prng_key: Optional[jnp.ndarray] = None, + logits_processor: Optional[FlaxLogitsProcessorList] = None, + logits_warper: Optional[FlaxLogitsProcessorList] = None, + trace: bool = True, + params: Optional[Dict[str, jnp.ndarray]] = None, + model_kwargs: Optional[Dict[str, jnp.ndarray]] = None, + ): + # init values + max_length = max_length if max_length is not None else self.generation_config.max_length + pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id + eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id + prng_key = prng_key if prng_key is not None else jax.random.PRNGKey(0) + + batch_size, cur_len = input_ids.shape + + eos_token_id = jnp.array(eos_token_id, dtype=jnp.int32 if eos_token_id is not None else None) + pad_token_id = jnp.array(pad_token_id, dtype=jnp.int32) + cur_len = jnp.array(cur_len) + + # per batch-item holding current token in loop. + sequences = jnp.full((batch_size, max_length), pad_token_id, dtype=jnp.int32) + sequences = lax.dynamic_update_slice(sequences, input_ids, (0, 0)) + + # per batch-item state bit indicating if sentence has finished. + is_sent_finished = jnp.zeros((batch_size,), dtype=jnp.bool_) + + # For Seq2Seq generation, we only need to use the decoder instead of the whole model in generation loop + # and pass it the `encoder_outputs`, which are part of the `model_kwargs`. + model = self.decode if self.config.is_encoder_decoder else self + + # initialize model specific kwargs + model_kwargs = self.prepare_inputs_for_generation(input_ids, max_length, **model_kwargs) + + # initialize state + state = SampleState( + cur_len=cur_len, + sequences=sequences, + running_token=input_ids, + is_sent_finished=is_sent_finished, + prng_key=prng_key, + model_kwargs=model_kwargs, + ) + + def sample_search_cond_fn(state): + """state termination condition fn.""" + has_reached_max_length = state.cur_len == max_length + all_sequence_finished = jnp.all(state.is_sent_finished) + finish_generation = jnp.logical_or(has_reached_max_length, all_sequence_finished) + return ~finish_generation + + def sample_search_body_fn(state): + """state update fn.""" + prng_key, prng_key_next = jax.random.split(state.prng_key) + model_outputs = model(state.running_token, params=params, **state.model_kwargs) + + logits = model_outputs.logits[:, -1] + + # apply min_length, ... + logits = logits_processor(state.sequences, logits, state.cur_len) + # apply top_p, top_k, temperature + logits = logits_warper(logits, logits, state.cur_len) + + next_token = jax.random.categorical(prng_key, logits, axis=-1) + + next_token = next_token * ~state.is_sent_finished + pad_token_id * state.is_sent_finished + next_is_sent_finished = state.is_sent_finished | (next_token == eos_token_id) + next_token = next_token[:, None] + + next_sequences = lax.dynamic_update_slice(state.sequences, next_token, (0, state.cur_len)) + next_model_kwargs = self.update_inputs_for_generation(model_outputs, state.model_kwargs) + + return SampleState( + cur_len=state.cur_len + 1, + sequences=next_sequences, + running_token=next_token, + is_sent_finished=next_is_sent_finished, + model_kwargs=next_model_kwargs, + prng_key=prng_key_next, + ) + + # The very first prompt often has sequence length > 1, so run outside of `lax.while_loop` to comply with TPU + if input_ids.shape[1] > 1: + state = sample_search_body_fn(state) + + if not trace: + state = self._run_loop_in_debug(sample_search_cond_fn, sample_search_body_fn, state) + else: + state = lax.while_loop(sample_search_cond_fn, sample_search_body_fn, state) + + return FlaxSampleOutput(sequences=state.sequences) + + def _beam_search( + self, + input_ids: None, + max_length: Optional[int] = None, + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + length_penalty: Optional[float] = None, + early_stopping: Optional[Union[bool, str]] = None, + logits_processor: Optional[FlaxLogitsProcessorList] = None, + trace: bool = True, + params: Optional[Dict[str, jnp.ndarray]] = None, + num_return_sequences: Optional[int] = None, + model_kwargs: Optional[Dict[str, jnp.ndarray]] = None, + ): + """ + This beam search function is heavily inspired by Flax's official example: + https://github.com/google/flax/blob/main/examples/wmt/decode.py + """ + + def flatten_beam_dim(tensor): + """Flattens the first two dimensions of a non-scalar array.""" + # ignore scalars (e.g. cache index) + if tensor.ndim == 0: + return tensor + return tensor.reshape((tensor.shape[0] * tensor.shape[1],) + tensor.shape[2:]) + + def unflatten_beam_dim(tensor, batch_size, num_beams): + """Unflattens the first, flat batch*beam dimension of a non-scalar array.""" + # ignore scalars (e.g. cache index) + if tensor.ndim == 0: + return tensor + return tensor.reshape((batch_size, num_beams) + tensor.shape[1:]) + + def gather_beams(nested, beam_indices, batch_size, new_num_beams): + """ + Gathers the beam slices indexed by beam_indices into new beam array. + """ + batch_indices = jnp.reshape( + jnp.arange(batch_size * new_num_beams) // new_num_beams, (batch_size, new_num_beams) + ) + + def gather_fn(tensor): + # ignore scalars (e.g. cache index) + if tensor.ndim == 0: + return tensor + else: + return tensor[batch_indices, beam_indices] + + return jax.tree_util.tree_map(gather_fn, nested) + + # init values + max_length = max_length if max_length is not None else self.generation_config.max_length + pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id + eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id + length_penalty = length_penalty if length_penalty is not None else self.generation_config.length_penalty + early_stopping = early_stopping if early_stopping is not None else self.generation_config.early_stopping + num_return_sequences = ( + num_return_sequences if num_return_sequences is not None else self.generation_config.num_return_sequences + ) + + batch_size, num_beams, cur_len = input_ids.shape + + eos_token_id = jnp.array(eos_token_id, dtype=jnp.int32 if eos_token_id is not None else None) + pad_token_id = jnp.array(pad_token_id, dtype=jnp.int32) + cur_len = jnp.array(cur_len) + + # record the prompt length of decoder + decoder_prompt_len = input_ids.shape[-1] + + # per batch,beam-item holding current token in loop. + sequences = jnp.full((batch_size, num_beams, max_length), pad_token_id, dtype=jnp.int32) + running_sequences = jnp.full((batch_size, num_beams, max_length), pad_token_id, dtype=jnp.int32) + running_sequences = lax.dynamic_update_slice(sequences, input_ids, (0, 0, 0)) + + # per batch,beam-item state bit indicating if sentence has finished. + is_sent_finished = jnp.zeros((batch_size, num_beams), dtype=jnp.bool_) + + # per batch,beam-item score, logprobs + running_scores = jnp.tile(jnp.array([0.0] + [np.array(-1.0e7)] * (num_beams - 1)), [batch_size, 1]) + scores = jnp.ones((batch_size, num_beams)) * np.array(-1.0e7) + + # For Seq2Seq generation, we only need to use the decoder instead of the whole model in generation loop + # and pass it the `encoder_outputs`, which are part of the `model_kwargs`. + model = self.decode if self.config.is_encoder_decoder else self + + # flatten beam dim + if "encoder_outputs" in model_kwargs: + model_kwargs["encoder_outputs"]["last_hidden_state"] = flatten_beam_dim( + model_kwargs["encoder_outputs"]["last_hidden_state"] + ) + for kwarg in ["attention_mask", "decoder_attention_mask"]: + if kwarg in model_kwargs: + model_kwargs[kwarg] = flatten_beam_dim(model_kwargs[kwarg]) + + # initialize model specific kwargs + model_kwargs = self.prepare_inputs_for_generation(flatten_beam_dim(input_ids), max_length, **model_kwargs) + + # initialize state + state = BeamSearchState( + cur_len=cur_len, + running_sequences=running_sequences, + running_scores=running_scores, + sequences=sequences, + scores=scores, + is_sent_finished=is_sent_finished, + model_kwargs=model_kwargs, + ) + + def beam_search_cond_fn(state): + """beam search state termination condition fn.""" + + # 1. is less than max length? + not_max_length_yet = state.cur_len < max_length + + # 2. can the new beams still improve? + # early_stopping == False -> apply heuristic = always get the best score from `cur_len`. See the discussion + # below for more details. + # https://github.com/huggingface/transformers/pull/20901#issuecomment-1369845565 + # early_stopping == "never" -> compute the best score from max_length or cur_len, depending on the sign of + # length_penalty. Positive length_penalty favors longer sequences, thus we use max_length there. + if early_stopping == "never" and length_penalty > 0.0: + best_running_score = state.running_scores[:, :1] / ( + (max_length - decoder_prompt_len) ** length_penalty + ) + else: + best_running_score = state.running_scores[:, :1] / ( + (state.cur_len - decoder_prompt_len) ** length_penalty + ) + worst_finished_score = jnp.where( + state.is_sent_finished, jnp.min(state.scores, axis=1, keepdims=True), np.array(-1.0e7) + ) + improvement_still_possible = jnp.any(best_running_score > worst_finished_score) + + # 3. is there still a beam that has not finished? + still_open_beam = ~(jnp.all(state.is_sent_finished) & (early_stopping is True)) + + return not_max_length_yet & still_open_beam & improvement_still_possible + + def beam_search_body_fn(state, input_ids_length=1): + """beam search state update fn.""" + # 1. Forward current tokens + # Collect the current position slice along length to feed the fast + # autoregressive decoder model. Flatten the beam dimension into batch + # dimension for feeding into the model. + # unflatten beam dimension + # Unflatten beam dimension in attention cache arrays + input_token = flatten_beam_dim( + lax.dynamic_slice( + state.running_sequences, + (0, 0, state.cur_len - input_ids_length), + (batch_size, num_beams, input_ids_length), + ) + ) + model_outputs = model(input_token, params=params, **state.model_kwargs) + + logits = unflatten_beam_dim(model_outputs.logits[:, -1], batch_size, num_beams) + cache = jax.tree_util.tree_map( + lambda tensor: unflatten_beam_dim(tensor, batch_size, num_beams), model_outputs.past_key_values + ) + + # adapt logits for FlaxMarianMTModel + logits = self._adapt_logits_for_beam_search(logits) + + # 2. Compute log probs + # get log probabilities from logits, + # process logits with processors (*e.g.* min_length, ...), and + # add new logprobs to existing running logprobs scores. + log_probs = jax.nn.log_softmax(logits) + log_probs = logits_processor( + flatten_beam_dim(state.running_sequences), flatten_beam_dim(log_probs), state.cur_len + ) + log_probs = unflatten_beam_dim(log_probs, batch_size, num_beams) + log_probs = log_probs + jnp.expand_dims(state.running_scores, axis=2) + vocab_size = log_probs.shape[2] + log_probs = log_probs.reshape((batch_size, num_beams * vocab_size)) + + # 3. Retrieve top-K + # Each item in batch has num_beams * vocab_size candidate sequences. + # For each item, get the top 2*k candidates with the highest log- + # probabilities. We gather the top 2*K beams here so that even if the best + # K sequences reach EOS simultaneously, we have another K sequences + # remaining to continue the live beam search. + # Gather the top 2*K scores from _all_ beams. + # Gather 2*k top beams. + # Recover the beam index by floor division. + # Recover token id by modulo division and expand Id array for broadcasting. + # Update sequences for the 2*K top-k new sequences. + beams_to_keep = 2 * num_beams + topk_log_probs, topk_indices = lax.top_k(log_probs, k=beams_to_keep) + topk_beam_indices = topk_indices // vocab_size + topk_running_sequences = gather_beams( + state.running_sequences, topk_beam_indices, batch_size, beams_to_keep + ) + topk_ids = jnp.expand_dims(topk_indices % vocab_size, axis=2) + topk_sequences = lax.dynamic_update_slice(topk_running_sequences, topk_ids, (0, 0, state.cur_len)) + + # 4. Check which sequences have ended + # Update current sequences: + # Did any of these sequences reach an end marker? + # To prevent these just finished sequences from being added to the current sequences + # set of active beam search sequences, set their log probs to a very large + # negative value. + did_topk_just_finished = topk_sequences[:, :, state.cur_len] == eos_token_id + running_topk_log_probs = topk_log_probs + did_topk_just_finished * np.array(-1.0e7) + # 5. Get running sequences scores for next + # Determine the top k beam indices (from top 2*k beams) from log probs + # and gather top k beams (from top 2*k beams). + next_topk_indices = lax.top_k(running_topk_log_probs, k=num_beams)[1] + next_running_sequences, next_running_scores = gather_beams( + [topk_sequences, running_topk_log_probs], next_topk_indices, batch_size, num_beams + ) + + # 6. Process topk logits + # Further process log probs: + # - add length penalty + # - make sure no scores can be added anymore if beam is full + # - make sure still running sequences cannot be chosen as finalized beam + topk_log_probs = topk_log_probs / ((state.cur_len + 1 - decoder_prompt_len) ** length_penalty) + beams_in_batch_are_full = jnp.broadcast_to( + state.is_sent_finished.all(axis=-1, keepdims=True), did_topk_just_finished.shape + ) & (early_stopping is True) + add_penalty = ~did_topk_just_finished | beams_in_batch_are_full + topk_log_probs += add_penalty * np.array(-1.0e7) + + # 7. Get scores, sequences, is sentence finished for next. + # Combine sequences, scores, and flags along the beam dimension and compare + # new finished sequence scores to existing finished scores and select the + # best from the new set of beams + merged_sequences = jnp.concatenate([state.sequences, topk_sequences], axis=1) + merged_scores = jnp.concatenate([state.scores, topk_log_probs], axis=1) + merged_is_sent_finished = jnp.concatenate([state.is_sent_finished, did_topk_just_finished], axis=1) + topk_merged_indices = lax.top_k(merged_scores, k=num_beams)[1] + next_sequences, next_scores, next_is_sent_finished = gather_beams( + [merged_sequences, merged_scores, merged_is_sent_finished], topk_merged_indices, batch_size, num_beams + ) + + # 8. Update model kwargs. + # Determine the top k beam indices from the original set of all beams. + # With these, gather the top k beam-associated caches. + next_running_indices = gather_beams(topk_beam_indices, next_topk_indices, batch_size, num_beams) + next_cache = gather_beams(cache, next_running_indices, batch_size, num_beams) + model_outputs["past_key_values"] = jax.tree_util.tree_map(lambda x: flatten_beam_dim(x), next_cache) + next_model_kwargs = self.update_inputs_for_generation(model_outputs, state.model_kwargs) + + return BeamSearchState( + cur_len=state.cur_len + 1, + running_scores=next_running_scores, + running_sequences=next_running_sequences, + scores=next_scores, + sequences=next_sequences, + is_sent_finished=next_is_sent_finished, + model_kwargs=next_model_kwargs, + ) + + # Always run first iteration outside of `lax.while_loop` to avoid calling `beam_search_cond_fn` + # when `state.cur_len` equals `decoder_prompt_len`. This also helps to comply with TPU when + # the very first prompt has sequence length > 1. + state = partial(beam_search_body_fn, input_ids_length=input_ids.shape[-1])(state) + + if not trace: + state = self._run_loop_in_debug(beam_search_cond_fn, beam_search_body_fn, state) + else: + state = lax.while_loop(beam_search_cond_fn, beam_search_body_fn, state) + + # Account for the edge-case where there are no finished sequences for a + # particular batch item. If so, return running sequences for that batch item. + none_finished = jnp.any(state.is_sent_finished, axis=1) + sequences = jnp.where(none_finished[:, None, None], state.sequences, state.running_sequences) + scores = jnp.where(none_finished[:, None], state.scores, state.running_scores) + + # Take best beams for each batch (the score is sorted in descending order) + sequences = flatten_beam_dim(sequences[:, :num_return_sequences, :]) + scores = flatten_beam_dim(scores[:, :num_return_sequences]) + + return FlaxBeamSearchOutput(sequences=sequences, scores=scores) diff --git a/lib/python3.12/site-packages/transformers/generation/logits_process.py b/lib/python3.12/site-packages/transformers/generation/logits_process.py new file mode 100644 index 0000000000000000000000000000000000000000..16c04478f08a64bbd1b65af5e720f1470940833b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/logits_process.py @@ -0,0 +1,2942 @@ +# coding=utf-8 +# Copyright 2024 The HuggingFace Inc. team and Google DeepMind. +# +# 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. + +import inspect +import math +from typing import Callable, Iterable, List, Optional, Tuple, Union + +import numpy as np +import torch + +from ..pytorch_utils import isin_mps_friendly +from ..utils import add_start_docstrings +from ..utils.logging import get_logger + + +logger = get_logger(__name__) + + +LOGITS_PROCESSOR_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`): + Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam + search or log softmax for each vocabulary token when using beam search + + Return: + `torch.FloatTensor` of shape `(batch_size, config.vocab_size)`: The processed prediction scores. + +""" + + +class LogitsProcessor: + """Abstract base class for all logit processors that can be applied during generation.""" + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + +class LogitsProcessorList(list): + """ + This class can be used to create a list of [`LogitsProcessor`] to subsequently process a `scores` input tensor. + This class inherits from list and adds a specific *__call__* method to apply each [`LogitsProcessor`] to the + inputs. + """ + + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.FloatTensor: + r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`): + Prediction scores of a language modeling head. These can be logits for each vocabulary when not using + beam search or log softmax for each vocabulary token when using beam search + kwargs (`Dict[str, Any]`, *optional*): + Additional kwargs that are specific to a logits processor. + + Return: + `torch.FloatTensor` of shape `(batch_size, config.vocab_size)`: + The processed prediction scores. + + """ + for processor in self: + function_args = inspect.signature(processor.__call__).parameters + if len(function_args) > 2: + if not all(arg in kwargs for arg in list(function_args.keys())[2:]): + raise ValueError( + f"Make sure that all the required parameters: {list(function_args.keys())} for " + f"{processor.__class__} are passed to the logits processor." + ) + scores = processor(input_ids, scores, **kwargs) + else: + scores = processor(input_ids, scores) + + return scores + + +class MinLengthLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] enforcing a min-length by setting EOS probability to 0. Note that, for decoder-only models + like most LLMs, the length includes the prompt. + + Args: + min_length (`int`): + The minimum length below which the score of `eos_token_id` is set to `-float("Inf")`. + eos_token_id (`Union[int, List[int], torch.Tensor]`): + The id(s) of the *end-of-sequence* token. + device (`str`, *optional*, defaults to `"cpu"`): + The device to allocate the tensors. + + Examples: + + ```python + >>> from transformers import AutoModelForCausalLM, AutoTokenizer + + >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m") + >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m") + + >>> inputs = tokenizer("A number:", return_tensors="pt") + >>> gen_out = model.generate(**inputs) + >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0]) + A number: one + + >>> # setting `min_length` to a value smaller than the uncontrolled output length has no impact + >>> gen_out = model.generate(**inputs, min_length=3) + >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0]) + A number: one + + >>> # setting a larger `min_length` will force the model to generate beyond its natural ending point, which is not + >>> # necessarily incorrect + >>> gen_out = model.generate(**inputs, min_length=10) + >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0]) + A number: one thousand, nine hundred and ninety-four + ``` + """ + + def __init__(self, min_length: int, eos_token_id: Union[int, List[int], torch.Tensor], device: str = "cpu"): + if not isinstance(min_length, int) or min_length < 0: + raise ValueError(f"`min_length` has to be a non-negative integer, but is {min_length}") + + if not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id, device=device) + + self.min_length = min_length + self.eos_token_id = eos_token_id + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + vocab_tensor = torch.arange(scores.shape[-1], device=scores.device) + eos_token_mask = isin_mps_friendly(vocab_tensor, self.eos_token_id) + scores_processed = scores.clone() + if input_ids.shape[-1] < self.min_length: + scores_processed = torch.where(eos_token_mask, -math.inf, scores) + return scores_processed + + +class MinNewTokensLengthLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] enforcing a min-length of new tokens by setting EOS (End-Of-Sequence) token probability to 0. + Contrarily to [`MinLengthLogitsProcessor`], this processor ignores the prompt. + + Args: + prompt_length_to_skip (`int`): + The input tokens length. Not a valid argument when used with `generate` as it will automatically assign the + input length. + min_new_tokens (`int`): + The minimum *new* tokens length below which the score of `eos_token_id` is set to `-float("Inf")`. + eos_token_id (`Union[int, List[int], torch.Tensor]`): + The id(s) of the *end-of-sequence* token. + device (`str`, *optional*, defaults to `"cpu"`): + The device to allocate the tensors. + + Examples: + + ```python + >>> from transformers import AutoModelForCausalLM, AutoTokenizer + + >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m") + >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m") + + >>> inputs = tokenizer(["A number:"], return_tensors="pt") + >>> gen_out = model.generate(**inputs) + >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0]) + A number: one + + >>> # setting `min_new_tokens` will force the model to generate beyond its natural ending point, which is not + >>> # necessarily incorrect + >>> gen_out = model.generate(**inputs, min_new_tokens=2) + >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0]) + A number: one thousand + ``` + """ + + def __init__( + self, + prompt_length_to_skip: int, + min_new_tokens: int, + eos_token_id: Union[int, List[int], torch.Tensor], + device: str = "cpu", + ): + for arg_name, arg_value in [ + ("prompt_length_to_skip", prompt_length_to_skip), + ("min_new_tokens", min_new_tokens), + ]: + if not isinstance(arg_value, int) or arg_value < 0: + raise ValueError(f"`{arg_name}` has to be a positive integer, but is {arg_value}") + + if not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id, device=device) + + self.prompt_length_to_skip = prompt_length_to_skip + self.min_new_tokens = min_new_tokens + self.eos_token_id = eos_token_id + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + new_tokens_length = input_ids.shape[-1] - self.prompt_length_to_skip + scores_processed = scores.clone() + vocab_tensor = torch.arange(scores.shape[-1], device=scores.device) + eos_token_mask = isin_mps_friendly(vocab_tensor, self.eos_token_id) + if new_tokens_length < self.min_new_tokens: + scores_processed = torch.where(eos_token_mask, -math.inf, scores) + + return scores_processed + + +class TemperatureLogitsWarper(LogitsProcessor): + r""" + [`LogitsProcessor`] for temperature (exponential scaling output probability distribution), which effectively means + that it can control the randomness of the predicted tokens. Often used together with [`TopPLogitsWarper`] and + [`TopKLogitsWarper`]. + + + + Make sure that `do_sample=True` is included in the `generate` arguments otherwise the temperature value won't have + any effect. + + + + Args: + temperature (`float`): + Strictly positive float value used to modulate the logits distribution. A value smaller than `1` decreases + randomness (and vice versa), with `0` being equivalent to shifting all probability mass to the most likely + token. + + Examples: + + ```python + >>> import torch + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed + + >>> set_seed(0) # for reproducibility + + >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> model.config.pad_token_id = model.config.eos_token_id + >>> inputs = tokenizer(["Hugging Face Company is"], return_tensors="pt") + + >>> # With temperature=1.0, the default, we consistently get random outputs due to random sampling. + >>> generate_kwargs = {"max_new_tokens": 10, "do_sample": True, "temperature": 1.0, "num_return_sequences": 2} + >>> outputs = model.generate(**inputs, **generate_kwargs) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)) + ['Hugging Face Company is one of these companies that is going to take a', + "Hugging Face Company is a brand created by Brian A. O'Neil"] + + >>> # However, with temperature close to 0, it approximates greedy decoding strategies (invariant) + >>> generate_kwargs["temperature"] = 0.0001 + >>> outputs = model.generate(**inputs, **generate_kwargs) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)) + ['Hugging Face Company is a company that has been around for over 20 years', + 'Hugging Face Company is a company that has been around for over 20 years'] + ``` + """ + + def __init__(self, temperature: float): + if not isinstance(temperature, float) or not (temperature > 0): + except_msg = ( + f"`temperature` (={temperature}) has to be a strictly positive float, otherwise your next token " + "scores will be invalid." + ) + if isinstance(temperature, float) and temperature == 0.0: + except_msg += " If you're looking for greedy decoding strategies, set `do_sample=False`." + raise ValueError(except_msg) + + self.temperature = temperature + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + scores_processed = scores / self.temperature + return scores_processed + + +class RepetitionPenaltyLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] that prevents the repetition of previous tokens through a penalty. This penalty is applied at + most once per token. Note that, for decoder-only models like most LLMs, the considered tokens include the prompt. + + In the original [paper](https://arxiv.org/pdf/1909.05858.pdf), the authors suggest the use of a penalty of around + 1.2 to achieve a good balance between truthful generation and lack of repetition. To penalize and reduce + repetition, use `penalty` values above 1.0, where a higher value penalizes more strongly. To reward and encourage + repetition, use `penalty` values between 0.0 and 1.0, where a lower value rewards more strongly. + + Args: + penalty (`float`): + The parameter for repetition penalty. 1.0 means no penalty. Above 1.0 penalizes previously generated + tokens. Between 0.0 and 1.0 rewards previously generated tokens. + + Examples: + + ```py + >>> from transformers import AutoTokenizer, AutoModelForCausalLM + + >>> # Initializing the model and tokenizer for it + >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2") + >>> inputs = tokenizer(["I'm not going to"], return_tensors="pt") + + >>> # This shows a normal generate without any specific parameters + >>> summary_ids = model.generate(**inputs) + >>> print(tokenizer.batch_decode(summary_ids, skip_special_tokens=True)[0]) + I'm not going to be able to do that. I'm going to be able to do that + + >>> # This generates a penalty for repeated tokens + >>> penalized_ids = model.generate(**inputs, repetition_penalty=1.1) + >>> print(tokenizer.batch_decode(penalized_ids, skip_special_tokens=True)[0]) + I'm not going to be able to do that. I'll just have to go out and play + ``` + """ + + def __init__(self, penalty: float): + if not isinstance(penalty, float) or not (penalty > 0): + raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}") + + self.penalty = penalty + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + score = torch.gather(scores, 1, input_ids) + + # if score < 0 then repetition penalty has to be multiplied to reduce the token probabilities + score = torch.where(score < 0, score * self.penalty, score / self.penalty) + + scores_processed = scores.scatter(1, input_ids, score) + return scores_processed + + +class EncoderRepetitionPenaltyLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] that works similarly to [`RepetitionPenaltyLogitsProcessor`], but with an *inverse* penalty + that is applied to the tokens present in the prompt. In other words, a penalty above 1.0 increases the odds of + selecting tokens that were present in the prompt. + + It was designed to avoid hallucination in input-grounded tasks, like summarization. Although originally intended + for encoder-decoder models, it can also be used with decoder-only models like LLMs. + + Args: + penalty (`float`): + The parameter for repetition penalty. 1.0 means no penalty. Above 1.0 rewards prompt tokens. Between 0.0 + and 1.0 penalizes prompt tokens. + encoder_input_ids (`torch.LongTensor`): + The encoder_input_ids that should be repeated within the decoder ids. + + Examples: + + ```python + >>> from transformers import AutoModelForCausalLM, AutoTokenizer + + >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m") + >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m") + + >>> inputs = tokenizer(["Alice and Bob. The third member's name was"], return_tensors="pt") + >>> gen_out = model.generate(**inputs) + >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0]) + Alice and Bob. The third member's name was not mentioned. + + >>> # With the `encoder_repetition_penalty` argument we can trigger this logits processor in `generate`, which can + >>> # promote the use of prompt tokens ("Bob" in this example) + >>> gen_out = model.generate(**inputs, encoder_repetition_penalty=1.2) + >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0]) + Alice and Bob. The third member's name was Bob. The third member's name was Bob. + ``` + """ + + def __init__(self, penalty: float, encoder_input_ids: torch.LongTensor): + if not isinstance(penalty, float) or not (penalty > 0): + raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}") + + self.penalty = 1 / penalty + self.encoder_input_ids = encoder_input_ids + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + score = torch.gather(scores, 1, self.encoder_input_ids) + + # if score < 0 then hallucination penalty has to be multiplied to increase the token probabilities + score = torch.where(score < 0, score * self.penalty, score / self.penalty) + + scores_processed = scores.scatter(1, self.encoder_input_ids, score) + return scores_processed + + +class TopPLogitsWarper(LogitsProcessor): + """ + [`LogitsProcessor`] that performs top-p, i.e. restricting to top tokens summing to prob_cut_off <= prob_cut_off. + Often used together with [`TemperatureLogitsWarper`] and [`TopKLogitsWarper`]. + + Args: + top_p (`float`): + If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or + higher are kept for generation. + filter_value (`float`, *optional*, defaults to -inf): + All filtered values will be set to this float value. + min_tokens_to_keep (`int`, *optional*, defaults to 1): + Minimum number of tokens that cannot be filtered. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed + + >>> set_seed(1) + >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2") + + >>> inputs = tokenizer("A sequence: 1, 2", return_tensors="pt") + + >>> # With sampling, the output is unexpected -- sometimes too unexpected. + >>> outputs = model.generate(**inputs, do_sample=True) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + A sequence: 1, 2, 3 | < 4 (left-hand pointer) ; + + + + >>> # With `top_p` sampling, the output gets restricted to high-probability tokens. + >>> # Pro tip: In practice, LLMs use `top_p` in the 0.9-0.95 range. + >>> outputs = model.generate(**inputs, do_sample=True, top_p=0.1) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + A sequence: 1, 2, 3, 4, 5, 6, 7, 8, 9 + ``` + """ + + def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1): + top_p = float(top_p) + if top_p < 0 or top_p > 1.0: + raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}") + if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1): + raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}") + + self.top_p = top_p + self.filter_value = filter_value + self.min_tokens_to_keep = min_tokens_to_keep + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + sorted_logits, sorted_indices = torch.sort(scores, descending=False) + cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1) + + # Remove tokens with cumulative top_p above the threshold (token with 0 are kept) + sorted_indices_to_remove = cumulative_probs <= (1 - self.top_p) + # Keep at least min_tokens_to_keep + sorted_indices_to_remove[..., -self.min_tokens_to_keep :] = 0 + + # scatter sorted tensors to original indexing + indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) + scores_processed = scores.masked_fill(indices_to_remove, self.filter_value) + return scores_processed + + +class TopKLogitsWarper(LogitsProcessor): + r""" + [`LogitsProcessor`] that performs top-k, i.e. restricting to the k highest probability elements. Often used + together with [`TemperatureLogitsWarper`] and [`TopPLogitsWarper`]. + + Args: + top_k (`int`): + The number of highest probability vocabulary tokens to keep for top-k-filtering. + filter_value (`float`, *optional*, defaults to -inf): + All filtered values will be set to this float value. + min_tokens_to_keep (`int`, *optional*, defaults to 1): + Minimum number of tokens that cannot be filtered. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed + + >>> set_seed(1) + >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2") + + >>> inputs = tokenizer("A sequence: A, B, C, D", return_tensors="pt") + + >>> # With sampling, the output is unexpected -- sometimes too unexpected. + >>> outputs = model.generate(**inputs, do_sample=True) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + A sequence: A, B, C, D, E — S — O, P — R + + >>> # With `top_k` sampling, the output gets restricted the k most likely tokens. + >>> # Pro tip: In practice, LLMs use `top_k` in the 5-50 range. + >>> outputs = model.generate(**inputs, do_sample=True, top_k=2) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + A sequence: A, B, C, D, E, F, G, H, I + ``` + """ + + def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1): + if not isinstance(top_k, int) or top_k <= 0: + raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}") + + self.top_k = max(top_k, min_tokens_to_keep) + self.filter_value = filter_value + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + top_k = min(self.top_k, scores.size(-1)) # Safety check + # Remove all tokens with a probability less than the last token of the top-k + indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None] + scores_processed = scores.masked_fill(indices_to_remove, self.filter_value) + return scores_processed + + +class MinPLogitsWarper(LogitsProcessor): + """ + [`LogitsProcessor`] that performs min-p, i.e. keeps all tokens that are above a minimum probability, scaled by the + probability of the most likely token. As a result, the filter becomes more agressive in the presence of + high-probability tokens, which is a sign of a confident output that we shouldn't deviate from. + + Often used together with [`TemperatureLogitsWarper`]. Used as an alternative to [`TopPLogitsWarper`] and + [`TopKLogitsWarper`]. + + Created by @menhguin and @kalomaze (github handles). Code adapted from [this external PR](https://github.com/oobabooga/text-generation-webui/pull/4449/files) + + Args: + min_p (`float`): + Minimum token probability, which will be scaled by the probability of the most likely token. It must be a + value between 0 and 1. Typical values are in the 0.01-0.2 range, comparably selective as setting `top_p` in + the 0.99-0.8 range (use the opposite of normal `top_p` values). + filter_value (`float`, *optional*, defaults to -inf): + All filtered values will be set to this float value. + min_tokens_to_keep (`int`, *optional*, defaults to 1): + Minimum number of tokens that cannot be filtered. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed + + >>> set_seed(1) + >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2") + + >>> inputs = tokenizer("A sequence: 1, 2", return_tensors="pt") + + >>> # With sampling, the output is unexpected -- sometimes too unexpected. + >>> outputs = model.generate(**inputs, do_sample=True) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + A sequence: 1, 2, 3 | < 4 (left-hand pointer) ; + + + + >>> # With `min_p` sampling, the output gets restricted to high-probability tokens. + >>> # Pro tip: In practice, LLMs use `min_p` in the 0.01-0.2 range. + >>> outputs = model.generate(**inputs, do_sample=True, min_p=0.1) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + A sequence: 1, 2, 3, 4, 5, 6, 7, 8, 9 + ``` + """ + + def __init__(self, min_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1): + if not (0 <= min_p <= 1.0): + raise ValueError(f"`min_p` has to be a float in the [0, 1] interval, but is {min_p}") + if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1): + raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}") + + self.min_p = min_p + self.filter_value = filter_value + self.min_tokens_to_keep = min_tokens_to_keep + + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + # Convert logits to probabilities + probs = torch.softmax(scores, dim=-1) + # Get the probability of the top token for each sequence in the batch + top_probs, _ = probs.max(dim=-1, keepdim=True) + # Calculate the actual min_p threshold by scaling min_p with the top token's probability + scaled_min_p = self.min_p * top_probs + # Create a mask for tokens that have a probability less than the scaled min_p + tokens_to_remove = probs < scaled_min_p + + sorted_indices = torch.argsort(scores, descending=True, dim=-1) + sorted_indices_to_remove = torch.gather(tokens_to_remove, dim=-1, index=sorted_indices) + # Keep at least min_tokens_to_keep + sorted_indices_to_remove[..., : self.min_tokens_to_keep] = False + + indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) + scores_processed = scores.masked_fill(indices_to_remove, self.filter_value) + return scores_processed + + +class TypicalLogitsWarper(LogitsProcessor): + r""" + [`LogitsProcessor`] that performs typical decoding. Inspired on how humans use language, it prioritizes tokens + whose log probability is close to the entropy of the token probability distribution. This means that the most + likely tokens may be discarded in the process. + + See [Typical Decoding for Natural Language Generation](https://arxiv.org/abs/2202.00666) for more information. + + Args: + mass (`float`, *optional*, defaults to 0.9): + Value of typical_p between 0 and 1 inclusive, defaults to 0.9. + filter_value (`float`, *optional*, defaults to -inf): + All filtered values will be set to this float value. + min_tokens_to_keep (`int`, *optional*, defaults to 1): + Minimum number of tokens that cannot be filtered. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed + + >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m") + >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m") + + >>> inputs = tokenizer("1, 2, 3", return_tensors="pt") + + >>> # We can see that greedy decoding produces a sequence of numbers + >>> outputs = model.generate(**inputs) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, + + >>> # For this particular seed, we can see that sampling produces nearly the same low-information (= low entropy) + >>> # sequence + >>> set_seed(18) + >>> outputs = model.generate(**inputs, do_sample=True) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + 1, 2, 3, 4, 5, 6, 7, 8, 9 and 10 + + >>> # With `typical_p` set, the most obvious sequence is no longer produced, which may be good for your problem + >>> set_seed(18) + >>> outputs = model.generate( + ... **inputs, do_sample=True, typical_p=0.1, return_dict_in_generate=True, output_scores=True + ... ) + >>> print(tokenizer.batch_decode(outputs.sequences, skip_special_tokens=True)[0]) + 1, 2, 3 and 5 + + >>> # We can see that the token corresponding to "4" (token 934) in the second position, the most likely token + >>> # as seen with greedy decoding, was entirely blocked out + >>> print(outputs.scores[1][0, 934]) + tensor(-inf) + ``` + """ + + def __init__(self, mass: float = 0.9, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1): + mass = float(mass) + if not (mass > 0 and mass < 1): + raise ValueError(f"`typical_p` has to be a float > 0 and < 1, but is {mass}") + if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1): + raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}") + + self.filter_value = filter_value + self.mass = mass + self.min_tokens_to_keep = min_tokens_to_keep + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + # calculate entropy + normalized = torch.nn.functional.log_softmax(scores, dim=-1) + p = torch.exp(normalized) + ent = -(normalized * p).nansum(-1, keepdim=True) + + # shift and sort + shifted_scores = torch.abs((-normalized) - ent) + sorted_scores, sorted_indices = torch.sort(shifted_scores, descending=False) + sorted_logits = scores.gather(-1, sorted_indices) + cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1) + + # Remove tokens with cumulative mass above the threshold + last_ind = (cumulative_probs < self.mass).sum(dim=1) + last_ind.clamp_(max=sorted_scores.shape[-1] - 1) + sorted_indices_to_remove = sorted_scores > sorted_scores.gather(1, last_ind.view(-1, 1)) + sorted_indices_to_remove[..., : self.min_tokens_to_keep] = 0 + indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) + + scores_processed = scores.masked_fill(indices_to_remove, self.filter_value) + return scores_processed + + +class EpsilonLogitsWarper(LogitsProcessor): + r""" + [`LogitsProcessor`] that performs epsilon-sampling, i.e. restricting to tokens with `prob >= epsilon`. Takes the + largest min_tokens_to_keep tokens if no tokens satisfy this constraint. See [Truncation Sampling as Language Model + Desmoothing](https://arxiv.org/abs/2210.15191) for more information. + + Args: + epsilon (`float`): + If set to > 0, only the most tokens with probabilities `epsilon` or higher are kept for generation. + filter_value (`float`, *optional*, defaults to -inf): + All filtered values will be set to this float value. + min_tokens_to_keep (`int`, *optional*, defaults to 1): + Minimum number of tokens that cannot be filtered. + + Examples: + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed + + >>> set_seed(1) + >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2") + + >>> inputs = tokenizer("A sequence: 1, 2", return_tensors="pt") + + >>> # With sampling, the output is unexpected -- sometimes too unexpected. + >>> outputs = model.generate(**inputs, do_sample=True) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + A sequence: 1, 2, 3 | < 4 (left-hand pointer) ; + + + + >>> # With epsilon sampling, the output gets restricted to high-probability tokens. Note that this is similar to + >>> # Top P sampling, which restricts tokens based on their cumulative probability. + >>> # Pro tip: The paper recomends using `epsilon_cutoff` values between 3e-4 and 9e-4 + >>> outputs = model.generate(**inputs, do_sample=True, epsilon_cutoff=0.1) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + A sequence: 1, 2, 3, 4, 5, 6, 7, 8, 9 + ``` + """ + + def __init__(self, epsilon: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1): + epsilon = float(epsilon) + if epsilon <= 0 or epsilon >= 1: + raise ValueError(f"`epsilon_cutoff` has to be a float > 0 and < 1, but is {epsilon}") + + min_tokens_to_keep = int(min_tokens_to_keep) + if min_tokens_to_keep < 1: + raise ValueError( + f"`min_tokens_to_keep` has to be a strictly positive integer, but is {min_tokens_to_keep}" + ) + + self.epsilon = epsilon + self.filter_value = filter_value + self.min_tokens_to_keep = min_tokens_to_keep + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + # Determine which indices to remove + probabilities = scores.softmax(dim=-1) + indices_to_remove = probabilities < self.epsilon + + # Keep the words with the 'min_tokens_to_keep'-highest probabilities + top_k = min(self.min_tokens_to_keep, scores.size(-1)) # Safety check + indices_to_remove = indices_to_remove & (scores < torch.topk(scores, top_k)[0][..., -1, None]) + + scores_processed = scores.masked_fill(indices_to_remove, self.filter_value) + return scores_processed + + +class EtaLogitsWarper(LogitsProcessor): + r""" + [`LogitsProcessor`] that performs eta-sampling, a technique to filter out tokens with probabilities below a dynamic + cutoff value, `eta`, which is calculated based on a combination of the hyperparameter `epsilon` and the entropy of + the token probabilities, i.e. `eta := min(epsilon, sqrt(epsilon * e^-entropy(probabilities)))`. Takes the largest + min_tokens_to_keep tokens if no tokens satisfy this constraint. It addresses the issue of poor quality in long + samples of text generated by neural language models leading to more coherent and fluent text. See [Truncation + Sampling as Language Model Desmoothing](https://arxiv.org/abs/2210.15191) for more information. Note: `do_sample` + must be set to `True` for this `LogitsProcessor` to work. + + + Args: + epsilon (`float`): + A float value in the range (0, 1). Hyperparameter used to calculate the dynamic cutoff value, `eta`. The + suggested values from the paper ranges from 3e-4 to 4e-3 depending on the size of the model. + filter_value (`float`, *optional*, defaults to -inf): + All values that are found to be below the dynamic cutoff value, `eta`, are set to this float value. This + parameter is useful when logits need to be modified for very low probability tokens that should be excluded + from generation entirely. + min_tokens_to_keep (`int`, *optional*, defaults to 1): + Specifies the minimum number of tokens that must be kept for generation, regardless of their probabilities. + For example, if `min_tokens_to_keep` is set to 1, at least one token will always be kept for generation, + even if all tokens have probabilities below the cutoff `eta`. + device (`str`, *optional*, defaults to `"cpu"`): + The device to allocate the tensors. + + Examples: + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed + + >>> set_seed(1) + >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2") + + >>> inputs = tokenizer("A sequence: 1, 2", return_tensors="pt") + + >>> # With sampling, the output is unexpected -- sometimes too unexpected. + >>> outputs = model.generate(**inputs, do_sample=True) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + A sequence: 1, 2, 3 | < 4 (left-hand pointer) ; + + + + >>> # With eta sampling, the output gets restricted to high-probability tokens. You can see it as a dynamic form of + >>> # epsilon sampling that adapts its cutoff probability based on the entropy (high entropy = lower cutoff). + >>> # Pro tip: The paper recomends using `eta_cutoff` values between 3e-4 to 4e-3 + >>> outputs = model.generate(**inputs, do_sample=True, eta_cutoff=0.1) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + A sequence: 1, 2, 3, 4, 5, 6, 7, 8, 9 + ``` + """ + + def __init__( + self, epsilon: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1, device: str = "cpu" + ): + epsilon = float(epsilon) + if epsilon <= 0 or epsilon >= 1: + raise ValueError(f"`eta_cutoff` has to be a float > 0 and < 1, but is {epsilon}") + + min_tokens_to_keep = int(min_tokens_to_keep) + if min_tokens_to_keep < 1: + raise ValueError( + f"`min_tokens_to_keep` has to be a strictly positive integer, but is {min_tokens_to_keep}" + ) + + self.epsilon = torch.tensor(epsilon, device=device) + self.filter_value = filter_value + self.min_tokens_to_keep = min_tokens_to_keep + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + probabilities = scores.softmax(dim=-1) + entropy = torch.distributions.Categorical(logits=scores).entropy() + eta = torch.min(self.epsilon, torch.sqrt(self.epsilon) * torch.exp(-entropy))[..., None] + indices_to_remove = probabilities < eta + + # Keep the words with the 'min_tokens_to_keep'-highest probabilities + top_k = min(self.min_tokens_to_keep, scores.size(-1)) # Safety check + indices_to_remove = indices_to_remove & (scores < torch.topk(scores, top_k)[0][..., -1, None]) + + scores_processed = scores.masked_fill(indices_to_remove, self.filter_value) + return scores_processed + + +def _get_ngrams(ngram_size: int, prev_input_ids: torch.Tensor, num_hypos: int): + """ + Assume ngram_size=2 and prev_input_ids=tensor([[40, 2883, 2712, 4346]]). The output of generated ngrams look like + this {(40,): [2883], (2883,): [2712], (2712,): [4346]}. + + Args: + ngram_size (`int`): + The number sequential tokens taken as a group which may only occur once before being banned. + prev_input_ids (`torch.Tensor`): + Generated token ids for the current hypothesis. + num_hypos (`int`): + The number of hypotheses for which n-grams need to be generated. + + Returns: + generated_ngrams (`dict`): + Dictionary of generated ngrams. + """ + # Initialize an empty list of dictionaries, one for each hypothesis (index) in the range of num_hypos + generated_ngrams = [{} for _ in range(num_hypos)] + for idx in range(num_hypos): + gen_tokens = prev_input_ids[idx].tolist() + generated_ngram = generated_ngrams[idx] + # Loop through each n-gram of size ngram_size in the list of tokens (gen_tokens) + for ngram in zip(*[gen_tokens[i:] for i in range(ngram_size)]): + prev_ngram_tuple = tuple(ngram[:-1]) + generated_ngram[prev_ngram_tuple] = generated_ngram.get(prev_ngram_tuple, []) + [ngram[-1]] + return generated_ngrams + + +def _get_generated_ngrams(banned_ngrams, prev_input_ids, ngram_size, cur_len): + """ + Determines the banned tokens for the current hypothesis based on previously generated n-grams. + + Args: + banned_ngrams (`dict`): + A dictionary containing previously generated n-grams for each hypothesis. + prev_input_ids (`torch.Tensor`): + Generated token ids for the current hypothesis. + ngram_size (`int`): + The number sequential tokens taken as a group which may only occur once before being banned. + cur_len (`int`): + The current length of the token sequences for which the n-grams are being checked. + + Returns: + List of tokens that are banned. + """ + # Before decoding the next token, prevent decoding of ngrams that have already appeared + start_idx = cur_len + 1 - ngram_size + ngram_idx = tuple(prev_input_ids[start_idx:cur_len].tolist()) + return banned_ngrams.get(ngram_idx, []) + + +def _calc_banned_ngram_tokens( + ngram_size: int, prev_input_ids: torch.Tensor, num_hypos: int, cur_len: int +) -> List[Iterable[int]]: + """Copied from fairseq for no_repeat_ngram in beam_search""" + if cur_len + 1 < ngram_size: + # return no banned tokens if we haven't generated no_repeat_ngram_size tokens yet + return [[] for _ in range(num_hypos)] + generated_ngrams = _get_ngrams(ngram_size, prev_input_ids, num_hypos) + banned_tokens = [ + _get_generated_ngrams(generated_ngrams[hypo_idx], prev_input_ids[hypo_idx], ngram_size, cur_len) + for hypo_idx in range(num_hypos) + ] + return banned_tokens + + +class NoRepeatNGramLogitsProcessor(LogitsProcessor): + r""" + N-grams are groups of "n" consecutive words, characters, or tokens taken from a sequence of text. Given the + sentence: "She runs fast", the bi-grams (n=2) would be ("she", "runs") and ("runs", "fast"). In text generation, + avoiding repetitions of word sequences provides a more diverse output. This [`LogitsProcessor`] enforces no + repetition of n-grams by setting the scores of banned tokens to negative infinity which eliminates those tokens + from consideration when further processing the scores. Note that, for decoder-only models like most LLMs, the + prompt is also considered to obtain the n-grams. + [Fairseq](https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345). + + + + Use n-gram penalties with care. For instance, penalizing 2-grams (bigrams) in an article about the city of New York + might lead to undesirable outcomes where the city's name appears only once in the entire text. + [Reference](https://huggingface.co/blog/how-to-generate) + + + + Args: + ngram_size (`int`): + All ngrams of size `ngram_size` can only occur once. + + Examples: + + ```py + >>> from transformers import AutoTokenizer, AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2") + >>> inputs = tokenizer(["Today I"], return_tensors="pt") + + >>> output = model.generate(**inputs) + >>> print(tokenizer.decode(output[0], skip_special_tokens=True)) + Today I’m not sure if I’m going to be able to do it. + + >>> # Now let's add ngram size using `no_repeat_ngram_size`. This stops the repetitions ("I’m") in the output. + >>> output = model.generate(**inputs, no_repeat_ngram_size=2) + >>> print(tokenizer.decode(output[0], skip_special_tokens=True)) + Today I’m not sure if I can get a better understanding of the nature of this issue + ``` + """ + + def __init__(self, ngram_size: int): + if not isinstance(ngram_size, int) or ngram_size <= 0: + raise ValueError(f"`ngram_size` has to be a strictly positive integer, but is {ngram_size}") + self.ngram_size = ngram_size + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + num_batch_hypotheses = scores.shape[0] + cur_len = input_ids.shape[-1] + scores_processed = scores.clone() + banned_batch_tokens = _calc_banned_ngram_tokens(self.ngram_size, input_ids, num_batch_hypotheses, cur_len) + for i, banned_tokens in enumerate(banned_batch_tokens): + scores_processed[i, banned_tokens] = -float("inf") + + return scores_processed + + +class EncoderNoRepeatNGramLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] that works similarly to [`NoRepeatNGramLogitsProcessor`], but applied exclusively to prevent + the repetition of n-grams present in the prompt. + + It was designed to promote chattiness in a language model, by preventing the generation of n-grams present in + previous conversation rounds. + + Args: + encoder_ngram_size (`int`): + All ngrams of size `ngram_size` can only occur within the encoder input ids. + encoder_input_ids (`int`): + The encoder_input_ids that should not be repeated within the decoder ids. + + Examples: + + ```py + >>> from transformers import AutoTokenizer, AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m") + >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m") + + >>> inputs = tokenizer("Alice: I love cats. What do you love?\nBob:", return_tensors="pt") + + >>> # With greedy decoding, we see Bob repeating Alice's opinion. If Bob was a chatbot, it would be a poor one. + >>> outputs = model.generate(**inputs) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + Alice: I love cats. What do you love? + Bob: I love cats. What do you + + >>> # With this logits processor, we can prevent Bob from repeating Alice's opinion. + >>> outputs = model.generate(**inputs, encoder_no_repeat_ngram_size=2) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + Alice: I love cats. What do you love? + Bob: My cats are very cute. + ``` + """ + + def __init__(self, encoder_ngram_size: int, encoder_input_ids: torch.LongTensor): + if not isinstance(encoder_ngram_size, int) or encoder_ngram_size <= 0: + raise ValueError( + f"`encoder_ngram_size` has to be a strictly positive integer, but is {encoder_ngram_size}" + ) + self.ngram_size = encoder_ngram_size + if len(encoder_input_ids.shape) == 1: + encoder_input_ids = encoder_input_ids.unsqueeze(0) + self.batch_size = encoder_input_ids.shape[0] + self.generated_ngrams = _get_ngrams(encoder_ngram_size, encoder_input_ids, self.batch_size) + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + # B x num_beams + num_hypos = scores.shape[0] + num_beams = num_hypos // self.batch_size + cur_len = input_ids.shape[-1] + scores_processed = scores.clone() + banned_batch_tokens = [ + _get_generated_ngrams( + self.generated_ngrams[hypo_idx // num_beams], input_ids[hypo_idx], self.ngram_size, cur_len + ) + for hypo_idx in range(num_hypos) + ] + + for i, banned_tokens in enumerate(banned_batch_tokens): + scores_processed[i, banned_tokens] = -float("inf") + + return scores_processed + + +class SequenceBiasLogitsProcessor(LogitsProcessor): + """ + [`LogitsProcessor`] that applies an additive bias on sequences. The bias is applied to the last token of a sequence + when the next generated token can complete it. Consequently, to take the most of biasing sequences with more than + one token, consider using beam methods (to gracefully work around partially completed sequences that have a + negative bias) and applying the bias to their prefixes (to ensure the bias is applied earlier). + + + + At a token-level, biasing a word is different from biasing a word with a space before it. If you want to bias + "foo" mid-sentence, you'll likely want to add a prefix space and bias " foo" instead. Check the tokenizer section + of our NLP course to find out why: https://huggingface.co/learn/nlp-course/chapter2/4?fw=pt + + + + Args: + sequence_bias (`List[List[Union[List[int], float]]]`): + List of lists that maps a sequence of tokens to its bias term (e.g. `[[[10, 45], -2.0], + [[64], -7.5]]`). Positive biases increase the odds of the + sequence being selected, while negative biases do the opposite. If a sequence has a length of 1, its bias + will always be applied. Otherwise, the bias will only be applied if the sequence in question is about to be + completed (in the token selection step after this processor is applied). + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct") + >>> inputs = tokenizer(["The full name of Donald is Donald"], return_tensors="pt") + + >>> summary_ids = model.generate(inputs["input_ids"], max_new_tokens=4, do_sample=False) + >>> print(tokenizer.batch_decode(summary_ids, skip_special_tokens=True)[0]) + The full name of Donald is Donald John Trump Sr. + + >>> def get_tokens(word): + ... return tokenizer([word], add_special_tokens=False).input_ids[0] + + >>> # IMPORTANT: Remember our tip about adding spaces before words to bias them correctly. + >>> sequence_bias = [[get_tokens("Trump"), -10.0],] # will fail to apply bias + >>> biased_ids = model.generate( + ... inputs["input_ids"], max_new_tokens=4, do_sample=False, sequence_bias=sequence_bias + ... ) + >>> print(tokenizer.batch_decode(biased_ids, skip_special_tokens=True)[0]) + The full name of Donald is Donald John Trump Sr. + + >>> sequence_bias = [[get_tokens(" Trump"), -10.0],] # will work + >>> biased_ids = model.generate( + ... inputs["input_ids"], max_new_tokens=4, do_sample=False, sequence_bias=sequence_bias + ... ) + >>> print(tokenizer.batch_decode(biased_ids, skip_special_tokens=True)[0]) + The full name of Donald is Donald John Harper. He + + >>> # We can also add a positive bias to nudge the model towards specific tokens or continuations. This technique + >>> # is also more effective when paired up with beam search. + >>> sequence_bias = [[get_tokens(" Donald Duck"), 10.0],] + >>> biased_ids = model.generate( + ... inputs["input_ids"], max_new_tokens=4, num_beams=4, do_sample=False, sequence_bias=sequence_bias + ... ) + >>> print(tokenizer.batch_decode(biased_ids, skip_special_tokens=True)[0]) + The full name of Donald is Donald Duck. He is + ``` + """ + + def __init__(self, sequence_bias: List[List[Union[List[int], float]]]): + self.sequence_bias = sequence_bias + self._validate_arguments() + self._convert_list_arguments_into_dict() + + # Bias variables that will be populated on the first call (for retrocompatibility purposes, the vocabulary size + # is inferred in the first usage, which inhibits initializing here) + self.length_1_bias = None + self.prepared_bias_variables = False + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + # 1 - Prepares the bias tensors. This is only needed the first time the logit processor is called. + if not self.prepared_bias_variables: + self._prepare_bias_variables(scores) + + # 2 - prepares an empty bias to add + bias = torch.zeros_like(scores) + + # 3 - include the bias from length = 1 + bias += self.length_1_bias + + # 4 - include the bias from length > 1, after determining which biased sequences may be completed. + for sequence_ids, sequence_bias in self.sequence_bias.items(): + if len(sequence_ids) == 1: # the sequence is of length 1, already applied + continue + if len(sequence_ids) > input_ids.shape[1]: # the sequence is longer than the context, ignore + continue + prefix_length = len(sequence_ids) - 1 + last_token = sequence_ids[-1] + matching_rows = torch.eq( + input_ids[:, -prefix_length:], + torch.tensor(sequence_ids[:-1], dtype=input_ids.dtype, device=input_ids.device), + ).prod(dim=1) + bias[:, last_token] += torch.where( + matching_rows.bool(), + torch.tensor(sequence_bias, device=input_ids.device), + torch.tensor(0.0, device=input_ids.device), + ) + + # 5 - apply the bias to the scores + scores_processed = scores + bias + return scores_processed + + def _prepare_bias_variables(self, scores: torch.FloatTensor): + vocabulary_size = scores.shape[-1] + + # Check biased tokens out of bounds + invalid_biases = [] + for sequence_ids in self.sequence_bias: + for token_id in sequence_ids: + if token_id >= vocabulary_size: + invalid_biases.append(token_id) + if len(invalid_biases) > 0: + raise ValueError( + f"The model vocabulary size is {vocabulary_size}, but the following tokens were being biased: " + f"{invalid_biases}" + ) + + # Precompute the bias tensors to be applied. Sequences of length 1 are kept separately, as they can be applied + # with simpler logic. + self.length_1_bias = torch.zeros((vocabulary_size,), dtype=torch.float, device=scores.device) + for sequence_ids, bias in self.sequence_bias.items(): + if len(sequence_ids) == 1: + self.length_1_bias[sequence_ids[-1]] = bias + + self.prepared_bias_variables = True + + def _validate_arguments(self): + sequence_bias = self.sequence_bias + if not isinstance(sequence_bias, dict) and not isinstance(sequence_bias, list) or len(sequence_bias) == 0: + raise ValueError( + f"`sequence_bias` has to be a non-empty dictionary, or non-empty list of lists but is {sequence_bias}." + ) + if isinstance(sequence_bias, dict) and any( + not isinstance(sequence_ids, tuple) for sequence_ids in sequence_bias.keys() + ): + raise ValueError(f"`sequence_bias` has to be a dict with tuples as keys, but is {sequence_bias}.") + if isinstance(sequence_bias, dict) and any( + any((not isinstance(token_id, (int, np.integer)) or token_id < 0) for token_id in sequence_ids) + or len(sequence_ids) == 0 + for sequence_ids in sequence_bias.keys() + ): + raise ValueError( + f"Each key in `sequence_bias` has to be a non-empty tuple of positive integers, but is " + f"{sequence_bias}." + ) + + def all_token_bias_pairs_are_valid(sequence): + return ( + isinstance(sequence[0], list) + and all(isinstance(token_id, (int, np.integer)) and token_id > 0 for token_id in sequence[0]) + and isinstance(sequence[1], float) + ) + + if isinstance(sequence_bias, list) and any( + (not all_token_bias_pairs_are_valid(sequence)) or len(sequence) == 0 for sequence in sequence_bias + ): + raise ValueError( + f"Each element in `sequence_bias` has to be a non-empty list of lists of positive integers and float, but is " + f"{sequence_bias}." + ) + if isinstance(sequence_bias, dict) and any(not isinstance(bias, float) for bias in sequence_bias.values()): + raise ValueError(f"`sequence_bias` has to be a dict with floats as values, but is {sequence_bias}.") + + def _convert_list_arguments_into_dict(self): + """BC: we used to accept `dict{tuple of tokens: float}` directly, now we expect a list""" + if isinstance(self.sequence_bias, list): + temp_sequence = self.sequence_bias + self.sequence_bias = {tuple(sublist[0]): sublist[1] for sublist in temp_sequence} + + +class NoBadWordsLogitsProcessor(SequenceBiasLogitsProcessor): + """ + [`LogitsProcessor`] that enforces that specified sequences will never be selected. + + + + In order to get the token ids of the words that should not appear in the generated text, make sure to set + `add_prefix_space=True` when initializing the tokenizer, and use `tokenizer(bad_words, + add_special_tokens=False).input_ids`. The `add_prefix_space` argument is only supported for some slow tokenizers, + as fast tokenizers' prefixing behaviours come from `pre tokenizers`. Read more + [here](https://huggingface.co/docs/tokenizers/api/pre-tokenizers). + + + + Args: + bad_words_ids (`List[List[int]]`): + List of list of token ids that are not allowed to be generated. + eos_token_id (`Union[int, List[int], torch.Tensor]`, *optional*): + The id(s) of the *end-of-sequence* token. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + >>> inputs = tokenizer(["In a word, the cake is a"], return_tensors="pt") + + >>> output_ids = model.generate(inputs["input_ids"], max_new_tokens=5, pad_token_id=tokenizer.eos_token_id) + >>> print(tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]) + In a word, the cake is a bit of a mess. + + >>> # Now let's take the bad words out. Please note that the tokenizer is initialized differently + >>> tokenizer_with_prefix_space = AutoTokenizer.from_pretrained("openai-community/gpt2", add_prefix_space=True) + + + >>> def get_tokens_as_list(word_list): + ... "Converts a sequence of words into a list of tokens" + ... tokens_list = [] + ... for word in word_list: + ... tokenized_word = tokenizer_with_prefix_space([word], add_special_tokens=False).input_ids[0] + ... tokens_list.append(tokenized_word) + ... return tokens_list + + + >>> bad_words_ids = get_tokens_as_list(word_list=["mess"]) + >>> output_ids = model.generate( + ... inputs["input_ids"], max_new_tokens=5, bad_words_ids=bad_words_ids, pad_token_id=tokenizer.eos_token_id + ... ) + >>> print(tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]) + In a word, the cake is a bit of a surprise. + ``` + """ + + def __init__( + self, bad_words_ids: List[List[int]], eos_token_id: Optional[Union[int, List[int], torch.Tensor]] = None + ): + self.bad_word_ids = bad_words_ids + self._validate_arguments() + + # Filter EOS token from bad_words_ids + if eos_token_id is not None: + if not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id) + + bad_words_ids = list( + filter(lambda bad_token_seq: all(bad_token_seq != [i] for i in eos_token_id), bad_words_ids) + ) + + # Forbidding a sequence is equivalent to setting its bias to -inf + sequence_bias = {tuple(sequence): float("-inf") for sequence in bad_words_ids} + super().__init__(sequence_bias=sequence_bias) + + def _validate_arguments(self): + bad_words_ids = self.bad_word_ids + if not isinstance(bad_words_ids, list) or len(bad_words_ids) == 0: + raise ValueError(f"`bad_words_ids` has to be a non-empty list, but is {bad_words_ids}.") + if any(not isinstance(bad_word_ids, list) for bad_word_ids in bad_words_ids): + raise ValueError(f"`bad_words_ids` has to be a list of lists, but is {bad_words_ids}.") + if any( + any((not isinstance(token_id, (int, np.integer)) or token_id < 0) for token_id in bad_word_ids) + for bad_word_ids in bad_words_ids + ): + raise ValueError( + f"Each list in `bad_words_ids` has to be a list of positive integers, but is {bad_words_ids}." + ) + + +class PrefixConstrainedLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] that enforces constrained generation and is useful for prefix-conditioned constrained + generation. See [Autoregressive Entity Retrieval](https://arxiv.org/abs/2010.00904) for more information. + + Args: + prefix_allowed_tokens_fn (`Callable[[int, torch.Tensor], List[int]]`): + This function constraints the beam search to allowed tokens only at each step. This function takes 2 + arguments `inputs_ids` and the batch ID `batch_id`. It has to return a list with the allowed tokens for the + next generation step conditioned on the previously generated tokens `inputs_ids` and the batch ID + `batch_id`. + + Examples: + + ```py + >>> from transformers import AutoTokenizer, AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("bigscience/bloomz-560m") + >>> tokenizer = AutoTokenizer.from_pretrained("bigscience/bloomz-560m") + + >>> inputs = tokenizer("Alice and Bob", return_tensors="pt") + + >>> # By default, it continues generating according to the model's logits + >>> outputs = model.generate(**inputs, max_new_tokens=5) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + Alice and Bob are friends + + >>> # We can contrain it with `prefix_allowed_tokens_fn` to force a certain behavior based on a prefix. + >>> # For instance, we can force an entire entity to be generated when its beginning is detected. + >>> entity = tokenizer(" Bob Marley", return_tensors="pt").input_ids[0] # 3 tokens + >>> def prefix_allowed_tokens_fn(batch_id, input_ids): + ... ''' + ... Attempts to generate 'Bob Marley' when 'Bob' is detected. + ... In this case, `batch_id` is not used, but you can set rules for each batch member. + ... ''' + ... if input_ids[-1] == entity[0]: + ... return [entity[1].item()] + ... elif input_ids[-2] == entity[0] and input_ids[-1] == entity[1]: + ... return [entity[2].item()] + ... return list(range(tokenizer.vocab_size)) # If no match, allow all tokens + + >>> outputs = model.generate(**inputs, max_new_tokens=5, prefix_allowed_tokens_fn=prefix_allowed_tokens_fn) + >>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]) + Alice and Bob Marley + ``` + """ + + def __init__(self, prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], List[int]], num_beams: int): + self._prefix_allowed_tokens_fn = prefix_allowed_tokens_fn + self._num_beams = num_beams + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + mask = torch.full_like(scores, -math.inf) + batch_size = input_ids.shape[0] // self._num_beams + + for batch_id in range(batch_size): + for beam_id in range(self._num_beams): + sent = input_ids[batch_id * self._num_beams + beam_id] + prefix_allowed_tokens = self._prefix_allowed_tokens_fn(batch_id, sent) + if len(prefix_allowed_tokens) == 0: + raise ValueError( + f"`prefix_allowed_tokens_fn` returned an empty list for batch ID {batch_id}." + f"This means that the constraint is unsatisfiable. Please check your implementation" + f"of `prefix_allowed_tokens_fn` " + ) + mask[batch_id * self._num_beams + beam_id, prefix_allowed_tokens] = 0 + + scores_processed = scores + mask + return scores_processed + + +class HammingDiversityLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] that enforces diverse beam search. + + Note that this logits processor is only effective for [`PreTrainedModel.group_beam_search`]. See [Diverse Beam + Search: Decoding Diverse Solutions from Neural Sequence Models](https://arxiv.org/pdf/1610.02424.pdf) for more + details. + + Traditional beam search often generates very similar sequences across different beams. + `HammingDiversityLogitsProcessor` addresses this by penalizing beams that generate tokens already chosen by other + beams in the same time step. + + Args: + diversity_penalty (`float`): + This value is subtracted from a beam's score if it generates a token same as any beam from other group at a + particular time. A higher `diversity_penalty` will enforce greater diversity among the beams. Adjusting + this value can help strike a balance between diversity and natural likelihood. + num_beams (`int`): + Number of beams for beam search. 1 means no beam search. + num_beam_groups (`int`): + Number of groups to divide `num_beams` into in order to ensure diversity among different groups of beams. + [this paper](https://arxiv.org/pdf/1610.02424.pdf) for more details. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForSeq2SeqLM + >>> import torch + + >>> # Initialize the model and tokenizer + >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base") + >>> model = AutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base") + + >>> # A long text about the solar system + >>> text = ( + ... "The Solar System is a gravitationally bound system comprising the Sun and the objects that orbit it, " + ... "either directly or indirectly. Of the objects that orbit the Sun directly, the largest are the eight " + ... "planets, with the remainder being smaller objects, such as the five dwarf planets and small Solar System " + ... "bodies. The Solar System formed 4.6 billion years ago from the gravitational collapse of a giant " + ... "interstellar molecular cloud." + ... ) + >>> inputs = tokenizer("summarize: " + text, return_tensors="pt") + + >>> # Generate diverse summary + >>> outputs_diverse = model.generate( + ... **inputs, + ... num_beam_groups=2, + ... diversity_penalty=10.0, + ... max_length=100, + ... num_beams=4, + ... num_return_sequences=2, + ... ) + >>> summaries_diverse = tokenizer.batch_decode(outputs_diverse, skip_special_tokens=True) + + >>> # Generate non-diverse summary + >>> outputs_non_diverse = model.generate( + ... **inputs, + ... max_length=100, + ... num_beams=4, + ... num_return_sequences=2, + ... ) + >>> summary_non_diverse = tokenizer.batch_decode(outputs_non_diverse, skip_special_tokens=True) + + >>> # With `diversity_penalty`, the resulting beams are much more diverse + >>> print(summary_non_diverse) + ['the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.', + 'the Solar System formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.'] + + >>> print(summaries_diverse) + ['the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets.', + 'the solar system formed 4.6 billion years ago from the collapse of a giant interstellar molecular cloud. of the objects that orbit the Sun directly, the largest are the eight planets. the rest of the objects are smaller objects, such as the five dwarf planets and small solar system bodies.'] + ``` + """ + + def __init__(self, diversity_penalty: float, num_beams: int, num_beam_groups: int): + if not isinstance(diversity_penalty, float) or (not diversity_penalty > 0.0): + raise ValueError("`diversity_penalty` should be a float strictly larger than 0.") + self._diversity_penalty = diversity_penalty + if not isinstance(num_beams, int) or num_beams < 2: + raise ValueError("`num_beams` should be an integer strictly larger than 1.") + self._num_beams = num_beams + if not isinstance(num_beam_groups, int) or num_beam_groups < 2: + raise ValueError("`num_beam_groups` should be an integer strictly larger than 1.") + if num_beam_groups > num_beams: + raise ValueError("`beam_groups` has to be smaller or equal to `num_beams`.") + self._num_sub_beams = num_beams // num_beam_groups + + def __call__( + self, + input_ids: torch.LongTensor, + scores: torch.FloatTensor, + current_tokens: torch.LongTensor, + beam_group_idx: int, + ) -> torch.FloatTensor: + r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids) + scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`): + Prediction scores of a language modeling head. These can be logits for each vocabulary when not using + beam search or log softmax for each vocabulary token when using beam search + current_tokens (`torch.LongTensor` of shape `(batch_size)`): + Indices of input sequence tokens in the vocabulary, corresponding to the tokens selected by the other + beam groups in the current generation step. + beam_group_idx (`int`): + The index of the beam group currently being processed. + + Return: + `torch.FloatTensor` of shape `(batch_size, config.vocab_size)`: + The processed prediction scores. + """ + # hamming diversity: penalise using same token in current group which was used in previous groups at + # the same time step + batch_size = current_tokens.shape[0] // self._num_beams + group_start_idx = beam_group_idx * self._num_sub_beams + group_end_idx = min(group_start_idx + self._num_sub_beams, self._num_beams) + group_size = group_end_idx - group_start_idx + vocab_size = scores.shape[-1] + + if group_start_idx == 0: + return scores + + scores_processed = scores.clone() + for batch_idx in range(batch_size): + # predicted tokens of last time step of previous groups + previous_group_tokens = current_tokens[ + batch_idx * self._num_beams : batch_idx * self._num_beams + group_start_idx + ] + token_frequency = torch.bincount(previous_group_tokens, minlength=vocab_size).to(scores.device) + scores_processed[batch_idx * group_size : (batch_idx + 1) * group_size] -= ( + self._diversity_penalty * token_frequency + ) + + return scores_processed + + +class ForcedBOSTokenLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] that enforces the specified token as the first generated token. Used with encoder-decoder + models. + + Args: + bos_token_id (`int`): + The id of the token to force as the first generated token. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForSeq2SeqLM + + >>> model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small") + >>> tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small") + + >>> inputs = tokenizer("Translate from English to German: I love cats.", return_tensors="pt") + + >>> # By default, it continues generating according to the model's logits + >>> outputs = model.generate(**inputs, max_new_tokens=10) + >>> print(tokenizer.batch_decode(outputs)[0]) + Ich liebe Kitty. + + >>> # We can use `forced_bos_token_id` to force the start of generation with an encoder-decoder model + >>> # (including forcing it to end straight away with an EOS token) + >>> outputs = model.generate(**inputs, max_new_tokens=10, forced_bos_token_id=tokenizer.eos_token_id) + >>> print(tokenizer.batch_decode(outputs)[0]) + + ``` + """ + + def __init__(self, bos_token_id: int): + self.bos_token_id = bos_token_id + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + cur_len = input_ids.shape[-1] + scores_processed = scores + if cur_len == 1: + scores_processed = torch.full_like(scores, -math.inf) + scores_processed[:, self.bos_token_id] = 0 + return scores_processed + + +class ForcedEOSTokenLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] that enforces the specified token as the last generated token when `max_length` is reached. + + Args: + max_length (`int`): + The maximum length of the sequence to be generated. + eos_token_id (`Union[int, List[int], torch.Tensor]`): + The id(s) of the *end-of-sequence* token. + device (`str`, *optional*, defaults to `"cpu"`): + The device to allocate the tensors. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2") + + >>> inputs = tokenizer("A sequence: 1, 2, 3", return_tensors="pt") + + >>> # By default, it continues generating according to the model's logits + >>> outputs = model.generate(**inputs, max_new_tokens=10) + >>> print(tokenizer.batch_decode(outputs)[0]) + A sequence: 1, 2, 3, 4, 5, 6, 7, 8 + + >>> # `forced_eos_token_id` ensures the generation ends with a EOS token + >>> outputs = model.generate(**inputs, max_new_tokens=10, forced_eos_token_id=tokenizer.eos_token_id) + >>> print(tokenizer.batch_decode(outputs)[0]) + A sequence: 1, 2, 3, 4, 5, 6, 7,<|endoftext|> + ``` + """ + + def __init__(self, max_length: int, eos_token_id: Union[int, List[int], torch.Tensor], device: str = "cpu"): + self.max_length = max_length + + if not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id, device=device) + self.eos_token_id = eos_token_id + + if torch.is_floating_point(eos_token_id) or (eos_token_id < 0).any(): + raise ValueError(f"`eos_token_id` has to be a list of positive integers, but is {eos_token_id}") + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + cur_len = input_ids.shape[-1] + scores_processed = scores + if cur_len == self.max_length - 1: + scores_processed = torch.full_like(scores, -math.inf) + scores_processed[:, self.eos_token_id] = 0 + return scores_processed + + +class InfNanRemoveLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] that removes all `nan` and `inf` values to avoid the generation method to fail. Note that using + the logits processor should only be used if necessary since it can slow down the generation method. + + This logits processor has no `generate` example, as there shouldn't be a correct combination of flags that warrants + its use. + """ + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + # set all nan values to 0.0 + scores_processed = torch.where(scores != scores, 0.0, scores) + + # set all +/-inf values to max/min possible value + scores_processed = torch.where(scores == float("inf"), torch.finfo(scores.dtype).max, scores_processed) + scores_processed = torch.where(scores == -float("inf"), torch.finfo(scores.dtype).min, scores_processed) + + return scores_processed + + +class ExponentialDecayLengthPenalty(LogitsProcessor): + r""" + [`LogitsProcessor`] that exponentially increases the score of the `eos_token_id` after `start_index` has been + reached. This allows generating shorter sequences without having a hard cutoff, allowing the `eos_token` to be + predicted in a meaningful position. + + Args: + exponential_decay_length_penalty (`tuple(int, float)`): + This tuple shall consist of: `(start_index, decay_factor)` where `start_index` indicates where penalty + starts and `decay_factor` represents the factor of exponential decay + eos_token_id (`Union[int, List[int], torch.Tensor]`): + The id(s) of the *end-of-sequence* token. + input_ids_seq_length (`int`): + The length of the input sequence. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed + + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + + >>> text = "Just wanted to let you know, I" + >>> inputs = tokenizer(text, return_tensors="pt") + + >>> # Let's consider that we want short sentences, so we limit `max_length=30`. However, we observe that the answer + >>> # tends to end abruptly. + >>> set_seed(1) + >>> outputs = model.generate(**inputs, do_sample=True, temperature=0.9, max_length=30, pad_token_id=50256) + >>> print(tokenizer.batch_decode(outputs)[0]) + Just wanted to let you know, I received a link to an ebook, the book How To Start A Social Network which was + published in 2010. Although + + >>> # To promote the appearance of the EOS token at the right time, we add the `exponential_decay_length_penalty = + >>> # (start_index, decay_factor)`. Instead of cutting at max_tokens, the output comes to an end before and usually + >>> # with more meaning. What happens is that starting from `start_index` the EOS token score will be increased + >>> # by `decay_factor` exponentially. However, if you set a high decay factor, you may also end up with abruptly + >>> # ending sequences. + >>> set_seed(1) + >>> outputs = model.generate( + ... **inputs, + ... do_sample=True, + ... temperature=0.9, + ... max_length=30, + ... pad_token_id=50256, + ... exponential_decay_length_penalty=(15, 1.6), + ... ) + >>> print(tokenizer.batch_decode(outputs)[0]) + Just wanted to let you know, I received a link to an ebook, the book How To Start A Social Network + which<|endoftext|> + + >>> # With a small decay factor, you will have a higher chance of getting a meaningful sequence. + >>> set_seed(1) + >>> outputs = model.generate( + ... **inputs, + ... do_sample=True, + ... temperature=0.9, + ... max_length=30, + ... pad_token_id=50256, + ... exponential_decay_length_penalty=(15, 1.01), + ... ) + >>> print(tokenizer.batch_decode(outputs)[0]) + Just wanted to let you know, I received a link to an ebook, the book How To Start A Social Network which was + published in 2010.<|endoftext|> + ``` + """ + + def __init__( + self, + exponential_decay_length_penalty: Tuple[int, float], + eos_token_id: Union[int, List[int], torch.Tensor], + input_ids_seq_length: int, + ): + self.regulation_start = exponential_decay_length_penalty[0] + input_ids_seq_length + self.regulation_factor = exponential_decay_length_penalty[1] + + if not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id) + self.eos_token_id = eos_token_id + + if torch.is_floating_point(eos_token_id) or (eos_token_id < 0).any(): + raise ValueError(f"`eos_token_id` has to be a list of positive integers, but is {eos_token_id}") + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + cur_len = input_ids.shape[-1] + self.eos_token_id = self.eos_token_id.to(scores.device) + penalties = torch.zeros_like(scores) + scores_processed = scores + if cur_len > self.regulation_start: + penalty_idx = cur_len - self.regulation_start + # To support negative logits we compute the penalty of the absolute value and add to the original logit + penalty = torch.abs(scores[:, self.eos_token_id]) * (pow(self.regulation_factor, penalty_idx) - 1) + penalties[:, self.eos_token_id] = penalty + scores_processed = scores + penalties + return scores_processed + + +class LogitNormalization(LogitsProcessor): + r""" + [`LogitsProcessor`] for normalizing the scores using log-softmax. It's important to normalize + the scores during beam search, after applying the logits processors or warpers, since the search algorithm used in + this library doesn't do it (it only does it before, but they may need re-normalization) but it still supposes that + the scores are normalized when comparing the hypotheses. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM + >>> import torch + + >>> model = AutoModelForCausalLM.from_pretrained("distilbert/distilgpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("distilbert/distilgpt2") + + >>> inputs = tokenizer("A sequence: 1, 2, 3", return_tensors="pt") + + >>> # By default, the scores are not normalized -- the sum of their exponentials is NOT a normalized probability + >>> # distribution, summing to 1 + >>> outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True) + >>> print(torch.allclose(torch.sum(torch.exp(outputs.scores[-1])), torch.Tensor((1.000,)), rtol=1e-4)) + False + + >>> # Normalizing them may have a positive impact on beam methods, or when using the scores on your application + >>> outputs = model.generate(**inputs, renormalize_logits=True, return_dict_in_generate=True, output_scores=True) + >>> print(torch.allclose(torch.sum(torch.exp(outputs.scores[-1])), torch.Tensor((1.000,)), rtol=1e-4)) + True + ``` + """ + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + scores_processed = scores.log_softmax(dim=-1) + return scores_processed + + +class SuppressTokensAtBeginLogitsProcessor(LogitsProcessor): + r""" + [`SuppressTokensAtBeginLogitsProcessor`] supresses a list of tokens as soon as the `generate` function starts + generating using `begin_index` tokens. This should ensure that the tokens defined by `begin_suppress_tokens` are + not generated at the beginning. Originally created for + [Whisper](https://huggingface.co/docs/transformers/model_doc/whisper). + + Examples: + + ```python + >>> from transformers import AutoProcessor, WhisperForConditionalGeneration + >>> from datasets import load_dataset + + >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en") + >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en") + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt") + + >>> # Whisper has `begin_suppress_tokens` set by default (= `[220, 50256]`). 50256 is the EOS token, so this means + >>> # it can't generate and EOS token in the first iteration, but it can in the others. + >>> outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True) + >>> print(outputs.scores[0][0, 50256]) + tensor(-inf) + >>> print(outputs.scores[-1][0, 50256]) # in other places we can see some probability mass for EOS + tensor(29.9010) + + >>> # If we disable `begin_suppress_tokens`, we can generate EOS in the first iteration. + >>> outputs = model.generate( + ... **inputs, return_dict_in_generate=True, output_scores=True, begin_suppress_tokens=None + ... ) + >>> print(outputs.scores[0][0, 50256]) + tensor(11.2027) + ``` + """ + + def __init__(self, begin_suppress_tokens, begin_index, device: str = "cpu"): + self.begin_suppress_tokens = torch.tensor(list(begin_suppress_tokens), device=device) + self.begin_index = begin_index + + def set_begin_index(self, begin_index): + self.begin_index = begin_index + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + vocab_tensor = torch.arange(scores.shape[-1], device=scores.device) + suppress_token_mask = isin_mps_friendly(vocab_tensor, self.begin_suppress_tokens) + scores_processed = scores + if input_ids.shape[-1] == self.begin_index: + scores_processed = torch.where(suppress_token_mask, -float("inf"), scores) + + return scores_processed + + +class SuppressTokensLogitsProcessor(LogitsProcessor): + r""" + This processor can be used to suppress a list of tokens. The processor will set their log probs to `-inf` so + that they are not generated. Originally created for + [Whisper](https://huggingface.co/docs/transformers/model_doc/whisper). + + Examples: + + ```python + >>> from transformers import AutoProcessor, WhisperForConditionalGeneration + >>> from datasets import load_dataset + + >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en") + >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en") + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt") + + >>> # Whisper has a long list of suppressed tokens. For instance, in this case, the token 1 is suppressed by default. + >>> outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True) + >>> print(outputs.scores[1][0, 1]) # 1 (and not 0) is the first freely generated token + tensor(-inf) + + >>> # If we disable `suppress_tokens`, we can generate it. + >>> outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True, suppress_tokens=None) + >>> print(outputs.scores[1][0, 1]) + tensor(6.0678) + ``` + """ + + def __init__(self, suppress_tokens, device: str = "cpu"): + self.suppress_tokens = torch.tensor(list(suppress_tokens), device=device) + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + vocab_tensor = torch.arange(scores.shape[-1], device=scores.device) + suppress_token_mask = isin_mps_friendly(vocab_tensor, self.suppress_tokens) + scores = torch.where(suppress_token_mask, -float("inf"), scores) + return scores + + +class WhisperTimeStampLogitsProcessor(LogitsProcessor): + r""" + + [`LogitsProcessor`] that modifies the logits for the generation of timestamps in the transcription. When the input + tokens are at a specific threshold, the processor sets the scores to negative infinity. The processor makes sure + that timestamp tokens appear in pairs, by masking out the logits that would break this pairing pattern. This is + done to maintain the consistency and structure of generated timestamps. It also ensures that when the predicted + probability of sampling any of the timestamp token is greater than any individual non-timestamp token, those + non-timestamp logits are set to negative infinity. This is done to ensure the generation of timestamps over other + potential tokens. + + + See [the paper](https://arxiv.org/abs/2212.04356) for more information. + + Args: + generate_config (`GenerateConfig`): + The generate config used to generate the output. The following parameters are required: + eos_token_id (`int`, *optional*, defaults to 50257): + The id of the *end-of-sequence* token. + no_timestamps_token_id (`int`, *optional*, defaults to 50363): + The id of the `"<|notimestamps|>"` token. + max_initial_timestamp_index (`int`, *optional*, defaults to 1): + Used to set the maximum value of the initial timestamp. This is used to prevent the model from + predicting timestamps that are too far in the future. + begin_index (`Optional`, *optional*): Token index of the first token that is generated by the model. + _detect_timestamp_from_logprob (`bool`, *optional*): Whether timestamps can be predicted from logprobs over all timestamps. + + Examples: + ``` python + >>> import torch + >>> from transformers import AutoProcessor, WhisperForConditionalGeneration, GenerationConfig + >>> from datasets import load_dataset + + >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en") + >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en") + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> inputs = processor(ds[3]["audio"]["array"], return_tensors="pt") + >>> input_features = inputs.input_features + + >>> #Displaying timestamps + >>> generated_ids = model.generate(inputs=input_features, return_timestamps=True) + >>> transcription = processor.batch_decode(generated_ids, decode_with_timestamps=True)[0] + >>> print("Transcription:", transcription) + Transcription: <|startoftranscript|><|0.00|> He has grave doubts whether Sir Frederick Layton's work is really Greek after all, and can<|6.44|><|6.44|> discover in it but little of rocky Ithaca.<|9.44|><|endoftext|> + + + >>> #No timestamps & change EOS: + >>> #This allows the user to select a specific token to terminate the sequence on, in this case it's the word "can"(460) + >>> model.generation_config.eos_token_id = 460 + >>> generated_ids = model.generate(inputs=input_features,return_timestamps=False) + >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] + >>> print("Transcription:", transcription) + Transcription: He has grave doubts whether Sir Frederick Layton's work is really Greek after all and can + ``` + """ + + def __init__( + self, + generate_config, + begin_index: Optional[int] = None, + _detect_timestamp_from_logprob: Optional[bool] = None, + ): # support for the kwargs + self.no_timestamps_token_id = generate_config.no_timestamps_token_id + self.timestamp_begin = generate_config.no_timestamps_token_id + 1 + self.eos_token_id = generate_config.eos_token_id or generate_config.bos_token_id + + # this variable is mostly just used for testing + self._detect_timestamp_from_logprob = ( + _detect_timestamp_from_logprob + if _detect_timestamp_from_logprob is not None + else getattr(generate_config, "_detect_timestamp_from_logprob", True) + ) + + num_forced_ids = ( + len(generate_config.forced_decoder_ids) if generate_config.forced_decoder_ids is not None else 0 + ) + self.begin_index = begin_index or (num_forced_ids + 1) + + self.max_initial_timestamp_index = getattr(generate_config, "max_initial_timestamp_index", None) + # TODO(Patrick): Make sure that official models have max_initial_timestamp_index set to 50 + # self.max_initial_timestamp_index = 50 + + def set_begin_index(self, begin_index): + self.begin_index = begin_index + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + # suppress <|notimestamps|> which is handled by without_timestamps + scores_processed = scores.clone() + scores_processed[:, self.no_timestamps_token_id] = -float("inf") + + # timestamps have to appear in pairs, except directly before eos_token; mask logits accordingly + for k in range(input_ids.shape[0]): + sampled_tokens = input_ids[k, self.begin_index :] + seq = list(sampled_tokens.tolist()) + + last_was_timestamp = len(seq) >= 1 and seq[-1] >= self.timestamp_begin + penultimate_was_timestamp = len(seq) < 2 or seq[-2] >= self.timestamp_begin + + if last_was_timestamp: + if penultimate_was_timestamp: # has to be non-timestamp + scores_processed[k, self.timestamp_begin :] = -float("inf") + else: # cannot be normal text tokens + scores_processed[k, : self.eos_token_id] = -float("inf") + + timestamps = sampled_tokens[sampled_tokens.ge(self.timestamp_begin)] + if timestamps.numel() > 0: + # `timestamps` shouldn't decrease; forbid timestamp tokens smaller than the last + # The following lines of code are copied from: https://github.com/openai/whisper/pull/914/files#r1137085090 + if last_was_timestamp and not penultimate_was_timestamp: + timestamp_last = timestamps[-1] + else: + # Avoid to emit <|0.00|> again + timestamp_last = timestamps[-1] + 1 + + scores_processed[k, self.timestamp_begin : timestamp_last] = -float("inf") + + # apply the `max_initial_timestamp` option + if input_ids.shape[1] == self.begin_index: + scores_processed[:, : self.timestamp_begin] = -float("inf") + + if self.max_initial_timestamp_index is not None: + last_allowed = self.timestamp_begin + self.max_initial_timestamp_index + scores_processed[:, last_allowed + 1 :] = -float("inf") + + # if sum of probability over timestamps is above any other token, sample timestamp + logprobs = torch.nn.functional.log_softmax(scores_processed.float(), dim=-1) + for k in range(input_ids.shape[0]): + timestamp_logprob = logprobs[k, self.timestamp_begin :].logsumexp(dim=-1) + max_text_token_logprob = logprobs[k, : self.timestamp_begin].max() + if timestamp_logprob > max_text_token_logprob and self._detect_timestamp_from_logprob: + scores_processed[k, : self.timestamp_begin] = -float("inf") + + return scores_processed + + +class WhisperNoSpeechDetection(LogitsProcessor): + r"""This processor can be used to detect silence when using Whisper. It should take as input unprocessed logits to follow the original implementation""" + + def __init__(self, no_speech_token: int, begin_index: int, scores_is_logprobs: bool = False): + self.no_speech_token = no_speech_token + # offset between token, , in paper and first generated token + # is equal to the position of the first generated token index + self.start_of_trans_offset = begin_index + + # `self.begin_index` is a running value that is changed on the fly + self.begin_index = begin_index + self._no_speech_prob = [0.0] + self.is_scores_logprobs = scores_is_logprobs + + # overwritten dynamically + self.model = None + self.inputs = None + + def set_model(self, model): + self.model = model + + def set_inputs(self, inputs): + self.inputs = {**self.model.prepare_inputs_for_generation(**inputs), **inputs} + self.inputs["input_features"] = self.inputs.pop("inputs") + + @property + def no_speech_prob(self): + return self._no_speech_prob + + def set_begin_index(self, begin_index): + self.begin_index = begin_index + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + is_scores_logprobs = self.is_scores_logprobs + + if input_ids.shape[1] == self.begin_index: + if self.start_of_trans_offset > 1: + with torch.no_grad(): + logits = self.model(**self.inputs).logits + + no_speech_index = self.begin_index - self.start_of_trans_offset + no_speech_scores = logits[:, no_speech_index] + is_scores_logprobs = False + else: + no_speech_scores = scores + + if is_scores_logprobs: + probs = no_speech_scores.exp() + else: + probs = no_speech_scores.float().softmax(dim=-1) + + self._no_speech_prob = probs[:, self.no_speech_token] + + return scores + + +class ClassifierFreeGuidanceLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] for classifier free guidance (CFG). The scores are split over the batch dimension, + where the first half correspond to the conditional logits (predicted from the input prompt) and the second half + correspond to the unconditional logits (predicted from an empty or 'null' prompt). The processor computes a + weighted average across the conditional and unconditional logits, parameterised by the `guidance_scale`. + + See [the paper](https://arxiv.org/abs/2306.05284) for more information. + + + + This logits processor is exclusively compatible with + [MusicGen](https://huggingface.co/docs/transformers/main/en/model_doc/musicgen) + + + + Args: + guidance_scale (float): + The guidance scale for classifier free guidance (CFG). CFG is enabled by setting `guidance_scale > 1`. + Higher guidance scale encourages the model to generate samples that are more closely linked to the input + prompt, usually at the expense of poorer quality. + + Examples: + + ```python + >>> from transformers import AutoProcessor, MusicgenForConditionalGeneration + + >>> processor = AutoProcessor.from_pretrained("facebook/musicgen-small") + >>> model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small") + + >>> inputs = processor( + ... text=["80s pop track with bassy drums and synth", "90s rock song with loud guitars and heavy drums"], + ... padding=True, + ... return_tensors="pt", + ... ) + >>> audio_values = model.generate(**inputs, do_sample=True, guidance_scale=3, max_new_tokens=256) + ``` + """ + + def __init__(self, guidance_scale): + if guidance_scale > 1: + self.guidance_scale = guidance_scale + else: + raise ValueError( + "Require guidance scale >1 to use the classifier free guidance processor, got guidance scale " + f"{guidance_scale}." + ) + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + # simple check to make sure we have compatible batch sizes between our + # logits scores (cond + uncond) and input ids (cond only) + if scores.shape[0] != 2 * input_ids.shape[0]: + raise ValueError( + f"Logits should have twice the batch size of the input ids, the first half of batches corresponding to " + f"the conditional inputs, and the second half of batches corresponding to the unconditional inputs. Got " + f"batch size {scores.shape[0]} for the logits and {input_ids.shape[0]} for the input ids." + ) + unguided_bsz = scores.shape[0] // 2 + cond_logits, uncond_logits = scores.split(unguided_bsz, dim=0) + scores_processed = uncond_logits + (cond_logits - uncond_logits) * self.guidance_scale + return scores_processed + + +class AlternatingCodebooksLogitsProcessor(LogitsProcessor): + r""" + [`LogitsProcessor`] enforcing alternated generation between the two codebooks of Bark. + + + + This logits processor is exclusively compatible with + [Bark](https://huggingface.co/docs/transformers/en/model_doc/bark)'s fine submodel. See the model documentation + for examples. + + + + Args: + input_start_len (`int`): + The length of the initial input sequence. + semantic_vocab_size (`int`): + Vocabulary size of the semantic part, i.e number of tokens associated to the semantic vocabulary. + codebook_size (`int`): + Number of tokens associated to the codebook. + """ + + def __init__(self, input_start_len: int, semantic_vocab_size: int, codebook_size: int): + if not isinstance(input_start_len, int) or input_start_len < 0: + raise ValueError(f"`input_starting_length` has to be a non-negative integer, but is {input_start_len}") + + self.input_start_len = input_start_len + self.semantic_vocab_size = semantic_vocab_size + self.codebook_size = codebook_size + + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + curr_len = input_ids.shape[-1] + + # even -> first codebook, odd -> second codebook + is_first_codebook = ((curr_len - self.input_start_len) % 2) == 0 + + scores_processed = scores.clone() + if is_first_codebook: + scores_processed[:, : self.semantic_vocab_size] = -float("inf") + scores_processed[:, self.semantic_vocab_size + self.codebook_size :] = -float("inf") + else: + scores_processed[:, : self.semantic_vocab_size + self.codebook_size] = -float("inf") + + return scores_processed + + +class UnbatchedClassifierFreeGuidanceLogitsProcessor(LogitsProcessor): + r""" + Logits processor for Classifier-Free Guidance (CFG). The processors computes a weighted average across scores + from prompt conditional and prompt unconditional (or negative) logits, parameterized by the `guidance_scale`. + The unconditional scores are computed internally by prompting `model` with the `unconditional_ids` branch. + + See [the paper](https://arxiv.org/abs/2306.17806) for more information. + + Args: + guidance_scale (`float`): + The guidance scale for classifier free guidance (CFG). CFG is enabled by setting `guidance_scale != 1`. + Higher guidance scale encourages the model to generate samples that are more closely linked to the input + prompt, usually at the expense of poorer quality. A value smaller than 1 has the opposite effect, while + making the negative prompt provided with negative_prompt_ids (if any) act as a positive prompt. + model (`PreTrainedModel`): + The model computing the unconditional scores. Supposedly the same as the one computing the conditional + scores. Both models must use the same tokenizer. + unconditional_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of input sequence tokens in the vocabulary for the unconditional branch. If unset, will default to + the last token of the prompt. + unconditional_attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Attention mask for unconditional_ids. + use_cache (`bool`, *optional*, defaults to `True`): + Whether to cache key/values during the negative prompt forward pass. + + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + >>> inputs = tokenizer(["Today, a dragon flew over Paris, France,"], return_tensors="pt") + >>> out = model.generate(inputs["input_ids"], guidance_scale=1.5) + >>> tokenizer.batch_decode(out, skip_special_tokens=True)[0] + 'Today, a dragon flew over Paris, France, killing at least 50 people and injuring more than 100' + + >>> # with a negative prompt + >>> neg_inputs = tokenizer(["A very happy event happened,"], return_tensors="pt") + >>> out = model.generate(inputs["input_ids"], guidance_scale=2, negative_prompt_ids=neg_inputs["input_ids"]) + >>> tokenizer.batch_decode(out, skip_special_tokens=True)[0] + 'Today, a dragon flew over Paris, France, killing at least 130 people. French media reported that' + + >>> # with a positive prompt + >>> neg_inputs = tokenizer(["A very happy event happened,"], return_tensors="pt") + >>> out = model.generate(inputs["input_ids"], guidance_scale=0, negative_prompt_ids=neg_inputs["input_ids"]) + >>> tokenizer.batch_decode(out, skip_special_tokens=True)[0] + "Today, a dragon flew over Paris, France, and I'm very happy to be here. I" + ``` + """ + + def __init__( + self, + guidance_scale: float, + model, + unconditional_ids: Optional[torch.LongTensor] = None, + unconditional_attention_mask: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = True, + ): + self.guidance_scale = guidance_scale + self.model = model + self.unconditional_context = { + "input_ids": unconditional_ids, + "attention_mask": unconditional_attention_mask, + "use_cache": use_cache, + "past_key_values": None, + "first_pass": True, + } + + def get_unconditional_logits(self, input_ids): + if self.unconditional_context["first_pass"]: + if self.unconditional_context["input_ids"] is None: + self.unconditional_context["input_ids"] = input_ids[:, -1:] + if self.unconditional_context["attention_mask"] is None: + self.unconditional_context["attention_mask"] = torch.ones_like( + self.unconditional_context["input_ids"], dtype=torch.long + ) + input_ids = self.unconditional_context["input_ids"] + attention_mask = self.unconditional_context["attention_mask"] + self.unconditional_context["first_pass"] = False + else: + attention_mask = torch.cat( + [ + self.unconditional_context["attention_mask"], + torch.ones_like(input_ids[:, -1:], dtype=torch.long), + ], + dim=1, + ) + if not self.unconditional_context["use_cache"]: + input_ids = torch.cat([self.unconditional_context["input_ids"], input_ids[:, -1:]], dim=1) + else: + input_ids = input_ids[:, -1:] + self.unconditional_context["input_ids"] = input_ids + self.unconditional_context["attention_mask"] = attention_mask + + out = self.model( + input_ids, + attention_mask=attention_mask, + use_cache=self.unconditional_context["use_cache"], + past_key_values=self.unconditional_context["past_key_values"], + ) + self.unconditional_context["past_key_values"] = out.get("past_key_values", None) + + return out.logits + + def __call__(self, input_ids, scores): + scores = torch.nn.functional.log_softmax(scores, dim=-1) + if self.guidance_scale == 1: + return scores + + logits = self.get_unconditional_logits(input_ids) + + unconditional_logits = torch.nn.functional.log_softmax(logits[:, -1], dim=-1) + scores_processed = self.guidance_scale * (scores - unconditional_logits) + unconditional_logits + return scores_processed + + +class BarkEosPrioritizerLogitsProcessor(LogitsProcessor): + r"""This processor ensures that the EOS token is selected if its probability is greater than the `min_eos_p`. + + + + This logits processor is exclusively compatible with + [Bark](https://huggingface.co/docs/transformers/en/model_doc/bark). See the model documentation for examples. + + + + Args: + eos_token_id (`Union[int, List[int], torch.Tensor]`): + The id(s) of the *end-of-sequence* token. + min_eos_p (`float`, *optional*): + Minimum end of speech threshold. + """ + + def __init__(self, eos_token_id: Union[int, List[int], torch.Tensor], min_eos_p: float, device: str = "cpu"): + if not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id, device=device) + self.eos_token_id = eos_token_id + + if torch.is_floating_point(eos_token_id) or (eos_token_id < 0).any(): + raise ValueError(f"`eos_token_id` has to be a list of positive integers, but is {eos_token_id}") + + if min_eos_p is not None and min_eos_p <= 0: + raise ValueError(f"`min_eos_p` has to be a positive float, but is {min_eos_p}") + self.min_eos_p = min_eos_p + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + scores_processed = scores + if self.min_eos_p: + probs = torch.nn.functional.softmax(scores.float(), dim=-1) + # create scores full of -inf except for the eos_token_id + early_stop_scores = torch.ones_like(scores) * -float("inf") + early_stop_scores[:, self.eos_token_id] = scores[:, self.eos_token_id] + + do_early_stop = probs[:, self.eos_token_id] > self.min_eos_p + do_early_stop = torch.any(do_early_stop, dim=1, keepdim=True) + scores_processed = torch.where(do_early_stop, early_stop_scores, scores) + + return scores_processed + + +class WatermarkLogitsProcessor(LogitsProcessor): + r""" + Logits processor for watermarking generated text. The processor modifies model output scores by adding a small bias to + randomized set of "green" tokens before generating the next token. "Green" tokens selection process depends on the + `seeding_scheme` used. The code was based on the [original repo](https://github.com/jwkirchenbauer/lm-watermarking/tree/main). + + The text generated by this `LogitsProcessor` can be detected using `WatermarkDetector`. See [`~WatermarkDetector.__call__`] for details, + + See [the paper](https://arxiv.org/abs/2306.04634) for more information. + + Args: + vocab_size (`int`): + The model tokenizer's vocab_size. Used to calculate "green" tokens ratio. + device (`str`): + The device where model is allocated. + greenlist_ratio (`float`, optional, *optional*, defaults to 0.25): + The ratio of "green" tokens used to the vocabulary size. Defaults to 0.25. + bias (`float`, optional, *optional*, defaults to 2.0): + The bias added to the selected "green" tokens' logits. Consider lowering the + `bias` if the text generation quality degrades. Recommended values are in the + range of [0.5, 2.0]. Defaults to 2.0. + hashing_key (`int`, optional, *optional*, defaults to 15485863): + Key used for hashing. If you deploy this watermark, we advise using another private key. + Defaults to 15485863 (the millionth prime). + seeding_scheme (`str`, optional, *optional*, defaults to `"lefthash"`): + The seeding scheme used for selecting "green" tokens. Accepts values: + - "lefthash" (default): "green" tokens selection depend on the last token (Algorithm 2 from paper) + - "selfhash": "green" tokens selection depends on the current token itself (Algorithm 3 from paper) + The downside of this scheme is that it considers all possible next tokens and can be slower than "lefthash". + The context length of previous tokens to use in seeding. Higher context length makes watermarking more robust. + context_width (`int`, *optional*, defaults to 1): + The number of previous tokens to use when setting the seed. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, WatermarkingConfig + + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + >>> inputs = tokenizer(["Alice and Bob are"], return_tensors="pt") + + >>> # normal generation + >>> out = model.generate(inputs["input_ids"], max_length=20, do_sample=False) + >>> tokenizer.batch_decode(out, skip_special_tokens=True)[0] + 'Alice and Bob are both in the same room.\n\n"I\'m not sure if you\'re' + + >>> # watermarked generation + >>> watermarking_config = WatermarkingConfig(bias=2.5, context_width=2, seeding_scheme="selfhash") + >>> out = model.generate(inputs["input_ids"], watermarking_config=watermarking_config, max_length=20, do_sample=False) + >>> tokenizer.batch_decode(out, skip_special_tokens=True)[0] + 'Alice and Bob are both still alive and well and the story is pretty much a one-hour adventure' + + >>> # to detect watermarked text use the WatermarkDetector class + >>> from transformers import WatermarkDetector + >>> detector = WatermarkDetector(model_config=model.config, device="cpu", watermarking_config= watermarking_config) + >>> detection_preds = detector(out) + >>> detection_preds + array([ True]) + ``` + """ + + def __init__( + self, + vocab_size, + device, + greenlist_ratio: float = 0.25, + bias: float = 2.0, + hashing_key: int = 15485863, + seeding_scheme: str = "lefthash", + context_width: int = 1, + ): + if seeding_scheme not in ["selfhash", "lefthash"]: + raise ValueError(f"seeding_scheme has to be one of [`selfhash`, `lefthash`], but found {seeding_scheme}") + if greenlist_ratio >= 1.0 or greenlist_ratio <= 0.0: + raise ValueError( + f"greenlist_ratio has be in range between 0.0 and 1.0, exclusively. but found {greenlist_ratio}" + ) + + self.vocab_size = vocab_size + self.greenlist_size = int(self.vocab_size * greenlist_ratio) + self.bias = bias + self.seeding_scheme = seeding_scheme + self.rng = torch.Generator(device=device) + self.hash_key = hashing_key + self.context_width = context_width + + self.rng.manual_seed(hashing_key) + self.table_size = 1_000_003 + self.fixed_table = torch.randperm(self.table_size, generator=self.rng, device=device) + + def set_seed(self, input_seq: torch.LongTensor): + input_seq = input_seq[-self.context_width :] + if self.seeding_scheme == "selfhash": + a = self.fixed_table[input_seq % self.table_size] + 1 + b = self.fixed_table[input_seq[-1] % self.table_size] + 1 + seed = (self.hash_key * a * b).min().item() + else: + seed = self.hash_key * input_seq[-1].item() + self.rng.manual_seed(seed % (2**64 - 1)) + + def _get_greenlist_ids(self, input_seq: torch.LongTensor) -> torch.LongTensor: + self.set_seed(input_seq) + vocab_permutation = torch.randperm(self.vocab_size, device=input_seq.device, generator=self.rng) + greenlist_ids = vocab_permutation[: self.greenlist_size] + return greenlist_ids + + def _score_rejection_sampling(self, input_seq: torch.LongTensor, scores: torch.FloatTensor) -> torch.LongTensor: + """ + Generate greenlist based on current candidate next token. Reject and move on if necessary. + Runs for a fixed number of steps only for efficiency, since the methods is not batched. + """ + final_greenlist = [] + _, greedy_predictions = scores.sort(dim=-1, descending=True) + + # 40 is an arbitrary number chosen to save compute and not run for long (taken from orig repo) + for i in range(40): + greenlist_ids = self._get_greenlist_ids(torch.cat([input_seq, greedy_predictions[i, None]], dim=-1)) + if greedy_predictions[i] in greenlist_ids: + final_greenlist.append(greedy_predictions[i]) + return torch.tensor(final_greenlist, device=input_seq.device) + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + if input_ids.shape[-1] < self.context_width: + logger.warning( + f"`input_ids` should have at least `{self.context_width}` tokens but has {input_ids.shape[-1]}. " + "The seeding will be skipped for this generation step!" + ) + return scores + + scores_processed = scores.clone() + for b_idx, input_seq in enumerate(input_ids): + if self.seeding_scheme == "selfhash": + greenlist_ids = self._score_rejection_sampling(input_seq, scores[b_idx]) + else: + greenlist_ids = self._get_greenlist_ids(input_seq) + scores_processed[b_idx, greenlist_ids] = scores_processed[b_idx, greenlist_ids] + self.bias + + return scores_processed + + +class SynthIDTextWatermarkState: + """SynthID watermarking state.""" + + def __init__( + self, + batch_size: int, + ngram_len: int, + context_history_size: int, + device: torch.device, + ): + """Initializes the state. + + Args: + batch_size (`int`): Batch size. + ngram_len (`int`): Ngram length. + context_history_size (`int`): Size of the tensor to keep track of seen contexts. + device (`int`): Device to use. + """ + self.context = torch.zeros( + (batch_size, ngram_len - 1), + dtype=torch.int64, + device=device, + ) + self.context_history = torch.zeros( + (batch_size, context_history_size), + dtype=torch.int64, + device=device, + ) + self.num_calls = 0 + + +class SynthIDTextWatermarkLogitsProcessor(LogitsProcessor): + r""" + Logits processor that implements watermarking techniques for text generation models. + This class facilitates the application of SynthID text watermarking, a method for embedding imperceptible signals + into generated text to aid in detecting synthetic content. It operates by subtly manipulating the probabilities of + token selection during text generation in a manner that can be reliably recovered later for verification. + + Key Features: + * **State Management:** Maintains internal state to track token sequences and generate watermarking keys + dynamically. + + * **Key Generation:** Computes hashes based on token sequences and watermarking parameters to create unique keys + for each position. + + * **G-Value Sampling:** Employs a pre-computed sampling table to sample watermarking values (g-values) based on + the generated keys. + + * **Score Adjustment:** Applies calculated g-values to modify token probabilities during generation, embedding the + watermark. + + * **Context Repetition Handling:** Incorporates logic to avoid watermarking tokens in repeated contexts, + preserving naturalness. + + * **EOS Token Masking:** Supports masking end-of-sentence tokens to prevent their inclusion in watermarking + calculations. + + * **Utility Functions:** Provides functions to compute g-values directly, check for context repetition, create + EOS token masks, and estimate expected mean g-values. + + Refer to paper url: https://www.nature.com/articles/s41586-024-08025-4 for more details around this. + + Args: + ngram_len (`int`): + Ngram length. + keys (`List[int]`): + A sequence of watermarking keys, one for each depth. + sampling_table_size (`int`): + Size of the sampling table. + sampling_table_seed (`int`): + Random seed to generate the sampling table. + context_history_size (`int`): + Size of the tensor to keep track of seen contexts. + device (`torch.device`): + Device to use. + skip_first_ngram_calls (`bool`, *optional*, defaults to `False`): + Whether to skip first ngram calls. + debug_mode (`bool`, optional, *optional*, defaults to `False`): + Logits are modified to uniform one got before watermarking modification is applied. This is to test the + implementation. + + Examples: + ```python + >>> from transformers import AutoModelForCausalLM, AutoTokenizer, SynthIDTextWatermarkingConfig + + >>> tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b', padding_side="left") + >>> model = AutoModelForCausalLM.from_pretrained('google/gemma-2-2b') + + >>> # SynthID Text configuration + >>> watermarking_config = SynthIDTextWatermarkingConfig( + ... keys=[654, 400, 836, 123, 340, 443, 597, 160, 57], + ... ngram_len=5, + ... ) + + >>> # Generation with watermarking + >>> tokenized_prompts = tokenizer(["Once upon a time, "], return_tensors="pt", padding=True) + >>> output_sequences = model.generate( + ... **tokenized_prompts, watermarking_config=watermarking_config, do_sample=True, max_new_tokens=10 + ... ) + >>> watermarked_text = tokenizer.batch_decode(output_sequences, skip_special_tokens=True) + ``` + """ + + def __init__( + self, + ngram_len: int, + keys: List[int], + sampling_table_size: int, + sampling_table_seed: int, + context_history_size: int, + device: torch.device, + skip_first_ngram_calls: bool = False, + debug_mode: bool = False, + ): + self.ngram_len = ngram_len + self.keys = torch.tensor(keys, device=device) + + generator = torch.Generator(device=device).manual_seed(sampling_table_seed) + # A random sampling table is pre-computed and modulo table size is applied to map from a hash of ngram keys to + # g values, this is similar to the hashtable implementation used in + # https://github.com/facebookresearch/three_bricks. We note that the hashing employed in this repository is + # different from that used to watermark the Gemini App, and hence the detectors trained based on the + # hashing in this repository will not transfer to text generated by the Gemini App. + self.sampling_table = torch.randint( + low=0, + high=2, + size=(sampling_table_size,), + generator=generator, + device=device, + ) + self.context_history_size = context_history_size + self.device = device + self.state = None + self.skip_first_ngram_calls = skip_first_ngram_calls + self.debug_mode = debug_mode + + def _init_state(self, batch_size: int): + """Initializes the state.""" + self.state = SynthIDTextWatermarkState( + batch_size=batch_size, + ngram_len=self.ngram_len, + context_history_size=self.context_history_size, + device=self.device, + ) + + def update_scores(self, scores: torch.FloatTensor, g_values: torch.FloatTensor) -> torch.FloatTensor: + """Updates scores using the g values. + + We assume that the scores are in the log space. + Args: + scores (`torch.FloatTensor`): Scores (batch_size, vocab_size). + g_values (`torch.FloatTensor`): G valus (batch_size, vocab_size, depth). + + Returns: + Updated scores (batch_size, vocab_size). + """ + _, _, depth = g_values.shape + + probs = torch.softmax(scores, dim=1) + + for i in range(depth): + g_values_at_depth = g_values[:, :, i] + g_mass_at_depth = (g_values_at_depth * probs).sum(axis=1, keepdims=True) + probs = probs * (1 + g_values_at_depth - g_mass_at_depth) + + log_probs = torch.log(probs) + log_probs = torch.where(torch.isfinite(log_probs), log_probs, torch.finfo(log_probs.dtype).min) + return log_probs + + @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + self._check_input_ids_shape(input_ids) + batch_size, vocab_size = scores.shape + + if self.debug_mode: + scores = torch.ones_like(scores) + + # Currently indices is just a arange to compute watermarking on the desnse logits. + all_indices = torch.stack([torch.arange(vocab_size, device=self.device) for _ in range(batch_size)]) + + if self.state is None: + # Initialize watermarking state if it does not exist. + self._init_state(batch_size) + else: + # Append last input id (which is the input id added in last call) to the + # previous context so we have the context to be used for current + # watermarking. + self.state.context = torch.concat( + (self.state.context, input_ids[:, -1:]), + dim=1, + ) + self.state.context = self.state.context[:, 1:] + + if self.state is None: + raise ValueError("self.state can't be None! Call `self._init_state` to initialize the state.") + + self.state.num_calls += 1 + + # Don't watermark the first ngram_len - 1 tokens if set. + if self.skip_first_ngram_calls and self.state.num_calls < self.ngram_len: + return scores + + # 2. Generate random keys for each ngram key combination. + ngram_keys, hash_result_with_just_context = self._compute_keys(self.state.context, all_indices) + # ngram_keys shape [batch_size, top_k, depth] + + # 3. Sample g values. + g_values = self.sample_g_values(ngram_keys) + # g_values shape [batch_size, top_k, depth] + + # 4. Modify scores. + updated_scores = self.update_scores(scores, g_values) + # updated scores shape [batch_size, top_k] + + # 5. Check if the current watermarking context was previously used, if yes skip watermarking. + hash_result_with_just_context = hash_result_with_just_context[:, None] + is_repeated_context = (self.state.context_history == hash_result_with_just_context).any( + dim=1, + keepdim=True, + ) + self.state.context_history = torch.concat( + (hash_result_with_just_context, self.state.context_history), + dim=1, + )[:, :-1] + + updated_watermarked_scores = torch.where( + is_repeated_context, + input=scores, + other=updated_scores, + ) + return updated_watermarked_scores + + def accumulate_hash( + self, + current_hash: torch.LongTensor, + data: torch.LongTensor, + multiplier: int = 6364136223846793005, + increment: int = 1, + ) -> torch.LongTensor: + """ + Accumulate hash of data on current hash. + + Method uses adapted linear congruential generator with newlib/musl parameters. + + This function has following property - + f(x, data[T]) = f(f(x, data[:T - 1]), data[T]) + + This function expects current_hash.shape and data.shape[:-1] to + match/broadcastable. + + Args: + current_hash (`torch.LongTensor`): + (shape,) + data (`torch.LongTensor`): + (shape, tensor_len) + multiplier (`int`, optional, *optional*, defaults to 6364136223846793005): + multiplier of linear congruential generator + increment (`int`, optional, *optional*, defaults to 1): + increment of linear congruential generator + + Returns: + updated hash (shape,) + """ + for i in range(data.shape[-1]): + current_hash = torch.add(current_hash, data[..., i]) + current_hash = torch.mul(current_hash, multiplier) + current_hash = torch.add(current_hash, increment) + return current_hash + + def compute_ngram_keys(self, ngrams: torch.LongTensor) -> torch.LongTensor: + """Computes random keys for each ngram and depth. + + Args: + ngrams (`torch.LongTensor`): + Ngrams (batch_size, num_ngrams, ngram_len). + + Returns: + ngram keys (batch_size, num_ngrams, depth). + """ + if len(ngrams.shape) != 3: + raise ValueError(f"Ngrams should be of shape (batch_size, num_ngrams, ngram_len), but is {ngrams.shape}") + if ngrams.shape[2] != self.ngram_len: + raise ValueError( + "Ngrams should be of shape (batch_size, num_ngrams, ngram_len)," + f" where ngram_len is {self.ngram_len}, but is {ngrams.shape}" + ) + batch_size, _, _ = ngrams.shape + + hash_result = torch.ones(batch_size, device=self.device, dtype=torch.long) + # hash_result shape [batch_size,] + # ngrams shape [batch_size, num_ngrams, ngram_len] + hash_result = torch.vmap(self.accumulate_hash, in_dims=(None, 1), out_dims=1)(hash_result, ngrams) + # hash_result shape [batch_size, num_ngrams] + + keys = self.keys[None, None, :, None] + # hash_result shape [batch_size, num_ngrams] + # keys shape [1, 1, depth, 1] + hash_result = torch.vmap(self.accumulate_hash, in_dims=(None, 2), out_dims=2)(hash_result, keys) + # hash_result shape [batch_size, num_ngrams, depth] + + return hash_result + + def _compute_keys( + self, n_minus_1_grams: torch.LongTensor, indices: torch.LongTensor + ) -> Tuple[torch.LongTensor, torch.LongTensor]: + """Computes random keys for each ngram and depth. + + Args: + n_minus_1_grams (`torch.LongTensor`): + Ngrams (batch_size, ngram_len - 1). + indices (`torch.LongTensor`): + indices of the continuations (batch_size, num_indices) + + Returns: + Ngram keys (batch_size, num_indices, depth). + """ + batch_size, _ = n_minus_1_grams.shape + + hash_result = torch.ones(batch_size, device=self.device, dtype=torch.long) + # First hash n_minus_1 gram, for each batch entry we have a single + # n_minus_1 gram context. + # hash_result shape [batch_size] + # n_minus_1_gram shape [batch_size, ngram_len - 1] + hash_result_with_just_context = self.accumulate_hash(hash_result, n_minus_1_grams) + # hash_result shape [batch_size,] + # Indices is of shape [batch_size, num_indices], so we make it + # [batch_size, num_indices, 1] so we can vmap over num_indices dim. + hash_result = torch.vmap(self.accumulate_hash, in_dims=(None, 1), out_dims=1)( + hash_result_with_just_context, indices[:, :, None] + ) + # hash_result shape [batch_size, num_indices] + # Basically we have a hash for each batch entry and each indices + # Now we add watermarking keys to this hash. + # keys are of shape [depth,] + # We add batch, num_indices and data dimension to this making it + # [1, 1, depth, 1]. + # So we can vmap over the depth dimension for compute_hash + keys = self.keys[None, None, :, None] + hash_result = torch.vmap(self.accumulate_hash, in_dims=(None, 2), out_dims=2)(hash_result, keys) + # hash_result shape should be [batch_size, num_indices, depth] + return hash_result, hash_result_with_just_context + + def sample_g_values(self, ngram_keys: torch.LongTensor) -> torch.LongTensor: + """ + Samples g values from Bernoulli distribution. + + It is not possible to pass random keys in a vectorized way in torch. Instead + we pre-compute a random sampling table, and use apply modulo table size to + map from ngram keys (int64) to g values. + + Args: + ngram_keys (`torch.LongTensor`): + Random keys (batch_size, num_ngrams, depth). + + Returns: + G values (batch_size, num_ngrams, depth). + """ + (sampling_table_size,) = self.sampling_table.shape + sampling_table = self.sampling_table.reshape((1, 1, sampling_table_size)) + ngram_keys = ngram_keys % sampling_table_size + return torch.take_along_dim(sampling_table, indices=ngram_keys, dim=2) + + def _check_input_ids_shape(self, input_ids: torch.LongTensor): + """Checks the shape of input ids.""" + if len(input_ids.shape) != 2: + raise ValueError(f"Input ids should be of shape (batch_size, input_len), but is {input_ids.shape}") + + def compute_g_values(self, input_ids: torch.LongTensor) -> torch.LongTensor: + """ + Computes g values for each ngram from the given sequence of tokens. + + Args: + input_ids (`torch.LongTensor`): + Input token ids (batch_size, input_len). + + Returns: + G values (batch_size, input_len - (ngram_len - 1), depth). + """ + self._check_input_ids_shape(input_ids) + ngrams = input_ids.unfold(dimension=1, size=self.ngram_len, step=1) + ngram_keys = self.compute_ngram_keys(ngrams) + return self.sample_g_values(ngram_keys) + + def compute_context_repetition_mask(self, input_ids: torch.LongTensor) -> torch.LongTensor: + """ + Computes repetition mask. + + 0 and 1 stand for repeated and not repeated context n-1 grams respectively. + + Args: + input_ids (`torch.LongTensor`): + Input token ids (batch_size, input_len). + + Returns: + Repetitions mask (batch_size, input_len - (ngram_len - 1)). + """ + self._check_input_ids_shape(input_ids) + batch_size, _ = input_ids.shape + state = SynthIDTextWatermarkState( + batch_size=batch_size, + ngram_len=self.ngram_len, + context_history_size=self.context_history_size, + device=self.device, + ) + contexts = input_ids[:, :-1].unfold( + dimension=1, + size=self.ngram_len - 1, + step=1, + ) + _, num_contexts, _ = contexts.shape + + are_repeated_contexts = [] + for i in range(num_contexts): + context = contexts[:, i, :] + hash_result = torch.ones(batch_size, device=self.device, dtype=torch.long) + context_hash = self.accumulate_hash(hash_result, context)[:, None] + is_repeated_context = (state.context_history == context_hash).any( + dim=1, + keepdim=True, + ) + are_repeated_contexts.append(is_repeated_context) + state.context_history = torch.concat( + (context_hash, state.context_history), + dim=1, + )[:, :-1] + are_repeated_contexts = torch.concat(are_repeated_contexts, dim=1) + + return torch.logical_not(are_repeated_contexts) + + def compute_eos_token_mask(self, input_ids: torch.LongTensor, eos_token_id: int) -> torch.LongTensor: + """ + Computes repetitions mask. + + 1 stands for ngrams that don't contain EOS tokens and vice versa. + + Args: + input_ids (`torch.LongTensor`): + Input token ids (batch_size, input_len). + eos_token_id (`int`): + EOS token ID. + + Returns: + EOS token mask (batch_size, input_len). + """ + self._check_input_ids_shape(input_ids) + noneos_masks = [] + all_eos_equated = input_ids == eos_token_id + for eos_equated in all_eos_equated: + nonzero_idx = torch.nonzero(eos_equated) + noneos_mask = torch.ones_like(eos_equated) + if nonzero_idx.shape[0] != 0: + noneos_mask[nonzero_idx[0][0] :] = 0 + noneos_masks.append(noneos_mask) + return torch.stack(noneos_masks, dim=0) + + def expected_mean_g_value(self, vocab_size: int, coinflip_prob: float = 0.5) -> float: + """ + Compute expected mean g-value after watermarking, assuming uniform LM dist. + + This is the theoretical expected value for single-layer watermarking. + + Args: + vocab_size (`int`): + The size of the vocabulary. + coinflip_prob arg_name (`float`, *optional*, defaults to 0.5): + Probability of 1 in boolean prf. + + Returns: + The expected mean g-value for watermarked text. + """ + return coinflip_prob + coinflip_prob * (1 - coinflip_prob) * (1 - (1 / vocab_size)) diff --git a/lib/python3.12/site-packages/transformers/generation/stopping_criteria.py b/lib/python3.12/site-packages/transformers/generation/stopping_criteria.py new file mode 100644 index 0000000000000000000000000000000000000000..4627aeb9702707602d1f09e2f2121ee47338fe39 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/stopping_criteria.py @@ -0,0 +1,519 @@ +import time +import warnings +from abc import ABC +from collections import OrderedDict +from copy import deepcopy +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import torch +from torch.nn import functional as F + +from ..pytorch_utils import isin_mps_friendly +from ..tokenization_utils_base import PreTrainedTokenizerBase +from ..utils import add_start_docstrings, logging + + +logger = logging.get_logger(__name__) +# We maintain a module-level cache of the embedding vectors for the stop string criterion +# because they are slow to compute +STOP_STRING_EMBEDDING_CACHE = OrderedDict() + + +STOPPING_CRITERIA_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`): + Prediction scores of a language modeling head. These can be scores for each vocabulary token before SoftMax + or scores for each vocabulary token after SoftMax. If this stopping criteria depends on the `scores` input, + make sure you pass `return_dict_in_generate=True, output_scores=True` to `generate`. + kwargs (`Dict[str, Any]`, *optional*): + Additional stopping criteria specific kwargs. + + Return: + `torch.BoolTensor`. (`torch.BoolTensor` of shape `(batch_size, 1)`), where `True` indicates we stop generation + for a particular row, `True` indicates we should continue. + +""" + + +class StoppingCriteria(ABC): + """Abstract base class for all stopping criteria that can be applied during generation. + + If your stopping criteria depends on the `scores` input, make sure you pass `return_dict_in_generate=True, + output_scores=True` to `generate`. + """ + + @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor: + raise NotImplementedError("StoppingCriteria needs to be subclassed") + + +class MaxLengthCriteria(StoppingCriteria): + """ + This class can be used to stop generation whenever the full generated number of tokens exceeds `max_length`. Keep + in mind for decoder-only type of transformers, this will include the initial prompted tokens. + + Args: + max_length (`int`): + The maximum length that the output sequence can have in number of tokens. + max_position_embeddings (`int`, *optional*): + The maximum model length, as defined by the model's `config.max_position_embeddings` attribute. + """ + + def __init__(self, max_length: int, max_position_embeddings: Optional[int] = None): + self.max_length = max_length + self.max_position_embeddings = max_position_embeddings + + @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor: + cur_len = input_ids.shape[-1] + is_done = cur_len >= self.max_length + if self.max_position_embeddings is not None and not is_done and cur_len >= self.max_position_embeddings: + logger.warning_once( + "This is a friendly reminder - the current text generation call will exceed the model's predefined " + f"maximum length ({self.max_position_embeddings}). Depending on the model, you may observe " + "exceptions, performance degradation, or nothing at all." + ) + return torch.full((input_ids.shape[0],), is_done, device=input_ids.device, dtype=torch.bool) + + +class MaxTimeCriteria(StoppingCriteria): + """ + This class can be used to stop generation whenever the full generation exceeds some amount of time. By default, the + time will start being counted when you initialize this function. You can override this by passing an + `initial_time`. + + Args: + max_time (`float`): + The maximum allowed time in seconds for the generation. + initial_time (`float`, *optional*, defaults to `time.time()`): + The start of the generation allowed time. + """ + + def __init__(self, max_time: float, initial_timestamp: Optional[float] = None): + self.max_time = max_time + self.initial_timestamp = time.time() if initial_timestamp is None else initial_timestamp + + @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor: + is_done = time.time() - self.initial_timestamp > self.max_time + return torch.full((input_ids.shape[0],), is_done, device=input_ids.device, dtype=torch.bool) + + +class StopStringCriteria(StoppingCriteria): + """ + This class can be used to stop generation whenever specific string sequences are generated. It preprocesses + the strings together with the tokenizer vocab to find positions where tokens can validly complete the stop strings. + + Generation is stopped as soon as a token is generated that completes any of the stop strings. + We want to catch any instance in which the stop string would be present in the decoded output, which means + we must also catch cases with "overhangs" off one or both ends. To make this more concrete, for the stop string + "stop", any of the following token sequences would trigger the match: + + - ["st", "op"] + - ["stop"] + - ["st", "opera"] + - ["sto", "pper"] + - ["las", "topper"] + - ["s", "to", "pped"] + + Note that a match will only be triggered if the stop string is at the end of the generated sequence. In other + words, these sequences will not trigger a match: + + - ["stop", "at"] + - ["st", "op", "at"] + - ["st", "opera", "tion"] + + The reason these are not a match is that the stop string does not overlap with the final token. If you can remove + one or more tokens from the end of the sequence without destroying the stop string, then this criterion will not + match that stop string. This is by design; because this check is run after each token is generated, we can't miss a + valid stop string if one is generated, but we don't want to halt generation just because the stop string exists + somewhere in the past input_ids. + + How is the match actually performed, though? We do it in quite a confusing way, because we want the entire match + process to be compilable with Torch or XLA, which means we cannot use standard string methods. However, it is possible, + with some work, to do string matching with pure tensor operations. We'll begin by describing the algorithm we use + with standard string operations, and then at the end we'll explain how this is converted to pure tensor operations. + + The key to the algorithm is an observation: Because the stop string must overlap with the end of the token sequence, we can start at + the end of the sequence and work backwards. Specifically, we check that there is an overlap between the start of + the final token and the end of the stop_string, or to put it another way, stop_string[-i:] == token[:i] for + some i > 0. If you look at the positive examples above, you'll see the last token in all of them fulfills this + property: + + - ["st", "op"] (overlap is "op", overlap length == 2) + - ["stop"] (overlap is "stop", overlap length == 4) + - ["st", "opera"] (overlap is "op", overlap length == 2) + - ["sto", "pper"] (overlap is "p", overlap length == 1) + - ["las", "topper"] (overlap is "top", overlap length == 3) + - ["s", "to", "pped"] (overlap is "p", overlap length == 1) + + It's impossible to construct a matching sequence that does not have this property (feel free to verify this + yourself). However, although this overlap between the start of the final token and the end of the stop string is + necessary for a match, it is not sufficient. We also need to check that the rest of the token sequence is + consistent with the stop string. + + How do we do that? Let's use ["s", "to", "pped"] as an example. We know that the final token, "pped", has an + overlap of 1 with the stop string, "stop". We then go back to the previous token, "to". Since we have already + matched 1 character from the stop string, the remainder to check is "sto". We check that the next token "to" + matches the end of the remainder, which it does. We have now matched 3 characters from the stop string, and the + remainder to match is "s". We go back to the previous token again, which is also "s". This is a match, and so + we have matched the entire stop string. + + How does it work when the tokens run off the start of the stop string, though? Let's consider the example of + ["las", "topper"]. The final token, "topper", has an overlap of 3 with the stop string, "stop". Therefore, + the remaining stop string to match is "s". We go back to the previous token, "las". Because the remainder to + match is just "s", with length 1, we consider only the final 1 character from the token, which is "s". This + matches the stop string, and so the entire string is matched. + + How do we compute these matches with tensor operations, though? Simply: we efficiently precompute the necessary + information for all tokens! For every token, we compute: + - Its overlap with the end of the stop string, if any + - The positions inside the stop string where the token matches, including matches that run off the start. + - The total length of the token + + For example, for the token "pped", we would compute an end overlap of 1, no internal matching positions, + and a length of 4. For the token "to", we would compute no end overlap, a single internal matching position + of 1 (counting from the end), and a length of 2. For the token "s", we would compute no end overlap, + a single internal matching position of 3 (again counting from the end) and a length of 1. + + As long as we have this information, we can execute the algorithm above without any string comparison + operations. We simply perform the following steps: + - Check if the final token has an end-overlap with the start string + - Continue backwards, keeping track of how much of the stop string we've matched so far + - At each point, check if the next token has the current position as one of its valid positions + - Continue until either a match fails, or we completely match the whole stop string + + Again, consider ["s", "to", "pped"] as an example. "pped" has an end overlap of 1, so we can begin a match. + We have matched 1 character so far, so we check that the next token "to", has 1 as a valid position (again, + counting from the end). It does, so we add the length of "to" to our position tracker. We have now matched + 3 characters, so we check that the next token "s" has 3 as a valid position. It does, so we add its length + to the position tracker. The position tracker is now 4, which is the length of the stop string. We have matched the + entire stop string. + + In the second case, ["las", "topper"], "topper" has an end overlap of 3, so we can begin a match. We have + matched 3 characters so far, so we check that the next token "las" has 3 as a valid position. It does, because we + allow tokens to match positions that run off the start of the stop string. We add its length to the position + tracker. The position tracker is now 6, which is greater than the length of the stop string! Don't panic, though - + this also counts as a match of the stop string. We have matched the entire stop string. + + + Args: + tokenizer (`PreTrainedTokenizer`): + The model's associated tokenizer (necessary to extract vocab and tokenize the termination sequences) + stop_strings (`Union[str, List[str]]`): + A list of strings that should end generation. If a string is passed, it will be treated like a + list with a single element. + + Examples: + + ```python + >>> from transformers import AutoModelForCausalLM, AutoTokenizer + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2") + >>> model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2") + >>> inputs = tokenizer("The biggest states in the USA by land area:", return_tensors="pt") + + >>> gen_out = model.generate(**inputs) + >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0]) + The biggest states in the USA by land area: + - Alaska + - Texas + - California + + >>> # Passing one or more stop strings will halt generation after those strings are emitted + >>> # Note that generating with stop strings requires you to pass the tokenizer too + >>> gen_out = model.generate(**inputs, stop_strings=["Texas"], tokenizer=tokenizer) + >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0]) + The biggest states in the USA by land area: + - Alaska + - Texas + ``` + """ + + def __init__(self, tokenizer: PreTrainedTokenizerBase, stop_strings: Union[str, List[str]]): + if isinstance(stop_strings, str): + stop_strings = [stop_strings] + self.stop_strings: Tuple[str, ...] = tuple(stop_strings) + vocab = tokenizer.get_vocab() + token_list, token_indices = tuple(vocab.keys()), tuple(vocab.values()) + self.embedding_vec, self.max_valid_positions, self.max_valid_end_lens = self.clean_and_embed_tokens_with_cache( + token_list, token_indices, tokenizer + ) + + self.maximum_token_len = max([len(stop_string) for stop_string in self.stop_strings]) + self.num_stop_strings = len(self.stop_strings) + self.target_lens = torch.tensor([len(stop_string) for stop_string in stop_strings], dtype=torch.int32) + + def clean_and_embed_tokens_with_cache(self, token_list, token_indices, tokenizer): + # We don't use the tokenizer in the cache key, because I don't trust it to have well-behaved equality + if (token_list, token_indices, self.stop_strings) in STOP_STRING_EMBEDDING_CACHE: + embedding_vec, max_valid_positions, max_valid_end_lens = STOP_STRING_EMBEDDING_CACHE[ + (token_list, token_indices, self.stop_strings) + ] + STOP_STRING_EMBEDDING_CACHE.move_to_end((token_list, token_indices, self.stop_strings)) + else: + clean_token_list, clean_token_indices = self.clean_tokenizer_vocab(tokenizer) + embedding_vec, max_valid_positions, max_valid_end_lens = self._stop_string_create_embedding_vec( + clean_token_list, clean_token_indices, self.stop_strings + ) + STOP_STRING_EMBEDDING_CACHE[(token_list, token_indices, self.stop_strings)] = ( + embedding_vec, + max_valid_positions, + max_valid_end_lens, + ) + if len(STOP_STRING_EMBEDDING_CACHE) > 8: + STOP_STRING_EMBEDDING_CACHE.popitem(last=False) # Pop from the start, the least recently used item + return embedding_vec, max_valid_positions, max_valid_end_lens + + @staticmethod + def clean_tokenizer_vocab(tokenizer, static_prefix="abcdef"): + """ + This method turns a tokenizer vocab into a "clean" vocab where each token represents the actual string + it will yield, without any special prefixes like "##" or "Ġ". This is trickier than it looks - the method + tokenizer.convert_tokens_to_string() does not always return the correct string because of issues with prefix + space addition/removal. To work around this, we add a static prefix to the start of the token, then remove + it (and any prefix that may have been introduced with it) after calling convert_tokens_to_string(). + """ + vocab = tokenizer.get_vocab() + clean_token_list = [] + clean_token_indices = [] + sentence_base = tokenizer(static_prefix, add_special_tokens=False)["input_ids"] + tokens_base = [tokenizer._convert_id_to_token(tok) for tok in sentence_base] + for token, token_idx in vocab.items(): + token_string = tokenizer.convert_tokens_to_string(tokens_base + [token]) + token_string = token_string[token_string.index(static_prefix) + len(static_prefix) :] + clean_token_list.append(token_string) + clean_token_indices.append(token_idx) + return tuple(clean_token_list), tuple(clean_token_indices) + + @staticmethod + def _stop_string_get_matching_positions( + token_list, token_indices, stop_strings + ) -> Tuple[Dict[str, Dict[str, List[int]]], Dict[str, Dict[str, List[int]]]]: + """This function preprocesses stop strings and the tokenizer vocabulary to determine where tokens can + validly appear in the stop strings. For each token, it computes a list of positions in the stop string where the + token appears, as well as a list of the possible "end overlaps" for that token - that is, the number of characters + from the end of the stop string that overlap with the start of the token, which can have more than one value. + + The reason for computing these may seem a bit cryptic - please see the docstring for StopStringCriteria for a full + explanation of what these values are for!""" + + token_valid_positions = {} + token_end_overlaps = {} + for stop_string in stop_strings: + reversed_stop_string = stop_string[::-1] + token_valid_positions[stop_string] = {} + token_end_overlaps[stop_string] = {} + for token, tok_idx in zip(token_list, token_indices): + reversed_token = token[::-1] + matching_positions = [] + possible_end_lengths = [] + for i in range(1 - len(token), len(stop_string)): + if i < 0: + tok = reversed_token[-i:] + i = 0 + else: + tok = reversed_token + stop = reversed_stop_string[i : i + len(tok)] + if tok.startswith(stop): + if i == 0: + possible_end_lengths.append(min(len(tok), len(stop))) + else: + matching_positions.append(i) + + if matching_positions: + token_valid_positions[stop_string][tok_idx] = matching_positions + if possible_end_lengths: + token_end_overlaps[stop_string][tok_idx] = possible_end_lengths + return token_valid_positions, token_end_overlaps + + @staticmethod + def _stop_string_create_embedding_vec(token_list, token_indices, stop_strings) -> Dict[str, torch.tensor]: + """This function precomputes everything needed for the run-time checks in StopStringCriteria, and packs + them into an embedding tensor that can be accessed with pure tensor operations. For the specifics of the values + that are precomputed and what they are used for, please refer to the StopStringCriteria docstring!""" + token_valid_positions, token_end_overlaps = StopStringCriteria._stop_string_get_matching_positions( + token_list, token_indices, stop_strings + ) + all_valid_positions = [len(val) for positions in token_valid_positions.values() for val in positions.values()] + # In some cases, tokens may have no valid internal positions (such as single-character stop strings), so + # we need a fallback to handle this case + max_valid_positions = max(all_valid_positions) if all_valid_positions else 1 + # There should always be at least one valid end_len, however, so no fallback needed here + valid_end_lens = [len(val) for positions in token_end_overlaps.values() for val in positions.values()] + if not valid_end_lens: + raise ValueError( + "Stop string preprocessing was unable to identify tokens matching one or more of the " + "supplied stop string(s). This is most often caused by the stop " + "strings containing unusual characters that are not in the tokenizer vocabulary." + ) + max_valid_end_lens = max(valid_end_lens) + vec_size = len(stop_strings) * (max_valid_positions + max_valid_end_lens) + 1 + # We use +2 instead of +1 so we can have a dummy entry at the end. We will clamp all token values + # over the max to this, ensuring they do not contribute to stop string matching. + gather_vec = np.full((max(token_indices) + 2, vec_size), dtype=np.int32, fill_value=-1) + + for i, stop_string in enumerate(stop_strings): + positions = token_valid_positions[stop_string] + end_lens = token_end_overlaps[stop_string] + + # Since this is lots of very small assignments of lists, we build it with numpy rather + # than torch for speed + simplicity, then convert to torch at the end + for token_idx, valid_positions in positions.items(): + gather_vec[token_idx, max_valid_positions * i : max_valid_positions * i + len(valid_positions)] = ( + valid_positions + ) + for token_idx, possible_end_lens in end_lens.items(): + gather_vec[ + token_idx, + max_valid_positions * len(stop_strings) + max_valid_end_lens * i : max_valid_positions + * len(stop_strings) + + max_valid_end_lens * i + + len(possible_end_lens), + ] = possible_end_lens + for token, token_idx in zip(token_list, token_indices): + gather_vec[token_idx, -1] = len(token) + + gather_vec = torch.tensor(gather_vec, dtype=torch.int32) + + return gather_vec, max_valid_positions, max_valid_end_lens + + @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.Tensor: + self.embedding_vec = self.embedding_vec.to(input_ids.device) + self.target_lens = self.target_lens.to(input_ids.device) + # The maximum length we need to consider is 1 token per character. Note that input_ids can also be + # *shorter* than the global max, and the code below should be ready for that + input_ids = input_ids[:, -self.maximum_token_len :] + + # Flip input_ids because we're only matching strings at the end of the generated sequence + flipped_ids = torch.flip(input_ids, (1,)) + + # Clip out-of-vocab values to the dummy value at the end of the embedding vector + flipped_ids = torch.clamp(flipped_ids, max=self.embedding_vec.size(0) - 1) + + # Size of the vector of positions a single token can match + max_valid_positions = self.max_valid_positions + + # The embedding vec contains the valid positions, end_lengths and total lengths for each token + embedded = F.embedding(flipped_ids, self.embedding_vec) + + # Now we split the embedding vector. valid_positions is the positions in the stop string the token can fit + valid_positions = embedded[:, 1:, : max_valid_positions * self.num_stop_strings].unflatten( + -1, (self.num_stop_strings, -1) + ) + # end_lengths is the number of characters from the string, counting from the end, that the token + # contains. It can have multiple values if the same token can overlap different end lengths + end_lengths = embedded[:, :1, max_valid_positions * self.num_stop_strings : -1].unflatten( + -1, (self.num_stop_strings, -1) + ) + # Lengths is the total length of each token. Unlike the others, it always has a single value + lengths = embedded[:, 1:, None, -1:] # Insert a dummy dimension for stop_strings even though lengths are const + + # Concatenate lengths onto each possible end_lengths value + lengths = lengths.expand((-1, -1, end_lengths.shape[-2], end_lengths.shape[-1])) + lengths_with_ends = torch.cat([end_lengths, lengths], dim=1) + + # cumsum() to get the number of matched characters in the stop string after each token + cumsum = lengths_with_ends.cumsum(dim=1) # B x maximum_token_len x num_stop_strings x max_valid_end_lens + + # The calculation above assumes that all tokens are in valid positions. Now we mask the ones that are not. + # First, tokens match the start of the string if they have a positive value in the end_lengths vector + initial_match = end_lengths > 0 + + # Tokens continue the string if the cumsum() so far is one of the valid positions for that token + # Note that we're actually tracking one cumsum() for for each possible end_length + later_match = torch.any(cumsum[:, :-1, :, None] == valid_positions[:, :, :, :, None], axis=-2) + + # The match vector is a boolean vector that indicates which positions have valid tokens + match = torch.cat([initial_match, later_match], dim=1) + + # Once a single position does not match, all positions following that position are masked + mask = (~match).cumsum(dim=1, dtype=torch.int32) + mask = mask == 0 + + # The string is matched if we reached a cumsum equal to or greater than the length of the string + # before hitting the mask + string_matches = torch.amax(cumsum * mask, dim=(1, -1)) >= self.target_lens[None, :] + + # We return a per-sample vector that is True if any stop string is matched for that sample + return torch.any(string_matches, dim=-1) + + +class EosTokenCriteria(StoppingCriteria): + """ + This class can be used to stop generation whenever the "end-of-sequence" token is generated. + By default, it uses the `model.generation_config.eos_token_id`. + + Args: + eos_token_id (`Union[int, List[int], torch.Tensor]`): + The id(s) of the *end-of-sequence* token. + """ + + def __init__(self, eos_token_id: Union[int, List[int], torch.Tensor]): + if not isinstance(eos_token_id, torch.Tensor): + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + eos_token_id = torch.tensor(eos_token_id) + self.eos_token_id = eos_token_id + + @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor: + self.eos_token_id = self.eos_token_id.to(input_ids.device) + is_done = isin_mps_friendly(input_ids[:, -1], self.eos_token_id) + return is_done + + +class ConfidenceCriteria(StoppingCriteria): + """ + This class can be used to stop generation whenever assistant model's confidence in its prediction for the current token is lower than the threshold + `model.generation_config.assistant_confidence_threshold` even if the number of speculative tokens (defined by `num_assistant_tokens`) is not yet reached. + + Args: + assistant_confidence_threshold (`float`): + The value of the threshold. + """ + + def __init__(self, assistant_confidence_threshold): + self.assistant_confidence_threshold = assistant_confidence_threshold + + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor: + probs = scores[-1].softmax(-1) + p = probs[0, input_ids[0, -1]].item() + if p < self.assistant_confidence_threshold: + return True + return False + + +class StoppingCriteriaList(list): + @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING) + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor: + is_done = torch.full((input_ids.shape[0],), False, device=input_ids.device, dtype=torch.bool) + for criteria in self: + is_done = is_done | criteria(input_ids, scores, **kwargs) + return is_done + + @property + def max_length(self) -> Optional[int]: + for stopping_criterium in self: + if isinstance(stopping_criterium, MaxLengthCriteria): + return stopping_criterium.max_length + return None + + +def validate_stopping_criteria(stopping_criteria: StoppingCriteriaList, max_length: int) -> StoppingCriteriaList: + stopping_max_length = stopping_criteria.max_length + new_stopping_criteria = deepcopy(stopping_criteria) + if stopping_max_length is not None and stopping_max_length != max_length: + warnings.warn("You set different `max_length` for stopping criteria and `max_length` parameter", UserWarning) + elif stopping_max_length is None: + new_stopping_criteria.append(MaxLengthCriteria(max_length=max_length)) + return new_stopping_criteria diff --git a/lib/python3.12/site-packages/transformers/generation/streamers.py b/lib/python3.12/site-packages/transformers/generation/streamers.py new file mode 100644 index 0000000000000000000000000000000000000000..c0fc26442b466fdaa09aa7c52e607a37dd0d78ab --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/streamers.py @@ -0,0 +1,318 @@ +# coding=utf-8 +# Copyright 2023 The HuggingFace Inc. team. +# +# 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 asyncio +from queue import Queue +from typing import TYPE_CHECKING, Optional + + +if TYPE_CHECKING: + from ..models.auto import AutoTokenizer + + +class BaseStreamer: + """ + Base class from which `.generate()` streamers should inherit. + """ + + def put(self, value): + """Function that is called by `.generate()` to push new tokens""" + raise NotImplementedError() + + def end(self): + """Function that is called by `.generate()` to signal the end of generation""" + raise NotImplementedError() + + +class TextStreamer(BaseStreamer): + """ + Simple text streamer that prints the token(s) to stdout as soon as entire words are formed. + + + + The API for the streamer classes is still under development and may change in the future. + + + + Parameters: + tokenizer (`AutoTokenizer`): + The tokenized used to decode the tokens. + skip_prompt (`bool`, *optional*, defaults to `False`): + Whether to skip the prompt to `.generate()` or not. Useful e.g. for chatbots. + decode_kwargs (`dict`, *optional*): + Additional keyword arguments to pass to the tokenizer's `decode` method. + + Examples: + + ```python + >>> from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer + + >>> tok = AutoTokenizer.from_pretrained("openai-community/gpt2") + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> inputs = tok(["An increasing sequence: one,"], return_tensors="pt") + >>> streamer = TextStreamer(tok) + + >>> # Despite returning the usual output, the streamer will also print the generated text to stdout. + >>> _ = model.generate(**inputs, streamer=streamer, max_new_tokens=20) + An increasing sequence: one, two, three, four, five, six, seven, eight, nine, ten, eleven, + ``` + """ + + def __init__(self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, **decode_kwargs): + self.tokenizer = tokenizer + self.skip_prompt = skip_prompt + self.decode_kwargs = decode_kwargs + + # variables used in the streaming process + self.token_cache = [] + self.print_len = 0 + self.next_tokens_are_prompt = True + + def put(self, value): + """ + Receives tokens, decodes them, and prints them to stdout as soon as they form entire words. + """ + if len(value.shape) > 1 and value.shape[0] > 1: + raise ValueError("TextStreamer only supports batch size 1") + elif len(value.shape) > 1: + value = value[0] + + if self.skip_prompt and self.next_tokens_are_prompt: + self.next_tokens_are_prompt = False + return + + # Add the new token to the cache and decodes the entire thing. + self.token_cache.extend(value.tolist()) + text = self.tokenizer.decode(self.token_cache, **self.decode_kwargs) + + # After the symbol for a new line, we flush the cache. + if text.endswith("\n"): + printable_text = text[self.print_len :] + self.token_cache = [] + self.print_len = 0 + # If the last token is a CJK character, we print the characters. + elif len(text) > 0 and self._is_chinese_char(ord(text[-1])): + printable_text = text[self.print_len :] + self.print_len += len(printable_text) + # Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words, + # which may change with the subsequent token -- there are probably smarter ways to do this!) + else: + printable_text = text[self.print_len : text.rfind(" ") + 1] + self.print_len += len(printable_text) + + self.on_finalized_text(printable_text) + + def end(self): + """Flushes any remaining cache and prints a newline to stdout.""" + # Flush the cache, if it exists + if len(self.token_cache) > 0: + text = self.tokenizer.decode(self.token_cache, **self.decode_kwargs) + printable_text = text[self.print_len :] + self.token_cache = [] + self.print_len = 0 + else: + printable_text = "" + + self.next_tokens_are_prompt = True + self.on_finalized_text(printable_text, stream_end=True) + + def on_finalized_text(self, text: str, stream_end: bool = False): + """Prints the new text to stdout. If the stream is ending, also prints a newline.""" + print(text, flush=True, end="" if not stream_end else None) + + def _is_chinese_char(self, cp): + """Checks whether CP is the codepoint of a CJK character.""" + # This defines a "chinese character" as anything in the CJK Unicode block: + # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block) + # + # Note that the CJK Unicode block is NOT all Japanese and Korean characters, + # despite its name. The modern Korean Hangul alphabet is a different block, + # as is Japanese Hiragana and Katakana. Those alphabets are used to write + # space-separated words, so they are not treated specially and handled + # like the all of the other languages. + if ( + (cp >= 0x4E00 and cp <= 0x9FFF) + or (cp >= 0x3400 and cp <= 0x4DBF) # + or (cp >= 0x20000 and cp <= 0x2A6DF) # + or (cp >= 0x2A700 and cp <= 0x2B73F) # + or (cp >= 0x2B740 and cp <= 0x2B81F) # + or (cp >= 0x2B820 and cp <= 0x2CEAF) # + or (cp >= 0xF900 and cp <= 0xFAFF) + or (cp >= 0x2F800 and cp <= 0x2FA1F) # + ): # + return True + + return False + + +class TextIteratorStreamer(TextStreamer): + """ + Streamer that stores print-ready text in a queue, to be used by a downstream application as an iterator. This is + useful for applications that benefit from acessing the generated text in a non-blocking way (e.g. in an interactive + Gradio demo). + + + + The API for the streamer classes is still under development and may change in the future. + + + + Parameters: + tokenizer (`AutoTokenizer`): + The tokenized used to decode the tokens. + skip_prompt (`bool`, *optional*, defaults to `False`): + Whether to skip the prompt to `.generate()` or not. Useful e.g. for chatbots. + timeout (`float`, *optional*): + The timeout for the text queue. If `None`, the queue will block indefinitely. Useful to handle exceptions + in `.generate()`, when it is called in a separate thread. + decode_kwargs (`dict`, *optional*): + Additional keyword arguments to pass to the tokenizer's `decode` method. + + Examples: + + ```python + >>> from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer + >>> from threading import Thread + + >>> tok = AutoTokenizer.from_pretrained("openai-community/gpt2") + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> inputs = tok(["An increasing sequence: one,"], return_tensors="pt") + >>> streamer = TextIteratorStreamer(tok) + + >>> # Run the generation in a separate thread, so that we can fetch the generated text in a non-blocking way. + >>> generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=20) + >>> thread = Thread(target=model.generate, kwargs=generation_kwargs) + >>> thread.start() + >>> generated_text = "" + >>> for new_text in streamer: + ... generated_text += new_text + >>> generated_text + 'An increasing sequence: one, two, three, four, five, six, seven, eight, nine, ten, eleven,' + ``` + """ + + def __init__( + self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, timeout: Optional[float] = None, **decode_kwargs + ): + super().__init__(tokenizer, skip_prompt, **decode_kwargs) + self.text_queue = Queue() + self.stop_signal = None + self.timeout = timeout + + def on_finalized_text(self, text: str, stream_end: bool = False): + """Put the new text in the queue. If the stream is ending, also put a stop signal in the queue.""" + self.text_queue.put(text, timeout=self.timeout) + if stream_end: + self.text_queue.put(self.stop_signal, timeout=self.timeout) + + def __iter__(self): + return self + + def __next__(self): + value = self.text_queue.get(timeout=self.timeout) + if value == self.stop_signal: + raise StopIteration() + else: + return value + + +class AsyncTextIteratorStreamer(TextStreamer): + """ + Streamer that stores print-ready text in a queue, to be used by a downstream application as an async iterator. + This is useful for applications that benefit from acessing the generated text asynchronously (e.g. in an + interactive Gradio demo). + + + + The API for the streamer classes is still under development and may change in the future. + + + + Parameters: + tokenizer (`AutoTokenizer`): + The tokenized used to decode the tokens. + skip_prompt (`bool`, *optional*, defaults to `False`): + Whether to skip the prompt to `.generate()` or not. Useful e.g. for chatbots. + timeout (`float`, *optional*): + The timeout for the text queue. If `None`, the queue will block indefinitely. Useful to handle exceptions + in `.generate()`, when it is called in a separate thread. + decode_kwargs (`dict`, *optional*): + Additional keyword arguments to pass to the tokenizer's `decode` method. + + Raises: + TimeoutError: If token generation time exceeds timeout value. + + Examples: + + ```python + >>> from transformers import AutoModelForCausalLM, AutoTokenizer, AsyncTextIteratorStreamer + >>> from threading import Thread + >>> import asyncio + + >>> tok = AutoTokenizer.from_pretrained("openai-community/gpt2") + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> inputs = tok(["An increasing sequence: one,"], return_tensors="pt") + + >>> # Run the generation in a separate thread, so that we can fetch the generated text in a non-blocking way. + >>> async def main(): + ... # Important: AsyncTextIteratorStreamer must be initialized inside a coroutine! + ... streamer = AsyncTextIteratorStreamer(tok) + ... generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=20) + ... thread = Thread(target=model.generate, kwargs=generation_kwargs) + ... thread.start() + ... generated_text = "" + ... async for new_text in streamer: + ... generated_text += new_text + >>> print(generated_text) + >>> asyncio.run(main()) + An increasing sequence: one, two, three, four, five, six, seven, eight, nine, ten, eleven, + ``` + """ + + def __init__( + self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, timeout: Optional[float] = None, **decode_kwargs + ): + super().__init__(tokenizer, skip_prompt, **decode_kwargs) + self.text_queue = asyncio.Queue() + self.stop_signal = None + self.timeout = timeout + self.loop = asyncio.get_running_loop() + self.has_asyncio_timeout = hasattr(asyncio, "timeout") + + def on_finalized_text(self, text: str, stream_end: bool = False): + """Put the new text in the queue. If the stream is ending, also put a stop signal in the queue.""" + self.loop.call_soon_threadsafe(self.text_queue.put_nowait, text) + if stream_end: + self.loop.call_soon_threadsafe(self.text_queue.put_nowait, self.stop_signal) + + def __aiter__(self): + return self + + async def __anext__(self): + try: + if self.has_asyncio_timeout: + async with asyncio.timeout(self.timeout): + value = await self.text_queue.get() + else: + value = await asyncio.wait_for(self.text_queue.get(), timeout=self.timeout) + except asyncio.TimeoutError: + raise TimeoutError() + else: + if value == self.stop_signal: + raise StopAsyncIteration() + else: + return value diff --git a/lib/python3.12/site-packages/transformers/generation/tf_logits_process.py b/lib/python3.12/site-packages/transformers/generation/tf_logits_process.py new file mode 100644 index 0000000000000000000000000000000000000000..f70655fb7c1d0b502dfc8f1c4f2c491867aa7e8c --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/tf_logits_process.py @@ -0,0 +1,603 @@ +# coding=utf-8 +# Copyright 2022 The HuggingFace Inc. team +# +# 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. + +import inspect +from typing import List, Tuple + +import numpy as np +import tensorflow as tf + +from ..tf_utils import stable_softmax +from ..utils import add_start_docstrings +from ..utils.logging import get_logger + + +logger = get_logger(__name__) + + +TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING = r""" + Args: + input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + + Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + scores (`tf.Tensor` of shape `(batch_size, config.vocab_size)`): + Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam + search or log softmax for each vocabulary token when using beam search. + cur_len (`int`): + The current length of valid input sequence tokens. In the TF implementation, the input_ids' sequence length + is the maximum length generate can produce, and we need to know which of its tokens are valid. + kwargs (`Dict[str, Any]`, *optional*): + Additional logits processor specific kwargs. + + Return: + `tf.Tensor` of shape `(batch_size, config.vocab_size)`: The processed prediction scores. +""" + + +class TFLogitsProcessor: + """Abstract base class for all logit processors that can be applied during generation.""" + + @add_start_docstrings(TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + """TF method for processing logits.""" + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + +class TFLogitsWarper: + """Abstract base class for all logit warpers that can be applied during generation with multinomial sampling.""" + + @add_start_docstrings(TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + """TF method for warping logits.""" + raise NotImplementedError( + f"{self.__class__} is an abstract class. Only classes inheriting this class can be called." + ) + + +class TFLogitsProcessorList(list): + """ + This class can be used to create a list of [`TFLogitsProcessor`] to subsequently process a `scores` input tensor. + This class inherits from list and adds a specific *__call__* method to apply each [`TFLogitsProcessor`] to the + inputs. + """ + + @add_start_docstrings(TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING) + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int, **kwargs) -> tf.Tensor: + for processor in self: + function_args = inspect.signature(processor.__call__).parameters + if len(function_args) > 3: + if not all(arg in kwargs for arg in list(function_args.keys())[2:]): + raise ValueError( + f"Make sure that all the required parameters: {list(function_args.keys())} for " + f"{processor.__class__} are passed to the logits processor." + ) + scores = processor(input_ids, scores, cur_len, **kwargs) + else: + scores = processor(input_ids, scores, cur_len) + return scores + + +class TFTemperatureLogitsWarper(TFLogitsWarper): + r""" + [`TFLogitsWarper`] for temperature (exponential scaling output probability distribution). + + Args: + temperature (`float`): + The value used to module the logits distribution. + """ + + def __init__(self, temperature: float): + if not isinstance(temperature, float) or not (temperature > 0): + raise ValueError(f"`temperature` has to be a strictly positive float, but is {temperature}") + + self.temperature = temperature + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + scores = scores / self.temperature + return scores + + +class TFTopKLogitsWarper(TFLogitsWarper): + r""" + [`TFLogitsWarper`] that performs top-k, i.e. restricting to the k highest probability elements. + + Args: + top_k (`int`): + The number of highest probability vocabulary tokens to keep for top-k-filtering. + filter_value (`float`, *optional*, defaults to -inf): + All filtered values will be set to this float value. + min_tokens_to_keep (`int`, *optional*, defaults to 1): + Minimum number of tokens that cannot be filtered. + """ + + def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1): + if not isinstance(top_k, int) or top_k <= 0: + raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}") + + self.top_k = max(top_k, min_tokens_to_keep) + self.filter_value = filter_value + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + top_k = min(self.top_k, scores.shape[-1]) # Safety check + # Boolean mask containing all tokens with a probability less than the last token of the top-k + indices_to_remove = scores < tf.math.top_k(scores, k=top_k)[0][..., -1:] + next_scores = tf.where(indices_to_remove, self.filter_value, scores) + return next_scores + + +class TFTopPLogitsWarper(TFLogitsWarper): + """ + [`TFLogitsWarper`] that performs top-p, i.e. restricting to top tokens summing to <= prob_cut_off. + + Args: + top_p (`float`): + If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or + higher are kept for generation. + filter_value (`float`, *optional*, defaults to -inf): + All filtered values will be set to this float value. + min_tokens_to_keep (`int`, *optional*, defaults to 1): + Minimum number of tokens that cannot be filtered. + """ + + def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1): + if not isinstance(top_p, float) or (top_p < 0 or top_p > 1.0): + raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}") + if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1): + raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}") + + self.top_p = top_p + self.filter_value = filter_value + self.min_tokens_to_keep = min_tokens_to_keep + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + topk_scores, topk_indices = tf.math.top_k(scores, scores.shape[-1]) + + mask_scores = tf.fill(scores.shape, self.filter_value) + cumulative_probs = tf.math.cumsum(stable_softmax(topk_scores, axis=-1), axis=-1) + score_mask = cumulative_probs < self.top_p + + # Also include the token that is higher than top_p (the first false = shift and insert a True on the left) + score_mask = tf.concat((tf.ones([score_mask.shape[0], 1], dtype=tf.bool), score_mask[:, :-1]), axis=-1) + + # Ensure min tokens to keep + score_mask = tf.concat( + ( + tf.ones([score_mask.shape[0], self.min_tokens_to_keep], dtype=tf.bool), + score_mask[:, self.min_tokens_to_keep :], + ), + axis=-1, + ) + + # Mask the values that do not fit the criteria + topk_next_scores = tf.where(score_mask, topk_scores, mask_scores) + + # Undo the topk sorting: converts the 2D matrix of per-row original indices of shape (batch_size, vocab_size) + # to a 3D tensor of shape (batch_size, vocab_size, 2) containing the original score coordinate, from which we + # can scatter (i.e. `scatter_indices[row, col, :]` is a tensor containing `[row, topk_indices[row, col]]`) + scatter_rows = tf.tile(tf.expand_dims(tf.range(topk_indices.shape[0]), axis=-1), [1, topk_indices.shape[-1]]) + scatter_indices = tf.stack((scatter_rows, topk_indices), axis=-1) + next_scores = tf.scatter_nd(scatter_indices, topk_next_scores, shape=topk_next_scores.shape) + + return next_scores + + +class TFMinLengthLogitsProcessor(TFLogitsProcessor): + r""" + [`TFLogitsProcessor`] enforcing a min-length by setting EOS probability to 0. + + Args: + min_length (`int`): + The minimum length below which the score of `eos_token_id` is set to `-float("Inf")`. + eos_token_id (`int`): + The id of the *end-of-sequence* token. + """ + + def __init__(self, min_length: int, eos_token_id: int): + if not isinstance(min_length, int) or min_length < 0: + raise ValueError(f"`min_length` has to be a positive integer, but is {min_length}") + + if not isinstance(eos_token_id, int) or eos_token_id < 0: + raise ValueError(f"`eos_token_id` has to be a positive integer, but is {eos_token_id}") + + self.min_length = min_length + self.eos_token_id = eos_token_id + + def _apply_eos_token_mask(self, scores: tf.Tensor) -> tf.Tensor: + eos_token_id_mask = tf.range(scores.shape[-1]) == self.eos_token_id + scores = tf.where(eos_token_id_mask, float("-inf"), scores) + return scores + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + # applies eos token masking if the first argument is true + scores = tf.cond( + tf.less(cur_len, self.min_length), + lambda: self._apply_eos_token_mask(scores), + lambda: tf.identity(scores), + ) + return scores + + +class TFRepetitionPenaltyLogitsProcessor(TFLogitsProcessor): + r""" + [`TFLogitsProcessor`] enforcing an exponential penalty on repeated sequences. + + Args: + repetition_penalty (`float`): + The parameter for repetition penalty. 1.0 means no penalty. See [this + paper](https://arxiv.org/pdf/1909.05858.pdf) for more details. + """ + + def __init__(self, penalty: float): + if not isinstance(penalty, float) or not (penalty > 0): + raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}") + + self.penalty = penalty + + def _create_score_penalties(self, input_ids: tf.Tensor, logits: tf.Tensor) -> tf.Tensor: + # We want to populate the penalties in the positions of `input_ids`. Since XLA can't handle shapes unknown + # before runtime, `tf.unique` can't be used. Therefore, we may have redundant updates, when a given row has + # the same token multiple times. + + # Gathers the penalties to apply + logit_penalties = tf.gather(logits, input_ids, axis=1, batch_dims=1) + logit_penalties = tf.where(logit_penalties > 0, 1 / self.penalty, logit_penalties) + logit_penalties = tf.where(logit_penalties < 0, self.penalty, logit_penalties) + + # Scatters the penalties + token_penalties = tf.ones(logits.shape) + batch_size = input_ids.shape[0] + seq_len = tf.shape(input_ids)[1] # the sequence length has dynamic size, hence the dynamic shape + indexable_prev_input_ids = tf.concat( + ( + tf.expand_dims(tf.repeat(tf.range(batch_size), seq_len), axis=-1), + tf.expand_dims(tf.reshape(input_ids, [-1]), axis=-1), + ), + axis=1, + ) + token_penalties = tf.tensor_scatter_nd_update( + token_penalties, indices=indexable_prev_input_ids, updates=tf.reshape(logit_penalties, [-1]) + ) + return token_penalties + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + score_penalties = self._create_score_penalties(input_ids[:, :cur_len], scores) + + scores = tf.math.multiply(scores, score_penalties) + + return scores + + +class TFNoBadWordsLogitsProcessor(TFLogitsProcessor): + """ + [`TFLogitsProcessor`] that enforces that specified sequences will never be sampled. + + Args: + bad_words_ids (`List[List[int]]`): + List of list of token ids that are not allowed to be generated. In order to get the tokens of the words + that should not appear in the generated text, make sure to set `add_prefix_space=True` when initializing + the tokenizer, and use `tokenizer(bad_words, add_special_tokens=False).input_ids`. The `add_prefix_space` + argument is only supported for some slow tokenizers, as fast tokenizers' prefixing behaviours come from + `pre tokenizers`. Read more [here](https://huggingface.co/docs/tokenizers/api/pre-tokenizers). + eos_token_id (`int`): + The id of the *end-of-sequence* token. + """ + + def __init__(self, bad_words_ids: List[List[int]], eos_token_id: int): + if not isinstance(bad_words_ids, List) or len(bad_words_ids) == 0: + raise ValueError(f"`bad_words_ids` has to be a non-empty list, but is {bad_words_ids}.") + if any(not isinstance(bad_word_ids, list) for bad_word_ids in bad_words_ids): + raise ValueError(f"`bad_words_ids` has to be a list of lists, but is {bad_words_ids}.") + if any( + any((not isinstance(token_id, (int, np.integer)) or token_id < 0) for token_id in bad_word_ids) + for bad_word_ids in bad_words_ids + ): + raise ValueError( + f"Each list in `bad_words_ids` has to be a list of positive integers, but is {bad_words_ids}." + ) + + # stores the information about bad words in three tensors: + # 1. a rectangular tensor with the forbidden sequences (padded with `-1`), for full data comparisons + self.bad_word_seqs_ids = tf.ragged.constant(bad_words_ids).to_tensor(default_value=-1) + # 2. a tensor with the unpadded length of each forbidden sequence, for quick length comparisons + bad_word_seqs_len = [len(bad_words) for bad_words in bad_words_ids] + if any(word_len == 0 for word_len in bad_word_seqs_len): + raise ValueError(f"Banned words token sequences {bad_words_ids} cannot have an empty list") + self.bad_word_seqs_len = tf.convert_to_tensor(bad_word_seqs_len, dtype=tf.int32) + # 3. a tensor containing the last token for each sequence, for easy access to the tokens that may be banned + self.seq_forbidden_tokens = tf.convert_to_tensor([bad_words[-1] for bad_words in bad_words_ids]) + + def _calc_row_banned_bad_tokens(self, row_input_ids: tf.Tensor) -> tf.Tensor: + def _tokens_match(bad_word_seq_number): + def _len_one(): + # If the bad sequence only has one token, always mask it + return tf.cond( + tf.math.equal(self.bad_word_seqs_len[bad_word_seq_number], 1), + lambda: tf.ones((), dtype=tf.bool), + _len_greater_than_cur_len, + ) + + def _len_greater_than_cur_len(): + # Otherwise, if the bad sequence is longer than the current length they can't ever match + return tf.cond( + tf.math.greater(self.bad_word_seqs_len[bad_word_seq_number], tf.shape(row_input_ids)[0]), + lambda: tf.zeros((), dtype=tf.bool), + _match_found, + ) + + def _match_found(): + # Finaly, runs the actual comparison. Can only be called if the previous comparisons do not yield + # an answer (otherwise we get indexing exceptions) + compare_len = self.bad_word_seqs_len[bad_word_seq_number] - 1 + return tf.cond( + tf.math.reduce_all( + tf.math.equal( + row_input_ids[-compare_len:], self.bad_word_seqs_ids[bad_word_seq_number, :compare_len] + ) + ), + lambda: tf.ones((), dtype=tf.bool), + lambda: tf.zeros((), dtype=tf.bool), + ) + + match = _len_one() + return match + + # Compares the current row against all bad word sequences, obtaining a mask with the matches. + match_mask = tf.map_fn(_tokens_match, tf.range(self.bad_word_seqs_ids.shape[0]), fn_output_signature=tf.bool) + row_banned_tokens = self.seq_forbidden_tokens[match_mask] + return row_banned_tokens + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + # We want to mask some banned tokens, at a score level. Since the banned tokens depend on the previous + # `input_ids`, they may have a different length for each row, and they may even be empty for some rows. + # To remain simple and XLA-compatible, we work on a per-row fashion. + # TODO (Joao): this function might trigger XLA retracing as `cur_len` increases. Fix it if it becomes + # a frequent choke point. (make `cur_len` a tensor?) + def _get_row_updated_score(row_inputs: Tuple[tf.Tensor]) -> tf.Tensor: + row_input_ids, row_score = row_inputs + banned_tokens = self._calc_row_banned_bad_tokens(row_input_ids[:cur_len]) + banned_tokens_mask = tf.scatter_nd( + indices=tf.expand_dims(banned_tokens, axis=-1), + updates=tf.ones_like(banned_tokens, dtype=tf.bool), + shape=row_score.shape, + ) + row_score = tf.where(banned_tokens_mask, -float("inf"), row_score) + return row_score + + scores = tf.map_fn(_get_row_updated_score, (input_ids, scores), fn_output_signature=tf.float32) + return scores + + +class TFNoRepeatNGramLogitsProcessor(TFLogitsProcessor): + r""" + [`TFLogitsProcessor`] that enforces no repetition of n-grams. See + [Fairseq](https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345). + + Args: + ngram_size (`int`): + All ngrams of size `ngram_size` can only occur once. + """ + + def __init__(self, ngram_size: int): + if not isinstance(ngram_size, int) or ngram_size <= 0: + raise ValueError(f"`ngram_size` has to be a strictly positive integer, but is {ngram_size}") + self.ngram_size = ngram_size + + def calc_banned_ngram_tokens(self, input_ids, num_hypos, cur_len): + # Copied from fairseq for no_repeat_ngram in beam_search + if cur_len + 1 < self.ngram_size: + # return no banned tokens if we haven't generated ngram_size tokens yet + return [[] for _ in range(num_hypos)] + generated_ngrams = [{} for _ in range(num_hypos)] + prev_input_ids = input_ids[:, :cur_len] + for idx in range(num_hypos): + gen_tokens = prev_input_ids[idx].numpy().tolist() + generated_ngram = generated_ngrams[idx] + for ngram in zip(*[gen_tokens[i:] for i in range(self.ngram_size)]): + prev_ngram_tuple = tuple(ngram[:-1]) + generated_ngram[prev_ngram_tuple] = generated_ngram.get(prev_ngram_tuple, []) + [ngram[-1]] + + def _get_generated_ngrams(hypo_idx): + # Before decoding the next token, prevent decoding of ngrams that have already appeared + start_idx = cur_len + 1 - self.ngram_size + ngram_idx = tuple(prev_input_ids[hypo_idx, start_idx:cur_len].numpy().tolist()) + return generated_ngrams[hypo_idx].get(ngram_idx, []) + + banned_tokens = [_get_generated_ngrams(hypo_idx) for hypo_idx in range(num_hypos)] + + return banned_tokens + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + # TODO (joao): enable XLA on this logits processor. See discussion and attempts in + # https://github.com/huggingface/transformers/pull/16974 + if not tf.executing_eagerly(): + raise NotImplementedError("TFNoRepeatNGramLogitsProcessor is only implemented for eager execution.") + + batch_size, vocab_size = scores.shape + banned_tokens = self.calc_banned_ngram_tokens(input_ids, batch_size, cur_len) + + # create banned_tokens boolean mask + banned_tokens_indices_mask = [] + for banned_tokens_slice in banned_tokens: + banned_tokens_indices_mask.append( + [True if token in banned_tokens_slice else False for token in range(vocab_size)] + ) + + scores = tf.where(tf.convert_to_tensor(banned_tokens_indices_mask, dtype=tf.bool), -float("inf"), scores) + + return scores + + +class TFForcedBOSTokenLogitsProcessor(TFLogitsProcessor): + r""" + [`TFLogitsProcessor`] that enforces the specified token as the first generated token. + + Args: + bos_token_id (`int`): + The id of the token to force as the first generated token. + """ + + def __init__(self, bos_token_id: int): + if bos_token_id < 0: + raise ValueError(f"The forced bos token id must be a non-negative integer, got {bos_token_id}") + self.bos_token_id = bos_token_id + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + if cur_len == 1: + batch_size, num_tokens = scores.shape + # sets the score to 0 in the bos_token_id column + scores = tf.zeros((batch_size, 1)) + # sets the score to -inf everywhere else + if self.bos_token_id > 0: + scores = tf.concat((tf.broadcast_to(-float("inf"), (batch_size, self.bos_token_id)), scores), axis=-1) + if self.bos_token_id < (num_tokens - 1): + scores = tf.concat( + (scores, tf.broadcast_to(-float("inf"), (batch_size, (num_tokens - 1) - self.bos_token_id))), + axis=-1, + ) + return scores + + +class TFForcedEOSTokenLogitsProcessor(TFLogitsProcessor): + r""" + [`TFLogitsProcessor`] that enforces the specified token as the last generated token when `max_length` is reached. + + Args: + max_length (`int`): + The maximum length of the sequence to be generated. + eos_token_id (`int`): + The id of the token to force as the last generated token when `max_length` is reached. + """ + + def __init__(self, max_length: int, eos_token_id: int): + self.max_length = max_length + if eos_token_id < 0: + raise ValueError(f"The forced eos token id must be a non-negative integer, got {eos_token_id}") + self.eos_token_id = eos_token_id + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + if cur_len == self.max_length - 1: + batch_size, num_tokens = scores.shape + # sets the score to 0 in the eos_token_id column + scores = tf.zeros((batch_size, 1)) + # sets the score to -inf everywhere else + if self.eos_token_id > 0: + scores = tf.concat((tf.broadcast_to(-float("inf"), (batch_size, self.eos_token_id)), scores), axis=-1) + if self.eos_token_id < (num_tokens - 1): + scores = tf.concat( + (scores, tf.broadcast_to(-float("inf"), (batch_size, (num_tokens - 1) - self.eos_token_id))), + axis=-1, + ) + return scores + + +class TFSuppressTokensAtBeginLogitsProcessor(TFLogitsProcessor): + r""" + [`TFSuppressTokensAtBeginLogitsProcessor`] suppresses a list of tokens as soon as the `generate` function starts + generating using `begin_index` tokens. This should ensure that the tokens defined by `begin_suppress_tokens` at not + sampled at the beginning of the generation. + """ + + def __init__(self, begin_suppress_tokens, begin_index): + self.begin_suppress_tokens = list(begin_suppress_tokens) + self.begin_index = begin_index + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + suppressed_indices = [] + for token in self.begin_suppress_tokens: + if token < scores.shape[-1]: # to ensure we don't go beyond the vocab size + suppressed_indices.extend([[i, token] for i in range(scores.shape[0])]) + + if len(suppressed_indices) > 0: + scores = tf.cond( + tf.equal(cur_len, self.begin_index), + lambda: tf.tensor_scatter_nd_update( + scores, + indices=suppressed_indices, + updates=[-float("inf") for _ in range(scores.shape[0] * len(self.begin_suppress_tokens))], + ), + lambda: scores, + ) + return scores + + +class TFSuppressTokensLogitsProcessor(TFLogitsProcessor): + r"""This processor can be used to suppress a list of tokens. The processor will set their log probs to `-inf` so that they + are not sampled.""" + + def __init__(self, suppress_tokens): + self.suppress_tokens = list(suppress_tokens) + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + suppressed_indices = [] + for token in self.suppress_tokens: + if token < scores.shape[-1]: # to ensure we don't go beyond the vocab size + suppressed_indices.extend([[i, token] for i in range(scores.shape[0])]) + + if len(suppressed_indices) > 0: + scores = tf.tensor_scatter_nd_update( + scores, + indices=[[i, token] for i in range(scores.shape[0]) for token in self.suppress_tokens], + updates=[-float("inf") for _ in range(scores.shape[0] * len(self.suppress_tokens))], + ) + return scores + + +class TFForceTokensLogitsProcessor(TFLogitsProcessor): + r"""This processor takes a list of pairs of integers which indicates a mapping from generation indices to token + indices that will be forced before sampling. The processor will set their log probs to `0` and all other tokens to + `-inf` so that they are sampled at their corresponding index.""" + + def __init__(self, force_token_map: List[List[int]]): + force_token_map = dict(force_token_map) + # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the + # index of the array corresponds to the index of the token to be forced, for XLA compatibility. + # Indexes without forced tokens will have an negative value. + force_token_array = np.ones((max(force_token_map.keys()) + 1), dtype=np.int32) * -1 + for index, token in force_token_map.items(): + if token is not None: + force_token_array[index] = token + self.force_token_array = tf.convert_to_tensor(force_token_array, dtype=tf.int32) + + def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor: + def _force_token(generation_idx): + batch_size = scores.shape[0] + current_token = self.force_token_array[generation_idx] + + new_scores = tf.zeros_like(scores, dtype=scores.dtype) + tf.constant([scores.dtype.min]) + indices = tf.stack((tf.range(batch_size), tf.tile([current_token], [batch_size])), axis=1) + updates = tf.zeros((batch_size,), dtype=scores.dtype) + new_scores = tf.tensor_scatter_nd_update(new_scores, indices, updates) + return new_scores + + scores = tf.cond( + tf.greater_equal(cur_len, tf.shape(self.force_token_array)[0]), + # If the current length is geq than the length of force_token_array, the processor does nothing. + lambda: tf.identity(scores), + # Otherwise, it may force a certain token. + lambda: tf.cond( + tf.greater_equal(self.force_token_array[cur_len], 0), + # Only valid (positive) tokens are forced + lambda: _force_token(cur_len), + # Otherwise, the processor does nothing. + lambda: scores, + ), + ) + return scores diff --git a/lib/python3.12/site-packages/transformers/generation/tf_utils.py b/lib/python3.12/site-packages/transformers/generation/tf_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..344147e6e34c67b1072449885dc81dcc24621809 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/tf_utils.py @@ -0,0 +1,3132 @@ +# coding=utf-8 +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. 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. + +import copy +import inspect +import warnings +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple, Union + +import numpy as np +import tensorflow as tf +from tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice + +from ..modeling_tf_outputs import TFCausalLMOutputWithPast, TFSeq2SeqLMOutput +from ..models.auto import ( + TF_MODEL_FOR_CAUSAL_LM_MAPPING, + TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, + TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, + TF_MODEL_FOR_VISION_2_SEQ_MAPPING, +) +from ..tf_utils import shape_list, stable_softmax +from ..utils import ModelOutput, logging +from .configuration_utils import GenerationConfig +from .tf_logits_process import ( + TFForcedBOSTokenLogitsProcessor, + TFForcedEOSTokenLogitsProcessor, + TFForceTokensLogitsProcessor, + TFLogitsProcessorList, + TFMinLengthLogitsProcessor, + TFNoBadWordsLogitsProcessor, + TFNoRepeatNGramLogitsProcessor, + TFRepetitionPenaltyLogitsProcessor, + TFSuppressTokensAtBeginLogitsProcessor, + TFSuppressTokensLogitsProcessor, + TFTemperatureLogitsWarper, + TFTopKLogitsWarper, + TFTopPLogitsWarper, +) + + +logger = logging.get_logger(__name__) + + +@dataclass +class TFGreedySearchDecoderOnlyOutput(ModelOutput): + """ + Base class for outputs of decoder-only generation models using greedy search. + + + Args: + sequences (`tf.Tensor` of shape `(batch_size, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each + generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, generated_length, hidden_size)`. + """ + + sequences: Optional[tf.Tensor] = None + scores: Optional[Tuple[tf.Tensor]] = None + attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None + + +@dataclass +class TFGreedySearchEncoderDecoderOutput(ModelOutput): + """ + Base class for outputs of encoder-decoder generation models using greedy search. Hidden states and attention + weights of the decoder (respectively the encoder) can be accessed via the encoder_attentions and the + encoder_hidden_states attributes (respectively the decoder_attentions and the decoder_hidden_states attributes) + + + Args: + sequences (`tf.Tensor` of shape `(batch_size, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each + generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, generated_length, hidden_size)`. + """ + + sequences: Optional[tf.Tensor] = None + scores: Optional[Tuple[tf.Tensor]] = None + encoder_attentions: Optional[Tuple[tf.Tensor]] = None + encoder_hidden_states: Optional[Tuple[tf.Tensor]] = None + decoder_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + cross_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + decoder_hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None + + +@dataclass +class TFSampleDecoderOnlyOutput(ModelOutput): + """ + Base class for outputs of decoder-only generation models using sampling. + + + Args: + sequences (`tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each + generated token), with each tensor of shape `(batch_size*num_return_sequences, config.vocab_size)`. + attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(num_return_sequences*batch_size, num_heads, generated_length, sequence_length)`. + hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(num_return_sequences*batch_size, generated_length, hidden_size)`. + """ + + sequences: Optional[tf.Tensor] = None + scores: Optional[Tuple[tf.Tensor]] = None + attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None + + +@dataclass +class TFSampleEncoderDecoderOutput(ModelOutput): + """ + Base class for outputs of encoder-decoder generation models using sampling. Hidden states and attention weights of + the decoder (respectively the encoder) can be accessed via the encoder_attentions and the encoder_hidden_states + attributes (respectively the decoder_attentions and the decoder_hidden_states attributes) + + + Args: + sequences (`tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each + generated token), with each tensor of shape `(batch_size*num_return_sequences, config.vocab_size)`. + encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size*num_return_sequences, + num_heads, sequence_length, sequence_length)`. + encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size*num_return_sequences, sequence_length, hidden_size)`. + decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size*num_return_sequences, num_heads, generated_length, sequence_length)`. + cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size*num_return_sequences, generated_length, hidden_size)`. + """ + + sequences: Optional[tf.Tensor] = None + scores: Optional[Tuple[tf.Tensor]] = None + encoder_attentions: Optional[Tuple[tf.Tensor]] = None + encoder_hidden_states: Optional[Tuple[tf.Tensor]] = None + decoder_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + cross_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + decoder_hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None + + +@dataclass +class TFBeamSearchDecoderOnlyOutput(ModelOutput): + """ + Base class for outputs of decoder-only generation models using beam search. + + Args: + sequences (`tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + sequences_scores (`tf.Tensor` of shape `(batch_size*num_return_sequences)`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Final beam scores of the generated `sequences`. + scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log + softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this + beam. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each generated token), + with each tensor of shape `(batch_size*num_beams*num_return_sequences, config.vocab_size)`. + beam_indices (`tf.Tensor`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Beam indices of generated token id at each generation step. `tf.Tensor` of shape + `(batch_size*num_return_sequences, sequence_length)`. + attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size*num_beams, num_heads, generated_length, sequence_length)`. + hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size*num_beams*num_return_sequences, generated_length, hidden_size)`. + """ + + sequences: Optional[tf.Tensor] = None + sequences_scores: Optional[tf.Tensor] = None + scores: Optional[Tuple[tf.Tensor]] = None + beam_indices: Optional[tf.Tensor] = None + attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None + + +@dataclass +class TFBeamSearchEncoderDecoderOutput(ModelOutput): + """ + Base class for outputs of encoder-decoder generation models using beam search. Hidden states and attention weights + of the decoder (respectively the encoder) can be accessed via the encoder_attentions and the encoder_hidden_states + attributes (respectively the decoder_attentions and the decoder_hidden_states attributes) + + Args: + sequences (`tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + sequences_scores (`tf.Tensor` of shape `(batch_size*num_return_sequences)`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Final beam scores of the generated `sequences`. + scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log + softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this + beam. `Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each generated token), + with each tensor of shape `(batch_size*num_beams, config.vocab_size)`. + beam_indices (`tf.Tensor`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Beam indices of generated token id at each generation step. `tf.Tensor` of shape + `(batch_size*num_return_sequences, sequence_length)`. + encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size*num_beams*num_return_sequences, sequence_length, hidden_size)`. + decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size*num_beams*num_return_sequences, num_heads, generated_length, + sequence_length)`. + cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size*num_beams*num_return_sequences, generated_length, hidden_size)`. + """ + + sequences: Optional[tf.Tensor] = None + sequences_scores: Optional[tf.Tensor] = None + scores: Optional[Tuple[tf.Tensor]] = None + beam_indices: Optional[tf.Tensor] = None + encoder_attentions: Optional[Tuple[tf.Tensor]] = None + encoder_hidden_states: Optional[Tuple[tf.Tensor]] = None + decoder_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + cross_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + decoder_hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None + + +@dataclass +class TFBeamSampleDecoderOnlyOutput(ModelOutput): + """ + Base class for outputs of decoder-only generation models using beam sample. + + Args: + sequences (`tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + sequences_scores (`tf.Tensor` of shape `(batch_size * num_return_sequence)`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Final beam scores of the generated `sequences`. + scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log + softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this + beam. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each generated token), + with each tensor of shape `(batch_size*num_beams*num_return_sequences, config.vocab_size)`. + beam_indices (`tf.Tensor`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Beam indices of generated token id at each generation step. `tf.Tensor` of shape + `(batch_size*num_return_sequences, sequence_length)`. + attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size*num_beams, num_heads, generated_length, sequence_length)`. + hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size*num_beams, generated_length, hidden_size)`. + """ + + sequences: Optional[tf.Tensor] = None + sequences_scores: Optional[tf.Tensor] = None + scores: Optional[Tuple[tf.Tensor]] = None + beam_indices: Optional[tf.Tensor] = None + attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None + + +@dataclass +class TFBeamSampleEncoderDecoderOutput(ModelOutput): + """ + Base class for outputs of encoder-decoder generation models using beam sampling. Hidden states and attention + weights of the decoder (respectively the encoder) can be accessed via the encoder_attentions and the + encoder_hidden_states attributes (respectively the decoder_attentions and the decoder_hidden_states attributes) + + Args: + sequences (`tf.Tensor` of shape `(batch_size*num_beams, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + sequences_scores (`tf.Tensor` of shape `(batch_size * num_return_sequence)`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Final beam scores of the generated `sequences`. + scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Processed beam scores for each vocabulary token at each generation step. Beam scores consisting of log + softmax scores for each vocabulary token and sum of log softmax of previously generated tokens in this + beam. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each generated token), + with each tensor of shape `(batch_size*num_beams, config.vocab_size)`. + beam_indices (`tf.Tensor`, *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Beam indices of generated token id at each generation step. `tf.Tensor` of shape + `(batch_size*num_return_sequences, sequence_length)`. + encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size*num_beams, sequence_length, hidden_size)`. + decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size*num_beams, num_heads, generated_length, sequence_length)`. + cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size*num_beams, generated_length, hidden_size)`. + """ + + sequences: Optional[tf.Tensor] = None + sequences_scores: Optional[tf.Tensor] = None + scores: Optional[Tuple[tf.Tensor]] = None + beam_indices: Optional[tf.Tensor] = None + encoder_attentions: Optional[Tuple[tf.Tensor]] = None + encoder_hidden_states: Optional[Tuple[tf.Tensor]] = None + decoder_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + cross_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + decoder_hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None + + +@dataclass +class TFContrastiveSearchDecoderOnlyOutput(ModelOutput): + """ + Base class for outputs of decoder-only generation models using contrastive search. + + Args: + sequences (`tf.Tensor` of shape `(batch_size, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each + generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, generated_length, hidden_size)`. + """ + + sequences: Optional[tf.Tensor] = None + scores: Optional[Tuple[tf.Tensor]] = None + attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None + + +@dataclass +class TFContrastiveSearchEncoderDecoderOutput(ModelOutput): + """ + Base class for outputs of encoder-decoder generation models using contrastive search. Hidden states and attention + weights of the decoder (respectively the encoder) can be accessed via the encoder_attentions and the + encoder_hidden_states attributes (respectively the decoder_attentions and the decoder_hidden_states attributes) + + Args: + sequences (`tf.Tensor` of shape `(batch_size, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + scores (`tuple(tf.Tensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`): + Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `tf.Tensor` with up to `max_new_tokens` elements (one element for each + generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + decoder_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + cross_attentions (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + decoder_hidden_states (`tuple(tuple(tf.Tensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `tf.Tensor` of shape `(batch_size, generated_length, hidden_size)`. + """ + + sequences: Optional[tf.Tensor] = None + scores: Optional[Tuple[tf.Tensor]] = None + encoder_attentions: Optional[Tuple[tf.Tensor]] = None + encoder_hidden_states: Optional[Tuple[tf.Tensor]] = None + decoder_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + cross_attentions: Optional[Tuple[Tuple[tf.Tensor]]] = None + decoder_hidden_states: Optional[Tuple[Tuple[tf.Tensor]]] = None + + +TFGreedySearchOutput = Union[TFGreedySearchEncoderDecoderOutput, TFGreedySearchDecoderOnlyOutput] +TFSampleOutput = Union[TFSampleEncoderDecoderOutput, TFSampleDecoderOnlyOutput] +TFBeamSearchOutput = Union[TFBeamSearchEncoderDecoderOutput, TFBeamSearchDecoderOnlyOutput] +TFBeamSampleOutput = Union[TFBeamSampleEncoderDecoderOutput, TFBeamSampleDecoderOnlyOutput] +TFContrastiveSearchOutput = Union[TFContrastiveSearchEncoderDecoderOutput, TFContrastiveSearchDecoderOnlyOutput] +TFGenerateOutput = Union[ + TFGreedySearchOutput, TFSampleOutput, TFBeamSearchOutput, TFBeamSampleOutput, TFContrastiveSearchOutput +] + + +class TFGenerationMixin: + """ + A class containing all of the functions supporting generation, to be used as a mixin in [`TFPreTrainedModel`]. + + The class exposes [`~generation.TFGenerationMixin.generate`], which can be used for: + - *greedy decoding* by calling [`~generation.TFGenerationMixin.greedy_search`] if `num_beams=1` and + `do_sample=False` + - *contrastive search* by calling [`~generation.TFGenerationMixin.contrastive_search`] if `penalty_alpha>0` and + `top_k>1` + - *multinomial sampling* by calling [`~generation.TFGenerationMixin.sample`] if `num_beams=1` and + `do_sample=True` + - *beam-search decoding* by calling [`~generation.TFGenerationMixin.beam_search`] if `num_beams>1` + + You do not need to call any of the above methods directly. Pass custom parameter values to 'generate' instead. To + learn more about decoding strategies refer to the [text generation strategies guide](../generation_strategies). + """ + + _seed_generator = None + + @property + def seed_generator(self): + warnings.warn("`seed_generator` is deprecated and will be removed in a future version.", UserWarning) + if self._seed_generator is None: + self._seed_generator = tf.random.Generator.from_non_deterministic_state() + return self._seed_generator + + supports_xla_generation = True + + def prepare_inputs_for_generation(self, *args, **kwargs): + raise NotImplementedError( + "A model class needs to define a `prepare_inputs_for_generation` method in order to use `generate`." + ) + + def compute_transition_scores( + self, + sequences: tf.Tensor, + scores: Tuple[tf.Tensor], + beam_indices: Optional[tf.Tensor] = None, + normalize_logits: bool = False, + ) -> tf.Tensor: + """ + Computes the transition scores of sequences given the generation scores (and beam indices, if beam search was + used). This is a convenient method to quickly obtain the scores of the selected tokens at generation time. + + Parameters: + sequences (`tf.Tensor`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or + shorter if all batches finished early due to the `eos_token_id`. + scores (`tuple(tf.Tensor)`): + Transition scores for each vocabulary token at each generation step. Beam transition scores consisting + of log probabilities of tokens conditioned on log softmax of previously generated tokens Tuple of + `tf.Tensor` with up to `max_new_tokens` elements (one element for each generated token), with each + tensor of shape `(batch_size*num_beams, config.vocab_size)`. + beam_indices (`tf.Tensor`, *optional*): + Beam indices of generated token id at each generation step. `tf.Tensor` of shape + `(batch_size*num_return_sequences, sequence_length)`. Only required if a `num_beams>1` at + generate-time. + normalize_logits (`bool`, *optional*, defaults to `False`): + Whether to normalize the logits (which, for legacy reasons, may be unnormalized). + + Return: + `tf.Tensor`: A `tf.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)` containing + the transition scores (logits) + + Examples: + + ```python + >>> from transformers import GPT2Tokenizer, TFAutoModelForCausalLM + >>> import numpy as np + + >>> tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2") + >>> model = TFAutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> tokenizer.pad_token_id = tokenizer.eos_token_id + >>> inputs = tokenizer(["Today is"], return_tensors="tf") + + >>> # Example 1: Print the scores for each token generated with Greedy Search + >>> outputs = model.generate(**inputs, max_new_tokens=5, return_dict_in_generate=True, output_scores=True) + >>> transition_scores = model.compute_transition_scores( + ... outputs.sequences, outputs.scores, normalize_logits=True + ... ) + >>> # input_length is the length of the input prompt for decoder-only models, like the GPT family, and 1 for + >>> # encoder-decoder models, like BART or T5. + >>> input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1] + >>> generated_tokens = outputs.sequences[:, input_length:] + >>> for tok, score in zip(generated_tokens[0], transition_scores[0]): + ... # | token | token string | logits | probability + ... print(f"| {tok:5d} | {tokenizer.decode(tok):8s} | {score.numpy():.3f} | {np.exp(score.numpy()):.2%}") + | 262 | the | -1.414 | 24.33% + | 1110 | day | -2.609 | 7.36% + | 618 | when | -2.010 | 13.40% + | 356 | we | -1.859 | 15.58% + | 460 | can | -2.508 | 8.14% + + >>> # Example 2: Reconstruct the sequence scores from Beam Search + >>> outputs = model.generate( + ... **inputs, + ... max_new_tokens=5, + ... num_beams=4, + ... num_return_sequences=4, + ... return_dict_in_generate=True, + ... output_scores=True, + ... ) + >>> transition_scores = model.compute_transition_scores( + ... outputs.sequences, outputs.scores, outputs.beam_indices, normalize_logits=False + ... ) + >>> # If you sum the generated tokens' scores and apply the length penalty, you'll get the sequence scores. + >>> # Tip: recomputing the scores is only guaranteed to match with `normalize_logits=False`. Depending on the + >>> # use case, you might want to recompute it with `normalize_logits=True`. + >>> output_length = np.sum(transition_scores.numpy() < 0, axis=1) + >>> length_penalty = model.generation_config.length_penalty + >>> reconstructed_scores = np.sum(transition_scores, axis=1) / (output_length**length_penalty) + >>> print(np.allclose(outputs.sequences_scores, reconstructed_scores)) + True + ```""" + # 1. In absence of `beam_indices`, we can assume that we come from e.g. greedy search, which is equivalent + # to a beam search approach were the first (and only) beam is always selected + if beam_indices is None: + beam_indices = tf.tile(tf.expand_dims(tf.range(scores[0].shape[0]), axis=1), [1, len(scores)]) + + # 2. reshape scores as [batch_size, vocab_size, # generation steps] with # generation steps being + # seq_len - input_length + scores = tf.transpose(tf.reshape(tf.stack(scores), (len(scores), -1)), (1, 0)) + scores = tf.reshape(scores, (-1, self.config.vocab_size, scores.shape[-1])) + + # 3. Optionally normalize the logits (across the vocab dimension) + if normalize_logits: + scores = tf.nn.log_softmax(scores, axis=1) + + # 4. cut beam_indices to longest beam length + beam_indices_mask = beam_indices < 0 + max_beam_length = tf.math.reduce_max( + tf.math.reduce_sum((1 - tf.cast(beam_indices_mask, dtype=tf.int32)), axis=-1) + ) + beam_indices = beam_indices[:, -max_beam_length:] + beam_indices_mask = beam_indices_mask[:, -max_beam_length:] + + # 5. Set indices of beams that finished early to 0; such indices will be masked correctly afterwards + beam_indices = tf.where(beam_indices_mask, 0, beam_indices) + + # 6. Define which indices contributed to scores + cut_idx = sequences.shape[-1] - max_beam_length + token_indices = sequences[:, cut_idx:] + gen_step_idx = tf.broadcast_to(tf.range(scores.shape[-1]), token_indices.shape) + indices = tf.stack([beam_indices, token_indices, gen_step_idx], axis=-1) + + # 7. Compute scores + transition_scores = tf.gather_nd(scores, indices) + + # 8. Mask out transition_scores of beams that stopped early + transition_scores = tf.where(beam_indices_mask, 0, transition_scores) + + return transition_scores + + def _validate_model_class(self): + """ + Confirms that the model class is compatible with generation. If not, raises an exception that points to the + right class to use. + """ + if not self.can_generate(): + generate_compatible_mappings = [ + TF_MODEL_FOR_CAUSAL_LM_MAPPING, + TF_MODEL_FOR_VISION_2_SEQ_MAPPING, + TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, + TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING, + ] + generate_compatible_classes = set() + for model_mapping in generate_compatible_mappings: + supported_models = model_mapping.get(type(self.config), default=None) + if supported_models is not None: + generate_compatible_classes.add(supported_models.__name__) + exception_message = ( + f"The current model class ({self.__class__.__name__}) is not compatible with `.generate()`, as " + "it doesn't have a language model head." + ) + if generate_compatible_classes: + exception_message += f" Please use one of the following classes instead: {generate_compatible_classes}" + raise TypeError(exception_message) + + def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]): + """Validates model kwargs for generation. Generate argument typos will also be caught here.""" + # Excludes arguments that are handled before calling any model function + if self.config.is_encoder_decoder: + for key in ["decoder_input_ids"]: + model_kwargs.pop(key, None) + + unused_model_args = [] + model_args = set(inspect.signature(self.prepare_inputs_for_generation).parameters) + # `kwargs`/`model_kwargs` is often used to handle optional forward pass inputs like `attention_mask`. If + # `prepare_inputs_for_generation` doesn't accept them, then a stricter check can be made ;) + if "kwargs" in model_args or "model_kwargs" in model_args: + model_args |= set(inspect.signature(self.call).parameters) + for key, value in model_kwargs.items(): + if value is not None and key not in model_args: + unused_model_args.append(key) + + if unused_model_args: + raise ValueError( + f"The following `model_kwargs` are not used by the model: {unused_model_args} (note: typos in the" + " generate arguments will also show up in this list)" + ) + + def generate( + self, + inputs: Optional[tf.Tensor] = None, + generation_config: Optional[GenerationConfig] = None, + logits_processor: Optional[TFLogitsProcessorList] = None, + seed=None, + **kwargs, + ) -> Union[TFGenerateOutput, tf.Tensor]: + r""" + Generates sequences of token ids for models with a language modeling head. + + + + Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the + model's default generation configuration. You can override any `generation_config` by passing the corresponding + parameters to generate, e.g. `.generate(inputs, num_beams=4, do_sample=True)`. + + For an overview of generation strategies and code examples, check out the [following + guide](../generation_strategies). + + + + Parameters: + inputs (`tf.Tensor` of varying shape depending on the modality, *optional*): + The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the + method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs` + should of in the format of `input_ids`. For encoder-decoder models *inputs* can represent any of + `input_ids`, `input_values`, `input_features`, or `pixel_values`. + generation_config (`~generation.GenerationConfig`, *optional*): + The generation configuration to be used as base parametrization for the generation call. `**kwargs` + passed to generate matching the attributes of `generation_config` will override them. If + `generation_config` is not provided, the default will be used, which had the following loading + priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model + configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s + default values, whose documentation should be checked to parameterize generation. + logits_processor (`LogitsProcessorList`, *optional*): + Custom logits processors that complement the default logits processors built from arguments and + generation config. If a logit processor is passed that is already created with the arguments or a + generation config an error is thrown. This feature is intended for advanced users. + seed (`List[int]`, *optional*): + Random seed to control sampling, containing two integers, used when `do_sample` is `True`. See the + `seed` argument from stateless functions in `tf.random`. + kwargs (`Dict[str, Any]`, *optional*): + Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be + forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder + specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*. + + Return: + [`~utils.ModelOutput`] or `tf.Tensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True` or when + `config.return_dict_in_generate=True`) or a `tf.Tensor`. + + If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible + [`~utils.ModelOutput`] types are: + + - [`~generation.TFGreedySearchDecoderOnlyOutput`], + - [`~generation.TFSampleDecoderOnlyOutput`], + - [`~generation.TFBeamSearchDecoderOnlyOutput`], + - [`~generation.TFBeamSampleDecoderOnlyOutput`] + + If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible + [`~utils.ModelOutput`] types are: + + - [`~generation.TFGreedySearchEncoderDecoderOutput`], + - [`~generation.TFSampleEncoderDecoderOutput`], + - [`~generation.TFBeamSearchEncoderDecoderOutput`], + - [`~generation.TFBeamSampleEncoderDecoderOutput`] + + """ + + # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call + self._validate_model_class() + + # priority: `generation_config` argument > `model.generation_config` (the default generation config) + if generation_config is None: + # legacy: users may modify the model configuration to control generation. To trigger this legacy behavior, + # two conditions must be met + # 1) the generation config must have been created from the model config (`_from_model_config` field); + # 2) the generation config must have seen no modification since its creation (the hash is the same). + if self.generation_config._from_model_config and self.generation_config._original_object_hash == hash( + self.generation_config + ): + new_generation_config = GenerationConfig.from_model_config(self.config) + if new_generation_config != self.generation_config: + warnings.warn( + "You have modified the pretrained model configuration to control generation. This is a" + " deprecated strategy to control generation and will be removed soon, in a future version." + " Please use and modify the model generation configuration (see" + " https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )" + ) + self.generation_config = new_generation_config + generation_config = self.generation_config + + generation_config = copy.deepcopy(generation_config) + model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs + self._validate_model_kwargs(model_kwargs.copy()) + + # 2. Cast input dtypes to tf.int32 unless they're floats (which happens for some image models) + if inputs is not None: + if isinstance(inputs, tf.Tensor) and inputs.dtype.is_floating: + pass + elif isinstance(inputs, np.ndarray) and np.issubdtype(inputs.dtype, np.floating): + pass + else: + inputs = tf.cast(inputs, tf.int32) + if model_kwargs.get("attention_mask") is not None: + model_kwargs["attention_mask"] = tf.cast(model_kwargs["attention_mask"], tf.int32) + if "decoder_input_ids" in model_kwargs: + if ( + isinstance(model_kwargs["decoder_input_ids"], tf.Tensor) + and model_kwargs["decoder_input_ids"].dtype.is_floating + ): + pass + elif isinstance(model_kwargs["decoder_input_ids"], np.ndarray) and np.issubdtype( + model_kwargs["decoder_input_ids"].dtype, np.floating + ): + pass + else: + model_kwargs["decoder_input_ids"] = tf.cast(model_kwargs["decoder_input_ids"], tf.int32) + + # 3. Set generation parameters if not already defined + logits_processor = logits_processor if logits_processor is not None else TFLogitsProcessorList() + + if generation_config.pad_token_id is None and generation_config.eos_token_id is not None: + if model_kwargs.get("attention_mask") is None: + logger.warning( + "The attention mask and the pad token id were not set. As a consequence, you may observe " + "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results." + ) + eos_token_id = generation_config.eos_token_id + if isinstance(eos_token_id, list): + eos_token_id = eos_token_id[0] + generation_config.pad_token_id = eos_token_id + + use_xla = not tf.executing_eagerly() + if use_xla and not self.supports_xla_generation: + raise ValueError( + "The selected model does not support Graph mode nor XLA generation (e.g. from tf.function())" + ) + + # 4. Define model inputs + inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs( + inputs, generation_config.bos_token_id, model_kwargs + ) + # inputs_ids now has to be defined and cannot be None anymore + batch_size = shape_list(inputs_tensor)[0] + + # 5. Prepare other model kwargs + model_kwargs["output_attentions"] = generation_config.output_attentions + model_kwargs["output_hidden_states"] = generation_config.output_hidden_states + model_kwargs["use_cache"] = generation_config.use_cache + + accepts_attention_mask = "attention_mask" in set(inspect.signature(self.call).parameters.keys()) + requires_attention_mask = "encoder_outputs" not in model_kwargs + + if model_kwargs.get("attention_mask", None) is None and requires_attention_mask and accepts_attention_mask: + model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation( + inputs_tensor, generation_config.pad_token_id, generation_config.eos_token_id + ) + + # decoder-only models should use left-padding for generation + if not self.config.is_encoder_decoder: + if generation_config.pad_token_id is not None and tf.math.reduce_any( + inputs_tensor[:, -1] == generation_config.pad_token_id + ): + logger.warning( + "A decoder-only architecture is being used, but right-padding was detected! For correct " + "generation results, please set `padding_side='left'` when initializing the tokenizer." + ) + if self.config.is_encoder_decoder and "encoder_outputs" not in model_kwargs: + # if model is encoder decoder encoder_outputs are created and added to `model_kwargs` + model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation( + inputs_tensor, model_kwargs, model_input_name + ) + + # 6. Prepare model inputs which will be used for auto-regressive generation + if self.config.is_encoder_decoder: + input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation( + batch_size=batch_size, + model_input_name=model_input_name, + model_kwargs=model_kwargs, + decoder_start_token_id=generation_config.decoder_start_token_id, + bos_token_id=generation_config.bos_token_id, + ) + else: + input_ids = inputs_tensor if model_input_name == "input_ids" else model_kwargs.pop("input_ids") + + # 7. Prepare `max_length` depending on other stopping criteria. + input_ids_seq_length = shape_list(input_ids)[-1] + has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None + if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20: + # 20 is the default max_length of the generation config + warnings.warn( + f"Using the model-agnostic default `max_length` (={generation_config.max_length}) " + "to control the generation length. recommend setting `max_new_tokens` to control the maximum length of the generation.", + UserWarning, + ) + elif generation_config.max_new_tokens is not None: + if not has_default_max_length and generation_config.max_length is not None: + logger.warning( + f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(=" + f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. " + "Please refer to the documentation for more information. " + "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)" + ) + generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length + + # If the input length is a tensor (i.e. dynamic length), skip length checks + if not isinstance(input_ids_seq_length, tf.Tensor): + if ( + generation_config.min_length is not None + and generation_config.min_length > generation_config.max_length + ): + raise ValueError( + f"Unfeasable length constraints: the minimum length ({generation_config.min_length}) is larger" + f" than the maximum length ({generation_config.max_length})" + ) + if input_ids_seq_length >= generation_config.max_length: + input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids" + logger.warning( + f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to" + f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider" + " increasing`max_new_tokens`." + ) + + # 8. determine generation mode + is_contrastive_search_gen_mode = ( + generation_config.top_k is not None + and generation_config.top_k > 1 + and generation_config.do_sample is False + and generation_config.penalty_alpha is not None + and generation_config.penalty_alpha > 0 + ) + is_greedy_gen_mode = ( + not is_contrastive_search_gen_mode + and (generation_config.num_beams == 1) + and generation_config.do_sample is False + ) + is_beam_gen_mode = ( + not is_contrastive_search_gen_mode + and (generation_config.num_beams > 1) + and generation_config.do_sample is False + ) + is_sample_gen_mode = (generation_config.num_beams == 1) and generation_config.do_sample is True + is_beam_sample_gen_mode = (generation_config.num_beams > 1) and generation_config.do_sample is True + + # 9. prepare distribution pre_processing samplers + logits_processor = self._get_logits_processor( + generation_config=generation_config, + input_ids_seq_length=input_ids_seq_length, + logits_processor=logits_processor, + ) + + # 10. go into different generation modes + if is_greedy_gen_mode: + if generation_config.num_return_sequences > 1: + raise ValueError( + f"num_return_sequences has to be 1, but is {generation_config.num_return_sequences} when doing" + " greedy search." + ) + # 11. run greedy search + return self.greedy_search( + input_ids, + max_length=generation_config.max_length, + pad_token_id=generation_config.pad_token_id, + eos_token_id=generation_config.eos_token_id, + logits_processor=logits_processor, + output_scores=generation_config.output_scores, + return_dict_in_generate=generation_config.return_dict_in_generate, + **model_kwargs, + ) + elif is_contrastive_search_gen_mode: + if generation_config.num_return_sequences > 1: + raise ValueError( + f"num_return_sequences has to be 1, but is {generation_config.num_return_sequences} when doing" + " contrastive search." + ) + # 11. run contrastive search + return self.contrastive_search( + input_ids, + top_k=generation_config.top_k, + penalty_alpha=generation_config.penalty_alpha, + logits_processor=logits_processor, + max_length=generation_config.max_length, + pad_token_id=generation_config.pad_token_id, + eos_token_id=generation_config.eos_token_id, + output_scores=generation_config.output_scores, + return_dict_in_generate=generation_config.return_dict_in_generate, + **model_kwargs, + ) + elif is_sample_gen_mode: + # 11. prepare logits warper + logits_warper = self._get_logits_warper(generation_config=generation_config) + + # 12. expand input_ids with `num_return_sequences` additional sequences per batch + input_ids, model_kwargs = self._expand_inputs_for_generation( + input_ids=input_ids, + expand_size=generation_config.num_return_sequences, + is_encoder_decoder=self.config.is_encoder_decoder, + **model_kwargs, + ) + + # 13. run sample + return self.sample( + input_ids, + logits_processor=logits_processor, + logits_warper=logits_warper, + max_length=generation_config.max_length, + pad_token_id=generation_config.pad_token_id, + eos_token_id=generation_config.eos_token_id, + seed=seed, + output_scores=generation_config.output_scores, + return_dict_in_generate=generation_config.return_dict_in_generate, + **model_kwargs, + ) + + elif is_beam_gen_mode: + if generation_config.num_beams < generation_config.num_return_sequences: + raise ValueError( + "Beam search decoding cannot return more sequences than it has beams. Please set num_beams >=" + f" num_return_sequences, got {generation_config.num_beams} and" + f" {generation_config.num_return_sequences} (respectivelly)" + ) + + # 11. broadcast inputs to the desired number of beams + input_ids, model_kwargs = self._expand_inputs_for_generation( + input_ids=input_ids, + expand_size=generation_config.num_beams, + is_encoder_decoder=self.config.is_encoder_decoder, + expand_in_new_axis=True, + **model_kwargs, + ) + + # 12. run beam search + return self.beam_search( + input_ids, + max_length=generation_config.max_length, + pad_token_id=generation_config.pad_token_id, + eos_token_id=generation_config.eos_token_id, + length_penalty=generation_config.length_penalty, + early_stopping=generation_config.early_stopping, + logits_processor=logits_processor, + output_scores=generation_config.output_scores, + return_dict_in_generate=generation_config.return_dict_in_generate, + num_return_sequences=generation_config.num_return_sequences, + **model_kwargs, + ) + + elif is_beam_sample_gen_mode: + if generation_config.num_beams < generation_config.num_return_sequences: + raise ValueError( + "Beam search decoding cannot return more sequences than it has beams. Please set num_beams >=" + f" num_return_sequences, got {generation_config.num_beams} and" + f" {generation_config.num_return_sequences} (respectivelly)" + ) + + # 11. prepare logits warper + logits_warper = self._get_logits_warper(generation_config=generation_config) + + # 12. broadcast inputs to the desired number of beams + input_ids, model_kwargs = self._expand_inputs_for_generation( + input_ids=input_ids, + expand_size=generation_config.num_beams, + is_encoder_decoder=self.config.is_encoder_decoder, + expand_in_new_axis=True, + **model_kwargs, + ) + + # 13. run beam sample (beam search with sampling) + return self.beam_search( + input_ids, + do_sample=True, + max_length=generation_config.max_length, + pad_token_id=generation_config.pad_token_id, + eos_token_id=generation_config.eos_token_id, + length_penalty=generation_config.length_penalty, + early_stopping=generation_config.early_stopping, + logits_processor=logits_processor, + logits_warper=logits_warper, + output_scores=generation_config.output_scores, + return_dict_in_generate=generation_config.return_dict_in_generate, + num_return_sequences=generation_config.num_return_sequences, + **model_kwargs, + ) + + def _prepare_attention_mask_for_generation( + self, + inputs: tf.Tensor, + pad_token_id: Optional[int], + eos_token_id: Optional[int], + ) -> tf.Tensor: + is_input_ids = len(inputs.shape) == 2 and inputs.dtype in (tf.int32, tf.int64) + is_pad_token_in_inputs = (pad_token_id is not None) and tf.math.reduce_any(inputs == pad_token_id) + is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or (pad_token_id != eos_token_id) + + # Check if input is input_ids and padded -> only then is attention_mask defined + if is_input_ids and is_pad_token_in_inputs and is_pad_token_not_equal_to_eos_token_id: + return tf.cast(tf.math.not_equal(inputs, pad_token_id), dtype=tf.int32) + else: + return tf.ones(inputs.shape[:2], dtype=tf.int32) + + def _prepare_encoder_decoder_kwargs_for_generation( + self, inputs_tensor: tf.Tensor, model_kwargs, model_input_name: Optional[str] = None + ) -> Dict[str, Any]: + # 1. get encoder and store encoder outputs + encoder = self.get_encoder() + + # 2. prepare encoder args and encoder kwargs from model kwargs + irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"] + encoder_kwargs = { + argument: value + for argument, value in model_kwargs.items() + if not any(argument.startswith(p) for p in irrelevant_prefix) + } + encoder_signature = set(inspect.signature(encoder.call).parameters) + encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature + if not encoder_accepts_wildcard: + encoder_kwargs = { + argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature + } + + # 3. vision models don't use `attention_mask`. + encoder_kwargs["return_dict"] = True + encoder_kwargs[model_input_name] = inputs_tensor + if model_input_name != self.main_input_name: # in Keras, the first input must always be passed + encoder_kwargs[self.main_input_name] = None + encoder_outputs = encoder(**encoder_kwargs) + model_kwargs["encoder_outputs"] = encoder_outputs + + return model_kwargs + + def _prepare_decoder_input_ids_for_generation( + self, + batch_size: int, + model_input_name: str, + model_kwargs: Dict[str, tf.Tensor], + decoder_start_token_id: Optional[int] = None, + bos_token_id: Optional[int] = None, + ) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]: + """Prepares `decoder_input_ids` for generation with encoder-decoder models""" + # 1. Check whether the user has defined `decoder_input_ids` manually. To facilitate in terms of input naming, + # we also allow the user to pass it under `input_ids`, if the encoder does not use it as the main input. + if model_kwargs is not None and "decoder_input_ids" in model_kwargs: + decoder_input_ids = model_kwargs.pop("decoder_input_ids") + elif "input_ids" in model_kwargs and model_input_name != "input_ids": + decoder_input_ids = model_kwargs.pop("input_ids") + else: + decoder_input_ids = None + + # 2. Encoder-decoder models expect the `decoder_input_ids` to start with a special token. Let's ensure that. + decoder_start_token_id = self._get_decoder_start_token_id(decoder_start_token_id, bos_token_id) + decoder_input_ids_start = tf.ones((batch_size, 1), dtype=tf.int32) * decoder_start_token_id + + # no user input -> use decoder_start_token_id as decoder_input_ids + if decoder_input_ids is None: + decoder_input_ids = decoder_input_ids_start + # user input but doesn't start with decoder_start_token_id -> prepend decoder_start_token_id (and adjust + # decoder_attention_mask if provided) + elif tf.reduce_all(decoder_input_ids[:, 0] != decoder_start_token_id): + decoder_input_ids = tf.concat([decoder_input_ids_start, decoder_input_ids], axis=-1) + if "decoder_attention_mask" in model_kwargs: + decoder_attention_mask = model_kwargs["decoder_attention_mask"] + decoder_attention_mask = tf.concat( + (tf.ones_like(decoder_attention_mask)[:, :1], decoder_attention_mask), + axis=-1, + ) + model_kwargs["decoder_attention_mask"] = decoder_attention_mask + + return decoder_input_ids, model_kwargs + + def _get_decoder_start_token_id( + self, decoder_start_token_id: Optional[int] = None, bos_token_id: Optional[int] = None + ) -> int: + # retrieve decoder_start_token_id for encoder-decoder models + # fall back to bos_token_id if necessary + decoder_start_token_id = ( + decoder_start_token_id + if decoder_start_token_id is not None + else self.generation_config.decoder_start_token_id + ) + bos_token_id = bos_token_id if bos_token_id is not None else self.generation_config.bos_token_id + + if decoder_start_token_id is not None: + return decoder_start_token_id + elif bos_token_id is not None: + return bos_token_id + raise ValueError( + "`decoder_start_token_id` or `bos_token_id` has to be defined for encoder-decoder generation." + ) + + @staticmethod + def _expand_inputs_for_generation( + expand_size: int = 1, + is_encoder_decoder: bool = False, + input_ids: Optional[tf.Tensor] = None, + expand_in_new_axis: bool = False, + **model_kwargs, + ) -> Tuple[tf.Tensor, Dict[str, Any]]: + """ + Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...] or [batch_size, expand_size, ...], + depending on `expand_in_new_axis`. Beam-based approaches expect this function to be used with + `expand_in_new_axis=True` + """ + + def _expand_tensor(tensor: tf.Tensor): + if expand_in_new_axis: + shape = shape_list(tensor) + return tf.broadcast_to(tensor[:, None], (shape[0], expand_size) + tuple(shape[1:])) + else: + return tf.repeat(tensor, expand_size, axis=0) + + def _expand_dict_for_generation(dict_to_expand): + for key in dict_to_expand: + if dict_to_expand[key] is not None and isinstance(dict_to_expand[key], tf.Tensor): + dict_to_expand[key] = _expand_tensor(dict_to_expand[key]) + return dict_to_expand + + if input_ids is not None: + input_ids = _expand_tensor(input_ids) + + model_kwargs = _expand_dict_for_generation(model_kwargs) + + if is_encoder_decoder: + if model_kwargs.get("encoder_outputs") is None: + raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.") + model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"]) + + return input_ids, model_kwargs + + def _prepare_model_inputs( + self, + inputs: Optional[tf.Tensor] = None, + bos_token_id: Optional[int] = None, + model_kwargs: Optional[Dict[str, tf.Tensor]] = None, + ) -> Tuple[tf.Tensor, Optional[str], Dict[str, tf.Tensor]]: + """ + This function extracts the model-specific `inputs` for generation. + """ + # 1. retrieve all kwargs that are non-None or non-model input related. + # some encoder-decoder models have different names for model and encoder + if ( + self.config.is_encoder_decoder + and hasattr(self, "encoder") + and hasattr(self.encoder, "main_input_name") + and self.encoder.main_input_name != self.main_input_name + ): + input_name = self.encoder.main_input_name + else: + input_name = self.main_input_name + + model_kwargs = {k: v for k, v in model_kwargs.items() if v is not None or k != input_name} + + # 2. check whether model_input_name is passed as kwarg + # if yes and `inputs` is None use kwarg inputs + inputs_kwarg = model_kwargs.pop(input_name, None) + if inputs_kwarg is not None and inputs is not None: + raise ValueError( + f"`inputs`: {inputs}` were passed alongside {input_name} which is not allowed. " + f"Make sure to either pass {inputs} or {input_name}=..." + ) + elif inputs_kwarg is not None: + inputs = inputs_kwarg + + # 3. In the presence of `inputs_embeds` for text models: + # - decoder-only models should complain if the user attempts to pass `inputs_embeds`, but the model + # doesn't have its forwarding implemented. `inputs_embeds` is kept in `model_kwargs` and can coexist with + # input_ids (`inputs_embeds` will be used in the 1st generation step, as opposed to `input_ids`) + # - encoder-decoder models should complain if the user attempts to pass `inputs_embeds` and `input_ids`, and + # pull the former to inputs. It will be used in place of `input_ids` to get the encoder hidden states. + if input_name == "input_ids" and "inputs_embeds" in model_kwargs: + if not self.config.is_encoder_decoder: + has_inputs_embeds_forwarding = "inputs_embeds" in set( + inspect.signature(self.prepare_inputs_for_generation).parameters.keys() + ) + if not has_inputs_embeds_forwarding: + raise ValueError( + f"You passed `inputs_embeds` to `.generate()`, but the model class {self.__class__.__name__} " + "doesn't have its forwarding implemented. See the GPT2 implementation for an example " + "(https://github.com/huggingface/transformers/pull/21405), and feel free to open a PR with it!" + ) + # In this case, `input_ids` is moved to the `model_kwargs`, so a few automations (like the creation of + # the attention mask) can rely on the actual model input. + model_kwargs["input_ids"] = self._maybe_initialize_input_ids_for_generation( + inputs, bos_token_id, model_kwargs=model_kwargs + ) + else: + if inputs is not None: + raise ValueError("You passed `inputs_embeds` and `input_ids` to `.generate()`. Please pick one.") + inputs, input_name = model_kwargs["inputs_embeds"], "inputs_embeds" + + # 4. if `inputs` is still None, try to create `input_ids` from BOS token + inputs = self._maybe_initialize_input_ids_for_generation(inputs, bos_token_id, model_kwargs) + + return inputs, input_name, model_kwargs + + def _maybe_initialize_input_ids_for_generation( + self, + inputs: Optional[tf.Tensor] = None, + bos_token_id: Optional[int] = None, + model_kwargs: Optional[Dict[str, tf.Tensor]] = None, + ) -> tf.Tensor: + """Initializes input ids for generation, if necessary.""" + if inputs is not None: + return inputs + + encoder_outputs = model_kwargs.get("encoder_outputs") + if self.config.is_encoder_decoder and encoder_outputs is not None: + # make dummy input_ids with value -100, as a sanity check ensuring that they won't be used for encoding + shape = encoder_outputs.last_hidden_state.shape[:-1] + return tf.ones(shape, dtype=tf.int32) * -100 + + if bos_token_id is None: + raise ValueError("`bos_token_id` has to be defined when no `input_ids` are provided.") + + # If there is some tensor in `model_kwargs`, we can infer the batch size from it. This is helpful with + # soft-prompting or in multimodal implementations built on top of decoder-only language models. + batch_size = 1 + for value in model_kwargs.values(): + if isinstance(value, tf.Tensor): + batch_size = value.shape[0] + break + return tf.ones((batch_size, 1), dtype=tf.int32) * bos_token_id + + @staticmethod + def _extract_past_from_model_output(outputs: ModelOutput): + past_key_values = None + if "past_key_values" in outputs: + past_key_values = outputs.past_key_values + elif "mems" in outputs: + past_key_values = outputs.mems + elif "past_buckets_states" in outputs: + past_key_values = outputs.past_buckets_states + return past_key_values + + def _update_model_kwargs_for_generation( + self, outputs: ModelOutput, model_kwargs: Dict[str, Any], is_encoder_decoder: bool = False + ) -> Dict[str, Any]: + # update past_key_values + model_kwargs["past_key_values"] = self._extract_past_from_model_output(outputs) + + # update attention mask + if not is_encoder_decoder: + if "attention_mask" in model_kwargs: + attention_mask = model_kwargs["attention_mask"] + model_kwargs["attention_mask"] = tf.concat( + [attention_mask, tf.ones((shape_list(attention_mask)[0], 1), dtype=tf.int32)], axis=-1 + ) + + return model_kwargs + + def _update_model_kwargs_for_xla_generation( + self, + model_outputs: ModelOutput, + model_kwargs: Dict[str, Any], + cur_len: int, + max_length: int, + batch_size: int, + is_encoder_decoder: bool = False, + batch_axis: int = 0, + ): + def _initialize_attention(model_kwargs, num_padding_values, is_encoder_decoder): + """initializes the appropriate attention mask -- encoder-decoder models use `decoder_attention_mask`""" + if is_encoder_decoder: + # One 1 for decoder_start_token_id, 0s for the currently-unfilled locations in the past_key_values tensor, + # 1s for the actual input_ids + decoder_attention_mask = tf.concat( + [ + tf.ones((batch_size, 1), dtype=tf.int32), + tf.zeros((batch_size, num_padding_values), dtype=tf.int32), + tf.ones((batch_size, 1), dtype=tf.int32), + ], + axis=1, + ) + mask = {"decoder_attention_mask": decoder_attention_mask} + else: + attention_mask = model_kwargs.pop("attention_mask") + # 0s for the currently-unfilled locations in the past_key_values tensor, 1s for the actual input_ids + attention_mask = tf.concat( + [ + attention_mask, + tf.zeros((batch_size, num_padding_values), dtype=attention_mask.dtype), + tf.ones((batch_size, 1), dtype=attention_mask.dtype), + ], + axis=1, + ) + mask = {"attention_mask": attention_mask} + return mask + + def _update_attention(model_kwargs, new_past_index, is_encoder_decoder): + """updates the appropriate attention mask -- encoder-decoder models use `decoder_attention_mask`""" + update_start = tf.constant([0, 1], dtype=tf.int32) * new_past_index + if is_encoder_decoder: + decoder_attention_mask = model_kwargs.pop("decoder_attention_mask") + decoder_attention_mask_update_slice = tf.ones((batch_size, 1), dtype=decoder_attention_mask.dtype) + decoder_attention_mask = dynamic_update_slice( + decoder_attention_mask, decoder_attention_mask_update_slice, update_start + ) + mask = {"decoder_attention_mask": decoder_attention_mask} + else: + attention_mask = model_kwargs.pop("attention_mask") + attention_mask_update_slice = tf.ones((batch_size, 1), dtype=attention_mask.dtype) + attention_mask = dynamic_update_slice(attention_mask, attention_mask_update_slice, update_start) + mask = {"attention_mask": attention_mask} + return mask + + def _initialize_past(past_key_values, num_padding_values, batch_axis): + """initialize past_key_values with zeros -- the structure depends on `batch_axis`""" + if batch_axis == 0: + padding_values = tf.constant([[0, 0], [0, 0], [0, num_padding_values], [0, 0]], dtype=tf.int32) + new_past = () + for past_layer in past_key_values: + new_past_layer = list(past_layer) + for i in range(len(new_past_layer[:2])): + new_past_layer[i] = tf.pad(past_layer[i], padding_values) + new_past += (tuple(new_past_layer),) + else: + padding_values = tf.scatter_nd(indices=[[3, 1]], updates=[num_padding_values], shape=(5, 2)) + new_past = list(past_key_values) + for i in range(len(past_key_values)): + new_past[i] = tf.pad(past_key_values[i], padding_values) + return new_past + + def _update_past(past_key_values, new_past_index, batch_axis): + if batch_axis == 0: + slice_start_base = tf.constant([0, 0, 1, 0]) + new_past = () + for past_layer in past_key_values: + new_past_layer = list(past_layer) + for i in range(len(new_past_layer[:2])): + update_slice = past_layer[i][:, :, -1:] + # Write the last slice to the first open location in the padded past_key_values array + # and then truncate the last slice off the array + new_past_layer[i] = dynamic_update_slice( + past_layer[i][:, :, :-1], update_slice, slice_start_base * new_past_index + ) + new_past += (tuple(new_past_layer),) + else: + slice_start_base = tf.constant([0, 0, 0, 1, 0]) + new_past = [None for _ in range(len(past_key_values))] + for i in range(len(past_key_values)): + update_slice = past_key_values[i][:, :, :, -1:] + # Write the last slice to the first open location in the padded past_key_values array + # and then truncate the last slice off the array + new_past[i] = dynamic_update_slice( + past_key_values[i][:, :, :, :-1], update_slice, slice_start_base * new_past_index + ) + return new_past + + past_key_values = self._extract_past_from_model_output(model_outputs) + if past_key_values is None: + raise ValueError( + "No known `past_key_values variable` found in model outputs (model outputs keys:" + f" {list(model_outputs.keys())})" + ) + is_past_initialized = model_kwargs.pop("past_key_values", None) is not None + + if not is_past_initialized: + # The padded version of `past_key_values` has a length of `max_length - 1`, as `past_key_values` holds information relative to + # previous autoregressive generation steps (step 0 has no past_key_values, step 1 has 1 past_key_values value, ..., the last step + # has `max_length - 1` past_key_values values). + num_padding_values = max_length - cur_len - 1 + mask = _initialize_attention(model_kwargs, num_padding_values, is_encoder_decoder) + new_past = _initialize_past(past_key_values, num_padding_values, batch_axis) + else: + # The new index of past_key_values to be filled corresponds to the current length of the sequence, with two + # subtractions: -1 because past_key_values holds information regarding previous generation steps (read comment above) + # and -1 again because in an array the index is the length of the array minus 1. + new_past_index = cur_len - 2 + mask = _update_attention(model_kwargs, new_past_index, is_encoder_decoder) + new_past = _update_past(past_key_values, new_past_index, batch_axis) + + # sets the updated variables (mask and past_key_values) + model_kwargs.update(mask) + model_kwargs["past_key_values"] = tuple(new_past) + + return model_kwargs + + def _get_logits_warper( + self, + generation_config: GenerationConfig, + ) -> TFLogitsProcessorList: + """ + This class returns a [`TFLogitsProcessorList`] list object that contains all relevant [`TFLogitsWarper`] + instances used for multinomial sampling. + """ + + # instantiate warpers list + warpers = TFLogitsProcessorList() + + # In beam methods, we need to keep at least one non-eos token to explore continuations that might have a + # better score (i.e. keep len(generation_config.eos_token_id) + 1) + if generation_config.num_beams > 1: + if isinstance(generation_config.eos_token_id, list): + min_tokens_to_keep = len(generation_config.eos_token_id) + 1 + else: + min_tokens_to_keep = 2 + else: + min_tokens_to_keep = 1 + + if generation_config.temperature is not None and generation_config.temperature != 1.0: + warpers.append(TFTemperatureLogitsWarper(generation_config.temperature)) + if generation_config.top_k is not None and generation_config.top_k != 0: + warpers.append(TFTopKLogitsWarper(top_k=generation_config.top_k, min_tokens_to_keep=min_tokens_to_keep)) + if generation_config.top_p is not None and generation_config.top_p < 1.0: + warpers.append(TFTopPLogitsWarper(top_p=generation_config.top_p, min_tokens_to_keep=min_tokens_to_keep)) + return warpers + + def _get_logits_processor( + self, + generation_config: GenerationConfig, + input_ids_seq_length: int, + logits_processor: Optional[TFLogitsProcessorList], + ) -> TFLogitsProcessorList: + """ + This class returns a [`TFLogitsProcessorList`] list object that contains all relevant [`TFLogitsProcessor`] + instances used to modify the scores of the language model head. + """ + processors = TFLogitsProcessorList() + + # instantiate processors list + if generation_config.repetition_penalty is not None and generation_config.repetition_penalty != 1.0: + processors.append(TFRepetitionPenaltyLogitsProcessor(penalty=generation_config.repetition_penalty)) + if generation_config.no_repeat_ngram_size is not None and generation_config.no_repeat_ngram_size > 0: + processors.append(TFNoRepeatNGramLogitsProcessor(generation_config.no_repeat_ngram_size)) + if generation_config.bad_words_ids is not None: + processors.append( + TFNoBadWordsLogitsProcessor(generation_config.bad_words_ids, generation_config.eos_token_id) + ) + if ( + generation_config.min_length is not None + and generation_config.eos_token_id is not None + and generation_config.min_length > 0 + ): + processors.append(TFMinLengthLogitsProcessor(generation_config.min_length, generation_config.eos_token_id)) + if generation_config.forced_bos_token_id is not None: + processors.append(TFForcedBOSTokenLogitsProcessor(generation_config.forced_bos_token_id)) + if generation_config.forced_eos_token_id is not None: + processors.append( + TFForcedEOSTokenLogitsProcessor(generation_config.max_length, generation_config.forced_eos_token_id) + ) + if generation_config.suppress_tokens is not None: + processors.append(TFSuppressTokensLogitsProcessor(generation_config.suppress_tokens)) + if generation_config.begin_suppress_tokens is not None: + begin_index = input_ids_seq_length + begin_index = ( + begin_index + if (input_ids_seq_length > 1 or generation_config.forced_bos_token_id is None) + else begin_index + 1 + ) + if generation_config.forced_decoder_ids is not None: + begin_index += generation_config.forced_decoder_ids[-1][ + 0 + ] # generation starts after the last token that is forced + processors.append( + TFSuppressTokensAtBeginLogitsProcessor(generation_config.begin_suppress_tokens, begin_index) + ) + if generation_config.forced_decoder_ids is not None: + processors.append(TFForceTokensLogitsProcessor(generation_config.forced_decoder_ids)) + + processors = self._merge_criteria_processor_list(processors, logits_processor) + return processors + + def _merge_criteria_processor_list( + self, + default_list: TFLogitsProcessorList, + custom_list: TFLogitsProcessorList, + ) -> TFLogitsProcessorList: + if len(custom_list) == 0: + return default_list + for default in default_list: + for custom in custom_list: + if type(custom) is type(default): + object_type = "logits processor" + raise ValueError( + f"A custom {object_type} of type {type(custom)} with values {custom} has been passed to" + f" `generate`, but it has already been created with the values {default}. {default} has been" + " created by passing the corresponding arguments to generate or by the model's config default" + f" values. If you just want to change the default values of {object_type} consider passing" + f" them as arguments to `generate` instead of using a custom {object_type}." + ) + default_list.extend(custom_list) + return default_list + + def greedy_search( + self, + input_ids: tf.Tensor, + max_length: Optional[int] = None, + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + logits_processor: Optional[TFLogitsProcessorList] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_scores: Optional[bool] = None, + return_dict_in_generate: Optional[bool] = None, + **model_kwargs, + ) -> Union[TFGreedySearchOutput, tf.Tensor]: + r""" + Generates sequences for models with a language modeling head using greedy decoding. + + Parameters: + input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + logits_processor (`TFLogitsProcessorList`, *optional*): + An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + max_length (`int`, *optional*, defaults to 20): + The maximum length of the sequence to be generated. + pad_token_id (`int`, *optional*): + The id of the *padding* token. + eos_token_id (`Union[int, List[int]]`, *optional*): + The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens. + output_attentions (`bool`, *optional*, defaults to `False`): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more details. + output_hidden_states (`bool`, *optional*, defaults to `False`): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more details. + output_scores (`bool`, *optional*, defaults to `False`): + Whether or not to return the prediction scores. See `scores` under returned tensors for more details. + return_dict_in_generate (`bool`, *optional*, defaults to `False`): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + model_kwargs: + Additional model specific keyword arguments will be forwarded to the `call` function of the model. If + model is an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.TFGreedySearchDecoderOnlyOutput`], [`~generation.TFGreedySearchEncoderDecoderOutput`] or + `tf.Tensor`: A `tf.Tensor` containing the generated tokens (default behaviour) or a + [`~generation.TFGreedySearchDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.TFGreedySearchEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + + Examples: + + ```python + >>> from transformers import ( + ... AutoTokenizer, + ... TFAutoModelForCausalLM, + ... TFLogitsProcessorList, + ... TFMinLengthLogitsProcessor, + ... ) + + >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + >>> model = TFAutoModelForCausalLM.from_pretrained("openai-community/gpt2") + + >>> # set pad_token_id to eos_token_id because GPT2 does not have a PAD token + >>> model.generation_config.pad_token_id = model.generation_config.eos_token_id + + >>> input_prompt = "Today is a beautiful day, and" + >>> input_ids = tokenizer(input_prompt, return_tensors="tf").input_ids + + >>> # instantiate logits processors + >>> logits_processor = TFLogitsProcessorList( + ... [ + ... TFMinLengthLogitsProcessor(15, eos_token_id=model.generation_config.eos_token_id), + ... ] + ... ) + + >>> outputs = model.greedy_search(input_ids, logits_processor=logits_processor) + >>> tokenizer.batch_decode(outputs, skip_special_tokens=True) + ["Today is a beautiful day, and I'm so happy to be here. I'm so happy to"] + ```""" + + # 1. init greedy_search values + logits_processor = logits_processor if logits_processor is not None else TFLogitsProcessorList() + + max_length = max_length if max_length is not None else self.generation_config.max_length + pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id + eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + output_scores = output_scores if output_scores is not None else self.generation_config.output_scores + output_attentions = ( + output_attentions if output_attentions is not None else self.generation_config.output_attentions + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states + ) + return_dict_in_generate = ( + return_dict_in_generate + if return_dict_in_generate is not None + else self.generation_config.return_dict_in_generate + ) + use_cache = model_kwargs.pop("use_cache", self.generation_config.use_cache) + use_xla = not tf.executing_eagerly() + # TODO (Joao): fix cache format or find programatic way to detect cache index + # GPT2 and other models has a slightly different cache structure, with a different batch axis + model_name = str(self.decoder) if "EncoderDecoder" in str(self) else str(self) + cache_batch_axis = 1 if any(model_prefix in model_name for model_prefix in ("TFGPT2", "TFCTRL")) else 0 + # some models, like XLNet, need more than the last token in the presence of past_key_values + needs_full_input = "use_mems" in set(inspect.signature(self.prepare_inputs_for_generation).parameters.keys()) + + # 2. init `attentions`, `hidden_states`, and `scores` tuples + scores = [] if (return_dict_in_generate and output_scores) else None + decoder_attentions = [] if (return_dict_in_generate and output_attentions) else None + cross_attentions = [] if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = [] if (return_dict_in_generate and output_hidden_states) else None + + # 3. init tensors to use for "xla-compileable" generate function + batch_size, cur_len = shape_list(input_ids) + + # initialize `generated` (`input_ids` padded with `pad_token_id`), `finished_sequences` + input_ids_padding = tf.ones((batch_size, max_length - cur_len), dtype=tf.int32) * (pad_token_id or 0) + generated = tf.concat([input_ids, input_ids_padding], axis=-1) + finished_sequences = tf.zeros((batch_size,), dtype=tf.bool) + + # 4. define "xla-compile-able" stop-condition and auto-regressive function + # define condition fn + def greedy_search_cond_fn(generated, finished_sequences, cur_len, model_kwargs): + """state termination condition fn.""" + return ~tf.reduce_all(finished_sequences) + + # define condition fn + def greedy_search_body_fn(generated, finished_sequences, cur_len, model_kwargs): + """state update fn.""" + if model_kwargs.get("past_key_values") is None or needs_full_input: + input_ids = generated[:, :cur_len] + else: + input_ids = tf.expand_dims(generated[:, cur_len - 1], -1) + model_inputs = self.prepare_inputs_for_generation(input_ids, use_cache=use_cache, **model_kwargs) + # forward pass to get next token logits + model_outputs = self( + **model_inputs, + return_dict=True, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + next_token_logits = model_outputs.logits[:, -1] + + # pre-process distribution + next_tokens_scores = logits_processor(generated, next_token_logits, cur_len) + + # Store scores, attentions and hidden_states when required + if not use_xla and return_dict_in_generate: + if output_scores: + scores.append(next_tokens_scores) + if output_attentions and self.config.is_encoder_decoder: + decoder_attentions.append(model_outputs.decoder_attentions) + elif output_attentions and not self.config.is_encoder_decoder: + decoder_attentions.append(model_outputs.attentions) + if self.config.is_encoder_decoder: + cross_attentions.append(model_outputs.cross_attentions) + + if output_hidden_states and self.config.is_encoder_decoder: + decoder_hidden_states.append(model_outputs.decoder_hidden_states) + elif output_hidden_states and self.config.is_encoder_decoder: + decoder_hidden_states.append(model_outputs.hidden_states) + + # argmax + next_tokens = tf.argmax(next_tokens_scores, axis=-1, output_type=tf.int32) + + if eos_token_id is not None: + if pad_token_id is None: + raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.") + unfinished_seq = 1 - tf.cast(finished_sequences, tf.int32) + next_tokens = next_tokens * unfinished_seq + pad_token_id * (1 - unfinished_seq) + next_token_is_eos = tf.math.reduce_any( + tf.equal( + tf.broadcast_to(next_tokens, (len(eos_token_id), batch_size)), tf.expand_dims(eos_token_id, -1) + ), + axis=0, + ) + finished_sequences = finished_sequences | next_token_is_eos + + # update `generated` and `cur_len` + update_indices = tf.stack([tf.range(batch_size), tf.broadcast_to(cur_len, [batch_size])], axis=-1) + generated = tf.tensor_scatter_nd_update(tensor=generated, indices=update_indices, updates=next_tokens) + cur_len += 1 + + # update model_kwargs + if use_xla: + model_kwargs = self._update_model_kwargs_for_xla_generation( + model_outputs=model_outputs, + model_kwargs=model_kwargs, + cur_len=cur_len, + max_length=max_length, + batch_size=batch_size, + is_encoder_decoder=self.config.is_encoder_decoder, + batch_axis=cache_batch_axis, + ) + else: + model_kwargs = self._update_model_kwargs_for_generation( + model_outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder + ) + # if we don't cache past_key_values key values we need the whole input + if model_kwargs.get("past_key_values", None) is None: + # let's throw out `past_key_values` since we don't want `None` tensors + model_kwargs.pop("past_key_values", None) + + return generated, finished_sequences, cur_len, model_kwargs + + # 5. run generation + # 1st generation step has to be run before to initialize `past_key_values` + generated, finished_sequences, cur_len, model_kwargs = greedy_search_body_fn( + generated, finished_sequences, cur_len, model_kwargs + ) + + # 2-to-n generation steps can then be run in autoregressive fashion + # only in case 1st generation step does NOT yield EOS token though + maximum_iterations = max_length - cur_len + generated, _, cur_len, _ = tf.while_loop( + greedy_search_cond_fn, + greedy_search_body_fn, + (generated, finished_sequences, cur_len, model_kwargs), + maximum_iterations=maximum_iterations, + ) + + # 6. prepare outputs + if not use_xla: + # cut for backward compatibility + generated = generated[:, :cur_len] + + if return_dict_in_generate: + if self.config.is_encoder_decoder: + # if model is an encoder-decoder, retrieve encoder attention weights + # and hidden states + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + scores = tuple(scores) if scores is not None else None + decoder_attentions = tuple(decoder_attentions) if decoder_attentions is not None else None + cross_attentions = tuple(cross_attentions) if cross_attentions is not None else None + decoder_hidden_states = tuple(decoder_hidden_states) if decoder_hidden_states is not None else None + + return TFGreedySearchEncoderDecoderOutput( + sequences=generated, + scores=scores, + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + ) + else: + return TFGreedySearchDecoderOnlyOutput( + sequences=generated, + scores=scores, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + ) + else: + return generated + + def sample( + self, + input_ids: tf.Tensor, + logits_processor: Optional[TFLogitsProcessorList] = None, + logits_warper: Optional[TFLogitsProcessorList] = None, + max_length: Optional[int] = None, + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + seed: Optional[Tuple[int, int]] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_scores: Optional[bool] = None, + return_dict_in_generate: Optional[bool] = None, + **model_kwargs, + ) -> Union[TFSampleOutput, tf.Tensor]: + r""" + Generates sequences for models with a language modeling head using multinomial sampling. + + Parameters: + input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + logits_processor (`TFLogitsProcessorList`, *optional*): + An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + logits_warper (`TFLogitsProcessorList`, *optional*): + An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsWarper`] + used to warp the prediction score distribution of the language modeling head applied before multinomial + sampling at each generation step. + max_length (`int`, *optional*, defaults to 20): + The maximum length of the sequence to be generated. + pad_token_id (`int`, *optional*): + The id of the *padding* token. + eos_token_id (`Union[int, List[int]]`, *optional*): + The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens. + seed (`List[int]`, *optional*): + Random seed to control sampling, containing two integers, used when `do_sample` is `True`. See the + `seed` argument from stateless functions in `tf.random`. + output_attentions (`bool`, *optional*, defaults to `False`): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more details. + output_hidden_states (`bool`, *optional*, defaults to `False`): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more details. + output_scores (`bool`, *optional*, defaults to `False`): + Whether or not to return the prediction scores. See `scores` under returned tensors for more details. + return_dict_in_generate (`bool`, *optional*, defaults to `False`): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + model_kwargs: + Additional model specific kwargs will be forwarded to the `call` function of the model. If model is an + encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.TFSampleDecoderOnlyOutput`], [`~generation.TFSampleEncoderDecoderOutput`] or `tf.Tensor`: A + `tf.Tensor` containing the generated tokens (default behaviour) or a + [`~generation.TFSampleDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.TFSampleEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + + Examples: + + ```python + >>> import tensorflow as tf + >>> from transformers import ( + ... AutoTokenizer, + ... TFAutoModelForCausalLM, + ... TFLogitsProcessorList, + ... TFMinLengthLogitsProcessor, + ... TFTopKLogitsWarper, + ... TFTemperatureLogitsWarper, + ... ) + + >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + >>> model = TFAutoModelForCausalLM.from_pretrained("openai-community/gpt2") + + >>> # set pad_token_id to eos_token_id because GPT2 does not have a EOS token + >>> model.generation_config.pad_token_id = model.generation_config.eos_token_id + + >>> input_prompt = "Today is a beautiful day, and" + >>> input_ids = tokenizer(input_prompt, return_tensors="tf").input_ids + + >>> # instantiate logits processors + >>> logits_processor = TFLogitsProcessorList( + ... [ + ... TFMinLengthLogitsProcessor(15, eos_token_id=model.generation_config.eos_token_id), + ... ] + ... ) + >>> # instantiate logits processors + >>> logits_warper = TFLogitsProcessorList( + ... [ + ... TFTopKLogitsWarper(50), + ... TFTemperatureLogitsWarper(0.7), + ... ] + ... ) + + >>> tf.random.set_seed(0) + >>> outputs = model.sample(input_ids, logits_processor=logits_processor, logits_warper=logits_warper) + + >>> tokenizer.batch_decode(outputs, skip_special_tokens=True) + ['Today is a beautiful day, and I love my country. But when I look at Donald Trump,'] + ```""" + + # 1. init greedy_search values + logits_processor = logits_processor if logits_processor is not None else TFLogitsProcessorList() + logits_warper = logits_warper if logits_warper is not None else TFLogitsProcessorList() + + max_length = max_length if max_length is not None else self.generation_config.max_length + pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id + eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + output_scores = output_scores if output_scores is not None else self.generation_config.output_scores + output_attentions = ( + output_attentions if output_attentions is not None else self.generation_config.output_attentions + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states + ) + return_dict_in_generate = ( + return_dict_in_generate + if return_dict_in_generate is not None + else self.generation_config.return_dict_in_generate + ) + use_cache = model_kwargs.pop("use_cache", self.generation_config.use_cache) + use_xla = not tf.executing_eagerly() + # TODO (Joao): fix cache format or find programatic way to detect cache index + # GPT2 and other models has a slightly different cache structure, with a different batch axis + model_name = str(self.decoder) if "EncoderDecoder" in str(self) else str(self) + cache_batch_axis = 1 if any(model_prefix in model_name for model_prefix in ("TFGPT2", "TFCTRL")) else 0 + # some models, like XLNet, need more than the last token in the presence of past_key_values + needs_full_input = "use_mems" in set(inspect.signature(self.prepare_inputs_for_generation).parameters.keys()) + + # 2. init `attentions`, `hidden_states`, and `scores` tuples + scores = [] if (return_dict_in_generate and output_scores) else None + decoder_attentions = [] if (return_dict_in_generate and output_attentions) else None + cross_attentions = [] if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = [] if (return_dict_in_generate and output_hidden_states) else None + + # 3. init tensors to use for "xla-compileable" generate function + batch_size, cur_len = shape_list(input_ids) + + # initialize `generated` (pre-populated with `pad_token_id`), `finished_sequences` + input_ids_padding = tf.ones((batch_size, max_length - cur_len), dtype=tf.int32) * (pad_token_id or 0) + generated = tf.concat([input_ids, input_ids_padding], axis=-1) + finished_sequences = tf.zeros((batch_size,), dtype=tf.bool) + + # 4. define "xla-compile-able" stop-condition and auto-regressive function + def sample_cond_fn(generated, finished_sequences, cur_len, model_kwargs): + return ~tf.reduce_all(finished_sequences) + + def sample_body_fn(generated, finished_sequences, cur_len, model_kwargs): + if model_kwargs.get("past_key_values") is None or needs_full_input: + input_ids = generated[:, :cur_len] + else: + input_ids = tf.expand_dims(generated[:, cur_len - 1], -1) + model_inputs = self.prepare_inputs_for_generation(input_ids, use_cache=use_cache, **model_kwargs) + # forward pass to get next token logits + model_outputs = self( + **model_inputs, + return_dict=True, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + next_token_logits = model_outputs.logits[:, -1] + + # pre-process distribution + next_tokens_scores = logits_processor(generated, next_token_logits, cur_len) + next_tokens_scores = logits_warper(generated, next_tokens_scores, cur_len) + + # Store scores, attentions and hidden_states when required + if not use_xla and return_dict_in_generate: + if output_scores: + scores.append(next_tokens_scores) + if output_attentions and self.config.is_encoder_decoder: + decoder_attentions.append(model_outputs.decoder_attentions) + elif output_attentions and not self.config.is_encoder_decoder: + decoder_attentions.append(model_outputs.attentions) + if self.config.is_encoder_decoder: + cross_attentions.append(model_outputs.cross_attentions) + + if output_hidden_states and self.config.is_encoder_decoder: + decoder_hidden_states.append(model_outputs.decoder_hidden_states) + elif output_hidden_states and self.config.is_encoder_decoder: + decoder_hidden_states.append(model_outputs.hidden_states) + + # sample + if seed is not None: + sample_seed = seed + else: + sample_seed = tf.experimental.numpy.random.randint(tf.int32.min, tf.int32.max, (2,), dtype=tf.int32) + next_tokens = tf.squeeze( + tf.random.stateless_categorical( + logits=next_tokens_scores, num_samples=1, seed=sample_seed, dtype=tf.int32 + ), + axis=1, + ) + + if eos_token_id is not None: + if pad_token_id is None: + raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.") + unfinished_seq = 1 - tf.cast(finished_sequences, tf.int32) + next_tokens = next_tokens * unfinished_seq + pad_token_id * (1 - unfinished_seq) + next_token_is_eos = tf.math.reduce_any( + tf.equal( + tf.broadcast_to(next_tokens, (len(eos_token_id), batch_size)), tf.expand_dims(eos_token_id, -1) + ), + axis=0, + ) + finished_sequences = finished_sequences | next_token_is_eos + + # update `generated` and `cur_len` + update_indices = tf.stack([tf.range(batch_size), tf.broadcast_to(cur_len, [batch_size])], axis=-1) + generated = tf.tensor_scatter_nd_update(tensor=generated, indices=update_indices, updates=next_tokens) + cur_len += 1 + + # update model_kwargs + if use_xla: + model_kwargs = self._update_model_kwargs_for_xla_generation( + model_outputs=model_outputs, + model_kwargs=model_kwargs, + cur_len=cur_len, + max_length=max_length, + batch_size=batch_size, + is_encoder_decoder=self.config.is_encoder_decoder, + batch_axis=cache_batch_axis, + ) + else: + model_kwargs = self._update_model_kwargs_for_generation( + model_outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder + ) + # if we don't cache past_key_values key values we need the whole input + if model_kwargs.get("past_key_values", None) is None: + # let's throw out `past_key_values` since we don't want `None` tensors + model_kwargs.pop("past_key_values", None) + + return generated, finished_sequences, cur_len, model_kwargs + + # 5. run generation + # 1st generation step has to be run before to initialize `past_key_values` + generated, finished_sequences, cur_len, model_kwargs = sample_body_fn( + generated, finished_sequences, cur_len, model_kwargs + ) + + # 2-to-n generation steps can then be run in autoregressive fashion + # only in case 1st generation step does NOT yield EOS token though + maximum_iterations = max_length - cur_len + generated, _, cur_len, _ = tf.while_loop( + sample_cond_fn, + sample_body_fn, + (generated, finished_sequences, cur_len, model_kwargs), + maximum_iterations=maximum_iterations, + ) + + # 6. prepare outputs + if not use_xla: + # cut for backward compatibility + generated = generated[:, :cur_len] + + if return_dict_in_generate: + if self.config.is_encoder_decoder: + # if model is an encoder-decoder, retrieve encoder attention weights + # and hidden states + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + scores = tuple(scores) if scores is not None else None + decoder_attentions = tuple(decoder_attentions) if decoder_attentions is not None else None + cross_attentions = tuple(cross_attentions) if cross_attentions is not None else None + decoder_hidden_states = tuple(decoder_hidden_states) if decoder_hidden_states is not None else None + + return TFSampleEncoderDecoderOutput( + sequences=generated, + scores=scores, + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + ) + else: + return TFSampleDecoderOnlyOutput( + sequences=generated, + scores=scores, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + ) + else: + return generated + + @staticmethod + def _gather_beams(nested, beam_indices, batch_axis=0): + """Gathers the beam slices indexed by beam_indices into new beam array.""" + + def gather_fn(tensor): + if batch_axis > 0: + # pushes all dimentions before the batch to the end, so we get (batch, beam_id, ...) + perm = tf.concat((tf.range(tf.rank(tensor))[batch_axis:], tf.range(batch_axis)), axis=0) + tensor = tf.transpose(tensor, perm=perm) + + gathered_tensor = tf.gather(params=tensor, indices=beam_indices, axis=1, batch_dims=1) + if batch_axis > 0: + # transposes back to the original dimensions + perm = tf.concat((tf.range(tf.rank(tensor))[batch_axis:], tf.range(batch_axis)), axis=0) + perm = tf.math.invert_permutation(perm) + gathered_tensor = tf.transpose(gathered_tensor, perm=perm) + + return gathered_tensor + + return tf.nest.map_structure(gather_fn, nested) + + def beam_search( + self, + input_ids: tf.Tensor, + do_sample: bool = False, + max_length: Optional[int] = None, + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + length_penalty: Optional[float] = None, + early_stopping: Optional[Union[bool, str]] = None, + logits_processor: Optional[TFLogitsProcessorList] = None, + logits_warper: Optional[TFLogitsProcessorList] = None, + num_return_sequences: Optional[int] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_scores: Optional[bool] = None, + return_dict_in_generate: Optional[bool] = None, + **model_kwargs, + ) -> Union[TFBeamSearchOutput, TFBeamSampleOutput, tf.Tensor]: + r""" + Generates sequences for models with a language modeling head using beam search. If `do_sample` is `False`, uses + a greedy approach, otherwise does multinomial sampling without replacement. + + Parameters: + input_ids (`tf.Tensor` of shape `(batch_size, num_beams, sequence_length)`): + The sequence used as a prompt for the generation. + do_sample (`bool`, *optional*, defaults to `False`): + Whether or not to use sampling ; use greedy decoding otherwise. + max_length (`int`, *optional*, defaults to 20): + The maximum length of the sequence to be generated. + pad_token_id (`int`, *optional*): + The id of the *padding* token. + eos_token_id (`Union[int, List[int]]`, *optional*): + The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens. + length_penalty (`float`, *optional*, defaults to 1.0): + Exponential penalty to the length that is used with beam-based generation. It is applied as an exponent + to the sequence length, which in turn is used to divide the score of the sequence. Since the score is + the log likelihood of the sequence (i.e. negative), `length_penalty` > 0.0 promotes longer sequences, + while `length_penalty` < 0.0 encourages shorter sequences. + early_stopping (`bool` or `str`, *optional*, defaults to `False`): + Controls the stopping condition for beam-based methods, like beam-search. It accepts the following + values: `True`, where the generation stops as soon as there are `num_beams` complete candidates; + `False`, where an heuristic is applied and the generation stops when is it very unlikely to find better + candidates; `"never"`, where the beam search procedure only stops when there cannot be better + candidates (canonical beam search algorithm). + logits_processor (`[TFLogitsProcessorList]`, *optional*): + An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + logits_warper (`TFLogitsProcessorList`, *optional*): + An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsWarper`] + used to warp the prediction score distribution of the language modeling head applied before multinomial + sampling at each generation step. + num_return_sequences(`int`, *optional*, defaults to 1): + The number of independently computed returned sequences for each element in the batch. + output_attentions (`bool`, *optional*, defaults to `False`): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more details. + output_hidden_states (`bool`, *optional*, defaults to `False`): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more details. + return_dict_in_generate (`bool`, *optional*, defaults to `False`): + Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. + model_kwargs: + Additional model specific kwargs will be forwarded to the `call` function of the model. If model is an + encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.TFBeamSearchDecoderOnlyOutput`], [`~generation.TFBeamSearchEncoderDecoderOutput`] or + `tf.Tensor`: A `tf.Tensor` containing the generated tokens (default behaviour) or a + [`~generation.TFBeamSearchDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.TFBeamSearchEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + + Examples: + + ```python + >>> from transformers import ( + ... AutoTokenizer, + ... TFAutoModelForSeq2SeqLM, + ... TFLogitsProcessorList, + ... TFMinLengthLogitsProcessor, + ... ) + >>> import tensorflow as tf + + >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-base") + >>> model = TFAutoModelForSeq2SeqLM.from_pretrained("google-t5/t5-base") + + >>> encoder_input_str = "translate English to German: How old are you?" + >>> encoder_input_ids = tokenizer(encoder_input_str, return_tensors="tf").input_ids + + >>> # lets run beam search using 3 beams + >>> num_beams = 3 + >>> # define decoder start token ids + >>> input_ids = tf.ones((1, num_beams, 1), dtype=tf.int32) + >>> input_ids = input_ids * model.generation_config.decoder_start_token_id + + >>> # add encoder_outputs to model keyword arguments + >>> encoder_outputs = model.get_encoder()(encoder_input_ids, return_dict=True) + >>> encoder_outputs.last_hidden_state = tf.repeat( + ... tf.expand_dims(encoder_outputs.last_hidden_state, axis=0), num_beams, axis=1 + ... ) + >>> model_kwargs = {"encoder_outputs": encoder_outputs} + + >>> # instantiate logits processors + >>> logits_processor = TFLogitsProcessorList( + ... [TFMinLengthLogitsProcessor(5, eos_token_id=model.generation_config.eos_token_id)] + ... ) + + >>> outputs = model.beam_search(input_ids, logits_processor=logits_processor, **model_kwargs) + >>> tokenizer.batch_decode(outputs, skip_special_tokens=True) + ['Wie alt bist du?'] + ```""" + + def flatten_beam_dim(tensor, batch_axis=0): + """Flattens the first two dimensions of a non-scalar array.""" + shape = shape_list(tensor) + return tf.reshape( + tensor, + shape[:batch_axis] + [shape[batch_axis] * shape[batch_axis + 1]] + shape[batch_axis + 2 :], + ) + + def unflatten_beam_dim(tensor, num_beams, batch_axis=0): + """Unflattens the first, flat batch*beam dimension of a non-scalar array.""" + shape = shape_list(tensor) + return tf.reshape(tensor, shape[:batch_axis] + [-1, num_beams] + shape[batch_axis + 1 :]) + + # 1. init beam_search values + logits_processor = logits_processor if logits_processor is not None else TFLogitsProcessorList() + logits_warper = logits_warper if logits_warper is not None else TFLogitsProcessorList() + + max_length = max_length if max_length is not None else self.generation_config.max_length + pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id + eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + num_return_sequences = ( + num_return_sequences if num_return_sequences is not None else self.generation_config.num_return_sequences + ) + + output_attentions = ( + output_attentions if output_attentions is not None else self.generation_config.output_attentions + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states + ) + output_scores = output_scores if output_scores is not None else self.generation_config.output_scores + return_dict_in_generate = ( + return_dict_in_generate + if return_dict_in_generate is not None + else self.generation_config.return_dict_in_generate + ) + + length_penalty = length_penalty if length_penalty is not None else self.generation_config.length_penalty + early_stopping = early_stopping if early_stopping is not None else self.generation_config.early_stopping + + use_cache = model_kwargs.pop("use_cache", self.generation_config.use_cache) + use_xla = not tf.executing_eagerly() + # TODO (Joao): fix cache format or find programatic way to detect cache index + # GPT2 and other models has a slightly different cache structure, with a different batch axis + model_name = str(self.decoder) if "EncoderDecoder" in str(self) else str(self) + cache_batch_axis = 1 if any(model_prefix in model_name for model_prefix in ("TFGPT2", "TFCTRL")) else 0 + # some models, like XLNet, need more than the last token in the presence of past_key_values + needs_full_input = "use_mems" in set(inspect.signature(self.prepare_inputs_for_generation).parameters.keys()) + + # 2. init `attentions`, `hidden_states`, and `scores` tuples + all_scores = [] if (return_dict_in_generate and output_scores) else None + decoder_attentions = [] if (return_dict_in_generate and output_attentions) else None + cross_attentions = [] if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = [] if (return_dict_in_generate and output_hidden_states) else None + + # 3. init tensors to use for "xla-compileable" generate function + batch_size, num_beams, cur_len = shape_list(input_ids) + # store the prompt length of decoder + decoder_prompt_len = cur_len + + # per batch, beam-item holding current token in loop, pre-populated with `pad_token_id` + input_ids_padding = tf.ones((batch_size, num_beams, max_length - cur_len), dtype=tf.int32) * ( + pad_token_id or 0 + ) + running_sequences = tf.concat([input_ids, input_ids_padding], axis=-1) + sequences = tf.ones((batch_size, num_beams, max_length), dtype=tf.int32) * (pad_token_id or 0) + + # per batch,beam-item state bit indicating if sentence has finished. + is_sent_finished = tf.zeros((batch_size, num_beams), dtype=tf.bool) + + # per batch, beam-item score, logprobs + running_scores = tf.tile( + tf.expand_dims(tf.convert_to_tensor([0.0] + [-1.0e9] * (num_beams - 1)), axis=0), [batch_size, 1] + ) + scores = tf.ones((batch_size, num_beams)) * -1.0e9 + + # per batch beam indices + running_beam_indices = tf.ones((batch_size, num_beams, max_length - decoder_prompt_len), dtype=tf.int32) * -1 + beam_indices = tf.ones((batch_size, num_beams, max_length - decoder_prompt_len), dtype=tf.int32) * -1 + + # flatten beam dim + if "encoder_outputs" in model_kwargs: + model_kwargs["encoder_outputs"]["last_hidden_state"] = flatten_beam_dim( + model_kwargs["encoder_outputs"]["last_hidden_state"] + ) + if "attention_mask" in model_kwargs: + model_kwargs["attention_mask"] = flatten_beam_dim(model_kwargs["attention_mask"]) + + # 4. define "xla-compile-able" stop-condition and auto-regressive function + # define stop-condition and auto-regressive function + def beam_search_cond_fn( + cur_len, + running_sequences, + running_scores, + running_beam_indices, + sequences, + scores, + beam_indices, + is_sent_finished, + decoder_prompt_len, + model_kwargs, + ): + """ + Beam Search termination condition function -- halts the generation loop if any of these conditions becomes + False + """ + # 1. is less than max length? + not_max_length_yet = cur_len < max_length + + # 2. can the new beams still improve? + # early_stopping == False -> apply heuristic = always get the best score from `cur_len - decoder_prompt_len`. See the discussion + # below for more details. + # https://github.com/huggingface/transformers/pull/20901#issuecomment-1369845565 + # early_stopping == "never" -> compute the best score from max_length or cur_len, depending on the sign of + # length_penalty. Positive length_penalty favors longer sequences, thus we use max_length there. + if early_stopping == "never" and length_penalty > 0.0: + best_running_score = running_scores[:, :1] / ((max_length - decoder_prompt_len) ** length_penalty) + else: + best_running_score = running_scores[:, :1] / ( + tf.cast(cur_len - decoder_prompt_len, dtype=tf.float32) ** length_penalty + ) + worst_finished_score = tf.where( + is_sent_finished, tf.math.reduce_min(scores, axis=1, keepdims=True), -1.0e9 + ) + improvement_still_possible = tf.math.reduce_any(best_running_score > worst_finished_score) + + # 3. is there still a beam that has not finished? + still_open_beam = ~(tf.math.reduce_all(is_sent_finished) & (early_stopping is True)) + + return not_max_length_yet & still_open_beam & improvement_still_possible + + def beam_search_body_fn( + cur_len, + running_sequences, + running_scores, + running_beam_indices, + sequences, + scores, + beam_indices, + is_sent_finished, + decoder_prompt_len, + model_kwargs, + ): + """ + Beam Search iterative update function -- each iteration adds a new token and updates the best sequences + seen so far + """ + # 1. Forward current tokens + if model_kwargs.get("past_key_values") is None or needs_full_input: + input_ids = running_sequences[:, :, :cur_len] + else: + input_ids = tf.expand_dims(running_sequences[:, :, cur_len - 1], -1) + model_inputs = self.prepare_inputs_for_generation( + flatten_beam_dim(input_ids), use_cache=use_cache, **model_kwargs + ) + model_outputs = self( + **model_inputs, + return_dict=True, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + logits = unflatten_beam_dim(model_outputs.logits[:, -1], num_beams) + + # 2. Compute log probs + # get log probabilities from logits, process logits with processors (*e.g.* min_length, ...), and + # add new logprobs to existing running logprobs scores. + log_probs = tf.nn.log_softmax(logits) + log_probs = logits_processor(flatten_beam_dim(running_sequences), flatten_beam_dim(log_probs), cur_len) + log_probs = unflatten_beam_dim(log_probs, num_beams) + if do_sample: + log_probs = logits_warper(flatten_beam_dim(running_sequences), flatten_beam_dim(log_probs), cur_len) + log_probs = unflatten_beam_dim(log_probs, num_beams) + log_probs_processed = log_probs + log_probs = log_probs + tf.expand_dims(running_scores, axis=2) + vocab_size = log_probs.shape[2] + log_probs = tf.reshape(log_probs, (batch_size, num_beams * vocab_size)) + + # Store scores, attentions and hidden_states when required + if not use_xla and return_dict_in_generate: + if output_scores: + all_scores.append( + logits_warper( + flatten_beam_dim(running_sequences), + flatten_beam_dim(log_probs_processed), + cur_len, + ) + ) + if output_attentions and self.config.is_encoder_decoder: + decoder_attentions.append(model_outputs.decoder_attentions) + elif output_attentions and not self.config.is_encoder_decoder: + decoder_attentions.append(model_outputs.attentions) + if self.config.is_encoder_decoder: + cross_attentions.append(model_outputs.cross_attentions) + + if output_hidden_states and self.config.is_encoder_decoder: + decoder_hidden_states.append(model_outputs.decoder_hidden_states) + elif output_hidden_states and self.config.is_encoder_decoder: + decoder_hidden_states.append(model_outputs.hidden_states) + + # 3. Retrieve top-K + # Each item in batch has num_beams * vocab_size candidate sequences. For each item, get the top 2*k + # candidates with the highest log-probabilities. We gather the top 2*K beams here so that even if the + # best K sequences reach EOS simultaneously, we have another K sequences remaining to continue the live + # beam search. + # Gather the top 2*K scores from _all_ beams. + # Gather 2*k top beams. + # Recover the beam index by floor division. + # Recover token id by modulo division and expand Id array for broadcasting. + # Update sequences for the 2*K top-k new sequences. + beams_to_keep = 2 * num_beams + if do_sample: + topk_indices = sample_without_replacement(log_probs, beams_to_keep) + topk_log_probs = tf.gather(log_probs, topk_indices, axis=1, batch_dims=1) + else: + topk_log_probs, topk_indices = tf.math.top_k(log_probs, k=beams_to_keep) + topk_current_beam_indices = topk_indices // vocab_size + topk_running_beam_indices = self._gather_beams(running_beam_indices, topk_current_beam_indices) + topk_running_sequences = self._gather_beams(running_sequences, topk_current_beam_indices) + topk_ids = topk_indices % vocab_size + + # writes the new token + indices_batch = tf.repeat(tf.range(batch_size), [beams_to_keep]) + indices_beam = tf.tile(tf.range(beams_to_keep), [batch_size]) + update_indices = tf.stack( + [indices_batch, indices_beam, tf.broadcast_to(cur_len, [batch_size * beams_to_keep])], axis=-1 + ) + topk_sequences = tf.tensor_scatter_nd_update( + tensor=topk_running_sequences, + indices=update_indices, + updates=tf.reshape(topk_ids, [batch_size * beams_to_keep]), + ) + + # we want to store the beam indices with batch information -> real beam index = beam index % num beams + batch_modified_indices = topk_current_beam_indices + tf.broadcast_to( + tf.expand_dims(tf.range(batch_size) * num_beams, axis=1), topk_current_beam_indices.shape + ) + update_indices = tf.stack( + [ + indices_batch, + indices_beam, + tf.broadcast_to(cur_len - decoder_prompt_len, [batch_size * beams_to_keep]), + ], + axis=-1, + ) + topk_beam_indices = tf.tensor_scatter_nd_update( + tensor=topk_running_beam_indices, + indices=update_indices, + updates=tf.reshape(batch_modified_indices, [batch_size * beams_to_keep]), + ) + + # 4. Check which sequences have ended + # Update current sequences: Did the top `num_beams` sequences reach an end marker? + # To prevent these just finished sequences from being added to the current sequences + # set of active beam search sequences, set their log probs to a very large negative value. + if eos_token_id is None: + eos_in_next_token = tf.zeros(topk_sequences[:, :, cur_len].shape, dtype=tf.bool) + else: + eos_in_next_token = tf.math.reduce_any( + tf.equal( + tf.broadcast_to( + topk_sequences[:, :, cur_len], + [len(eos_token_id)] + topk_sequences[:, :, cur_len].shape, + ), + tf.expand_dims(tf.expand_dims(eos_token_id, -1), -1), + ), + axis=0, + ) + did_topk_just_finished = eos_in_next_token & tf.broadcast_to( + tf.concat((tf.ones((num_beams), dtype=tf.bool), tf.zeros((num_beams), dtype=tf.bool)), axis=0), + shape_list(eos_in_next_token), + ) + + # non-top `num_beams` eos tokens can't be used to finish a beam, but the others can't be used in the next + # running sentences either + running_topk_log_probs = topk_log_probs + tf.cast(eos_in_next_token, tf.float32) * -1.0e9 + + # 5. Get running sequences scores for next + # Determine the top k beam indices (from top 2*k beams) from log probs and gather top k beams + # (from top 2*k beams). + next_topk_indices = tf.math.top_k(running_topk_log_probs, k=num_beams)[1] + next_running_sequences, next_running_scores, next_running_beam_indices = self._gather_beams( + [topk_sequences, running_topk_log_probs, topk_beam_indices], next_topk_indices + ) + + # 6. Process topk logits + # Further process log probs: + # - add length penalty + # - make sure no scores can be added anymore if beam is full + # - make sure still running sequences cannot be chosen as finalized beam + topk_log_probs = topk_log_probs / ( + tf.cast(cur_len + 1 - decoder_prompt_len, dtype=tf.float32) ** length_penalty + ) + beams_in_batch_are_full = tf.broadcast_to( + tf.math.reduce_all(is_sent_finished, axis=-1, keepdims=True), shape_list(did_topk_just_finished) + ) & (early_stopping is True) + add_penalty = ~did_topk_just_finished | beams_in_batch_are_full + topk_log_probs += tf.cast(add_penalty, tf.float32) * -1.0e9 + + # 7. Get scores, sequences, is sentence finished for next. + # Combine sequences, scores, and flags along the beam dimension and compare new finished sequence scores + # to existing finished scores and select the best from the new set of beams + merged_sequences = tf.concat([sequences, topk_sequences], axis=1) + merged_scores = tf.concat([scores, topk_log_probs], axis=1) + merged_beams = tf.concat([beam_indices, topk_beam_indices], axis=1) + merged_is_sent_finished = tf.concat([is_sent_finished, did_topk_just_finished], axis=1) + topk_merged_indices = tf.math.top_k(merged_scores, k=num_beams)[1] + next_sequences, next_scores, next_beam_indices, next_is_sent_finished = self._gather_beams( + [merged_sequences, merged_scores, merged_beams, merged_is_sent_finished], topk_merged_indices + ) + + # 8. Prepare data for the next iteration + # Determine the top k beam indices from the original set of all beams. With these, gather the top k + # beam-associated caches. + cur_len = cur_len + 1 + if "past_key_values" in model_outputs: + cache = tf.nest.map_structure( + lambda tensor: unflatten_beam_dim(tensor, num_beams, batch_axis=cache_batch_axis), + model_outputs.past_key_values, + ) + next_running_indices = self._gather_beams(topk_current_beam_indices, next_topk_indices) + next_cache = self._gather_beams(cache, next_running_indices, batch_axis=cache_batch_axis) + model_outputs["past_key_values"] = tf.nest.map_structure( + lambda tensor: flatten_beam_dim(tensor, batch_axis=cache_batch_axis), next_cache + ) + + if use_xla: + next_model_kwargs = self._update_model_kwargs_for_xla_generation( + model_outputs=model_outputs, + model_kwargs=model_kwargs, + cur_len=cur_len, + max_length=max_length, + batch_size=(batch_size * num_beams), + is_encoder_decoder=self.config.is_encoder_decoder, + batch_axis=cache_batch_axis, + ) + else: + next_model_kwargs = self._update_model_kwargs_for_generation( + model_outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder + ) + + # if we don't cache past_key_values key values we need the whole input + if model_kwargs.get("past_key_values", None) is None: + # let's throw out `past_key_values` since we don't want `None` tensors + model_kwargs.pop("past_key_values", None) + + return ( + cur_len, + next_running_sequences, + next_running_scores, + next_running_beam_indices, + next_sequences, + next_scores, + next_beam_indices, + next_is_sent_finished, + decoder_prompt_len, + next_model_kwargs, + ) + + # 5. run generation + # 1st generation step has to be run before to initialize `past_key_values` (if active) + ( + cur_len, + running_sequences, + running_scores, + running_beam_indices, + sequences, + scores, + beam_indices, + is_sent_finished, + decoder_prompt_len, + model_kwargs, + ) = beam_search_body_fn( + cur_len, + running_sequences, + running_scores, + running_beam_indices, + sequences, + scores, + beam_indices, + is_sent_finished, + decoder_prompt_len, + model_kwargs, + ) + + # 2-to-n generation steps can then be run in autoregressive fashion (only in case 1st generation step does + # NOT yield EOS token though) + maximum_iterations = max_length - cur_len + ( + cur_len, + running_sequences, + running_scores, + running_beam_indices, + sequences, + scores, + beam_indices, + is_sent_finished, + decoder_prompt_len, + _, + ) = tf.while_loop( + beam_search_cond_fn, + beam_search_body_fn, + ( + cur_len, + running_sequences, + running_scores, + running_beam_indices, + sequences, + scores, + beam_indices, + is_sent_finished, + decoder_prompt_len, + model_kwargs, + ), + maximum_iterations=maximum_iterations, + ) + + # 6. prepare outputs + # Account for the edge-case where there are no finished sequences for a particular batch item. If so, return + # running sequences for that batch item. + none_finished = tf.math.reduce_any(is_sent_finished, axis=1) + sequences = tf.where(none_finished[:, None, None], sequences, running_sequences) + beam_indices = tf.where(none_finished[:, None, None], beam_indices, running_beam_indices) + + # Apply the length penalty so that running scores match the finalized scores if they are used + running_scores = running_scores / (tf.cast(cur_len - decoder_prompt_len, dtype=tf.float32) ** length_penalty) + scores = tf.where(none_finished[:, None], scores, running_scores) + + # Take best beams for each batch (the score is sorted in descending order) + sequences = flatten_beam_dim(sequences[:, :num_return_sequences, :]) + scores = flatten_beam_dim(scores[:, :num_return_sequences]) + beam_indices = flatten_beam_dim(beam_indices[:, :num_return_sequences, :]) + + if not use_xla: + # Cut for backward compatibility + sequences = sequences[:, :cur_len] + beam_indices = beam_indices[:, : cur_len - decoder_prompt_len] + + if return_dict_in_generate: + if self.config.is_encoder_decoder: + # if model is an encoder-decoder, retrieve encoder attention weights and hidden states + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + output_cls = TFBeamSampleEncoderDecoderOutput if do_sample else TFBeamSearchEncoderDecoderOutput + return output_cls( + sequences=sequences, + sequences_scores=scores, + scores=all_scores, + beam_indices=beam_indices, + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + ) + else: + output_cls = TFBeamSampleDecoderOnlyOutput if do_sample else TFBeamSearchDecoderOnlyOutput + return output_cls( + sequences=sequences, + sequences_scores=scores, + scores=all_scores, + beam_indices=beam_indices, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + ) + else: + return sequences + + def contrastive_search( + self, + input_ids: tf.Tensor, + top_k: Optional[int] = 1, + penalty_alpha: Optional[float] = 0, + logits_processor: Optional[TFLogitsProcessorList] = None, + logits_warper: Optional[TFLogitsProcessorList] = None, + max_length: Optional[int] = None, + pad_token_id: Optional[int] = None, + eos_token_id: Optional[int] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_scores: Optional[bool] = None, + return_dict_in_generate: Optional[bool] = None, + **model_kwargs, + ) -> Union[TFContrastiveSearchOutput, tf.Tensor]: + r""" + Generates sequences of token ids for models with a language modeling head using **contrastive search** and can + be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. + + Parameters: + input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + top_k (`int`, *optional*, defaults to 1): + The size of the candidate set that is used to re-rank for contrastive search + penalty_alpha (`float`, *optional*, defaults to 0): + The degeneration penalty for contrastive search; activate when it is larger than 0 + logits_processor (`TFLogitsProcessorList`, *optional*): + An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + logits_warper (`TFLogitsProcessorList`, *optional*): + An instance of [`TFLogitsProcessorList`]. List of instances of class derived from [`TFLogitsWarper`] + used to warp the prediction score distribution of the language modeling head applied before multinomial + sampling at each generation step. + max_length (`int`, *optional*, defaults to 20): + The maximum length of the sequence to be generated. + pad_token_id (`int`, *optional*): + The id of the *padding* token. + eos_token_id (`Union[int, List[int]]`, *optional*): + The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens. + output_attentions (`bool`, *optional*, defaults to `False`): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more details. + output_hidden_states (`bool`, *optional*, defaults to `False`): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more details. + output_scores (`bool`, *optional*, defaults to `False`): + Whether or not to return the prediction scores. See `scores` under returned tensors for more details. + return_dict_in_generate (`bool`, *optional*, defaults to `False`): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + model_kwargs: + Additional model specific keyword arguments will be forwarded to the `call` function of the model. If + model is an encoder-decoder model the kwargs should include `encoder_outputs`. + Return: + [`~generation.TFContrastiveSearchDecoderOnlyOutput`], + [`~generation.TFContrastiveSearchEncoderDecoderOutput`] or `tf.Tensor`: A `tf.Tensor` containing the + generated tokens (default behaviour) or a [`~generation.TFContrastiveySearchDecoderOnlyOutput`] if + `model.config.is_encoder_decoder=False` and `return_dict_in_generate=True` or a + [`~generation.TFContrastiveSearchEncoderDecoderOutput`] if `model.config.is_encoder_decoder=True`. + Examples: + ```python + >>> from transformers import AutoTokenizer, TFAutoModelForCausalLM + + >>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m") + >>> model = TFAutoModelForCausalLM.from_pretrained("facebook/opt-125m") + >>> # set pad_token_id to eos_token_id because OPT does not have a PAD token + >>> model.config.pad_token_id = model.config.eos_token_id + >>> input_prompt = "DeepMind Company is" + >>> input_ids = tokenizer(input_prompt, return_tensors="tf") + >>> outputs = model.contrastive_search(**input_ids, penalty_alpha=0.6, top_k=4, max_length=64) + >>> tokenizer.batch_decode(outputs, skip_special_tokens=True) + ['DeepMind Company is a company that focuses on the development and commercialization of artificial intelligence (AI). DeepMind’s mission is to help people understand and solve problems that are difficult to solve in the world today.\n\nIn this post, we talk about the benefits of deep learning in business and how it'] + ```""" + + def gather_best_candidate(nested, selected_idx_stacked, batch_axis=0): + """Gathers the slices indexed by selected_idx_stacked from a potentially nested structure of tensors.""" + + def gather_fn(tensor): + gathered_tensor = tf.gather(params=tensor, indices=selected_idx_stacked, axis=batch_axis) + return gathered_tensor + + return tf.nest.map_structure(gather_fn, nested) + + # 1. init greedy_search values + logits_processor = logits_processor if logits_processor is not None else TFLogitsProcessorList() + logits_warper = logits_warper if logits_warper is not None else TFLogitsProcessorList() + max_length = max_length if max_length is not None else self.generation_config.max_length + pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id + eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id + if isinstance(eos_token_id, int): + eos_token_id = [eos_token_id] + output_scores = output_scores if output_scores is not None else self.generation_config.output_scores + output_attentions = ( + output_attentions if output_attentions is not None else self.generation_config.output_attentions + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states + ) + return_dict_in_generate = ( + return_dict_in_generate + if return_dict_in_generate is not None + else self.generation_config.return_dict_in_generate + ) + use_cache = True # In contrastive search, we always use cache + model_kwargs.pop("use_cache", None) + + use_xla = not tf.executing_eagerly() + # TODO (Joao): fix cache format or find programatic way to detect cache index + # GPT2 and other models has a slightly different cache structure, with a different batch axis + model_name = str(self.decoder) if "EncoderDecoder" in str(self) else str(self) + cache_batch_axis = 1 if any(model_prefix in model_name for model_prefix in ("TFGPT2", "TFCTRL")) else 0 + + # 2. init `attentions`, `hidden_states`, and `scores` tuples + scores = [] if (return_dict_in_generate and output_scores) else None + decoder_attentions = [] if (return_dict_in_generate and output_attentions) else None + cross_attentions = [] if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = [] if (return_dict_in_generate and output_hidden_states) else None + + # 3. init tensors to use for "xla-compileable" generate function + batch_size, cur_len = shape_list(input_ids) + + # initialize `generated` (`input_ids` padded with `pad_token_id`), `finished_sequences` + input_ids_padding = tf.ones((batch_size, max_length - cur_len), dtype=tf.int32) * (pad_token_id or 0) + generated = tf.concat([input_ids, input_ids_padding], axis=-1) + finished_sequences = tf.zeros((batch_size,), dtype=tf.bool) + + # 4. define "xla-compile-able" stop-condition and auto-regressive function + # define condition fn + def contrastive_search_cond_fn( + generated, finished_sequences, cur_len, model_kwargs, next_step_cached_variables + ): + """state termination condition fn.""" + return ~tf.reduce_all(finished_sequences) + + # define condition fn + def contrastive_search_body_fn( + generated, finished_sequences, cur_len, model_kwargs, next_step_cached_variables + ): + """state update fn.""" + + # if the first step in the loop, encode all the prefix and obtain: (1) past_key_values; + # (2) last_hidden_states; (3) logit_for_next_step; (4) update model kwargs for the next step + if model_kwargs.get("past_key_values") is None: + # prepare inputs + model_inputs = self.prepare_inputs_for_generation( + generated[:, :cur_len], use_cache=use_cache, **model_kwargs + ) + + # encode the given prefix and prepare model inputs; encoder-decoder model process the prefix and save + # the `encoder_outputs` + outputs = self( + **model_inputs, return_dict=True, output_hidden_states=True, output_attentions=output_attentions + ) + + # last decoder hidden states will be used to compute the degeneration penalty (cosine similarity with + # previous tokens) + if self.config.is_encoder_decoder: + last_hidden_states = outputs.decoder_hidden_states[-1] + else: + last_hidden_states = outputs.hidden_states[-1] + + # XLA: last_hidden_states normally grows at each step, but in XLA it is padded so as to be used across + # iterations (with fixed shapes) + if use_xla: + last_hidden_states = tf.pad(last_hidden_states, [[0, 0], [0, max_length - cur_len], [0, 0]]) + + # next logit for contrastive search to select top-k candidate tokens + logit_for_next_step = outputs.logits[:, -1, :] + + if use_xla: + model_kwargs = self._update_model_kwargs_for_xla_generation( + model_outputs=outputs, + model_kwargs=model_kwargs, + cur_len=cur_len, + max_length=max_length, + batch_size=batch_size, + is_encoder_decoder=self.config.is_encoder_decoder, + batch_axis=cache_batch_axis, + ) + else: + model_kwargs = self._update_model_kwargs_for_generation( + outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder + ) + + # Expands model inputs top_k times, for batched forward passes (akin to beam search). + _, model_kwargs = self._expand_inputs_for_generation( + expand_size=top_k, is_encoder_decoder=self.config.is_encoder_decoder, **model_kwargs + ) + + past_key_values = model_kwargs.get("past_key_values") + if past_key_values is None: + raise ValueError( + f"{self.__class__.__name__} does not support caching and therefore **can't** be used " + "for contrastive search." + ) + elif ( + not isinstance(past_key_values[0], (tuple, tf.Tensor)) + or past_key_values[0][0].shape[0] != batch_size + ): + raise ValueError( + f"{self.__class__.__name__} does not have a standard cache format and therefore **can't** be " + "used for contrastive search without further modifications." + ) + else: + logit_for_next_step = next_step_cached_variables["logit_for_next_step"] + last_hidden_states = next_step_cached_variables["last_hidden_states"] + outputs = next_step_cached_variables["outputs"] + + # contrastive_search main logic start: + # contrastive search decoding consists of two steps: (1) candidate tokens recall; (2) candidate re-rank by + # degeneration penalty + + logit_for_next_step = logits_processor(generated, logit_for_next_step, cur_len) + logit_for_next_step = logits_warper(generated, logit_for_next_step, cur_len) + next_probs = stable_softmax(logit_for_next_step, axis=-1) + top_k_probs, top_k_ids = tf.math.top_k(next_probs, k=top_k) + + # Store scores, attentions and hidden_states when required + if not use_xla and return_dict_in_generate: + if output_scores: + scores.append(logit_for_next_step) + if output_attentions and self.config.is_encoder_decoder: + decoder_attentions.append(outputs.decoder_attentions) + elif output_attentions and not self.config.is_encoder_decoder: + decoder_attentions.append(outputs.attentions) + if self.config.is_encoder_decoder: + cross_attentions.append(outputs.cross_attentions) + + if output_hidden_states and self.config.is_encoder_decoder: + decoder_hidden_states.append(outputs.decoder_hidden_states) + elif output_hidden_states and self.config.is_encoder_decoder: + decoder_hidden_states.append(outputs.hidden_states) + + # Replicates the new past_key_values to match the `top_k` candidates + model_kwargs["past_key_values"] = tf.nest.map_structure( + lambda tensor: tf.repeat(tensor, top_k, axis=cache_batch_axis), model_kwargs["past_key_values"] + ) + + # compute the candidate tokens by the language model and collects their hidden_states + next_model_inputs = self.prepare_inputs_for_generation( + tf.reshape(top_k_ids, [-1, 1]), use_cache=use_cache, **model_kwargs + ) + outputs = self( + **next_model_inputs, return_dict=True, output_hidden_states=True, output_attentions=output_attentions + ) + next_past_key_values = self._extract_past_from_model_output(outputs) + + logits = outputs.logits[:, -1, :] + # name is different for encoder-decoder and decoder-only models + if self.config.is_encoder_decoder: + next_hidden = outputs.decoder_hidden_states[-1] + full_hidden_states = outputs.decoder_hidden_states + else: + next_hidden = outputs.hidden_states[-1] + full_hidden_states = outputs.hidden_states + context_hidden = tf.repeat(last_hidden_states[:, :cur_len, :], top_k, axis=0) + + # compute the degeneration penalty and re-rank the candidates based on the degeneration penalty and the + # model confidence + selected_idx = _ranking_fast(context_hidden, next_hidden, top_k_probs, penalty_alpha, top_k) + + # converts indices to a dimension of top_k to the stacked top_k * batch_size dimension, for indexing + # without a need to reshape on tensors that have these two dimensions stacked + selected_idx_stacked = selected_idx + tf.range(selected_idx.shape[0], dtype=tf.int64) * top_k + + # prepare for the next step: (1) next token_id; (2) past_key_values; (3) last_hidden_states for computing + # the degeneration penalty; (4) logits for selecting next top-k candidates; (5) selected tokens scores + # (model confidence minus degeneration penalty); (6) decoder hidden_states + next_tokens = tf.gather(top_k_ids, selected_idx, axis=1, batch_dims=1) + next_hidden = gather_best_candidate(next_hidden, selected_idx_stacked) + + # XLA: last_hidden_states normally grows at each step, but in XLA it is padded so as to be used across + # iterations (with fixed shapes) + if use_xla: + last_hidden_states = dynamic_update_slice(last_hidden_states, next_hidden, [0, cur_len, 0]) + else: + last_hidden_states = tf.concat([last_hidden_states, next_hidden], axis=1) + + next_decoder_hidden_states = gather_best_candidate(full_hidden_states, selected_idx_stacked) + next_past_key_values = gather_best_candidate( + next_past_key_values, selected_idx_stacked, batch_axis=cache_batch_axis + ) + logit_for_next_step = gather_best_candidate(logits, selected_idx_stacked) + + # Rebuilds the relevant parts of the model output for the selected token, for use in the next iteration + if self.config.is_encoder_decoder: + next_step_cross_attentions = () + next_step_decoder_attentions = () + if output_attentions: + next_step_cross_attentions = gather_best_candidate(outputs.cross_attentions, selected_idx_stacked) + next_step_decoder_attentions = gather_best_candidate( + outputs.decoder_attentions, selected_idx_stacked + ) + outputs = TFSeq2SeqLMOutput( + past_key_values=next_past_key_values, + decoder_hidden_states=next_decoder_hidden_states, + decoder_attentions=next_step_decoder_attentions or None, + cross_attentions=next_step_cross_attentions or None, + ) + else: + next_step_attentions = () + if output_attentions: + next_step_attentions = gather_best_candidate(outputs.attentions, selected_idx_stacked) + outputs = TFCausalLMOutputWithPast( + past_key_values=next_past_key_values, + hidden_states=next_decoder_hidden_states, + attentions=next_step_attentions or None, + ) + # contrastive_search main logic end + + if eos_token_id is not None: + if pad_token_id is None: + raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.") + unfinished_seq = 1 - tf.cast(finished_sequences, tf.int32) + next_tokens = next_tokens * unfinished_seq + pad_token_id * (1 - unfinished_seq) + next_token_is_eos = tf.math.reduce_any( + tf.equal( + tf.broadcast_to(next_tokens, (len(eos_token_id), batch_size)), tf.expand_dims(eos_token_id, -1) + ), + axis=0, + ) + finished_sequences = finished_sequences | next_token_is_eos + + # update `generated` and `cur_len` + update_indices = tf.stack([tf.range(batch_size), tf.broadcast_to(cur_len, [batch_size])], axis=-1) + generated = tf.tensor_scatter_nd_update(tensor=generated, indices=update_indices, updates=next_tokens) + cur_len += 1 + + if use_xla: + # NOTE: 1) relative to other generation strategies, contrastive search is always running forward + # passes one step ahead -- hence the `cur_len=cur_len + 1`; 2) the attention mask here is expanded from + # [batch_size, ...] to [batch_size*top_k, ...] -- hence the `batch_size=batch_size * top_k` + model_kwargs = self._update_model_kwargs_for_xla_generation( + model_outputs=outputs, + model_kwargs=model_kwargs, + cur_len=cur_len + 1, + max_length=max_length, + batch_size=batch_size * top_k, + is_encoder_decoder=self.config.is_encoder_decoder, + batch_axis=cache_batch_axis, + ) + else: + model_kwargs = self._update_model_kwargs_for_generation( + outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder + ) + + next_step_cached_variables = { + "logit_for_next_step": logit_for_next_step, + "last_hidden_states": last_hidden_states, + "outputs": outputs, + } + return generated, finished_sequences, cur_len, model_kwargs, next_step_cached_variables + + # 5. run generation + # 1st generation step has to be run before to initialize `past_key_values` + generated, finished_sequences, cur_len, model_kwargs, next_step_cached_variables = contrastive_search_body_fn( + generated, finished_sequences, cur_len, model_kwargs, None + ) + + # 2-to-n generation steps can then be run in autoregressive fashion + # only in case 1st generation step does NOT yield EOS token though + maximum_iterations = max_length - cur_len + generated, _, cur_len, _, _ = tf.while_loop( + contrastive_search_cond_fn, + contrastive_search_body_fn, + (generated, finished_sequences, cur_len, model_kwargs, next_step_cached_variables), + maximum_iterations=maximum_iterations, + ) + + # 6. prepare outputs + if not use_xla: + # cut for backward compatibility + generated = generated[:, :cur_len] + + if return_dict_in_generate: + if self.config.is_encoder_decoder: + # if model is an encoder-decoder, retrieve encoder attention weights + # and hidden states + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + scores = tuple(scores) if scores is not None else None + decoder_attentions = tuple(decoder_attentions) if decoder_attentions is not None else None + cross_attentions = tuple(cross_attentions) if cross_attentions is not None else None + decoder_hidden_states = tuple(decoder_hidden_states) if decoder_hidden_states is not None else None + + return TFContrastiveSearchEncoderDecoderOutput( + sequences=generated, + scores=scores, + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + ) + else: + return TFContrastiveSearchDecoderOnlyOutput( + sequences=generated, + scores=scores, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + ) + else: + return generated + + +def scatter_values_on_batch_indices(values, batch_indices): + shape = shape_list(batch_indices) + # broadcast batch dim to shape + broad_casted_batch_dims = tf.reshape(tf.broadcast_to(tf.expand_dims(tf.range(shape[0]), axis=-1), shape), [1, -1]) + # transform batch_indices to pair_indices + pair_indices = tf.transpose(tf.concat([broad_casted_batch_dims, tf.reshape(batch_indices, [1, -1])], 0)) + # scatter values to pair indices + return tf.scatter_nd(pair_indices, tf.reshape(values, [-1]), shape) + + +def sample_without_replacement(logits, num_samples): + """ + categorical sampling without replacement is currently not implemented the gumbel-max trick will do for now see + https://github.com/tensorflow/tensorflow/issues/9260 for more info + """ + z = -tf.math.log(-tf.math.log(tf.random.uniform(shape_list(logits), 0, 1))) + _, indices = tf.nn.top_k(logits + z, num_samples) + return indices + + +def _ranking_fast( + context_hidden: tf.Tensor, + next_hidden: tf.Tensor, + next_top_k_probs: tf.Tensor, + alpha: float, + beam_width: int, +) -> tf.Tensor: + """ + Reranks the top_k candidates based on a degeneration penalty (cosine similarity with previous tokens), as described + in the paper "A Contrastive Framework for Neural Text Generation". Returns the index of the best candidate for each + row in the batch. + """ + norm_context_hidden = context_hidden / tf.norm(context_hidden, axis=2, keepdims=True) + norm_next_hidden = next_hidden / tf.norm(next_hidden, axis=2, keepdims=True) + cosine_matrix = tf.squeeze(tf.linalg.matmul(norm_context_hidden, norm_next_hidden, transpose_b=True), axis=-1) + degeneration_penalty = tf.reduce_max(cosine_matrix, axis=-1) + next_top_k_probs = tf.reshape(next_top_k_probs, shape=[-1]) + contrastive_score = (1.0 - alpha) * next_top_k_probs - alpha * degeneration_penalty + contrastive_score = tf.reshape(contrastive_score, shape=[-1, beam_width]) + selected_idx = tf.argmax(contrastive_score, axis=1) + return selected_idx diff --git a/lib/python3.12/site-packages/transformers/generation/utils.py b/lib/python3.12/site-packages/transformers/generation/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b68de89f66afb64ecc2d8d979e2ce315f70f76b8 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/utils.py @@ -0,0 +1,5245 @@ +# coding=utf-8 +# Copyright 2020 The Google AI Language Team Authors, Facebook AI Research authors and The HuggingFace Inc. team. +# Copyright (c) 2020, NVIDIA CORPORATION. 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. +import copy +import inspect +import os +import warnings +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.distributed as dist +from packaging import version +from torch import nn +from torch.nn import functional as F + +from transformers.generation.candidate_generator import AssistantVocabTranslatorCache + +from ..cache_utils import ( + Cache, + DynamicCache, + EncoderDecoderCache, + OffloadedCache, + QuantizedCacheConfig, + StaticCache, +) +from ..configuration_utils import PretrainedConfig +from ..integrations.deepspeed import is_deepspeed_zero3_enabled +from ..integrations.fsdp import is_fsdp_managed_module +from ..modeling_outputs import CausalLMOutputWithPast, Seq2SeqLMOutput +from ..pytorch_utils import isin_mps_friendly +from ..tokenization_utils import ExtensionsTrie +from ..utils import ( + ModelOutput, + is_accelerate_available, + is_hqq_available, + is_optimum_quanto_available, + is_torchdynamo_exporting, + logging, +) +from .beam_constraints import DisjunctiveConstraint, PhrasalConstraint +from .beam_search import BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer +from .candidate_generator import ( + AssistedCandidateGenerator, + AssistedCandidateGeneratorDifferentTokenizers, + CandidateGenerator, + EarlyExitCandidateGenerator, + PromptLookupCandidateGenerator, + UniversalSpeculativeDecodingGenerator, + _crop_past_key_values, + _prepare_attention_mask, + _prepare_token_type_ids, +) +from .configuration_utils import ( + NEED_SETUP_CACHE_CLASSES_MAPPING, + QUANT_BACKEND_CLASSES_MAPPING, + GenerationConfig, + GenerationMode, +) +from .logits_process import ( + EncoderNoRepeatNGramLogitsProcessor, + EncoderRepetitionPenaltyLogitsProcessor, + EpsilonLogitsWarper, + EtaLogitsWarper, + ExponentialDecayLengthPenalty, + ForcedBOSTokenLogitsProcessor, + ForcedEOSTokenLogitsProcessor, + HammingDiversityLogitsProcessor, + InfNanRemoveLogitsProcessor, + LogitNormalization, + LogitsProcessorList, + MinLengthLogitsProcessor, + MinNewTokensLengthLogitsProcessor, + MinPLogitsWarper, + NoBadWordsLogitsProcessor, + NoRepeatNGramLogitsProcessor, + PrefixConstrainedLogitsProcessor, + RepetitionPenaltyLogitsProcessor, + SequenceBiasLogitsProcessor, + SuppressTokensAtBeginLogitsProcessor, + SuppressTokensLogitsProcessor, + TemperatureLogitsWarper, + TopKLogitsWarper, + TopPLogitsWarper, + TypicalLogitsWarper, + UnbatchedClassifierFreeGuidanceLogitsProcessor, +) +from .stopping_criteria import ( + ConfidenceCriteria, + EosTokenCriteria, + MaxLengthCriteria, + MaxTimeCriteria, + StoppingCriteria, + StoppingCriteriaList, + StopStringCriteria, +) + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + from ..tokenization_utils_base import PreTrainedTokenizerBase + from .streamers import BaseStreamer + +logger = logging.get_logger(__name__) + +if is_accelerate_available(): + from accelerate.hooks import AlignDevicesHook, add_hook_to_module + + +# Variable names used to hold the cache at generation time +ALL_CACHE_NAMES = [ + "past_key_values", # default + "cache_params", # mamba-based models + "state", # rwkv + "mems", # xlnet + "past_buckets_states", # reformer +] + + +@dataclass +class GenerateDecoderOnlyOutput(ModelOutput): + """ + Outputs of decoder-only generation models, when using non-beam methods. + + Args: + sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`): + Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for + each generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`): + Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for + each generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`. + past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True`): + Returns the model cache, used to speed up decoding. Different models have a different cache format, check + the model's documentation. Usually, a [`~cache_utils.Cache`] instance. + """ + + sequences: torch.LongTensor + scores: Optional[Tuple[torch.FloatTensor]] = None + logits: Optional[Tuple[torch.FloatTensor]] = None + attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + past_key_values: Optional[Tuple[Tuple[Tuple[torch.FloatTensor]]]] = None + + +@dataclass +class GenerateEncoderDecoderOutput(ModelOutput): + """ + Outputs of encoder-decoder generation models, when using non-beam methods. + + Args: + sequences (`torch.LongTensor` of shape `(batch_size*num_return_sequences, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`): + Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for + each generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`): + Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for + each generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, + sequence_length, sequence_length)`. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of + shape `(batch_size, sequence_length, hidden_size)`. + decoder_attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + cross_attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + decoder_hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`. + past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Returns the model cache, used to speed up decoding. Different models have a different cache format, check + the model's documentation. Usually, a [`~cache_utils.Cache`] instance. + """ + + sequences: torch.LongTensor + scores: Optional[Tuple[torch.FloatTensor]] = None + logits: Optional[Tuple[torch.FloatTensor]] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None + decoder_attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + cross_attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + past_key_values: Optional[Tuple[Tuple[Tuple[torch.FloatTensor]]]] = None + + +@dataclass +class GenerateBeamDecoderOnlyOutput(ModelOutput): + """ + Outputs of decoder-only generation models, when using beam methods. + + Args: + sequences (`torch.LongTensor` of shape `(batch_size*num_return_sequences, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + sequences_scores (`torch.FloatTensor` of shape `(batch_size*num_return_sequences)`, *optional*, returned when `output_scores=True`): + Final beam scores of the generated `sequences`. + scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`): + Beam transition scores for each vocabulary token at each generation step. Beam transition scores consisting + of log probabilities of tokens conditioned on log softmax of previously generated tokens in this beam. + Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for each generated token), + with each tensor of shape `(batch_size*num_beams, config.vocab_size)`. + logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`): + Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for + each generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + beam_indices (`torch.LongTensor`, *optional*, returned when `output_scores=True`): + Beam indices of generated token id at each generation step. `torch.LongTensor` of shape + `(batch_size*num_return_sequences, sequence_length)`. + attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size*num_beams, num_heads, generated_length, sequence_length)`. + hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size*num_beams*num_return_sequences, generated_length, hidden_size)`. + past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True`): + Returns the model cache, used to speed up decoding. Different models have a different cache format, check + the model's documentation. Usually, a [`~cache_utils.Cache`] instance. + """ + + sequences: torch.LongTensor + sequences_scores: Optional[torch.FloatTensor] = None + scores: Optional[Tuple[torch.FloatTensor]] = None + logits: Optional[Tuple[torch.FloatTensor]] = None + beam_indices: Optional[torch.LongTensor] = None + attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + past_key_values: Optional[Tuple[Tuple[Tuple[torch.FloatTensor]]]] = None + + +@dataclass +class GenerateBeamEncoderDecoderOutput(ModelOutput): + """ + Outputs of encoder-decoder generation models, when using beam methods. + + Args: + sequences (`torch.LongTensor` of shape `(batch_size*num_return_sequences, sequence_length)`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter + if all batches finished early due to the `eos_token_id`. + sequences_scores (`torch.FloatTensor` of shape `(batch_size*num_return_sequences)`, *optional*, returned when `output_scores=True`): + Final beam scores of the generated `sequences`. + scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`): + Beam transition scores for each vocabulary token at each generation step. Beam transition scores consisting + of log probabilities of tokens conditioned on log softmax of previously generated tokens in this beam. + Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for each generated token), + with each tensor of shape `(batch_size*num_beams, config.vocab_size)`. + logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`): + Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax) + at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for + each generated token), with each tensor of shape `(batch_size, config.vocab_size)`. + beam_indices (`torch.LongTensor`, *optional*, returned when `output_scores=True`): + Beam indices of generated token id at each generation step. `torch.LongTensor` of shape + `(batch_size*num_return_sequences, sequence_length)`. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer of the decoder) of shape `(batch_size, num_heads, + sequence_length, sequence_length)`. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of + shape `(batch_size*num_beams*num_return_sequences, sequence_length, hidden_size)`. + decoder_attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size*num_beams*num_return_sequences, num_heads, generated_length, + sequence_length)`. + cross_attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`. + decoder_hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`): + Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of + `torch.FloatTensor` of shape `(batch_size*num_beams*num_return_sequences, generated_length, hidden_size)`. + past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True`): + Returns the model cache, used to speed up decoding. Different models have a different cache format, check + the model's documentation. Usually, a [`~cache_utils.Cache`] instance. + """ + + sequences: torch.LongTensor + sequences_scores: Optional[torch.FloatTensor] = None + scores: Optional[Tuple[torch.FloatTensor]] = None + logits: Optional[Tuple[torch.FloatTensor]] = None + beam_indices: Optional[torch.LongTensor] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None + decoder_attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + cross_attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + past_key_values: Optional[Tuple[Tuple[Tuple[torch.FloatTensor]]]] = None + + +# TODO (joao): remove the equivalent classes and typing shortcuts below in v5 +# Equivalent classes (kept for retrocompatibility purposes) +GreedySearchDecoderOnlyOutput = GenerateDecoderOnlyOutput +ContrastiveSearchDecoderOnlyOutput = GenerateDecoderOnlyOutput +SampleDecoderOnlyOutput = GenerateDecoderOnlyOutput + +ContrastiveSearchEncoderDecoderOutput = GenerateEncoderDecoderOutput +GreedySearchEncoderDecoderOutput = GenerateEncoderDecoderOutput +SampleEncoderDecoderOutput = GenerateEncoderDecoderOutput + +BeamSearchDecoderOnlyOutput = GenerateBeamDecoderOnlyOutput +BeamSampleDecoderOnlyOutput = GenerateBeamDecoderOnlyOutput + +BeamSearchEncoderDecoderOutput = GenerateBeamEncoderDecoderOutput +BeamSampleEncoderDecoderOutput = GenerateBeamEncoderDecoderOutput + +GreedySearchOutput = Union[GreedySearchEncoderDecoderOutput, GreedySearchDecoderOnlyOutput] +SampleOutput = Union[SampleEncoderDecoderOutput, SampleDecoderOnlyOutput] +BeamSearchOutput = Union[BeamSearchEncoderDecoderOutput, BeamSearchDecoderOnlyOutput] +BeamSampleOutput = Union[BeamSampleEncoderDecoderOutput, BeamSampleDecoderOnlyOutput] +ContrastiveSearchOutput = Union[ContrastiveSearchEncoderDecoderOutput, ContrastiveSearchDecoderOnlyOutput] + +# Typing shortcuts +GenerateNonBeamOutput = Union[GenerateDecoderOnlyOutput, GenerateEncoderDecoderOutput] +GenerateBeamOutput = Union[GenerateBeamDecoderOnlyOutput, GenerateBeamEncoderDecoderOutput] +GenerateOutput = Union[GenerateNonBeamOutput, GenerateBeamOutput] + + +class GenerationMixin: + """ + A class containing all functions for auto-regressive text generation, to be used as a mixin in model classes. + Inheriting from this class causes the model to have special generation-related behavior, such as loading a + `GenerationConfig` at initialization time or ensuring `generate`-related tests are run in `transformers` CI. + + A model class should inherit from `GenerationMixin` to enable calling methods like `generate`, or when it + has defined a custom `generate` method that relies on `GenerationMixin`, directly or indirectly, which + approximately shares the same interface to public methods like `generate`. Three examples: + - `LlamaForCausalLM` should inherit from `GenerationMixin` to enable calling `generate` and other public + methods in the mixin; + - `BlipForQuestionAnswering` has a custom `generate` method that approximately shares the same interface as + `GenerationMixin.generate` (it has a few extra arguments, and the same output). That function also calls + `GenerationMixin.generate` indirectly, through an inner model. As such, `BlipForQuestionAnswering` should + inherit from `GenerationMixin` to benefit from all generation-related automation in our codebase; + - `BarkModel` has a custom `generate` method and one of its inner models calls `GenerationMixin.generate`. + However, its `generate` does not share the same interface as `GenerationMixin.generate`. In this case, + `BarkModel` shoud NOT inherit from `GenerationMixin`, as it breaks the `generate` interface. + + The class exposes [`~generation.GenerationMixin.generate`], which can be used for: + - *greedy decoding* if `num_beams=1` and `do_sample=False` + - *contrastive search* if `penalty_alpha>0` and `top_k>1` + - *multinomial sampling* if `num_beams=1` and `do_sample=True` + - *beam-search decoding* if `num_beams>1` and `do_sample=False` + - *beam-search multinomial sampling* if `num_beams>1` and `do_sample=True` + - *diverse beam-search decoding* if `num_beams>1` and `num_beam_groups>1` + - *constrained beam-search decoding* if `constraints!=None` or `force_words_ids!=None` + - *assisted decoding* if `assistant_model` or `prompt_lookup_num_tokens` is passed to `.generate()` + + To learn more about decoding strategies refer to the [text generation strategies guide](../generation_strategies). + """ + + def _cache_dependant_input_preparation( + self, + input_ids: torch.LongTensor, + inputs_embeds: Optional[torch.FloatTensor], + cache_position: Optional[torch.LongTensor], + ) -> Tuple[torch.FloatTensor, torch.LongTensor]: + """ + Generic cache-dependent input preparation + The code is put in a separate function to allow granular unit testing + as it needs a different implementation to be exportable. + + If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens + - Exception 1: when passing input_embeds, input_ids may be missing entries + - Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here + - Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case. + - Excpetion 4: If input_embeds are passed then slice it through `cache_position`, to keep only the unprocessed tokens and + generate the first token for each sequence. Later use the generated Input ids for continuation. + + The current implementation does not rely on ``self`` and could be + a class method. It is left as a standard method to be easily rewritten. + """ + if is_torchdynamo_exporting(): + return self._cache_dependant_input_preparation_exporting(input_ids, inputs_embeds, cache_position) + if inputs_embeds is not None and input_ids.shape[1] == 0: # Exception 4 + inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] + elif ( + inputs_embeds is not None # Exception 1 + or (cache_position[-1] >= input_ids.shape[1]) # Exception 3 + ): + input_ids = input_ids[:, -cache_position.shape[0] :] + elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) + input_ids = input_ids[:, cache_position] + return inputs_embeds, input_ids + + def _cache_dependant_input_preparation_exporting( + self, + input_ids: torch.LongTensor, + inputs_embeds: Optional[torch.FloatTensor], + cache_position: Optional[torch.LongTensor], + ) -> Tuple[torch.FloatTensor, torch.LongTensor]: + """ + This method implements method ``_cache_dependant_input_preparation`` + with :func:`torch.cond` to make it exportable with :func:`torch.export.export`. + The code is put in a separate function to allow granular unit testing. + """ + if inputs_embeds is None: + input_ids = input_ids[:, cache_position] + else: + # This is the code we need to implemented with torch.cond. + # if input_ids.shape[1] == 0: + # inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :] + # else: + # if cache_position[-1] >= input_ids.shape[1]: + # input_ids = input_ids[:, -cache_position.shape[0] :] + # else: + # if input_ids.shape[1] != cache_position.shape[0]: + # input_ids = input_ids[:, cache_position] + def branch_1(inputs_embeds, cache_position): + return inputs_embeds[:, -cache_position.shape[0] :] + + def branch_2(input_ids, cache_position): + return input_ids[:, -cache_position.shape[0] :] + + def branch_3(input_ids, cache_position): + return input_ids[:, cache_position] + + inputs_embeds, input_ids = torch.cond( + input_ids.shape[1] == 0, + ( + lambda input_ids, inputs_embeds, cache_position: ( + branch_1(inputs_embeds, cache_position), + input_ids, + ) + ), + ( + lambda input_ids, inputs_embeds, cache_position: ( + inputs_embeds, + torch.cond( + cache_position[-1] >= input_ids.shape[1], + branch_2, + lambda input_ids, cache_position: ( + torch.cond( + input_ids.shape[1] != cache_position.shape[0], + branch_3, + (lambda input_ids, cache_position: input_ids), + [input_ids, cache_position], + ) + ), + [input_ids, cache_position], + ), + ) + ), + [input_ids, inputs_embeds, cache_position], + ) + return inputs_embeds, input_ids + + def prepare_inputs_for_generation( + self, + input_ids: torch.LongTensor, + past_key_values: Optional[Cache] = None, + attention_mask: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ): + """ + Prepare the model inputs for generation. In includes operations like computing the 4D attention mask or + slicing inputs given the existing cache. + + See the forward pass in the model documentation for expected arguments (different models might have different + requirements for e.g. `past_key_values`). This function should work as is for most LLMs. + """ + + # 1. Handle BC: + model_inputs = {} + # - some models don't have `Cache` support (which implies they don't expect `cache_position` in `forward`) + if self._supports_cache_class: + model_inputs["cache_position"] = cache_position + # - `cache_position` was not a mandatory input in `prepare_inputs_for_generation` for those models, and this + # function may be called outside of `generate`. Handle most use cases by creating `cache_position` on the fly + # (this alternative is not as robust as calling `generate` and letting it create `cache_position`) + elif cache_position is None: + past_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 + cache_position = torch.arange(past_length, input_ids.shape[1], dtype=torch.long, device=input_ids.device) + + # 2. Generic cache-dependent input preparation + if past_key_values is not None: + model_inputs["past_key_values"] = past_key_values + inputs_embeds, input_ids = self._cache_dependant_input_preparation( + input_ids, inputs_embeds, cache_position + ) + + # 3. Prepare base model inputs + input_ids_key = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids" + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step for every prompt. + if not self.config.is_encoder_decoder: + if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]: + model_inputs[input_ids_key] = None + model_inputs["inputs_embeds"] = inputs_embeds + else: + # `clone` calls in this function ensure a consistent stride. See #32227 + model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format) + model_inputs["inputs_embeds"] = None + else: + model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format) + + # 4. Create missing `position_ids` on the fly + encoder_attention_mask = attention_mask if self.config.is_encoder_decoder else None + attention_mask = ( + kwargs.pop("decoder_attention_mask", None) if self.config.is_encoder_decoder else attention_mask + ) + attention_mask_key = "decoder_attention_mask" if self.config.is_encoder_decoder else "attention_mask" + position_ids_key = "decoder_position_ids" if self.config.is_encoder_decoder else "position_ids" + if ( + attention_mask is not None + and kwargs.get(position_ids_key) is None + and position_ids_key in set(inspect.signature(self.forward).parameters.keys()) + ): + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + kwargs[position_ids_key] = position_ids # placed in kwargs for further processing (see below) + + # 5. Slice model inputs if it's an input that should have the same length as `input_ids` + for model_input_name in ["position_ids", "token_type_ids", "decoder_position_ids"]: + model_input = kwargs.get(model_input_name) + if model_input is not None: + if past_key_values is not None: + current_input_length = ( + model_inputs["inputs_embeds"].shape[1] + if model_inputs.get("inputs_embeds") is not None + else model_inputs[input_ids_key].shape[1] + ) + model_input = model_input[:, -current_input_length:] + model_input = model_input.clone(memory_format=torch.contiguous_format) + model_inputs[model_input_name] = model_input + + # 6. Create 4D attention mask is we are using a `StaticCache` (important for performant compiled forward pass) + if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2: + if model_inputs["inputs_embeds"] is not None: + batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape + device = model_inputs["inputs_embeds"].device + else: + batch_size, sequence_length = model_inputs[input_ids_key].shape + device = model_inputs[input_ids_key].device + + # Create the causal mask with fixed shape in advance, to reduce recompilations. If the function to create + # the 4D causal mask exists, it should be present in the base model (XXXModel class). + base_model = getattr(self, self.base_model_prefix, None) + if base_model is None: + causal_mask_creation_function = getattr( + self, "_prepare_4d_causal_attention_mask_with_cache_position", None + ) + else: + causal_mask_creation_function = getattr( + base_model, "_prepare_4d_causal_attention_mask_with_cache_position", None + ) + if causal_mask_creation_function is None: + logger.warning_once( + f"{self.__class__.__name__} has no `_prepare_4d_causal_attention_mask_with_cache_position` method " + "defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're " + "writing code, see Llama for an example implementation. If you're a user, please report this " + "issue on GitHub." + ) + else: + attention_mask = causal_mask_creation_function( + attention_mask, + sequence_length=sequence_length, + target_length=past_key_values.get_max_cache_shape(), + dtype=self.dtype, + device=device, + cache_position=cache_position, + batch_size=batch_size, + config=self.config, + past_key_values=past_key_values, + ) + if attention_mask is not None: + model_inputs[attention_mask_key] = attention_mask + + if encoder_attention_mask is not None: + model_inputs["attention_mask"] = encoder_attention_mask + + # 7. Forward ALL kwargs that are uninitialized (e.g. `use_cache`). + for key, value in kwargs.items(): + if key not in model_inputs: + model_inputs[key] = value + + # 8. Remove unexpected `generate` inputs (TODO @joao: fix trainer and examples) + model_inputs.pop("labels", None) + return model_inputs + + def _prepare_model_inputs( + self, + inputs: Optional[torch.Tensor] = None, + bos_token_id: Optional[torch.Tensor] = None, + model_kwargs: Optional[Dict[str, torch.Tensor]] = None, + ) -> Tuple[torch.Tensor, Optional[str], Dict[str, torch.Tensor]]: + """ + This function extracts the model-specific `inputs` for generation. + """ + # 1. retrieve all kwargs that are non-None or non-model input related. + # some encoder-decoder models have different names for model and encoder + if ( + self.config.is_encoder_decoder + and hasattr(self, "encoder") + and self.encoder.main_input_name != self.main_input_name + ): + input_name = self.encoder.main_input_name + else: + input_name = self.main_input_name + + model_kwargs = {k: v for k, v in model_kwargs.items() if v is not None or k != input_name} + + # 2. check whether model_input_name is passed as kwarg + # if yes and `inputs` is None use kwarg inputs + inputs_kwarg = model_kwargs.pop(input_name, None) + if inputs_kwarg is not None and inputs is not None: + raise ValueError( + f"`inputs`: {inputs}` were passed alongside {input_name} which is not allowed. " + f"Make sure to either pass {inputs} or {input_name}=..." + ) + elif inputs_kwarg is not None: + inputs = inputs_kwarg + + # 3. In the presence of `inputs_embeds` for text models: + # - decoder-only models should complain if the user attempts to pass `inputs_embeds`, but the model + # doesn't have its forwarding implemented. `inputs_embeds` is kept in `model_kwargs` and can coexist with + # input_ids (`inputs_embeds` will be used in the 1st generation step, as opposed to `input_ids`) + # - encoder-decoder models should complain if the user attempts to pass `inputs_embeds` and `input_ids`, and + # pull the former to inputs. It will be used in place of `input_ids` to get the encoder hidden states. + if input_name == "input_ids" and "inputs_embeds" in model_kwargs: + if not self.config.is_encoder_decoder: + has_inputs_embeds_forwarding = "inputs_embeds" in set( + inspect.signature(self.prepare_inputs_for_generation).parameters.keys() + ) + if not has_inputs_embeds_forwarding: + raise ValueError( + f"You passed `inputs_embeds` to `.generate()`, but the model class {self.__class__.__name__} " + "doesn't have its forwarding implemented. See the GPT2 implementation for an example " + "(https://github.com/huggingface/transformers/pull/21405), and feel free to open a PR with it!" + ) + # In this case, `input_ids` is moved to the `model_kwargs`, so a few automations (like the creation of + # the attention mask) can rely on the actual model input. + model_kwargs["input_ids"] = self._maybe_initialize_input_ids_for_generation( + inputs, bos_token_id, model_kwargs=model_kwargs + ) + else: + if inputs is not None: + raise ValueError("You passed `inputs_embeds` and `input_ids` to `.generate()`. Please pick one.") + inputs, input_name = model_kwargs["inputs_embeds"], "inputs_embeds" + + # 4. if `inputs` is still None, try to create `input_ids` from BOS token + inputs = self._maybe_initialize_input_ids_for_generation(inputs, bos_token_id, model_kwargs) + return inputs, input_name, model_kwargs + + def _maybe_initialize_input_ids_for_generation( + self, + inputs: Optional[torch.Tensor] = None, + bos_token_id: Optional[torch.Tensor] = None, + model_kwargs: Optional[Dict[str, torch.Tensor]] = None, + ) -> torch.LongTensor: + """Initializes input ids for generation, if necessary.""" + if inputs is not None: + return inputs + + encoder_outputs = model_kwargs.get("encoder_outputs") + if self.config.is_encoder_decoder and encoder_outputs is not None: + # make dummy input_ids with value -100, as a sanity check ensuring that they won't be used for encoding + shape = encoder_outputs.last_hidden_state.size()[:-1] + return torch.ones(shape, dtype=torch.long, device=self.device) * -100 + + # If there is some tensor in `model_kwargs`, we can infer the batch size from it. This is helpful with + # soft-prompting or in multimodal implementations built on top of decoder-only language models. + batch_size = 1 + for value in model_kwargs.values(): + if isinstance(value, torch.Tensor): + batch_size = value.shape[0] + break + + if "inputs_embeds" in model_kwargs: + return torch.ones((batch_size, 0), dtype=torch.long, device=self.device) + + if bos_token_id is None: + raise ValueError("`bos_token_id` has to be defined when no `input_ids` are provided.") + + return torch.ones((batch_size, 1), dtype=torch.long, device=self.device) * bos_token_id + + def _prepare_attention_mask_for_generation( + self, + inputs_tensor: torch.Tensor, + generation_config: GenerationConfig, + model_kwargs: Dict[str, Any], + ) -> torch.LongTensor: + pad_token_id = generation_config._pad_token_tensor + eos_token_id = generation_config._eos_token_tensor + + # `input_ids` may be present in the model kwargs, instead of being the main input (e.g. multimodal model) + if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0: + inputs_tensor = model_kwargs["input_ids"] + + # No information for attention mask inference -> return default attention mask + default_attention_mask = torch.ones(inputs_tensor.shape[:2], dtype=torch.long, device=inputs_tensor.device) + if pad_token_id is None: + return default_attention_mask + + is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long] + if not is_input_ids: + return default_attention_mask + + is_pad_token_in_inputs = (pad_token_id is not None) and ( + isin_mps_friendly(elements=inputs_tensor, test_elements=pad_token_id).any() + ) + is_pad_token_not_equal_to_eos_token_id = (eos_token_id is None) or ~( + isin_mps_friendly(elements=eos_token_id, test_elements=pad_token_id).any() + ) + can_infer_attention_mask = is_pad_token_in_inputs * is_pad_token_not_equal_to_eos_token_id + attention_mask_from_padding = inputs_tensor.ne(pad_token_id).long() + + attention_mask = ( + attention_mask_from_padding * can_infer_attention_mask + default_attention_mask * ~can_infer_attention_mask + ) + return attention_mask + + def _prepare_encoder_decoder_kwargs_for_generation( + self, + inputs_tensor: torch.Tensor, + model_kwargs, + model_input_name: Optional[str], + generation_config: GenerationConfig, + ) -> Dict[str, Any]: + # 1. get encoder + encoder = self.get_encoder() + # Compatibility with Accelerate big model inference: we need the encoder to outputs stuff on the same device + # as the inputs. + if hasattr(self, "hf_device_map"): + if hasattr(encoder, "_hf_hook"): + encoder._hf_hook.io_same_device = True + else: + add_hook_to_module(encoder, AlignDevicesHook(io_same_device=True)) + + # 2. Prepare encoder args and encoder kwargs from model kwargs and generation config. + irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"] + encoder_kwargs = { + argument: value + for argument, value in model_kwargs.items() + if not any(argument.startswith(p) for p in irrelevant_prefix) + } + encoder_signature = set(inspect.signature(encoder.forward).parameters) + encoder_accepts_wildcard = "kwargs" in encoder_signature or "model_kwargs" in encoder_signature + if not encoder_accepts_wildcard: + encoder_kwargs = { + argument: value for argument, value in encoder_kwargs.items() if argument in encoder_signature + } + encoder_kwargs["output_attentions"] = generation_config.output_attentions + encoder_kwargs["output_hidden_states"] = generation_config.output_hidden_states + + # 3. make sure that encoder returns `ModelOutput` + model_input_name = model_input_name if model_input_name is not None else self.main_input_name + encoder_kwargs["return_dict"] = True + encoder_kwargs[model_input_name] = inputs_tensor + model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs) # type: ignore + + return model_kwargs + + def _prepare_decoder_input_ids_for_generation( + self, + batch_size: int, + model_input_name: str, + model_kwargs: Dict[str, torch.Tensor], + decoder_start_token_id: torch.Tensor, + device: Optional[torch.device] = None, + ) -> Tuple[torch.LongTensor, Dict[str, torch.Tensor]]: + """Prepares `decoder_input_ids` for generation with encoder-decoder models""" + # 1. Check whether the user has defined `decoder_input_ids` manually. To facilitate in terms of input naming, + # we also allow the user to pass it under `input_ids`, if the encoder does not use it as the main input. + if model_kwargs is not None and "decoder_input_ids" in model_kwargs: + decoder_input_ids = model_kwargs.pop("decoder_input_ids") + elif "input_ids" in model_kwargs and model_input_name != "input_ids": + decoder_input_ids = model_kwargs.pop("input_ids") + else: + decoder_input_ids = None + + # 2. `decoder_start_token_id` must have shape (batch_size, 1) + if device is None: + device = self.device + if decoder_start_token_id.ndim == 1: + if decoder_start_token_id.shape[0] != batch_size: + raise ValueError( + f"`decoder_start_token_id` expected to have length {batch_size} but got {decoder_start_token_id.shape[0]}" + ) + decoder_start_token_id = decoder_start_token_id.view(-1, 1) + else: + decoder_start_token_id = ( + torch.ones((batch_size, 1), dtype=torch.long, device=device) * decoder_start_token_id + ) + + # 3. Encoder-decoder models expect the `decoder_input_ids` to start with a special token. Let's ensure that. + # no user input -> use decoder_start_token_id as decoder_input_ids + if decoder_input_ids is None: + decoder_input_ids = decoder_start_token_id + # exception: Donut checkpoints have task-specific decoder starts and don't expect a BOS token. Note that the + # original checkpoints can't be detected through `self.__class__.__name__.lower()`, needing custom logic. + # See: https://github.com/huggingface/transformers/pull/31470 + elif "donut" in self.__class__.__name__.lower() or ( + self.config.model_type == "vision-encoder-decoder" and "donut" in self.config.encoder.model_type.lower() + ): + pass + elif self.config.model_type in ["whisper"]: + pass + # user input but doesn't start with decoder_start_token_id -> prepend decoder_start_token_id (and adjust + # decoder_attention_mask if provided) + elif (decoder_input_ids[:, 0] != decoder_start_token_id[:, 0]).all().item(): + decoder_input_ids = torch.cat([decoder_start_token_id, decoder_input_ids], dim=-1) + if "decoder_attention_mask" in model_kwargs: + decoder_attention_mask = model_kwargs["decoder_attention_mask"] + decoder_attention_mask = torch.cat( + (torch.ones_like(decoder_attention_mask)[:, :1], decoder_attention_mask), + dim=-1, + ) + model_kwargs["decoder_attention_mask"] = decoder_attention_mask + + return decoder_input_ids, model_kwargs + + @staticmethod + def _expand_inputs_for_generation( + expand_size: int = 1, + is_encoder_decoder: bool = False, + input_ids: Optional[torch.LongTensor] = None, + **model_kwargs, + ) -> Tuple[torch.LongTensor, Dict[str, Any]]: + """Expands tensors from [batch_size, ...] to [batch_size * expand_size, ...]""" + # Do not call torch.repeat_interleave if expand_size is 1 because it clones + # the input tensor and thus requires more memory although no change is applied + if expand_size == 1: + return input_ids, model_kwargs + + def _expand_dict_for_generation(dict_to_expand): + for key in dict_to_expand: + if ( + key != "cache_position" + and dict_to_expand[key] is not None + and isinstance(dict_to_expand[key], torch.Tensor) + ): + dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0) + return dict_to_expand + + if input_ids is not None: + input_ids = input_ids.repeat_interleave(expand_size, dim=0) + + model_kwargs = _expand_dict_for_generation(model_kwargs) + + if is_encoder_decoder: + if model_kwargs.get("encoder_outputs") is None: + raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.") + model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"]) + + return input_ids, model_kwargs + + def _update_model_kwargs_for_generation( + self, + outputs: ModelOutput, + model_kwargs: Dict[str, Any], + is_encoder_decoder: bool = False, + num_new_tokens: int = 1, + ) -> Dict[str, Any]: + # update past_key_values keeping its naming used in model code + for possible_cache_name in ALL_CACHE_NAMES: + if possible_cache_name in outputs: + # TODO (joao): remove output/input mismatch when these old models (xlnet, reformer) are deprecated + if possible_cache_name in ("past_buckets_states", "mems"): + cache_name = "past_key_values" + else: + cache_name = possible_cache_name + model_kwargs[cache_name] = getattr(outputs, possible_cache_name) + break + + # update token_type_ids with last value + if "token_type_ids" in model_kwargs: + token_type_ids = model_kwargs["token_type_ids"] + model_kwargs["token_type_ids"] = torch.cat([token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1) + + if not is_encoder_decoder: + # update attention mask + if "attention_mask" in model_kwargs: + attention_mask = model_kwargs["attention_mask"] + model_kwargs["attention_mask"] = torch.cat( + [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1 + ) + else: + # update decoder attention mask + if "decoder_attention_mask" in model_kwargs: + decoder_attention_mask = model_kwargs["decoder_attention_mask"] + model_kwargs["decoder_attention_mask"] = torch.cat( + [decoder_attention_mask, decoder_attention_mask.new_ones((decoder_attention_mask.shape[0], 1))], + dim=-1, + ) + + if model_kwargs.get("use_cache", True): + model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens + else: + past_positions = model_kwargs.pop("cache_position") + new_positions = torch.arange( + past_positions[-1] + 1, past_positions[-1] + num_new_tokens + 1, dtype=past_positions.dtype + ).to(past_positions.device) + model_kwargs["cache_position"] = torch.cat((past_positions, new_positions)) + return model_kwargs + + def _reorder_cache(self, past_key_values, beam_idx): + raise NotImplementedError( + f"Make sure that a `_reorder_cache` function is correctly implemented in {self.__class__.__module__} to" + f" enable beam search for {self.__class__}" + ) + + def _get_candidate_generator( + self, + generation_config: GenerationConfig, + input_ids: torch.LongTensor, + inputs_tensor: torch.Tensor, + assistant_model: "PreTrainedModel", + logits_processor: LogitsProcessorList, + target_tokenizer: "PreTrainedTokenizerBase", + assistant_tokenizer: "PreTrainedTokenizerBase", + model_kwargs: Dict, + ) -> CandidateGenerator: + """ + Returns the candidate generator to be used in `assisted_generation` + """ + different_tokenizers = all(v is not None for v in (assistant_model, target_tokenizer, assistant_tokenizer)) + + if generation_config.assistant_early_exit is not None: + candidate_generator = EarlyExitCandidateGenerator( + input_ids=input_ids, + assistant_model=self, + generation_config=generation_config, + model_kwargs=model_kwargs, + inputs_tensor=inputs_tensor, + logits_processor=logits_processor, + ) + elif generation_config.prompt_lookup_num_tokens is not None: + candidate_generator = PromptLookupCandidateGenerator( + eos_token_id=generation_config._eos_token_tensor, + num_output_tokens=generation_config.prompt_lookup_num_tokens, + max_matching_ngram_size=generation_config.max_matching_ngram_size, + max_length=generation_config.max_length, + ) + elif different_tokenizers: + if generation_config.do_sample is True: + atm_translator = AssistantVocabTranslatorCache.get_translator( + target_tokenizer, assistant_tokenizer, self.config.vocab_size, assistant_model.device + ) + candidate_generator = UniversalSpeculativeDecodingGenerator( + input_ids=input_ids, + assistant_model=assistant_model, + generation_config=generation_config, + model_kwargs=model_kwargs, + inputs_tensor=inputs_tensor, + logits_processor=logits_processor, + target_tokenizer=target_tokenizer, + assistant_tokenizer=assistant_tokenizer, + atm_translator=atm_translator, + ) + elif generation_config.do_sample is False: + candidate_generator = AssistedCandidateGeneratorDifferentTokenizers( + input_ids=input_ids, + assistant_model=assistant_model, + generation_config=generation_config, + model_kwargs=model_kwargs, + inputs_tensor=inputs_tensor, + logits_processor=logits_processor, + target_tokenizer=target_tokenizer, + assistant_tokenizer=assistant_tokenizer, + ) + else: + raise ValueError( + f"Invalid value for `do_sample`: expected a boolean, got {type(generation_config.do_sample).__name__}" + ) + else: + candidate_generator = AssistedCandidateGenerator( + input_ids=input_ids, + assistant_model=assistant_model, + generation_config=generation_config, + model_kwargs=model_kwargs, + inputs_tensor=inputs_tensor, + logits_processor=logits_processor, + ) + return candidate_generator + + def _get_logits_processor( + self, + generation_config: GenerationConfig, + input_ids_seq_length: int, + encoder_input_ids: torch.LongTensor, + prefix_allowed_tokens_fn: Callable[[int, torch.Tensor], List[int]], + logits_processor: Optional[LogitsProcessorList], + device: Optional[str] = None, + model_kwargs: Optional[Dict[str, Any]] = None, + negative_prompt_ids: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + ) -> LogitsProcessorList: + """ + This class returns a [`LogitsProcessorList`] list object that contains all relevant [`LogitsProcessor`] + instances used to modify the scores of the language model head. + """ + # instantiate processors list + processors = LogitsProcessorList() + + if generation_config.guidance_scale is not None and generation_config.guidance_scale != 1: + processors.append( + UnbatchedClassifierFreeGuidanceLogitsProcessor( + generation_config.guidance_scale, + self, + unconditional_ids=negative_prompt_ids, + unconditional_attention_mask=negative_prompt_attention_mask, + use_cache=generation_config.use_cache, + ) + ) + if generation_config.sequence_bias is not None: + processors.append(SequenceBiasLogitsProcessor(sequence_bias=generation_config.sequence_bias)) + + if generation_config.diversity_penalty is not None and generation_config.diversity_penalty > 0.0: + processors.append( + HammingDiversityLogitsProcessor( + diversity_penalty=generation_config.diversity_penalty, + num_beams=generation_config.num_beams, + num_beam_groups=generation_config.num_beam_groups, + ) + ) + if ( + generation_config.encoder_repetition_penalty is not None + and generation_config.encoder_repetition_penalty != 1.0 + ): + if len(encoder_input_ids.shape) == 2: + processors.append( + EncoderRepetitionPenaltyLogitsProcessor( + penalty=generation_config.encoder_repetition_penalty, + encoder_input_ids=encoder_input_ids, + ) + ) + else: + warnings.warn( + "Passing `encoder_repetition_penalty` requires some form of `input_ids` to be passed to " + "`generate`, ignoring the argument.", + UserWarning, + ) + if generation_config.repetition_penalty is not None and generation_config.repetition_penalty != 1.0: + processors.append(RepetitionPenaltyLogitsProcessor(penalty=generation_config.repetition_penalty)) + if generation_config.no_repeat_ngram_size is not None and generation_config.no_repeat_ngram_size > 0: + processors.append(NoRepeatNGramLogitsProcessor(generation_config.no_repeat_ngram_size)) + if ( + generation_config.encoder_no_repeat_ngram_size is not None + and generation_config.encoder_no_repeat_ngram_size > 0 + ): + if len(encoder_input_ids.shape) == 2: + processors.append( + EncoderNoRepeatNGramLogitsProcessor( + generation_config.encoder_no_repeat_ngram_size, + encoder_input_ids, + ) + ) + else: + warnings.warn( + "Passing `encoder_no_repeat_ngram_size` requires some form of `input_ids` to be passed to " + "`generate`, ignoring the argument.", + UserWarning, + ) + if generation_config.bad_words_ids is not None: + processors.append( + NoBadWordsLogitsProcessor( + generation_config.bad_words_ids, + generation_config._eos_token_tensor, + ) + ) + if ( + generation_config.min_length is not None + and generation_config._eos_token_tensor is not None + and generation_config.min_length > 0 + ): + processors.append( + MinLengthLogitsProcessor( + generation_config.min_length, + generation_config._eos_token_tensor, + device=device, + ) + ) + if ( + generation_config.min_new_tokens is not None + and generation_config._eos_token_tensor is not None + and generation_config.min_new_tokens > 0 + ): + processors.append( + MinNewTokensLengthLogitsProcessor( + input_ids_seq_length, + generation_config.min_new_tokens, + generation_config._eos_token_tensor, + device=device, + ) + ) + if prefix_allowed_tokens_fn is not None: + processors.append( + PrefixConstrainedLogitsProcessor( + prefix_allowed_tokens_fn, + generation_config.num_beams // generation_config.num_beam_groups, + ) + ) + if generation_config.forced_bos_token_id is not None: + processors.append( + ForcedBOSTokenLogitsProcessor( + generation_config.forced_bos_token_id, + ) + ) + if generation_config.forced_eos_token_id is not None: + processors.append( + ForcedEOSTokenLogitsProcessor( + generation_config.max_length, + generation_config.forced_eos_token_id, + device=device, + ) + ) + if generation_config.remove_invalid_values is True: + processors.append(InfNanRemoveLogitsProcessor()) + if generation_config.exponential_decay_length_penalty is not None: + processors.append( + ExponentialDecayLengthPenalty( + generation_config.exponential_decay_length_penalty, + generation_config._eos_token_tensor, + input_ids_seq_length, + ) + ) + if generation_config.suppress_tokens is not None: + processors.append( + SuppressTokensLogitsProcessor( + generation_config.suppress_tokens, + device=device, + ) + ) + if generation_config.begin_suppress_tokens is not None: + begin_index = input_ids_seq_length + begin_index = ( + begin_index + if (input_ids_seq_length > 1 or generation_config.forced_bos_token_id is None) + else begin_index + 1 + ) + processors.append( + SuppressTokensAtBeginLogitsProcessor( + generation_config.begin_suppress_tokens, + begin_index, + device=device, + ) + ) + if generation_config.forced_decoder_ids is not None: + # TODO (sanchit): move this exception to GenerationConfig.validate() when TF & FLAX are aligned with PT + raise ValueError( + "You have explicitly specified `forced_decoder_ids`. Please remove the `forced_decoder_ids` argument " + "in favour of `input_ids` or `decoder_input_ids` respectively.", + ) + + # TODO (joao): find a strategy to specify the order of the processors + processors = self._merge_criteria_processor_list(processors, logits_processor) + + # Processors previously known as `LogitsWarpers`, only applied with sampling strategies + if generation_config.do_sample: + # In beam methods, we need to keep at least one non-eos token to explore continuations that might have a + # better score (i.e. keep len(list(generation_config._eos_token_tensor)) + 1) + if generation_config.num_beams > 1: + if isinstance(generation_config._eos_token_tensor, list): + min_tokens_to_keep = len(generation_config._eos_token_tensor) + 1 + elif isinstance(generation_config._eos_token_tensor, torch.Tensor): + min_tokens_to_keep = generation_config._eos_token_tensor.shape[0] + 1 + else: + min_tokens_to_keep = 2 + else: + min_tokens_to_keep = 1 + + # the following idea is largely copied from this PR: https://github.com/huggingface/transformers/pull/5420/files + # all samplers can be found in `generation_utils_samplers.py` + if generation_config.temperature is not None and generation_config.temperature != 1.0: + processors.append(TemperatureLogitsWarper(generation_config.temperature)) + if generation_config.top_k is not None and generation_config.top_k != 0: + processors.append( + TopKLogitsWarper(top_k=generation_config.top_k, min_tokens_to_keep=min_tokens_to_keep) + ) + if generation_config.top_p is not None and generation_config.top_p < 1.0: + processors.append( + TopPLogitsWarper(top_p=generation_config.top_p, min_tokens_to_keep=min_tokens_to_keep) + ) + if generation_config.min_p is not None: + # Applied after temperature scaling (see https://github.com/ggerganov/llama.cpp/pull/3841#issuecomment-2073826084) + processors.append( + MinPLogitsWarper(min_p=generation_config.min_p, min_tokens_to_keep=min_tokens_to_keep) + ) + if generation_config.typical_p is not None and generation_config.typical_p < 1.0: + processors.append( + TypicalLogitsWarper(mass=generation_config.typical_p, min_tokens_to_keep=min_tokens_to_keep) + ) + if generation_config.epsilon_cutoff is not None and 0.0 < generation_config.epsilon_cutoff < 1.0: + processors.append( + EpsilonLogitsWarper( + epsilon=generation_config.epsilon_cutoff, min_tokens_to_keep=min_tokens_to_keep + ) + ) + if generation_config.eta_cutoff is not None and 0.0 < generation_config.eta_cutoff < 1.0: + processors.append( + EtaLogitsWarper( + epsilon=generation_config.eta_cutoff, min_tokens_to_keep=min_tokens_to_keep, device=device + ) + ) + + # Watermarking should be after all logits processing is finished (see #34630) + if generation_config.watermarking_config is not None: + processors.append( + generation_config.watermarking_config.construct_processor(self.config.vocab_size, device) + ) + + # `LogitNormalization` should always be the last logit processor, when present + if generation_config.renormalize_logits is True: + processors.append(LogitNormalization()) + return processors + + def _get_stopping_criteria( + self, + generation_config: GenerationConfig, + stopping_criteria: Optional[StoppingCriteriaList], + tokenizer: Optional["PreTrainedTokenizerBase"] = None, + **kwargs, + ) -> StoppingCriteriaList: + criteria = StoppingCriteriaList() + if generation_config.max_length is not None: + max_position_embeddings = getattr(self.config, "max_position_embeddings", None) + criteria.append( + MaxLengthCriteria( + max_length=generation_config.max_length, + max_position_embeddings=max_position_embeddings, + ) + ) + if generation_config.max_time is not None: + criteria.append(MaxTimeCriteria(max_time=generation_config.max_time)) + if generation_config.stop_strings is not None: + if tokenizer is None: + raise ValueError( + "There are one or more stop strings, either in the arguments to `generate` or in the " + "model's generation config, but we could not locate a tokenizer. When generating with " + "stop strings, you must pass the model's tokenizer to the `tokenizer` argument of `generate`." + ) + criteria.append(StopStringCriteria(stop_strings=generation_config.stop_strings, tokenizer=tokenizer)) + if generation_config._eos_token_tensor is not None: + criteria.append(EosTokenCriteria(eos_token_id=generation_config._eos_token_tensor)) + if ( + generation_config.is_assistant + and generation_config.assistant_confidence_threshold is not None + and generation_config.assistant_confidence_threshold > 0 + ): + criteria.append( + ConfidenceCriteria(assistant_confidence_threshold=generation_config.assistant_confidence_threshold) + ) + criteria = self._merge_criteria_processor_list(criteria, stopping_criteria) + return criteria + + def _merge_criteria_processor_list( + self, + default_list: Union[LogitsProcessorList, StoppingCriteriaList], + custom_list: Union[LogitsProcessorList, StoppingCriteriaList], + ) -> Union[LogitsProcessorList, StoppingCriteriaList]: + """ + Merge user-defined processors/criteria with the ones instantiated inside `generate`. In case the same + processor/criteria is present on both lists, use the user-defined one. + + (Note: up to v4.49.0, this funtion threw an exception is the same logit processor was found twice.) + """ + if len(custom_list) == 0: + return default_list + + final_list = type(default_list)() + for default in default_list: + using_custom = False + for custom in custom_list: + if type(custom) is type(default): + object_type = "stopping criteria" if isinstance(custom, StoppingCriteria) else "logits processor" + logger.warning_once( + f"A custom {object_type} of type {type(custom)} has been passed to `.generate()`, but it " + f"was also created in `.generate()`, given its parameterization. The custom {type(custom)} " + f"will take precedence. Please check the docstring of {type(custom)} to see related " + "`.generate()` flags." + ) + final_list.append(custom) + using_custom = True + break + if not using_custom: + final_list.append(default) + + for custom in custom_list: + if custom not in final_list: + final_list.append(custom) + return final_list + + def compute_transition_scores( + self, + sequences: torch.Tensor, + scores: Tuple[torch.Tensor], + beam_indices: Optional[torch.Tensor] = None, + normalize_logits: bool = False, + ) -> torch.Tensor: + """ + Computes the transition scores of sequences given the generation scores (and beam indices, if beam search was + used). This is a convenient method to quickly obtain the scores of the selected tokens at generation time. + + Parameters: + sequences (`torch.LongTensor`): + The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or + shorter if all batches finished early due to the `eos_token_id`. + scores (`tuple(torch.FloatTensor)`): + Transition scores for each vocabulary token at each generation step. Beam transition scores consisting + of log probabilities of tokens conditioned on log softmax of previously generated tokens in this beam. + Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for each generated token), + with each tensor of shape `(batch_size*num_beams, config.vocab_size)`. + beam_indices (`torch.LongTensor`, *optional*): + Beam indices of generated token id at each generation step. `torch.LongTensor` of shape + `(batch_size*num_return_sequences, sequence_length)`. Only required if a `num_beams>1` at + generate-time. + normalize_logits (`bool`, *optional*, defaults to `False`): + Whether to normalize the logits (which, for legacy reasons, may be unnormalized). + + Return: + `torch.Tensor`: A `torch.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)` containing + the transition scores (logits) + + Examples: + + ```python + >>> from transformers import GPT2Tokenizer, AutoModelForCausalLM + >>> import numpy as np + + >>> tokenizer = GPT2Tokenizer.from_pretrained("gpt2") + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> tokenizer.pad_token_id = tokenizer.eos_token_id + >>> inputs = tokenizer(["Today is"], return_tensors="pt") + + >>> # Example 1: Print the scores for each token generated with Greedy Search + >>> outputs = model.generate(**inputs, max_new_tokens=5, return_dict_in_generate=True, output_scores=True) + >>> transition_scores = model.compute_transition_scores( + ... outputs.sequences, outputs.scores, normalize_logits=True + ... ) + >>> # input_length is the length of the input prompt for decoder-only models, like the GPT family, and 1 for + >>> # encoder-decoder models, like BART or T5. + >>> input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1] + >>> generated_tokens = outputs.sequences[:, input_length:] + >>> for tok, score in zip(generated_tokens[0], transition_scores[0]): + ... # | token | token string | log probability | probability + ... print(f"| {tok:5d} | {tokenizer.decode(tok):8s} | {score.numpy():.3f} | {np.exp(score.numpy()):.2%}") + | 262 | the | -1.414 | 24.33% + | 1110 | day | -2.609 | 7.36% + | 618 | when | -2.010 | 13.40% + | 356 | we | -1.859 | 15.58% + | 460 | can | -2.508 | 8.14% + + >>> # Example 2: Reconstruct the sequence scores from Beam Search + >>> outputs = model.generate( + ... **inputs, + ... max_new_tokens=5, + ... num_beams=4, + ... num_return_sequences=4, + ... return_dict_in_generate=True, + ... output_scores=True, + ... ) + >>> transition_scores = model.compute_transition_scores( + ... outputs.sequences, outputs.scores, outputs.beam_indices, normalize_logits=False + ... ) + >>> # If you sum the generated tokens' scores and apply the length penalty, you'll get the sequence scores. + >>> # Tip 1: recomputing the scores is only guaranteed to match with `normalize_logits=False`. Depending on the + >>> # use case, you might want to recompute it with `normalize_logits=True`. + >>> # Tip 2: the output length does NOT include the input length + >>> output_length = np.sum(transition_scores.numpy() < 0, axis=1) + >>> length_penalty = model.generation_config.length_penalty + >>> reconstructed_scores = transition_scores.sum(axis=1) / (output_length**length_penalty) + >>> print(np.allclose(outputs.sequences_scores, reconstructed_scores)) + True + ```""" + # 1. In absence of `beam_indices`, we can assume that we come from e.g. greedy search, which is equivalent + # to a beam search approach were the first (and only) beam is always selected + if beam_indices is None: + beam_indices = torch.arange(scores[0].shape[0]).view(-1, 1).to(sequences.device) + beam_indices = beam_indices.expand(-1, len(scores)) + + # 2. reshape scores as [batch_size*vocab_size, # generation steps] with # generation steps being + # seq_len - input_length + scores = torch.stack(scores).reshape(len(scores), -1).transpose(0, 1) + + # 3. Optionally normalize the logits (across the vocab dimension) + if normalize_logits: + scores = scores.reshape(-1, self.config.vocab_size, scores.shape[-1]) + scores = torch.nn.functional.log_softmax(scores, dim=1) + scores = scores.reshape(-1, scores.shape[-1]) + + # 4. cut beam_indices to longest beam length + beam_indices_mask = beam_indices < 0 + max_beam_length = (1 - beam_indices_mask.long()).sum(-1).max() + beam_indices = beam_indices.clone()[:, :max_beam_length] + beam_indices_mask = beam_indices_mask[:, :max_beam_length] + + # 5. Set indices of beams that finished early to 0; such indices will be masked correctly afterwards + beam_indices[beam_indices_mask] = 0 + + # 6. multiply beam_indices with vocab size to gather correctly from scores + beam_sequence_indices = beam_indices * self.config.vocab_size + + # 7. Define which indices contributed to scores + cut_idx = sequences.shape[-1] - max_beam_length + indices = sequences[:, cut_idx:] + beam_sequence_indices + + # 8. Compute scores + transition_scores = scores.gather(0, indices) + + # 9. Mask out transition_scores of beams that stopped early + transition_scores[beam_indices_mask] = 0 + + return transition_scores + + def _validate_model_class(self): + """ + Confirms that the model class is compatible with generation. If not, raises an exception that points to the + right class to use. + """ + # TODO(joao): remove this function in v4.50, i.e. when we remove the inheritance of `GenerationMixin` from + # `PreTrainedModel`. With that inheritance removed, all model classes inheriting from `GenerationMixin` can + # safely call `GenerationMixin.generate` + if not self.can_generate(): + terminations_with_generation_support = [ + "ForCausalLM", + "ForConditionalGeneration", + "ForSpeechSeq2Seq", + "ForVision2Seq", + ] + raise TypeError( + f"The current model class ({self.__class__.__name__}) is not compatible with `.generate()`, as " + "it doesn't have a language model head. Classes that support generation often end in one of these " + f"names: {terminations_with_generation_support}." + ) + + def _validate_assistant(self, assistant_model, tokenizer, assistant_tokenizer): + if assistant_model is None: + return + + if self.config.is_encoder_decoder and not assistant_model.config.is_encoder_decoder: + attributes_to_check = ["encoder_attention_heads", "encoder_ffn_dim", "encoder_layers"] + attributes_to_check = [attr for attr in dir(assistant_model.config) if attr in attributes_to_check] + are_equal = all( + getattr(self.config, attr) == getattr(assistant_model.config, attr) for attr in attributes_to_check + ) + if not are_equal: + raise ValueError( + "The main model and the assistant don't have compatible encoder-dependent input shapes. " + "Ensure you load the assistant with the correct encoder-decoder class, e.g. `AutoModelForSpeechSeq2Seq` for Whisper." + ) + + doc_reference = ( + "(see https://huggingface.co/docs/transformers/en/generation_strategies#universal-assisted-decoding)" + ) + if self.config.get_text_config().vocab_size == assistant_model.config.get_text_config().vocab_size: + if assistant_tokenizer is not None: + raise ValueError( + f"`assistant_tokenizer` is not required when the main and assistant models use the same tokenizer. Please omit `assistant_tokenizer` from `generate()` {doc_reference}." + ) + else: + if tokenizer is None or assistant_tokenizer is None: + raise ValueError( + f"The main and assistant moedels have different tokenizers. Please provide `tokenizer` and `assistant_tokenizer` to `generate()` {doc_reference}." + ) + + def _validate_model_kwargs(self, model_kwargs: Dict[str, Any]): + """Validates model kwargs for generation. Generate argument typos will also be caught here.""" + # If a `Cache` instance is passed, checks whether the model is compatible with it + if isinstance(model_kwargs.get("past_key_values", None), Cache) and not self._supports_cache_class: + raise ValueError( + f"{self.__class__.__name__} does not support an instance of `Cache` as `past_key_values`. Please " + "check the model documentation for supported cache formats." + ) + + # Excludes arguments that are handled before calling any model function + if self.config.is_encoder_decoder: + for key in ["decoder_input_ids"]: + model_kwargs.pop(key, None) + + unused_model_args = [] + model_args = set(inspect.signature(self.prepare_inputs_for_generation).parameters) + # `kwargs`/`model_kwargs` is often used to handle optional forward pass inputs like `attention_mask`. If + # `prepare_inputs_for_generation` doesn't accept them, then a stricter check can be made ;) + if "kwargs" in model_args or "model_kwargs" in model_args: + model_args |= set(inspect.signature(self.forward).parameters) + + # Encoder-Decoder models may also need Encoder arguments from `model_kwargs` + if self.config.is_encoder_decoder: + base_model = getattr(self, self.base_model_prefix, None) + + # allow encoder kwargs + encoder = getattr(self, "encoder", None) + # `MusicgenForConditionalGeneration` has `text_encoder` and `audio_encoder`. + # Also, it has `base_model_prefix = "encoder_decoder"` but there is no `self.encoder_decoder` + # TODO: A better way to handle this. + if encoder is None and base_model is not None: + encoder = getattr(base_model, "encoder", None) + + if encoder is not None: + encoder_model_args = set(inspect.signature(encoder.forward).parameters) + model_args |= encoder_model_args + + # allow decoder kwargs + decoder = getattr(self, "decoder", None) + if decoder is None and base_model is not None: + decoder = getattr(base_model, "decoder", None) + + if decoder is not None: + decoder_model_args = set(inspect.signature(decoder.forward).parameters) + model_args |= {f"decoder_{x}" for x in decoder_model_args} + + for key, value in model_kwargs.items(): + if value is not None and key not in model_args: + unused_model_args.append(key) + + if unused_model_args: + raise ValueError( + f"The following `model_kwargs` are not used by the model: {unused_model_args} (note: typos in the" + " generate arguments will also show up in this list)" + ) + + def _validate_generated_length(self, generation_config, input_ids_length, has_default_max_length): + """Performs validation related to the resulting generated length""" + # 1. Max length warnings related to poor parameterization + if has_default_max_length and generation_config.max_new_tokens is None and generation_config.max_length == 20: + # 20 is the default max_length of the generation config + warnings.warn( + f"Using the model-agnostic default `max_length` (={generation_config.max_length}) to control the " + "generation length. We recommend setting `max_new_tokens` to control the maximum length of the " + "generation.", + UserWarning, + ) + if input_ids_length >= generation_config.max_length: + input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids" + raise ValueError( + f"Input length of {input_ids_string} is {input_ids_length}, but `max_length` is set to" + f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider" + " increasing `max_length` or, better yet, setting `max_new_tokens`." + ) + + # 2. Min length warnings due to unfeasible parameter combinations + min_length_error_suffix = ( + " Generation will stop at the defined maximum length. You should decrease the minimum length and/or " + "increase the maximum length." + ) + if has_default_max_length: + min_length_error_suffix += ( + f" Note that `max_length` is set to {generation_config.max_length}, its default value." + ) + if generation_config.min_length is not None and generation_config.min_length > generation_config.max_length: + warnings.warn( + f"Unfeasible length constraints: `min_length` ({generation_config.min_length}) is larger than" + f" the maximum possible length ({generation_config.max_length})." + min_length_error_suffix, + UserWarning, + ) + if generation_config.min_new_tokens is not None: + min_length = generation_config.min_new_tokens + input_ids_length + if min_length > generation_config.max_length: + warnings.warn( + f"Unfeasible length constraints: `min_new_tokens` ({generation_config.min_new_tokens}), when " + f"added to the prompt length ({input_ids_length}), is larger than" + f" the maximum possible length ({generation_config.max_length})." + min_length_error_suffix, + UserWarning, + ) + + def _prepare_generated_length( + self, + generation_config, + has_default_max_length, + has_default_min_length, + model_input_name, + input_ids_length, + inputs_tensor, + ): + """Prepared max and min length in generation configs to avoid clashes between similar attributes""" + + if generation_config.max_new_tokens is not None: + if not has_default_max_length and generation_config.max_length is not None: + logger.warning( + f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(=" + f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. " + "Please refer to the documentation for more information. " + "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)" + ) + generation_config.max_length = generation_config.max_new_tokens + input_ids_length + + # if both `inputs_embeds` and `input_ids` are passed, we do not correct the length + # otherwise we need total length [inputs-embeds-len + new-tokens-len] to not go beyond indicated `max_length`` + elif ( + model_input_name == "inputs_embeds" + and input_ids_length != inputs_tensor.shape[1] + and not self.config.is_encoder_decoder + ): + generation_config.max_length -= inputs_tensor.shape[1] + elif has_default_max_length: # by default let's always generate 20 new tokens + if generation_config.max_length == GenerationConfig().max_length: + generation_config.max_length = generation_config.max_length + input_ids_length + max_position_embeddings = getattr(self.config, "max_position_embeddings", None) + if max_position_embeddings is not None: + generation_config.max_length = min(generation_config.max_length, max_position_embeddings) + + # same for min length + if generation_config.min_new_tokens is not None: + if not has_default_min_length: + logger.warning( + f"Both `min_new_tokens` (={generation_config.min_new_tokens}) and `min_length`(=" + f"{generation_config.min_length}) seem to have been set. `min_new_tokens` will take precedence. " + "Please refer to the documentation for more information. " + "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)" + ) + generation_config.min_length = generation_config.min_new_tokens + input_ids_length + + elif ( + model_input_name == "inputs_embeds" + and input_ids_length != inputs_tensor.shape[1] + and not self.config.is_encoder_decoder + ): + generation_config.min_length = max(generation_config.min_length - inputs_tensor.shape[1], 0) + + return generation_config + + def _prepare_generation_config( + self, generation_config: Optional[GenerationConfig], use_model_defaults: Optional[bool] = None, **kwargs: Dict + ) -> Tuple[GenerationConfig, Dict]: + """ + Prepares the base generation config, then applies any generation configuration options from kwargs. This + function handles retrocompatibility with respect to configuration files. + """ + # parameterization priority: + # kwargs > non-global default values in `generation_config` > `model.generation_config` > GenerationConfig() + # TODO (joao): per-model generation config classes. + + using_model_generation_config = False + if generation_config is None: + # legacy: users may modify the model configuration to control generation. To trigger this legacy behavior, + # the following conditions must be met + # 1) the generation config must have been created from the model config (`_from_model_config` field); + # 2) the generation config must have seen no modification since its creation (the hash is the same); + # 3) there are non-default generation parameters in the model config. + # 4) the user must have set new generation parameters in the model config. + if ( + self.generation_config._from_model_config # 1) + and self.generation_config._original_object_hash == hash(self.generation_config) # 2) + and len(self.config._get_non_default_generation_parameters()) > 0 # 3) + ): + new_generation_config = GenerationConfig.from_model_config(self.config) + if new_generation_config != self.generation_config: # 4) + warnings.warn( + "You have modified the pretrained model configuration to control generation. This is a" + " deprecated strategy to control generation and will be removed in v5." + " Please use and modify the model generation configuration (see" + " https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )", + UserWarning, + ) + self.generation_config = new_generation_config + + generation_config = self.generation_config + using_model_generation_config = True + + # `torch.export.export` usually raises an exception if it is called + # with ``strict=True``. deepcopy can only be processed if ``strict=False``. + generation_config = copy.deepcopy(generation_config) + + if not using_model_generation_config: + # If `generation_config` is provided: + # - `use_model_defaults`: let's fallback ALL default values to the model's generation config + # - otherwise: legacy behavior, let's just make sure we have the tokens defined + model_base_version = version.parse(version.parse(self.generation_config.transformers_version).base_version) + if use_model_defaults is True or ( + use_model_defaults is None and model_base_version >= version.parse("4.50.0") + ): + modified_values = {} + default_generation_config = GenerationConfig() + for key, default_value in default_generation_config.__dict__.items(): + if key.startswith("_") or key == "transformers_version": # metadata + continue + custom_gen_config_value = getattr(generation_config, key) + model_gen_config_value = getattr(self.generation_config, key) + if custom_gen_config_value == default_value and model_gen_config_value != default_value: + modified_values[key] = model_gen_config_value + setattr(generation_config, key, model_gen_config_value) + if len(modified_values) > 0: + logger.warning_once( + f"`generation_config` default values have been modified to match model-specific defaults: " + f"{modified_values}. If this is not desired, please set these values explicitly." + ) + else: + if generation_config.bos_token_id is None: + generation_config.bos_token_id = self.generation_config.bos_token_id + if generation_config.eos_token_id is None: + generation_config.eos_token_id = self.generation_config.eos_token_id + if generation_config.pad_token_id is None: + generation_config.pad_token_id = self.generation_config.pad_token_id + if generation_config.decoder_start_token_id is None: + generation_config.decoder_start_token_id = self.generation_config.decoder_start_token_id + + # Finally, apply any passed kwargs + model_kwargs = generation_config.update(**kwargs) + + return generation_config, model_kwargs + + def _get_initial_cache_position(self, input_ids, model_kwargs): + """Calculates `cache_position` for the pre-fill stage based on `input_ids` and optionally past length""" + # `torch.compile`-friendly `torch.arange` from a shape -- the lines below are equivalent to `torch.arange` + if "inputs_embeds" in model_kwargs and not self.config.is_encoder_decoder: + cache_position = torch.ones_like(model_kwargs["inputs_embeds"][0, :, 0], dtype=torch.int64).cumsum(0) - 1 + elif "decoder_inputs_embeds" in model_kwargs and self.config.is_encoder_decoder: + cache_position = ( + torch.ones_like(model_kwargs["decoder_inputs_embeds"][0, :, 0], dtype=torch.int64).cumsum(0) - 1 + ) + else: + cache_position = torch.ones_like(input_ids[0, :], dtype=torch.int64).cumsum(0) - 1 + + past_length = 0 + if model_kwargs.get("past_key_values") is not None: + cache = model_kwargs["past_key_values"] + past_length = 0 + if not isinstance(cache, Cache): + past_length = cache[0][0].shape[2] + elif hasattr(cache, "get_seq_length") and cache.get_seq_length() is not None: + past_length = cache.get_seq_length() + + cache_position = cache_position[past_length:] + + model_kwargs["cache_position"] = cache_position + return model_kwargs + + def _get_layer_device_map_for_cache_init(self) -> Optional[Dict[int, Union[str, int]]]: + """ + Returns the device map for each decoder layer, to allocate the cache on the right device. + Inspired from `dispatch_model` in accelerate. + """ + execution_device_map = None + + if hasattr(self, "hf_device_map"): + if set(self.hf_device_map.values()) == {"cpu"} or set(self.hf_device_map.values()) == {"cpu", "disk"}: + main_device = "cpu" + else: + main_device = [d for d in self.hf_device_map.values() if d not in ["cpu", "disk"]][0] + execution_device_map = { + name: main_device if device in ["cpu", "disk"] else device + for name, device in self.hf_device_map.items() + } + + # No `execution_device_map` -> rely on `self.device` to allocate the cache + if execution_device_map is None: + return None + + # Single device for all layers + num_hidden_layers = self.config.get_text_config().num_hidden_layers + if len(execution_device_map) == 1 and "" in execution_device_map: + return dict.fromkeys(range(num_hidden_layers), execution_device_map[""]) + + # Multiple devices in `execution_device_map` -> we need to map decoder layers to the correct device. + layer_device_map = {} + # Case 1: The model has a `get_decoder` method, we can use it to find the decoder name. + if hasattr(self, "get_decoder"): + decoder_name = None + for name, module in self.named_modules(): + if module is self.get_decoder(): + decoder_name = name + break + if decoder_name is None: + raise RuntimeError( + "`model.get_decoder()` is not returning a named module of the model. This is unexpected, please " + "open an issue on GitHub." + ) + + decoder_mapped_modules = [ + module_name for module_name in execution_device_map.keys() if decoder_name in module_name + ] + # The decoder name may be present in `execution_device_map` in two forms: + # a) each layer has a device mapping + if len(decoder_mapped_modules) >= num_hidden_layers: + for idx in range(num_hidden_layers): + for module_name in decoder_mapped_modules: + if f".{idx}." in f"{module_name}.": + layer_device_map[idx] = execution_device_map[module_name] + break + + # b) the whole module is mapped to a single device. If the decoder name is NOT present in the device map, + # then the mapping is done in a parent module + else: + while True: + if decoder_name in execution_device_map: + layer_device_map = dict.fromkeys(range(num_hidden_layers), execution_device_map[decoder_name]) + break + elif "." in decoder_name: + decoder_name = decoder_name.rsplit(".", 1)[0] # gets the name of the parent module + else: + raise RuntimeError(f"Decoder name {decoder_name} not found in execution device map") + + # Case 2: Legacy code path: assume the decoder layers are named as `(...).X` (X being the layer index) + else: + for layer in execution_device_map: + for idx in range(num_hidden_layers): + if f".{idx}." in f"{layer}.": + layer_device_map[idx] = execution_device_map[layer] + break + + for idx in range(num_hidden_layers): + if idx not in layer_device_map: + raise RuntimeError(f"layer {idx} has not been mapped to a device.") + return layer_device_map + + def _get_cache( + self, cache_implementation: str, batch_size: int, max_cache_len: int, device: torch.device, model_kwargs + ) -> Cache: + """ + Sets a cache for `generate`, that will persist across calls. A new cache will only be initialized a + new `generate` call requires a larger cache or uses a different batch size. + + Returns the resulting cache object. + """ + if cache_implementation == "hybrid" and "llama4" in getattr(self.config, "model_type", ""): + cache_implementation = "hybrid_chunked" + + cache_cls: Cache = NEED_SETUP_CACHE_CLASSES_MAPPING[cache_implementation] + requires_cross_attention_cache = ( + self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None + ) + + if hasattr(self, "_cache"): + cache_to_check = self._cache.self_attention_cache if requires_cross_attention_cache else self._cache + + if cache_implementation == "sliding_window": + max_cache_len = min(self.config.sliding_window, max_cache_len) + + need_new_cache = ( + not hasattr(self, "_cache") + or (not isinstance(cache_to_check, cache_cls)) + or cache_to_check.max_batch_size != batch_size + ) + if cache_implementation != "mamba": + need_new_cache = need_new_cache or cache_to_check.max_cache_len < max_cache_len + + if requires_cross_attention_cache and hasattr(self, "_cache"): + need_new_cache = ( + need_new_cache + or self._cache.cross_attention_cache.max_cache_len != model_kwargs["encoder_outputs"][0].shape[1] + ) + + if need_new_cache: + if hasattr(self.config, "_pre_quantization_dtype"): + cache_dtype = self.config._pre_quantization_dtype + else: + cache_dtype = self.dtype + + layer_device_map = self._get_layer_device_map_for_cache_init() + cache_kwargs = { + "config": self.config.get_text_config(), + "max_batch_size": batch_size, + "max_cache_len": max_cache_len, + "dtype": cache_dtype, + "device": device, + "layer_device_map": layer_device_map, + } + self._cache = cache_cls(**cache_kwargs) + if requires_cross_attention_cache: + encoder_kwargs = cache_kwargs.copy() + encoder_kwargs["max_cache_len"] = model_kwargs["encoder_outputs"][0].shape[1] + self._cache = EncoderDecoderCache(self._cache, cache_cls(**encoder_kwargs)) + else: + self._cache.reset() + return self._cache + + def _supports_default_dynamic_cache(self) -> bool: + """ + Return `True` if current model can use a `DynamicCache` instance when initializing the `past_key_values`. + This is mostly the same as `_supports_cache_class` attribute, but add exception for `Jamba` model which + uses its own `HybridMambaAttentionDynamicCache` and do not need to initialize the Cache in advance in + order to save memory (because no back and forth `to_legacy_cache` and `from_legacy_cache` will be performed + for `HybridMambaAttentionDynamicCache`). + """ + return ( + self._supports_cache_class + and "jamba" not in self.__class__.__name__.lower() + and "zamba" not in self.__class__.__name__.lower() + and "bamba" not in self.__class__.__name__.lower() + ) + + def _prepare_cache_for_generation( + self, + generation_config: GenerationConfig, + model_kwargs: Dict, + assistant_model: "PreTrainedModel", + batch_size: int, + max_cache_length: int, + device: torch.device, + ) -> bool: + """ + Prepares the cache for generation (if applicable), given `generate`'s parameterization. If a cache is + instantiated, writes it to `model_kwargs`, under the name expected by the model. + """ + + cache_name = "past_key_values" if "mamba" not in self.__class__.__name__.lower() else "cache_params" + requires_cross_attention_cache = ( + self.config.is_encoder_decoder or model_kwargs.get("encoder_outputs") is not None + ) + + # Quick escape route 1: if the user specifies a cache, we only need to: + # a) check for conflicting `generate` arguments + # b) convert to the new cache format (if the user passes a legacy cache and model supports it) + user_defined_cache = model_kwargs.get(cache_name) + if user_defined_cache is not None: + if generation_config.cache_implementation is not None: + raise ValueError( + f"Passing both `cache_implementation` (used to initialize certain caches) and `{cache_name}` (a " + "Cache object) is unsupported. Please use only one of the two." + ) + if isinstance(user_defined_cache, tuple) and self._supports_default_dynamic_cache(): + model_kwargs[cache_name] = ( + DynamicCache.from_legacy_cache(user_defined_cache) + if not requires_cross_attention_cache + else EncoderDecoderCache.from_legacy_cache(user_defined_cache) + ) + return + + # Quick escape route 2: if the user specifies no cache is to be used. (conflicting arguments are handled in + # `generation_config.validate()`) + if generation_config.use_cache is False: + return + + # Quick escape route 3: model that only supports legacy caches = nothing to prepare + if not self._supports_default_dynamic_cache(): + if generation_config.cache_implementation is not None: + warnings.warn( + "This model does not support `Cache` instances, it only supports the legacy cache format (tuple " + f"of tuples). `cache_implementation` (set to {generation_config.cache_implementation}) will be " + "ignored.", + UserWarning, + ) + return + + # Otherwise we NEED to prepare a cache, based on `generation_config.cache_implementation` + + # TODO(joao): support static caches in assisted generation. assisted generation needs to roll back caches, + # which is only supported in dynamic caches atm + if assistant_model is not None and generation_config.cache_implementation is not None: + logger.warning_once( + "An assistant model is provided, using a dynamic cache instead of a cache of type=" + f"'{generation_config.cache_implementation}'." + ) + generation_config.cache_implementation = None + + generation_config.cache_implementation = generation_config.cache_implementation or getattr( + self.config.get_text_config(), "cache_implementation", None + ) + if generation_config.cache_implementation is not None: + if generation_config.cache_implementation in NEED_SETUP_CACHE_CLASSES_MAPPING: + if generation_config.cache_implementation == "static" and not self._supports_static_cache: + raise ValueError( + "This model does not support `cache_implementation='static'`. Please check the following " + "issue: https://github.com/huggingface/transformers/issues/28981" + ) + model_kwargs[cache_name] = self._get_cache( + cache_implementation=generation_config.cache_implementation, + batch_size=max(generation_config.num_beams, generation_config.num_return_sequences) * batch_size, + max_cache_len=max_cache_length, + device=device, + model_kwargs=model_kwargs, + ) + elif generation_config.cache_implementation == "quantized": + if not self._supports_quantized_cache: + raise ValueError( + "This model does not support the quantized cache. If you want your model to support quantized " + "cache, please open an issue and tag @zucchini-nlp." + ) + + cache_config = ( + generation_config.cache_config + if generation_config.cache_config is not None + else QuantizedCacheConfig() + ) + cache_class = QUANT_BACKEND_CLASSES_MAPPING[cache_config.backend] + + if cache_config.backend == "quanto" and not is_optimum_quanto_available(): + raise ImportError( + "You need to install optimum-quanto in order to use KV cache quantization with optimum-quanto backend. " + "Please install it via with `pip install optimum-quanto`" + ) + elif cache_config.backend == "HQQ" and not is_hqq_available(): + raise ImportError( + "You need to install `HQQ` in order to use KV cache quantization with HQQ backend. " + "Please install it via with `pip install hqq`" + ) + + model_kwargs[cache_name] = cache_class(cache_config) + elif generation_config.cache_implementation == "offloaded": + model_kwargs[cache_name] = OffloadedCache() + elif generation_config.cache_implementation == "dynamic": + model_kwargs[cache_name] = DynamicCache() + + # Use DynamicCache() instance by default. This will avoid back and forth from legacy format that + # keeps copying the cache thus using much more memory + else: + model_kwargs[cache_name] = ( + DynamicCache() + if not requires_cross_attention_cache + else EncoderDecoderCache(DynamicCache(), DynamicCache()) + ) + + def _supports_logits_to_keep(self) -> bool: + """ + Return True if the current model supports the keyword argument `logits_to_keep` in forward() + to save memory. Checking it in this way allows to avoid using a new model attribute. + """ + return "logits_to_keep" in set(inspect.signature(self.forward).parameters.keys()) + + def _prepare_special_tokens( + self, + generation_config: GenerationConfig, + kwargs_has_attention_mask: Optional[bool] = None, + device: Optional[Union[torch.device, str]] = None, + ): + """ + Prepares the special tokens for generation, overwriting the generation config with their processed versions + converted to tensor. + + Note that `generation_config` is changed in place and stops being serializable after this method is called. + That is no problem if called within `generate` (`generation_config` is a local copy that doesn't leave the + function). However, if called outside `generate`, consider creating a copy of `generation_config` first. + """ + + # Convert special tokens to tensors + def _tensor_or_none(token, device=None): + if token is None: + return token + + device = device if device is not None else self.device + if isinstance(token, torch.Tensor): + return token.to(device) + return torch.tensor(token, device=device, dtype=torch.long) + + bos_token_tensor = _tensor_or_none(generation_config.bos_token_id, device=device) + eos_token_tensor = _tensor_or_none(generation_config.eos_token_id, device=device) + pad_token_tensor = _tensor_or_none(generation_config.pad_token_id, device=device) + decoder_start_token_tensor = _tensor_or_none(generation_config.decoder_start_token_id, device=device) + + # for BC we also try to get `decoder_start_token_id` or `bos_token_id` (#30892) + if self.config.is_encoder_decoder: + decoder_start_token_tensor = ( + decoder_start_token_tensor if decoder_start_token_tensor is not None else bos_token_tensor + ) + + # We can have more than one eos token. Always treat it as a 1D tensor (when it exists). + if eos_token_tensor is not None and eos_token_tensor.ndim == 0: + eos_token_tensor = eos_token_tensor.unsqueeze(0) + + # Set pad token if unset (and there are conditions to do so) + if pad_token_tensor is None and eos_token_tensor is not None: + if kwargs_has_attention_mask is not None and not kwargs_has_attention_mask: + logger.warning( + "The attention mask and the pad token id were not set. As a consequence, you may observe " + "unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results." + ) + pad_token_tensor = eos_token_tensor[0] + logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{pad_token_tensor} for open-end generation.") + + # Sanity checks/warnings + if self.config.is_encoder_decoder and decoder_start_token_tensor is None: + raise ValueError( + "`decoder_start_token_id` or `bos_token_id` has to be defined for encoder-decoder generation." + ) + if ( + eos_token_tensor is not None + and isin_mps_friendly(elements=eos_token_tensor, test_elements=pad_token_tensor).any() + ): + if kwargs_has_attention_mask is not None and not kwargs_has_attention_mask: + logger.warning_once( + "The attention mask is not set and cannot be inferred from input because pad token is same as " + "eos token. As a consequence, you may observe unexpected behavior. Please pass your input's " + "`attention_mask` to obtain reliable results." + ) + if eos_token_tensor is not None and ( + torch.is_floating_point(eos_token_tensor) or (eos_token_tensor < 0).any() + ): + logger.warning( + f"`eos_token_id` should consist of positive integers, but is {eos_token_tensor}. Your generation " + "will not stop until the maximum length is reached. Depending on other flags, it may even crash." + ) + + # Update generation config with the updated special tokens tensors + # NOTE: this must be written into a different attribute name than the one holding the original special tokens + # (in their non-tensor form), in order to enable end-to-end compilation. See + # https://pytorch.org/docs/stable/torch.compiler_cudagraph_trees.html#limitations + generation_config._bos_token_tensor = bos_token_tensor + generation_config._eos_token_tensor = eos_token_tensor + generation_config._pad_token_tensor = pad_token_tensor + generation_config._decoder_start_token_tensor = decoder_start_token_tensor + + @torch.no_grad() + def generate( + self, + inputs: Optional[torch.Tensor] = None, + generation_config: Optional[GenerationConfig] = None, + logits_processor: Optional[LogitsProcessorList] = None, + stopping_criteria: Optional[StoppingCriteriaList] = None, + prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None, + synced_gpus: Optional[bool] = None, + assistant_model: Optional["PreTrainedModel"] = None, + streamer: Optional["BaseStreamer"] = None, + negative_prompt_ids: Optional[torch.Tensor] = None, + negative_prompt_attention_mask: Optional[torch.Tensor] = None, + use_model_defaults: Optional[bool] = None, + **kwargs, + ) -> Union[GenerateOutput, torch.LongTensor]: + r""" + + Generates sequences of token ids for models with a language modeling head. + + + + Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the + model's default generation configuration. You can override any `generation_config` by passing the corresponding + parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`. + + For an overview of generation strategies and code examples, check out the [following + guide](../generation_strategies). + + + + Parameters: + inputs (`torch.Tensor` of varying shape depending on the modality, *optional*): + The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the + method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs` + should be in the format of `input_ids`. For encoder-decoder models *inputs* can represent any of + `input_ids`, `input_values`, `input_features`, or `pixel_values`. + generation_config ([`~generation.GenerationConfig`], *optional*): + The generation configuration to be used as base parametrization for the generation call. `**kwargs` + passed to generate matching the attributes of `generation_config` will override them. If + `generation_config` is not provided, the default will be used, which has the following loading + priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model + configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s + default values, whose documentation should be checked to parameterize generation. + logits_processor (`LogitsProcessorList`, *optional*): + Custom logits processors that complement the default logits processors built from arguments and + generation config. If a logit processor is passed that is already created with the arguments or a + generation config an error is thrown. This feature is intended for advanced users. + stopping_criteria (`StoppingCriteriaList`, *optional*): + Custom stopping criteria that complements the default stopping criteria built from arguments and a + generation config. If a stopping criteria is passed that is already created with the arguments or a + generation config an error is thrown. If your stopping criteria depends on the `scores` input, make + sure you pass `return_dict_in_generate=True, output_scores=True` to `generate`. This feature is + intended for advanced users. + prefix_allowed_tokens_fn (`Callable[[int, torch.Tensor], List[int]]`, *optional*): + If provided, this function constraints the beam search to allowed tokens only at each step. If not + provided no constraint is applied. This function takes 2 arguments: the batch ID `batch_id` and + `input_ids`. It has to return a list with the allowed tokens for the next generation step conditioned + on the batch ID `batch_id` and the previously generated tokens `inputs_ids`. This argument is useful + for constrained generation conditioned on the prefix, as described in [Autoregressive Entity + Retrieval](https://arxiv.org/abs/2010.00904). + synced_gpus (`bool`, *optional*): + Whether to continue running the while loop until max_length. Unless overridden, this flag will be set + to `True` if using `FullyShardedDataParallel` or DeepSpeed ZeRO Stage 3 with multiple GPUs to avoid + deadlocking if one GPU finishes generating before other GPUs. Otherwise, defaults to `False`. + assistant_model (`PreTrainedModel`, *optional*): + An assistant model that can be used to accelerate generation. The assistant model must have the exact + same tokenizer. The acceleration is achieved when forecasting candidate tokens with the assistant model + is much faster than running generation with the model you're calling generate from. As such, the + assistant model should be much smaller. + streamer (`BaseStreamer`, *optional*): + Streamer object that will be used to stream the generated sequences. Generated tokens are passed + through `streamer.put(token_ids)` and the streamer is responsible for any further processing. + negative_prompt_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + The negative prompt needed for some processors such as CFG. The batch size must match the input batch + size. This is an experimental feature, subject to breaking API changes in future versions. + negative_prompt_attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Attention_mask for `negative_prompt_ids`. + use_model_defaults (`bool`, *optional*): + When it is `True`, unset parameters in `generation_config` will be set to the model-specific default + generation configuration (`model.generation_config`), as opposed to the global defaults + (`GenerationConfig()`). If unset, models saved starting from `v4.50` will consider this flag to be + `True`. + kwargs (`Dict[str, Any]`, *optional*): + Ad hoc parametrization of `generation_config` and/or additional model-specific kwargs that will be + forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder + specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*. + + Return: + [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True` + or when `config.return_dict_in_generate=True`) or a `torch.LongTensor`. + + If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible + [`~utils.ModelOutput`] types are: + + - [`~generation.GenerateDecoderOnlyOutput`], + - [`~generation.GenerateBeamDecoderOnlyOutput`] + + If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible + [`~utils.ModelOutput`] types are: + + - [`~generation.GenerateEncoderDecoderOutput`], + - [`~generation.GenerateBeamEncoderDecoderOutput`] + """ + + # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call + self._validate_model_class() + tokenizer = kwargs.pop("tokenizer", None) # Pull this out first, we only use it for stopping criteria + assistant_tokenizer = kwargs.pop("assistant_tokenizer", None) # only used for assisted generation + + generation_config, model_kwargs = self._prepare_generation_config( + generation_config, use_model_defaults, **kwargs + ) + self._validate_model_kwargs(model_kwargs.copy()) + self._validate_assistant(assistant_model, tokenizer, assistant_tokenizer) + + # 2. Set generation parameters if not already defined + if synced_gpus is None: + synced_gpus = (is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self)) and dist.get_world_size() > 1 + + logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList() + stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList() + + accepts_attention_mask = "attention_mask" in set(inspect.signature(self.forward).parameters.keys()) + requires_attention_mask = "encoder_outputs" not in model_kwargs + kwargs_has_attention_mask = model_kwargs.get("attention_mask", None) is not None + + # 3. Define model inputs + inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs( + inputs, generation_config.bos_token_id, model_kwargs + ) + batch_size = inputs_tensor.shape[0] + + device = inputs_tensor.device + self._prepare_special_tokens(generation_config, kwargs_has_attention_mask, device=device) + + # decoder-only models must use left-padding for batched generation. + if not self.config.is_encoder_decoder: + # If `input_ids` was given, check if the last id in any sequence is `pad_token_id` + # Note: If using, `inputs_embeds` this check does not work, because we want to be more hands-off. + if ( + generation_config._pad_token_tensor is not None + and batch_size > 1 + and len(inputs_tensor.shape) == 2 + and torch.sum(inputs_tensor[:, -1] == generation_config._pad_token_tensor) > 0 + ): + logger.warning( + "A decoder-only architecture is being used, but right-padding was detected! For correct " + "generation results, please set `padding_side='left'` when initializing the tokenizer." + ) + + # 4. Define other model kwargs + # decoder-only models with inputs_embeds forwarding must use caching (otherwise we can't detect whether we are + # generating the first new token or not, and we only want to use the embeddings for the first new token) + if not self.config.is_encoder_decoder and model_input_name == "inputs_embeds": + generation_config.use_cache = True + + if not kwargs_has_attention_mask and requires_attention_mask and accepts_attention_mask: + model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation( + inputs_tensor, generation_config, model_kwargs + ) + elif kwargs_has_attention_mask: + # TODO (joao): generalize this check with other types of inputs + if model_input_name == "input_ids" and len(model_kwargs["attention_mask"].shape) > 2: + raise ValueError("`attention_mask` passed to `generate` must be 2D.") + + if self.config.is_encoder_decoder and "encoder_outputs" not in model_kwargs: + # if model is encoder decoder encoder_outputs are created and added to `model_kwargs` + model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation( + inputs_tensor, model_kwargs, model_input_name, generation_config + ) + + # 5. Prepare `input_ids` which will be used for auto-regressive generation + if self.config.is_encoder_decoder: + input_ids, model_kwargs = self._prepare_decoder_input_ids_for_generation( + batch_size=batch_size, + model_input_name=model_input_name, + model_kwargs=model_kwargs, + decoder_start_token_id=generation_config._decoder_start_token_tensor, + device=inputs_tensor.device, + ) + else: + input_ids = inputs_tensor if model_input_name == "input_ids" else model_kwargs.pop("input_ids") + + if generation_config.token_healing: + input_ids = self.heal_tokens(input_ids, tokenizer) + + if streamer is not None: + streamer.put(input_ids.cpu()) + + # 6. Prepare `max_length` depending on other stopping criteria. + input_ids_length = input_ids.shape[-1] + has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None + has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None + generation_config = self._prepare_generated_length( + generation_config=generation_config, + has_default_max_length=has_default_max_length, + has_default_min_length=has_default_min_length, + model_input_name=model_input_name, + inputs_tensor=inputs_tensor, + input_ids_length=input_ids_length, + ) + + # If the model supports `logits_to_keep` in forward(), set it to 1 to avoid computing the whole + # logit matrix. This can save a lot of memory during the first forward pass. Note that assisted decoding + # dynamically overrides this value as it can need more than the last token logits + if self._supports_logits_to_keep() and "logits_to_keep" not in model_kwargs: + model_kwargs["logits_to_keep"] = 1 + + self._validate_generated_length(generation_config, input_ids_length, has_default_max_length) + + # 7. Prepare the cache. + # - `model_kwargs` may be updated in place with a cache as defined by the parameters in `generation_config`. + # - different models have a different cache name expected by the model (default = "past_key_values") + # - `max_length`, prepared above, is used to determine the maximum cache length + max_cache_length = generation_config.max_length - 1 + if ( + inputs_tensor.shape[1] != input_ids_length + and model_input_name == "inputs_embeds" + and not self.config.is_encoder_decoder + ): + max_cache_length += inputs_tensor.shape[1] + self._prepare_cache_for_generation( + generation_config, model_kwargs, assistant_model, batch_size, max_cache_length, device + ) + + # 8. determine generation mode + generation_mode = generation_config.get_generation_mode(assistant_model) + + if streamer is not None and (generation_config.num_beams > 1): + raise ValueError( + "`streamer` cannot be used with beam search (yet!). Make sure that `num_beams` is set to 1." + ) + + if self.device.type != input_ids.device.type: + warnings.warn( + "You are calling .generate() with the `input_ids` being on a device type different" + f" than your model's device. `input_ids` is on {input_ids.device.type}, whereas the model" + f" is on {self.device.type}. You may experience unexpected behaviors or slower generation." + " Please make sure that you have put `input_ids` to the" + f" correct device by calling for example input_ids = input_ids.to('{self.device.type}') before" + " running `.generate()`.", + UserWarning, + ) + + # 9. prepare logits processors and stopping criteria + prepared_logits_processor = self._get_logits_processor( + generation_config=generation_config, + input_ids_seq_length=input_ids_length, + encoder_input_ids=inputs_tensor, + prefix_allowed_tokens_fn=prefix_allowed_tokens_fn, + logits_processor=logits_processor, + device=inputs_tensor.device, + model_kwargs=model_kwargs, + negative_prompt_ids=negative_prompt_ids, + negative_prompt_attention_mask=negative_prompt_attention_mask, + ) + prepared_stopping_criteria = self._get_stopping_criteria( + generation_config=generation_config, stopping_criteria=stopping_criteria, tokenizer=tokenizer, **kwargs + ) + + # Set model_kwargs `use_cache` so we can use it later in forward runs + model_kwargs["use_cache"] = generation_config.use_cache + + # 10. go into different generation modes + if generation_mode == GenerationMode.ASSISTED_GENERATION: + if generation_config.num_return_sequences > 1: + raise ValueError( + "num_return_sequences has to be 1 when doing assisted generate, " + f"but is {generation_config.num_return_sequences}." + ) + if batch_size > 1: + raise ValueError("assisted generate is only supported for batch_size = 1") + if not model_kwargs["use_cache"]: + raise ValueError("assisted generate requires `use_cache=True`") + if generation_config.cache_implementation in ["static", "hybrid", "sliding_window"]: + raise ValueError("assisted generate is not supported with Static cache classes`") + if self._is_stateful: + # In assisted generation we need the ability to confirm whether the model would pick certain tokens, + # which is not possible with stateful models (they can't reset to a previous subset of generated text) + raise ValueError( + f"assisted generation is not supported with stateful models, such as {self.__class__.__name__}" + ) + + # 11. Get the candidate generator, given the parameterization + candidate_generator = self._get_candidate_generator( + generation_config=generation_config, + input_ids=input_ids, + inputs_tensor=inputs_tensor, + assistant_model=assistant_model, + logits_processor=logits_processor, + target_tokenizer=tokenizer, + assistant_tokenizer=assistant_tokenizer, + model_kwargs=model_kwargs, + ) + + # 12. run assisted generate + result = self._assisted_decoding( + input_ids, + candidate_generator=candidate_generator, + logits_processor=prepared_logits_processor, + stopping_criteria=prepared_stopping_criteria, + generation_config=generation_config, + synced_gpus=synced_gpus, + streamer=streamer, + **model_kwargs, + ) + elif generation_mode == GenerationMode.DOLA_GENERATION: + if self._is_stateful: + # DoLa decoding was not designed for stateful models, and would require some changes + raise ValueError( + f"dola decoding is not supported with stateful models, such as {self.__class__.__name__}" + ) + result = self._dola_decoding( + input_ids, + dola_layers=generation_config.dola_layers, + logits_processor=prepared_logits_processor, + stopping_criteria=prepared_stopping_criteria, + generation_config=generation_config, + synced_gpus=synced_gpus, + streamer=streamer, + **model_kwargs, + ) + + elif generation_mode == GenerationMode.CONTRASTIVE_SEARCH: + if not model_kwargs["use_cache"]: + raise ValueError("Contrastive search requires `use_cache=True`") + if self._is_stateful: + # Just like assisted generation, we need to be able to rollback to a previous state (see comment above) + raise ValueError( + f"contrastive search is not supported with stateful models, such as {self.__class__.__name__}" + ) + + result = self._contrastive_search( + input_ids, + logits_processor=prepared_logits_processor, + stopping_criteria=prepared_stopping_criteria, + generation_config=generation_config, + synced_gpus=synced_gpus, + streamer=streamer, + **model_kwargs, + ) + + elif generation_mode in (GenerationMode.SAMPLE, GenerationMode.GREEDY_SEARCH): + # 11. expand input_ids with `num_return_sequences` additional sequences per batch + input_ids, model_kwargs = self._expand_inputs_for_generation( + input_ids=input_ids, + expand_size=generation_config.num_return_sequences, + is_encoder_decoder=self.config.is_encoder_decoder, + **model_kwargs, + ) + + # 12. run sample (it degenerates to greedy search when `generation_config.do_sample=False`) + result = self._sample( + input_ids, + logits_processor=prepared_logits_processor, + stopping_criteria=prepared_stopping_criteria, + generation_config=generation_config, + synced_gpus=synced_gpus, + streamer=streamer, + **model_kwargs, + ) + + elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH): + # 11. interleave input_ids with `num_beams` additional sequences per batch + input_ids, model_kwargs = self._expand_inputs_for_generation( + input_ids=input_ids, + expand_size=generation_config.num_beams, + is_encoder_decoder=self.config.is_encoder_decoder, + **model_kwargs, + ) + # 12. run beam sample + result = self._beam_search( + input_ids, + logits_processor=prepared_logits_processor, + stopping_criteria=prepared_stopping_criteria, + generation_config=generation_config, + synced_gpus=synced_gpus, + **model_kwargs, + ) + + elif generation_mode == GenerationMode.GROUP_BEAM_SEARCH: + # 11. prepare beam search scorer + beam_scorer = BeamSearchScorer( + batch_size=batch_size, + num_beams=generation_config.num_beams, + device=inputs_tensor.device, + length_penalty=generation_config.length_penalty, + do_early_stopping=generation_config.early_stopping, + num_beam_hyps_to_keep=generation_config.num_return_sequences, + num_beam_groups=generation_config.num_beam_groups, + max_length=generation_config.max_length, + ) + # 12. interleave input_ids with `num_beams` additional sequences per batch + input_ids, model_kwargs = self._expand_inputs_for_generation( + input_ids=input_ids, + expand_size=generation_config.num_beams, + is_encoder_decoder=self.config.is_encoder_decoder, + **model_kwargs, + ) + # 13. run beam search + result = self._group_beam_search( + input_ids, + beam_scorer, + logits_processor=prepared_logits_processor, + stopping_criteria=prepared_stopping_criteria, + generation_config=generation_config, + synced_gpus=synced_gpus, + **model_kwargs, + ) + + elif generation_mode == GenerationMode.CONSTRAINED_BEAM_SEARCH: + final_constraints = [] + if generation_config.constraints is not None: + final_constraints = generation_config.constraints + + if generation_config.force_words_ids is not None: + + def typeerror(): + raise ValueError( + "`force_words_ids` has to either be a `List[List[List[int]]]` or `List[List[int]]` " + f"of positive integers, but is {generation_config.force_words_ids}." + ) + + if ( + not isinstance(generation_config.force_words_ids, list) + or len(generation_config.force_words_ids) == 0 + ): + typeerror() + + for word_ids in generation_config.force_words_ids: + if isinstance(word_ids[0], list): + if not isinstance(word_ids, list) or len(word_ids) == 0: + typeerror() + if any(not isinstance(token_ids, list) for token_ids in word_ids): + typeerror() + if any( + any((not isinstance(token_id, int) or token_id < 0) for token_id in token_ids) + for token_ids in word_ids + ): + typeerror() + + constraint = DisjunctiveConstraint(word_ids) + else: + if not isinstance(word_ids, list) or len(word_ids) == 0: + typeerror() + if any((not isinstance(token_id, int) or token_id < 0) for token_id in word_ids): + typeerror() + + constraint = PhrasalConstraint(word_ids) + final_constraints.append(constraint) + + # 11. prepare beam search scorer + constrained_beam_scorer = ConstrainedBeamSearchScorer( + constraints=final_constraints, + batch_size=batch_size, + num_beams=generation_config.num_beams, + device=inputs_tensor.device, + length_penalty=generation_config.length_penalty, + do_early_stopping=generation_config.early_stopping, + num_beam_hyps_to_keep=generation_config.num_return_sequences, + max_length=generation_config.max_length, + ) + # 12. interleave input_ids with `num_beams` additional sequences per batch + input_ids, model_kwargs = self._expand_inputs_for_generation( + input_ids=input_ids, + expand_size=generation_config.num_beams, + is_encoder_decoder=self.config.is_encoder_decoder, + **model_kwargs, + ) + # 13. run beam search + result = self._constrained_beam_search( + input_ids, + constrained_beam_scorer=constrained_beam_scorer, + logits_processor=prepared_logits_processor, + stopping_criteria=prepared_stopping_criteria, + generation_config=generation_config, + synced_gpus=synced_gpus, + **model_kwargs, + ) + + # Convert to legacy cache format if requested + if ( + generation_config.return_legacy_cache is True + and hasattr(result, "past_key_values") + and getattr(result.past_key_values, "to_legacy_cache") is not None + ): + result.past_key_values = result.past_key_values.to_legacy_cache() + return result + + def _has_unfinished_sequences(self, this_peer_finished: bool, synced_gpus: bool, device: torch.device) -> bool: + """ + Returns whether there are still unfinished sequences in the device. The existence of unfinished sequences is + fed through `this_peer_finished`. ZeRO stage 3-friendly. + """ + if synced_gpus: + # Under synced_gpus the `forward` call must continue until all gpus complete their sequence. + # The following logic allows an early break if all peers finished generating their sequence + this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0, device=device) + # send 0.0 if we finished, 1.0 otherwise + dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) + # did all peers finish? the reduced sum will be 0.0 then + if this_peer_finished_flag.item() == 0.0: + return False + elif this_peer_finished: + return False + return True + + def heal_tokens( + self, input_ids: torch.LongTensor, tokenizer: Optional["PreTrainedTokenizerBase"] = None + ) -> torch.LongTensor: + r""" + Generates sequences of token ids for models with a language modeling head. + Parameters: + input_ids (`torch.LongTensor`): The sequence used as a prompt for the generation. + tokenizer (`PreTrainedTokenizerBase`, *optional*): The tokenizer used to decode the input ids. + Return: + `torch.LongTensor` where each sequence has its tail token replaced with its appropriate extension. + """ + if tokenizer is None: + raise ValueError( + " When generating with token healing, you must pass the model's tokenizer to the `tokenizer` " + "argument of `generate`." + ) + + bos_token_id, pad_token_id = tokenizer.bos_token_id, tokenizer.pad_token_id + vocab_trie = ExtensionsTrie(tokenizer.get_vocab()) + generation_config = GenerationConfig(max_new_tokens=1, pad_token_id=pad_token_id) + + # assumption: leading/trailing whitespace is not meaningful, so the prompts are + # stripped before re-tokenizing to desensitize generation to whitespace artefacts + prompts = [p.strip() for p in tokenizer.batch_decode(input_ids, skip_special_tokens=True)] + input_ids = tokenizer( + prompts, + return_tensors="pt", + padding=True, + ).input_ids.to(input_ids.device) + + # replace bos with pad to not condition healing on it + input_ids = torch.where(input_ids == bos_token_id, pad_token_id, input_ids) + + """ + the latter code assumes the input_ids is not empty, + input_id has to be checked if contains elements + """ + if input_ids.numel() == 0: + return input_ids + + tail_ids = input_ids[:, -1].tolist() + + space_tok = tokenizer.convert_ids_to_tokens(tokenizer.convert_tokens_to_ids(" "))[0] + # tail tokens are used for a prefix search, thus, whitespaces are replaced with + # their tokenization (e.g. 'Ġ') to enable search for tokens prefixed with a whitespace + tail_toks = (tokenizer.decode(t).replace(" ", space_tok) for t in tail_ids) + + for batch_idx, (tail_id, tail_tok) in enumerate(zip(tail_ids, tail_toks)): + batch_ids = input_ids[batch_idx] + if torch.all(batch_ids == pad_token_id).item(): + continue # skip empty sequences (all pad ids) + + # apply bias for alternatives (extensions) to the tail token + """ + seq_bias key has to be tuple with int so have to use + tokenizer function to convert str to int + """ + seq_bias = { + (tokenizer.convert_tokens_to_ids(alt_tok),): 10.0 for alt_tok in vocab_trie.extensions(prefix=tail_tok) + } + + if len(seq_bias) == 1: + continue # skip if there are no token alternatives to heal with + + # slightly favor original token to limit aggressive healing e.g. 'http' -> 'https' + seq_bias[(tail_id,)] += 1.0 + generation_config.update(sequence_bias=seq_bias) + + trimmed_ids = batch_ids[:-1] + + """ + the latter code assumes trimmed_ids is not empty + so have to check the its element count + """ + if trimmed_ids.numel() == 0: + continue + + # if the prompt is a single (non-pad) token, regenerate from bos + if len(batch_ids[batch_ids != pad_token_id]) == 1: + trimmed_ids[-1] = bos_token_id + + input_ids[batch_idx] = self.generate(trimmed_ids.unsqueeze(0), generation_config=generation_config) + + return input_ids + + def _dola_decoding( + self, + input_ids: torch.LongTensor, + dola_layers: Union[str, List[int]], + logits_processor: LogitsProcessorList, + stopping_criteria: StoppingCriteriaList, + generation_config: GenerationConfig, + synced_gpus: bool, + streamer: "BaseStreamer", + **model_kwargs, + ) -> Union[GenerateNonBeamOutput, torch.LongTensor]: + r""" + Generates sequences of token ids for models with a language modeling head using **dola decoding** and can be + used for decoder-only text models. + The method is based on the paper "DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language + Models" (https://arxiv.org/abs/2309.03883) in ICLR 2024. + + Parameters: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + dola_layers (`Union[str, List[int]]`): + The candidate layers used in contrasting layers of DoLa. It can be either 1) 'low' or 'high', which + means the lower part or higher part of the model layers, respectively, or 2) a list of layer indices + to be used for candidate layers. The 0-th layer is the word embedding layer of the model. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + stopping_criteria (`StoppingCriteriaList`, *optional*): + An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] + used to tell if the generation loop should stop. + generation_config ([`~generation.GenerationConfig`]): + The generation configuration to be used as parametrization of the decoding method. + synced_gpus (`bool`): + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). + streamer (`BaseStreamer`, *optional*): + Streamer object that will be used to stream the generated sequences. Generated tokens are passed + through `streamer.put(token_ids)` and the streamer is responsible for any further processing. + model_kwargs: + Additional model specific keyword arguments will be forwarded to the `forward` function of the model. + If model is an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] + or `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a + [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + """ + + if self.config.is_encoder_decoder: + raise ValueError("DoLa decoding is only available for decoder-only models.") + # init values + + pad_token_id = generation_config._pad_token_tensor + output_attentions = generation_config.output_attentions + output_hidden_states = generation_config.output_hidden_states + output_scores = generation_config.output_scores + output_logits = generation_config.output_logits + return_dict_in_generate = generation_config.return_dict_in_generate + has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria) + do_sample = generation_config.do_sample + + # init attention / hidden states / scores tuples + scores = () if (return_dict_in_generate and output_scores) else None + raw_logits = () if (return_dict_in_generate and output_logits) else None + decoder_attentions = () if (return_dict_in_generate and output_attentions) else None + cross_attentions = () if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None + + # keep track of which sequences are already finished + batch_size = input_ids.shape[0] + unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) + model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) + + this_peer_finished = False + + # prepare layers for DoLa decoding + final_layer = self.config.get_text_config().num_hidden_layers + # if the model has tied word embeddings, we skip the word embeddings (0-th) layer and start from the 2nd layer, + # as the early exit from word embeddings will become identity function + # if the model is really shallow (<=2 layers), we use the 1st layer if it's not the final layer and the 0-th + # layer otherwise. Notice that DoLa does not help shallow models much. + if not self.config.tie_word_embeddings: + start_layer = 0 + elif final_layer > 2: + start_layer = 2 + elif final_layer == 2: + start_layer = 1 + else: + start_layer = 0 + + # For `N`-layer models with `N <= 40` layers, the layers of `range(0, N // 2, 2)` and `range(N // 2, N, 2)` + # are used for `'low'` and `'high'` layers, respectively. + # For models with `N > 40` layers, the layers of `range(0, 20, 2)` and `range(N - 20, N, 2)` are used for + # `'low'` and `'high'` layers, respectively. + if isinstance(dola_layers, str) and dola_layers == "low": + if start_layer == final_layer // 2: + candidate_premature_layers = [start_layer] + else: + candidate_premature_layers = ( + list(range(start_layer, final_layer // 2, 2)) + if final_layer <= 40 + else list(range(start_layer, 20, 2)) + ) + elif isinstance(dola_layers, str) and dola_layers == "high": + candidate_premature_layers = ( + list(range(final_layer // 2, final_layer, 2)) + if final_layer <= 40 + else list(range(final_layer - 20, final_layer, 2)) + ) + # Set the `dola_layers` to a list of integers for layer indices to contrast manually specified layers. + elif isinstance(dola_layers, list): + candidate_premature_layers = [i for i in dola_layers if i < final_layer] + else: + raise ValueError("dola_layers must be either 'low', 'high' or a list of integers.") + + lm_head = self.get_output_embeddings() + if lm_head is None: + raise ValueError("DoLa is not supported for models that don't have output embeddings.") + + while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): + # prepare model inputs + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + + # forward pass to get next token + outputs = self( + **model_inputs, + return_dict=True, + output_attentions=output_attentions, + output_hidden_states=True, + ) + + # .float() is needed to retain precision for later logits manipulations + final_layer_next_token_logits = outputs.logits[:, -1, :].detach().to(copy=True, dtype=torch.float32) + final_logits = outputs.logits[:, -1, :].float() + candidate_premature_logits = {} + for candidate_premature_layer in candidate_premature_layers: + candidate_premature_logits[candidate_premature_layer] = lm_head( + outputs.hidden_states[candidate_premature_layer][:, -1, :] + ).to(final_logits.device) + + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + if synced_gpus and this_peer_finished: + continue + + next_token_logits = _dola_select_contrast( + candidate_premature_layers, candidate_premature_logits, final_logits + ) + next_token_logits = next_token_logits.to(input_ids.device) + # pre-process distribution + next_token_scores = logits_processor(input_ids, next_token_logits) + + # Store scores, attentions and hidden_states when required + if return_dict_in_generate: + if output_scores: + scores += (next_token_scores,) + if output_logits: + raw_logits += (final_layer_next_token_logits,) + if output_attentions: + decoder_attentions += ( + (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) + ) + if self.config.is_encoder_decoder: + cross_attentions += (outputs.cross_attentions,) + + if output_hidden_states: + decoder_hidden_states += ( + (outputs.decoder_hidden_states,) + if self.config.is_encoder_decoder + else (outputs.hidden_states,) + ) + + if do_sample: # sample + probs = nn.functional.softmax(next_token_scores, dim=-1) + next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) + else: # argmax + next_tokens = torch.argmax(next_token_scores, dim=-1) + + # finished sentences should have their next token be a padding token + if has_eos_stopping_criteria: + next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) + + # update generated ids, model inputs, and length for next step + input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) + if streamer is not None: + streamer.put(next_tokens.cpu()) + + # stop when each sentence is finished + unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores) + this_peer_finished = unfinished_sequences.max() == 0 + + if streamer is not None: + streamer.end() + + if return_dict_in_generate: + return GenerateDecoderOnlyOutput( + sequences=input_ids, + scores=scores, + logits=raw_logits, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return input_ids + + @torch.no_grad() + def _contrastive_search( + self, + input_ids: torch.LongTensor, + logits_processor: LogitsProcessorList, + stopping_criteria: StoppingCriteriaList, + generation_config: GenerationConfig, + synced_gpus: bool, + streamer: Optional["BaseStreamer"], + **model_kwargs, + ) -> Union[GenerateNonBeamOutput, torch.LongTensor]: + r""" + Generates sequences of token ids for models with a language modeling head using **contrastive search** and can + be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. + + Parameters: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + stopping_criteria (`StoppingCriteriaList`): + An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] + used to tell if the generation loop should stop. + generation_config ([`~generation.GenerationConfig`]): + The generation configuration to be used as parametrization of the decoding method. + synced_gpus (`bool`): + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). + streamer (`BaseStreamer`, *optional*): + Streamer object that will be used to stream the generated sequences. Generated tokens are passed + through `streamer.put(token_ids)` and the streamer is responsible for any further processing. + model_kwargs: + Additional model specific keyword arguments will be forwarded to the `forward` function of the model. + If model is an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] + or `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a + [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + """ + # init values + has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria) + top_k = generation_config.top_k + penalty_alpha = generation_config.penalty_alpha + pad_token_id = generation_config._pad_token_tensor + output_attentions = generation_config.output_attentions + output_hidden_states = generation_config.output_hidden_states + output_scores = generation_config.output_scores + output_logits = generation_config.output_logits + return_dict_in_generate = generation_config.return_dict_in_generate + sequential = generation_config.low_memory + + # init attention / hidden states / scores tuples + raw_logits = () if (return_dict_in_generate and output_logits) else None + scores = () if (return_dict_in_generate and output_scores) else None + decoder_attentions = () if (return_dict_in_generate and output_attentions) else None + cross_attentions = () if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None + + # if model is an encoder-decoder, retrieve encoder attention weights and hidden states + if return_dict_in_generate and self.config.is_encoder_decoder: + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + # keep track of which sequences are already finished + batch_size = input_ids.shape[0] + unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) + model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) + + # Create cosine_matrix_mask based on the attention_mask + cosine_matrix_mask = torch.ones_like(input_ids, dtype=torch.long) + if self.config.is_encoder_decoder: + if "decoder_attention_mask" in model_kwargs and model_kwargs["decoder_attention_mask"] is not None: + cosine_matrix_mask = model_kwargs["decoder_attention_mask"] + else: + cosine_matrix_mask = model_kwargs["attention_mask"] + cosine_matrix_mask = cosine_matrix_mask.repeat_interleave(top_k, dim=0) + + this_peer_finished = False + + while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): + # if the first step in the loop, encode all the prefix and obtain: (1) past_key_values; + # (2) last_hidden_states; (3) logit_for_next_step; (4) update model kwargs for the next step + if model_kwargs.get("past_key_values") is None or ( + isinstance(model_kwargs["past_key_values"], (Cache, EncoderDecoderCache)) + and model_kwargs["past_key_values"].get_seq_length() == 0 + ): + # prepare inputs + model_kwargs["use_cache"] = True + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + + # encode the given prefix and prepare model inputs; encoder-decoder model process the prefix and save + # the `encoder_outputs` + outputs = self( + **model_inputs, return_dict=True, output_hidden_states=True, output_attentions=output_attentions + ) + + # last decoder hidden states will be used to compute the degeneration penalty (cosine similarity with + # previous tokens) + if self.config.is_encoder_decoder: + last_hidden_states = outputs.decoder_hidden_states[-1] + else: + last_hidden_states = outputs.hidden_states[-1] + + # next logit for contrastive search to select top-k candidate tokens + # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for this first iteration + # (the clone itself is always small) + # torch.float32 is needed to retain precision for later logits manipulations + logit_for_next_step = outputs.logits[:, -1, :].to( + copy=True, dtype=torch.float32, device=input_ids.device + ) + + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + + if not sequential: + # Expands model inputs top_k times, for batched forward passes (akin to beam search). + # input_ids is required for expanding visual inputs in qwen2vl + _, model_kwargs = self._expand_inputs_for_generation( + input_ids=input_ids, + expand_size=top_k, + is_encoder_decoder=self.config.is_encoder_decoder, + **model_kwargs, + ) + + past_key_values = model_kwargs.get("past_key_values") + if past_key_values is None: + raise ValueError( + f"{self.__class__.__name__} does not support caching and therefore **can't** be used " + "for contrastive search." + ) + elif ( + not isinstance(past_key_values[0], (tuple, torch.Tensor)) + or past_key_values[0][0].shape[0] != batch_size + ): + raise ValueError( + f"{self.__class__.__name__} does not have a standard cache format and therefore **can't** be " + "used for contrastive search without further modifications." + ) + + # contrastive_search main logic start: + # contrastive search decoding consists of two steps: (1) candidate tokens recall; (2) candidate re-rank by + # degeneration penalty + processed_logit_for_next_step = logits_processor(input_ids, logit_for_next_step) + next_probs = nn.functional.softmax(processed_logit_for_next_step, dim=-1) + + top_k_probs, top_k_ids = torch.topk(next_probs, dim=-1, k=top_k) + + # Store scores, attentions and hidden_states when required + if return_dict_in_generate: + if output_logits: + raw_logits += (logit_for_next_step,) + if output_scores: + scores += (processed_logit_for_next_step,) + if output_attentions: + decoder_attentions += ( + (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) + ) + if self.config.is_encoder_decoder: + cross_attentions += (outputs.cross_attentions,) + + if output_hidden_states: + decoder_hidden_states += ( + (outputs.decoder_hidden_states,) + if self.config.is_encoder_decoder + else (outputs.hidden_states,) + ) + + # This is needed to properly delete outputs.logits which may be very large for this first iteration + # Otherwise a reference to outputs.logits is kept all along until after the next call to self.forward() + del outputs + + if not sequential: + # Replicates the new past_key_values to match the `top_k` candidates + past = model_kwargs["past_key_values"] + # If it is a static cache, modify it in-place layer after layer to save memory + if isinstance(past, DynamicCache) or ( + isinstance(past, EncoderDecoderCache) and isinstance(past.self_attention_cache, DynamicCache) + ): + past.batch_repeat_interleave(top_k) + else: + new_key_values = [] + for layer in past: + items = [] + # item is either the key or the value matrix + for item in layer: + items.append(item.repeat_interleave(top_k, dim=0)) + new_key_values.append(tuple(items)) + + past = tuple(new_key_values) + + model_kwargs["past_key_values"] = past + + if sequential: + all_outputs = [] + for i in range(top_k): + # compute the candidate tokens by the language model and collect their hidden_states + next_model_inputs = self.prepare_inputs_for_generation(top_k_ids[:, i].view(-1, 1), **model_kwargs) + + outputs = self( + **next_model_inputs, + return_dict=True, + output_hidden_states=True, + output_attentions=output_attentions, + ) + if isinstance(outputs["past_key_values"], DynamicCache) or ( + isinstance(outputs["past_key_values"], EncoderDecoderCache) + and isinstance(outputs["past_key_values"].self_attention_cache, DynamicCache) + ): + # Remove past K-V from output since we don't need to stack later + outputs["past_key_values"] = None + # Remove last token from past K-V since we don't want to append it at this point + model_kwargs["past_key_values"].crop(-1) + + all_outputs.append(outputs) + outputs = stack_model_outputs(all_outputs, self.config.get_text_config()) + + else: + # compute the candidate tokens by the language model and collect their hidden_states + # assembles top_k_ids into batch of size k + next_model_inputs = self.prepare_inputs_for_generation(top_k_ids.view(-1, 1), **model_kwargs) + + outputs = self( + **next_model_inputs, + return_dict=True, + output_hidden_states=True, + output_attentions=output_attentions, + ) + + # This is essential to avoid having a last reference to the big past K-V and double the necessary memory + # in the next loop + del next_model_inputs + + # name is different for encoder-decoder and decoder-only models + if self.config.is_encoder_decoder: + next_hidden = outputs.decoder_hidden_states[-1] + full_hidden_states = outputs.decoder_hidden_states + else: + next_hidden = outputs.hidden_states[-1] + full_hidden_states = outputs.hidden_states + + # .float() is needed to retain precision for later logits manipulations + logits = outputs.logits[:, -1, :].float() + context_hidden = last_hidden_states.repeat_interleave(top_k, dim=0) + + # compute the degeneration penalty and re-rank the candidates based on the degeneration penalty and the + # model confidence. Keeping `selected_idx` on CPU enables multi-device contrastive search and doesn't + # introduce (noticeable) slowdowns on single-device runs. + selected_idx = _ranking_fast( + context_hidden, next_hidden, top_k_probs, cosine_matrix_mask, penalty_alpha, top_k + ) + cosine_matrix_mask = torch.cat( + [cosine_matrix_mask, cosine_matrix_mask.new_ones((cosine_matrix_mask.shape[0], 1))], dim=-1 + ) + selected_idx = selected_idx.to("cpu") + + # This will be used instead of the previous inneficient torch.stack(torch.split()) + augmented_idx = torch.tensor([x + i * top_k for i, x in enumerate(selected_idx)]) + + # prepare for the next step: (1) next token_id; (2) past_key_values; (3) last_hidden_states for computing + # the degeneration penalty; (4) logits for selecting next top-k candidates; (5) selected tokens scores + # (model confidence minus degeneration penalty); (6) decoder hidden_states + next_tokens = top_k_ids[range(len(top_k_ids)), selected_idx] + next_hidden = torch.stack(torch.split(next_hidden.squeeze(dim=1), top_k)) + next_hidden = next_hidden[range(batch_size), selected_idx, :] + last_hidden_states = torch.cat([last_hidden_states, next_hidden.unsqueeze(1)], dim=1) + + next_decoder_hidden_states = () + for layer in full_hidden_states: + layer = torch.stack(torch.split(layer, top_k))[range(batch_size), selected_idx, :] + next_decoder_hidden_states += (layer,) + + # generate past_key_values cache of only the selected token + if sequential: + next_model_input = self.prepare_inputs_for_generation( + top_k_ids[:, selected_idx].view(-1, 1), **model_kwargs + ) + + selected_outputs = self( + **next_model_input, + return_dict=True, + output_hidden_states=False, + output_attentions=False, + ) + next_past_key_values = selected_outputs["past_key_values"] + + else: + next_past_key_values = None + for possible_cache_name in ALL_CACHE_NAMES: + next_past_key_values = next_past_key_values or getattr(outputs, possible_cache_name, None) + # Do it in-place layer per layer to save memory + if isinstance(next_past_key_values, DynamicCache) or ( + isinstance(next_past_key_values, EncoderDecoderCache) + and isinstance(next_past_key_values.self_attention_cache, DynamicCache) + ): + next_past_key_values.batch_select_indices(augmented_idx) + else: + new_key_values = [] + for layer in next_past_key_values: + items = [] + # item is either the key or the value matrix + for item in layer: + items.append(item[augmented_idx, ...]) + new_key_values.append(tuple(items)) + + next_past_key_values = tuple(new_key_values) + + logit_for_next_step = torch.stack(torch.split(logits, top_k))[range(batch_size), selected_idx, :] + logit_for_next_step = logit_for_next_step.to(input_ids.device) + + # Rebuilds the relevant parts of the model output for the selected token, for use in the next iteration + if self.config.is_encoder_decoder: + next_step_cross_attentions = () + next_step_decoder_attentions = () + if output_attentions: + for layer in outputs.cross_attentions: + layer = torch.stack(torch.split(layer, top_k, dim=0))[range(batch_size), selected_idx, ...] + next_step_cross_attentions += (layer,) + for layer in outputs.decoder_attentions: + layer = torch.stack(torch.split(layer, top_k, dim=0))[range(batch_size), selected_idx, ...] + next_step_decoder_attentions += (layer,) + outputs = Seq2SeqLMOutput( + past_key_values=next_past_key_values, + decoder_hidden_states=next_decoder_hidden_states, + decoder_attentions=next_step_decoder_attentions or None, + cross_attentions=next_step_cross_attentions or None, + ) + else: + next_step_attentions = () + if output_attentions: + for layer in outputs.attentions: + layer = torch.stack(torch.split(layer, top_k, dim=0))[range(batch_size), selected_idx, ...] + next_step_attentions += (layer,) + outputs = CausalLMOutputWithPast( + past_key_values=next_past_key_values, + hidden_states=next_decoder_hidden_states, + attentions=next_step_attentions or None, + ) + # contrastive_search main logic end + + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + if synced_gpus and this_peer_finished: + continue + + # finished sentences should have their next token be a padding token + if has_eos_stopping_criteria: + next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) + + # update generated ids, model inputs, and length for next step + input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) + if streamer is not None: + streamer.put(next_tokens.cpu()) + + # stop when each sentence is finished + unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores) + this_peer_finished = unfinished_sequences.max() == 0 + + if streamer is not None: + streamer.end() + + if return_dict_in_generate: + # Contrastive search works by forward looking at the next token, so we need to exclude it from + # `past_key_values` to be consistent with the other decoding methods + if model_kwargs.get("past_key_values") is not None: + if isinstance(model_kwargs["past_key_values"], DynamicCache) or ( + isinstance(model_kwargs["past_key_values"], EncoderDecoderCache) + and isinstance(model_kwargs["past_key_values"].self_attention_cache, DynamicCache) + ): + model_kwargs["past_key_values"].crop(-1) + else: + past_key_values = [] + for layer in model_kwargs["past_key_values"]: + layer_past_key_values = [] + for item in layer: + layer_past_key_values.append(item[..., :-1, :]) + past_key_values.append(tuple(layer_past_key_values)) + model_kwargs["past_key_values"] = tuple(past_key_values) + + if self.config.is_encoder_decoder: + return GenerateEncoderDecoderOutput( + sequences=input_ids, + scores=scores, + logits=raw_logits, + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return GenerateDecoderOnlyOutput( + sequences=input_ids, + scores=scores, + logits=raw_logits, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return input_ids + + def _sample( + self, + input_ids: torch.LongTensor, + logits_processor: LogitsProcessorList, + stopping_criteria: StoppingCriteriaList, + generation_config: GenerationConfig, + synced_gpus: bool, + streamer: Optional["BaseStreamer"], + **model_kwargs, + ) -> Union[GenerateNonBeamOutput, torch.LongTensor]: + r""" + Generates sequences of token ids for models with a language modeling head using **multinomial sampling** and + can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. + + Parameters: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + stopping_criteria (`StoppingCriteriaList`): + An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] + used to tell if the generation loop should stop. + generation_config ([`~generation.GenerationConfig`]): + The generation configuration to be used as parametrization of the decoding method. + synced_gpus (`bool`): + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). + streamer (`BaseStreamer`, *optional*): + Streamer object that will be used to stream the generated sequences. Generated tokens are passed + through `streamer.put(token_ids)` and the streamer is responsible for any further processing. + model_kwargs: + Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is + an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or `torch.LongTensor`: + A `torch.LongTensor` containing the generated tokens (default behaviour) or a + [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + """ + # init values + pad_token_id = generation_config._pad_token_tensor + output_attentions = generation_config.output_attentions + output_hidden_states = generation_config.output_hidden_states + output_scores = generation_config.output_scores + output_logits = generation_config.output_logits + return_dict_in_generate = generation_config.return_dict_in_generate + has_eos_stopping_criteria = any(hasattr(criteria, "eos_token_id") for criteria in stopping_criteria) + do_sample = generation_config.do_sample + + # init attention / hidden states / scores tuples + scores = () if (return_dict_in_generate and output_scores) else None + raw_logits = () if (return_dict_in_generate and output_logits) else None + decoder_attentions = () if (return_dict_in_generate and output_attentions) else None + cross_attentions = () if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None + + # if model is an encoder-decoder, retrieve encoder attention weights and hidden states + if return_dict_in_generate and self.config.is_encoder_decoder: + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + # keep track of which sequences are already finished + batch_size, cur_len = input_ids.shape + this_peer_finished = False + unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) + model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) + + model_forward = self.__call__ + if isinstance(model_kwargs.get("past_key_values"), Cache): + is_compileable = model_kwargs["past_key_values"].is_compileable and self._supports_static_cache + if getattr(self, "hf_quantizer", None) is not None: + is_compileable &= self.hf_quantizer.is_compileable + is_compileable = is_compileable and not generation_config.disable_compile + if is_compileable and ( + self.device.type == "cuda" or generation_config.compile_config._compile_all_devices + ): + os.environ["TOKENIZERS_PARALLELISM"] = "0" + model_forward = self.get_compiled_call(generation_config.compile_config) + + if generation_config.prefill_chunk_size is not None: + model_kwargs = self._prefill_chunking(input_ids, generation_config, **model_kwargs) + is_prefill = False + else: + is_prefill = True + + while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): + # prepare model inputs + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + + # prepare variable output controls (note: some models won't accept all output controls) + model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) + model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) + + if is_prefill: + outputs = self(**model_inputs, return_dict=True) + is_prefill = False + else: + outputs = model_forward(**model_inputs, return_dict=True) + + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + if synced_gpus and this_peer_finished: + continue + + # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration + # (the clone itself is always small) + next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device) + + # pre-process distribution + next_token_scores = logits_processor(input_ids, next_token_logits) + + # Store scores, attentions and hidden_states when required + if return_dict_in_generate: + if output_scores: + scores += (next_token_scores,) + if output_logits: + raw_logits += (next_token_logits,) + if output_attentions: + decoder_attentions += ( + (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) + ) + if self.config.is_encoder_decoder: + cross_attentions += (outputs.cross_attentions,) + + if output_hidden_states: + decoder_hidden_states += ( + (outputs.decoder_hidden_states,) + if self.config.is_encoder_decoder + else (outputs.hidden_states,) + ) + + # token selection + if do_sample: + probs = nn.functional.softmax(next_token_scores, dim=-1) + # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution + next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) + else: + next_tokens = torch.argmax(next_token_scores, dim=-1) + + # finished sentences should have their next token be a padding token + if has_eos_stopping_criteria: + next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) + + # update generated ids, model inputs, and length for next step + input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) + if streamer is not None: + streamer.put(next_tokens.cpu()) + + unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores) + this_peer_finished = unfinished_sequences.max() == 0 + cur_len += 1 + + # This is needed to properly delete outputs.logits which may be very large for first iteration + # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration + del outputs + + if streamer is not None: + streamer.end() + + if return_dict_in_generate: + if self.config.is_encoder_decoder: + return GenerateEncoderDecoderOutput( + sequences=input_ids, + scores=scores, + logits=raw_logits, + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return GenerateDecoderOnlyOutput( + sequences=input_ids, + scores=scores, + logits=raw_logits, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return input_ids + + # Auxiliary functions for beam search + def _temporary_reorder_cache(self, past_key_values, beam_idx): + """ + Temporary function to handle the different types of cache reordering processes while we roll out `Cache`. + + TODO: standardize cache formats and make all models compatible with `Cache`. It would remove the need + for this function, with `Cache.reorder_cache` being the sole remaining code path + """ + model_class = self.__class__.__name__.lower() + # Exception 1: code path for models using the legacy cache format + if isinstance(past_key_values, (tuple, list)): + past_key_values = self._reorder_cache(past_key_values, beam_idx) + # Exception 2: models with different cache formats. These are limited to `DynamicCache` until their + # cache format is standardized, to avoid adding complexity to the codebase. + elif "gptbigcode" in model_class: + if not isinstance(past_key_values, (DynamicCache, EncoderDecoderCache)): + raise ValueError( + f"Using an unsupported cache format with {model_class}. Currently, it only supports the " + "legacy tuple format or `DynamicCache`" + ) + past_key_values = self._reorder_cache(past_key_values, beam_idx) + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + # Standard code path: use the `Cache.reorder_cache` + else: + past_key_values.reorder_cache(beam_idx) + return past_key_values + + @staticmethod + def _flatten_beam_dim(tensor: torch.Tensor) -> torch.Tensor: + """[batch_size, num_beams, ...] -> [batch_size * num_beams, ...]""" + shape = list(tensor.shape) + return torch.reshape(tensor, [shape[0] * shape[1]] + shape[2:]) + + @staticmethod + def _unflatten_beam_dim(tensor: torch.Tensor, batch_size: int, num_beams: int) -> torch.Tensor: + """[batch_size * num_beams, ...] -> [batch_size, num_beams, ...]""" + shape = list(tensor.shape) + return torch.reshape(tensor, [batch_size, num_beams] + shape[1:]) + + @staticmethod + def _gather_beams(tensor: torch.Tensor, beam_indices: torch.Tensor) -> torch.Tensor: + """ + Gathers the beam slices indexed by beam_indices into new beam array. + + Args: + tensor (`torch.Tensor`): A tensor containing data to be gathered. The tensor is a 2D or a 3D tensor + with the two first dimensions depicting the batch and the beam dimensions. + beam_indices (`torch.Tensor` of shape `(batch_size, num_beams_to_select)`): The indices of the beams to + select . + + Returns: + A tensor with the selected beams + """ + # `take_along_dim` requires its indices arg to have the same number of dims as `input` + while len(beam_indices.shape) < len(tensor.shape): + beam_indices = beam_indices.unsqueeze(-1) + gathered_tensor = torch.take_along_dim(input=tensor, indices=beam_indices, dim=1) + return gathered_tensor + + @staticmethod + def _beam_search_has_unfinished_sequences( + running_beam_scores: torch.Tensor, + beam_scores: torch.Tensor, + is_sent_finished: torch.Tensor, + next_token_hits_stopping_criteria: torch.Tensor, + cur_len: int, + max_length: int, + decoder_prompt_len: int, + early_stopping: Union[bool, str], + length_penalty: float, + ): + """ + Beam Search stopping condition -- halts the generation loop if any of these conditions becomes False + """ + # a. Can the open beams improve the top completed scores? + # early_stopping == False -> apply heuristic = always get the best score from + # `cur_len - decoder_prompt_len`. See the discussion below for more details. + # https://github.com/huggingface/transformers/pull/20901#issuecomment-1369845565 + # early_stopping == "never" -> compute the best score from `max_length` or `cur_len`, depending on the + # sign of `length_penalty`. Positive `length_penalty` favors longer sequences, thus we use + # `max_length` there. + if early_stopping == "never" and length_penalty > 0.0: + best_hypothetical_length = max_length - decoder_prompt_len + else: + best_hypothetical_length = cur_len - decoder_prompt_len + best_possible_running_score = running_beam_scores[:, :1] / (best_hypothetical_length**length_penalty) + worst_finished_score = torch.where(is_sent_finished, torch.min(beam_scores, dim=1, keepdim=True)[0], -1.0e9) + improvement_possible = torch.any(best_possible_running_score > worst_finished_score) + + # b. Is there still a beam without fully completed sequences? This is only relevant if early_stopping is + # enabled, where we want to finish as soon as all beams have a completed sequence. + exists_open_beam = ~(torch.all(is_sent_finished) & (early_stopping is True)) + + # c. Have we hit a stopping criteria with all running sequences and have no way to continue? e.g. we have + # reached `max_length`` + valid_continuations = ~torch.all(next_token_hits_stopping_criteria) + + return improvement_possible & exists_open_beam & valid_continuations + + def _get_top_k_continuations( + self, + accumulated_log_probs: torch.Tensor, + running_sequences: torch.Tensor, + running_beam_indices: torch.Tensor, + cur_len: int, + decoder_prompt_len: int, + do_sample: bool, + beams_to_keep: int, + num_beams: int, + vocab_size: int, + batch_size: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Get top-K continuations given the accumulated log probs on the next token. + + A few notes to understand what's going on: + 1. Each item in batch has `num_beams` * `vocab_size` candidate continuations. For each item, get the + top K [K = (number of EOS tokens + 1) * `num_beams`] candidates with the highest accumulated + log-probabilities, or sample them without replacement using the accumulated scores + 2. We gather the top K (as opposed to `num_beams`, or any number lower than K) here so that we have at + least `num_beams` sequences remaining to continue the live beam search. + 3. Note that other stopping criteria might result in impossible to continue beams, i.e. all continuations + selected in this step hit the stopping criteria. + """ + # TODO (joao): This function should take an optional beam scorer function, to manipulate the scores after + # token selection. The function should be an argument exposed, so that custom scoring functions can be + # defined. + + # Gather the top K scores from _all_ beams. + if do_sample: + topk_indices = torch.multinomial( + nn.functional.softmax(accumulated_log_probs, dim=-1), num_samples=beams_to_keep + ) + topk_log_probs = torch.gather(input=accumulated_log_probs, dim=1, index=topk_indices) + else: + topk_log_probs, topk_indices = torch.topk(accumulated_log_probs, k=beams_to_keep) + + # Gather K top beams, recover the beam index by floor division and token id by modulo division + topk_current_beam_indices = topk_indices // vocab_size + topk_running_beam_indices = self._gather_beams(running_beam_indices, topk_current_beam_indices) + topk_running_sequences = self._gather_beams(running_sequences, topk_current_beam_indices) + topk_ids = topk_indices % vocab_size + + # Update sequences for the K top-k new sequences. + topk_running_sequences[:, :, cur_len] = topk_ids + + # we want to store the beam indices with batch information -> real beam index = beam index % num beams + batch_offset = torch.arange(batch_size, device=topk_ids.device).view(-1, 1) * num_beams + batch_modified_indices = topk_current_beam_indices + batch_offset + topk_running_beam_indices[:, :, cur_len - decoder_prompt_len] = batch_modified_indices + + return topk_log_probs, topk_running_sequences, topk_running_beam_indices + + def _get_running_beams_for_next_iteration( + self, + topk_log_probs: torch.Tensor, + topk_running_sequences: torch.Tensor, + topk_running_beam_indices: torch.Tensor, + next_token_hits_stopping_criteria: torch.Tensor, + num_beams: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Given the top-K continuations, their scores, and whether they hit a stopping criteria, select the + best non-finished beams to continue beam search in the next iteration. + """ + # To prevent these just finished sequences from being used in subsequent iterations, set their log probs + # to a very large negative value + topk_running_log_probs = topk_log_probs + next_token_hits_stopping_criteria.to(torch.float32) * -1.0e9 + + next_topk_indices = torch.topk(topk_running_log_probs, k=num_beams)[1] + running_sequences = self._gather_beams(topk_running_sequences, next_topk_indices) + running_beam_scores = self._gather_beams(topk_running_log_probs, next_topk_indices) + running_beam_indices = self._gather_beams(topk_running_beam_indices, next_topk_indices) + return running_sequences, running_beam_scores, running_beam_indices + + def _update_finished_beams( + self, + sequences: torch.Tensor, + topk_running_sequences: torch.Tensor, + beam_scores: torch.Tensor, + topk_log_probs: torch.Tensor, + beam_indices: torch.Tensor, + topk_running_beam_indices: torch.Tensor, + is_sent_finished: torch.Tensor, + next_token_hits_stopping_criteria: torch.Tensor, + top_num_beam_mask: torch.Tensor, + num_beams: int, + cur_len: int, + decoder_prompt_len: int, + length_penalty: float, + early_stopping: Union[bool, str], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Updates the finished beams if (and only if) there are new completed sequences that have a higher score than + the current finished sequences. + """ + # Only the top `num_beam` sequences can be considered for the final returned sequences. Remember: the + # remaining sequences only exist as a backup to ensure that we have at least `num_beams` sequences to + # continue. + did_top_num_beams_just_finished = next_token_hits_stopping_criteria & top_num_beam_mask[None, :] + + # Further process topk logits for the finished beams + # - add length penalty + topk_log_probs = topk_log_probs / ((cur_len + 1 - decoder_prompt_len) ** length_penalty) + # - make sure no scores can be added anymore if beam is full and early stopping is on + beams_in_batch_are_full = torch.all(is_sent_finished, axis=-1, keepdims=True) & (early_stopping is True) + topk_log_probs += beams_in_batch_are_full.to(torch.float32) * -1.0e9 + # - make sure still running sequences cannot be chosen as finalized beam + topk_log_probs += (~did_top_num_beams_just_finished) * -1.0e9 + + # Get finalized `num_beam` sequences for the next generation step -- combine the previous finalized + # data with the new finalized sequences (if any, non-finalized sequences have a very large negative score + # in this step), and keep the best `num_beams` sequences. + merged_sequences = torch.cat((sequences, topk_running_sequences), dim=1) + merged_scores = torch.cat((beam_scores, topk_log_probs), dim=1) + merged_beam_indices = torch.cat((beam_indices, topk_running_beam_indices), dim=1) + merged_is_sent_finished = torch.cat((is_sent_finished, did_top_num_beams_just_finished), dim=1) + topk_merged_indices = torch.topk(merged_scores, k=num_beams)[1] + sequences = self._gather_beams(merged_sequences, topk_merged_indices) + beam_scores = self._gather_beams(merged_scores, topk_merged_indices) + beam_indices = self._gather_beams(merged_beam_indices, topk_merged_indices) + is_sent_finished = self._gather_beams(merged_is_sent_finished, topk_merged_indices) + return sequences, beam_scores, beam_indices, is_sent_finished + + # end of auxiliary functions for beam search + + def _beam_search( + self, + input_ids: torch.LongTensor, + logits_processor: LogitsProcessorList, + stopping_criteria: StoppingCriteriaList, + generation_config: GenerationConfig, + synced_gpus: bool, + **model_kwargs, + ) -> Union[GenerateBeamOutput, torch.LongTensor]: + r""" + Generates sequences of token ids for models with a language modeling head using **beam search decoding** and + can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. + + If it's the first time you're diving into Beam Search, we recommend you read the following blog post: + https://huggingface.co/blog/how-to-generate (especially the beam search section). + + You can recompute the sequence scores from the individual scores using the `compute_transition_scores` function + (https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationMixin.compute_transition_scores) + + Parameters: + input_ids (`torch.LongTensor` of shape `(batch_size*num_beams, sequence_length)`): + The sequence used as a prompt for the generation. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + stopping_criteria (`StoppingCriteriaList`: + An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] + used to tell if the generation loop should stop. + generation_config ([`~generation.GenerationConfig`]): + The generation configuration to be used as parametrization of the decoding method. + synced_gpus (`bool`): + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). + model_kwargs: + Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is + an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`generation.GenerateBeamDecoderOnlyOutput`], [`~generation.GenerateBeamEncoderDecoderOutput`] or + `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a + [`~generation.GenerateBeamDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.GenerateBeamEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + """ + + # 1. init beam_search values + pad_token_id = generation_config._pad_token_tensor + eos_token_id = generation_config._eos_token_tensor + output_attentions = generation_config.output_attentions + output_hidden_states = generation_config.output_hidden_states + output_scores = generation_config.output_scores + output_logits = generation_config.output_logits + return_dict_in_generate = generation_config.return_dict_in_generate + do_sample = generation_config.do_sample + early_stopping = generation_config.early_stopping + length_penalty = generation_config.length_penalty + max_length = generation_config.max_length + num_beams = generation_config.num_beams + num_return_sequences = generation_config.num_return_sequences + + batch_size_unflattened, cur_len = input_ids.shape + batch_size = batch_size_unflattened // num_beams + # TODO (joao): standardize special cases + if self.__class__.__name__ == "MoshiDepthDecoder": + vocab_size = self.config.audio_vocab_size + elif self.__class__.__name__ == "ImageGPTForCausalImageModeling": + vocab_size = self.get_output_embeddings().out_features + else: + vocab_size = self.config.get_text_config().vocab_size + decoder_prompt_len = cur_len + this_peer_finished = False + + # At each beam search step, we want to keep top K [K = (number of EOS tokens + 1) * `num_beams`] candidates + # with the highest log-probabilities, or sample K continuations without replacement. We gather the top K + # (as opposed to `num_beams`, or any number lower than K) so that we have at least `num_beams` sequences + # non-finished to continue the live beam search, in case the top `num_beams` all select an EOS token. + n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0 + beams_to_keep = max(2, 1 + n_eos_tokens) * num_beams + top_num_beam_mask = torch.cat( + (torch.ones((num_beams), dtype=torch.bool), torch.zeros((beams_to_keep - num_beams), dtype=torch.bool)), + dim=0, + ).to(input_ids.device) + + model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) + + # (joao) feature lost in the refactor. Probably won't implement, hurts readbility with minimal gains (there + # are newer low-memory alternatives like the offloaded cache) + sequential = generation_config.low_memory + if sequential: + raise ValueError( + "`low_memory=True` is not supported after the beam search refactor. Please check the discussion in " + "#35802 *after the PR got merged*, and add a comment there if your questions are not yet answered." + ) + + # 2. init output tuples + all_scores = () if (return_dict_in_generate and output_scores) else None + raw_logits = () if (return_dict_in_generate and output_logits) else None + beam_indices = () if (return_dict_in_generate and output_logits) else None + decoder_attentions = () if (return_dict_in_generate and output_attentions) else None + cross_attentions = () if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None + + # if model is an encoder-decoder, retrieve encoder attention weights and hidden states + if return_dict_in_generate and self.config.is_encoder_decoder: + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + # 3. init running tensors and static-shaped placeholders + + # per batch, beam-item holding current token in loop and completed sequences + output_fill_value = pad_token_id or eos_token_id[0] if eos_token_id is not None else -1 + running_sequences = torch.full( + (batch_size, num_beams, max_length), + fill_value=output_fill_value, + dtype=torch.int64, + device=input_ids.device, + ) + running_sequences[:, :, :cur_len] = self._unflatten_beam_dim(input_ids, batch_size, num_beams) + sequences = running_sequences.detach().clone() + + # per batch, beam-item score, logprobs + # initialise score of first beam with 0 and the rest with -1e9. This makes sure that only tokens + # of the first beam are considered to avoid sampling the exact same tokens across all beams. + running_beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device) + running_beam_scores[:, 1:] = -1e9 + beam_scores = torch.full((batch_size, num_beams), fill_value=-1e9, dtype=torch.float, device=input_ids.device) + + # per batch, beam-item state bit indicating if sentence has finished. + is_sent_finished = torch.zeros((batch_size, num_beams), dtype=torch.bool, device=input_ids.device) + + # per batch, beam-item state bit indicating if there are valid continuations. + next_token_hits_stopping_criteria = torch.zeros( + (batch_size, num_beams), dtype=torch.bool, device=input_ids.device + ) + + # per batch selected beam indices + running_beam_indices = torch.full( + (batch_size, num_beams, max_length - cur_len), fill_value=-1, dtype=torch.int32, device=input_ids.device + ) + beam_indices = running_beam_indices.detach().clone() + + # 4. run the generation loop + while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): + # a. Forward current tokens, obtain the logits + flat_running_sequences = self._flatten_beam_dim(running_sequences[:, :, :cur_len]) + model_inputs = self.prepare_inputs_for_generation(flat_running_sequences, **model_kwargs) + + # prepare variable output controls (note: some models won't accept all output controls) + model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) + model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) + + model_outputs = self(**model_inputs, return_dict=True) + + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + model_outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + if synced_gpus and this_peer_finished: + continue + + # Copy is needed to avoid keeping a hanging ref + logits = model_outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device) + + # b. Compute log probs -- get log probabilities from logits, process logits with processors (*e.g.* + # `temperature`, ...), and add new logprobs to existing running logprobs scores. + log_probs = nn.functional.log_softmax(logits, dim=-1) + log_probs = logits_processor(flat_running_sequences, log_probs) + + # Store logits, attentions and hidden_states when required + if return_dict_in_generate: + if output_logits: + raw_logits += (logits.clone(),) + if return_dict_in_generate and output_scores: + all_scores += (log_probs.clone(),) + + if output_attentions: + decoder_attentions += ( + (model_outputs.decoder_attentions,) + if self.config.is_encoder_decoder + else (model_outputs.attentions,) + ) + if self.config.is_encoder_decoder: + cross_attentions += (model_outputs.cross_attentions,) + + if output_hidden_states: + decoder_hidden_states += ( + (model_outputs.decoder_hidden_states,) + if self.config.is_encoder_decoder + else (model_outputs.hidden_states,) + ) + + # This is needed to properly delete logits which may be very large for first iteration + # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration + del model_outputs + + log_probs = self._unflatten_beam_dim(log_probs, batch_size, num_beams) + log_probs = log_probs + running_beam_scores[:, :, None] + log_probs = torch.reshape(log_probs, (batch_size, num_beams * vocab_size)) + + # c. Retrieve top-K continuations, i.e. select the next token (greedy or sampling) and then keep the best + # continuations among all beams based on the accumulated scores. + topk_log_probs, topk_running_sequences, topk_running_beam_indices = self._get_top_k_continuations( + accumulated_log_probs=log_probs, + running_sequences=running_sequences, + running_beam_indices=running_beam_indices, + cur_len=cur_len, + decoder_prompt_len=decoder_prompt_len, + do_sample=do_sample, + beams_to_keep=beams_to_keep, + num_beams=num_beams, + vocab_size=vocab_size, + batch_size=batch_size, + ) + + # d. Check which running sequences have finished + next_token_hits_stopping_criteria = stopping_criteria( + self._flatten_beam_dim(topk_running_sequences[:, :, : cur_len + 1]), # remove unfilled token indexes + all_scores, + ) + next_token_hits_stopping_criteria = self._unflatten_beam_dim( + next_token_hits_stopping_criteria, batch_size, beams_to_keep + ) + + # e. Get the non-finished running `num_beams` sequences for the next generation step + running_sequences, running_beam_scores, running_beam_indices = self._get_running_beams_for_next_iteration( + topk_log_probs=topk_log_probs, + topk_running_sequences=topk_running_sequences, + topk_running_beam_indices=topk_running_beam_indices, + next_token_hits_stopping_criteria=next_token_hits_stopping_criteria, + num_beams=num_beams, + ) + + # f. Update the completed beams if a new high score in a finished sequence is found + sequences, beam_scores, beam_indices, is_sent_finished = self._update_finished_beams( + sequences=sequences, + topk_running_sequences=topk_running_sequences, + beam_scores=beam_scores, + topk_log_probs=topk_log_probs, + beam_indices=beam_indices, + topk_running_beam_indices=topk_running_beam_indices, + is_sent_finished=is_sent_finished, + next_token_hits_stopping_criteria=next_token_hits_stopping_criteria, + top_num_beam_mask=top_num_beam_mask, + num_beams=num_beams, + cur_len=cur_len, + decoder_prompt_len=decoder_prompt_len, + length_penalty=length_penalty, + early_stopping=early_stopping, + ) + + # g. Prepare remaining data for the next iteration, including computing the stopping condition for + # beam search as a whole (as opposed to individual beams, i.e. `stopping_criteria`) + + # pluck the cache from the beam indices that will be used in the next iteration + if model_kwargs.get("past_key_values", None) is not None: + model_kwargs["past_key_values"] = self._temporary_reorder_cache( + past_key_values=model_kwargs["past_key_values"], + beam_idx=self._flatten_beam_dim(running_beam_indices[..., cur_len - decoder_prompt_len]), + ) + + cur_len = cur_len + 1 + this_peer_finished = not self._beam_search_has_unfinished_sequences( + running_beam_scores, + beam_scores, + is_sent_finished, + next_token_hits_stopping_criteria, + cur_len, + max_length, + decoder_prompt_len, + early_stopping, + length_penalty, + ) + + # 5. prepare outputs + # Take best beams for each batch (the score is sorted in descending order) + sequences = self._flatten_beam_dim(sequences[:, :num_return_sequences, :]) + beam_scores = self._flatten_beam_dim(beam_scores[:, :num_return_sequences]) + beam_indices = self._flatten_beam_dim(beam_indices[:, :num_return_sequences, :]) + + # Crop the static-shaped tensors to the actual size. + # `beam_indices` is initialized with -1s, and is updated with the beam index of the generated token at each + # step. We can use it to detect the generated length, which may be != `cur_len` (e.g. selected beam is from a + # previous decoding iteration) + max_generated_length = ((beam_indices + 1).bool()).sum(dim=1).max() + output_length = decoder_prompt_len + max_generated_length + sequences = sequences[:, :output_length] + beam_indices = beam_indices[:, :max_generated_length] + + if return_dict_in_generate: + if not output_scores: + beam_scores = None + + if self.config.is_encoder_decoder: + return GenerateBeamEncoderDecoderOutput( + sequences=sequences, + sequences_scores=beam_scores, + scores=all_scores, + logits=raw_logits, + beam_indices=beam_indices, + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return GenerateBeamDecoderOnlyOutput( + sequences=sequences, + sequences_scores=beam_scores, + scores=all_scores, + logits=raw_logits, + beam_indices=beam_indices, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return sequences + + def _group_beam_search( + self, + input_ids: torch.LongTensor, + beam_scorer: BeamScorer, + logits_processor: LogitsProcessorList, + stopping_criteria: StoppingCriteriaList, + generation_config: GenerationConfig, + synced_gpus: bool, + **model_kwargs, + ): + r""" + Generates sequences of token ids for models with a language modeling head using **diverse beam search + decoding** and can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. + + Parameters: + input_ids (`torch.LongTensor` of shape `(batch_size*num_beams, sequence_length)`): + The sequence used as a prompt for the generation. + beam_scorer (`BeamScorer`): + An derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and + sorted during generation. For more information, the documentation of [`BeamScorer`] should be read. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + stopping_criteria (`StoppingCriteriaList`): + An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] + used to tell if the generation loop should stop. + generation_config ([`~generation.GenerationConfig`]): + The generation configuration to be used as parametrization of the decoding method. + synced_gpus (`bool`): + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). + model_kwargs: + Additional model specific kwargs that will be forwarded to the `forward` function of the model. If + model is an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.GenerateBeamDecoderOnlyOutput`], [`~generation.GenerateBeamEncoderDecoderOutput`] or + `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a + [`~generation.GenerateBeamDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.GenerateBeamEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + """ + # init values + pad_token_id = generation_config._pad_token_tensor + eos_token_id = generation_config._eos_token_tensor + output_attentions = generation_config.output_attentions + output_hidden_states = generation_config.output_hidden_states + output_scores = generation_config.output_scores + output_logits = generation_config.output_logits + return_dict_in_generate = generation_config.return_dict_in_generate + + num_beams = beam_scorer.num_beams + num_beam_groups = beam_scorer.num_beam_groups + num_sub_beams = num_beams // num_beam_groups + batch_size = len(beam_scorer._beam_hyps) // num_beam_groups + device = input_ids.device + + batch_beam_size, cur_len = input_ids.shape + model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) + + if return_dict_in_generate and output_scores: + beam_indices = [tuple(() for _ in range(num_sub_beams * batch_size)) for _ in range(num_beam_groups)] + else: + beam_indices = None + + if num_beams * batch_size != batch_beam_size: + raise ValueError( + f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}." + ) + + # init attention / hidden states / scores tuples + scores = () if (return_dict_in_generate and output_scores) else None + raw_logits = () if (return_dict_in_generate and output_logits) else None + decoder_attentions = () if (return_dict_in_generate and output_attentions) else None + cross_attentions = () if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None + + # if model is an encoder-decoder, retrieve encoder attention weights and hidden states + if return_dict_in_generate and self.config.is_encoder_decoder: + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + # initialise score of first beam of each group with 0 and the rest with -1e9. This ensures that the beams in + # the same group don't produce same tokens every time. + beam_scores = torch.full((batch_size, num_beams), -1e9, dtype=torch.float, device=device) + beam_scores[:, ::num_sub_beams] = 0 + beam_scores = beam_scores.view((batch_size * num_beams,)) + + this_peer_finished = False + + decoder_prompt_len = input_ids.shape[-1] # record the prompt length of decoder + while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): + # predicted tokens in cur_len step + current_tokens = torch.zeros(batch_size * num_beams, dtype=input_ids.dtype, device=device) + + # indices which will form the beams in the next time step + reordering_indices = torch.zeros(batch_size * num_beams, dtype=torch.long, device=device) + + # do one decoder step on all beams of all sentences in batch + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + + # prepare variable output controls (note: some models won't accept all output controls) + model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) + model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) + + outputs = self(**model_inputs, return_dict=True) + + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + if synced_gpus and this_peer_finished: + cur_len = cur_len + 1 + continue + + if output_scores: + processed_score = torch.zeros_like(outputs.logits[:, -1, :]) + if output_logits: + # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration + # (the clone itself is always small) + raw_logit_score = outputs.logits[:, -1, :].to(copy=True, device=input_ids.device) + + for beam_group_idx in range(num_beam_groups): + group_start_idx = beam_group_idx * num_sub_beams + group_end_idx = min(group_start_idx + num_sub_beams, num_beams) + group_size = group_end_idx - group_start_idx + + # indices of beams of current group among all sentences in batch + batch_group_indices = [] + + for batch_idx in range(batch_size): + batch_group_indices.extend( + [batch_idx * num_beams + idx for idx in range(group_start_idx, group_end_idx)] + ) + group_input_ids = input_ids[batch_group_indices] + + # select outputs of beams of current group only + # No need to clone() the logits here as they will not retain outputs.logits at the end of the loop + # .float() is needed to retain precision for later logits manipulations + next_token_logits = outputs.logits[batch_group_indices, -1, :].to( + dtype=torch.float32, device=input_ids.device + ) + + next_token_scores = nn.functional.log_softmax( + next_token_logits, dim=-1 + ) # (batch_size * group_size, vocab_size) + vocab_size = next_token_scores.shape[-1] + + next_token_scores_processed = logits_processor( + group_input_ids, next_token_scores, current_tokens=current_tokens, beam_group_idx=beam_group_idx + ) + next_token_scores = next_token_scores_processed + beam_scores[batch_group_indices].unsqueeze(-1) + next_token_scores = next_token_scores.expand_as(next_token_scores_processed) + + if output_scores: + processed_score[batch_group_indices] = next_token_scores_processed + + # reshape for beam search + next_token_scores = next_token_scores.view(batch_size, group_size * vocab_size) + + # Sample 1 + len(eos_token_id) next tokens for each beam so we have at least 1 non eos token per beam. + n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0 + next_token_scores, next_tokens = torch.topk( + next_token_scores, max(2, 1 + n_eos_tokens) * group_size, dim=1, largest=True, sorted=True + ) + + next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor") + next_tokens = next_tokens % vocab_size + + # stateless + process_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None + beam_outputs = beam_scorer.process( + group_input_ids, + next_token_scores, + next_tokens, + next_indices, + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + beam_indices=process_beam_indices, + group_index=beam_group_idx, + decoder_prompt_len=decoder_prompt_len, + ) + beam_scores[batch_group_indices] = beam_outputs["next_beam_scores"] + beam_next_tokens = beam_outputs["next_beam_tokens"] + beam_idx = beam_outputs["next_beam_indices"] + + if return_dict_in_generate and output_scores: + beam_indices[beam_group_idx] = tuple( + beam_indices[beam_group_idx][beam_idx[i]] + (beam_idx[i],) for i in range(len(beam_indices[0])) + ) + + input_ids[batch_group_indices] = group_input_ids[beam_idx] + group_input_ids = torch.cat([group_input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) + current_tokens[batch_group_indices] = group_input_ids[:, -1] + + # (beam_idx // group_size) -> batch_idx + # (beam_idx % group_size) -> offset of idx inside the group + reordering_indices[batch_group_indices] = ( + num_beams * torch.div(beam_idx, group_size, rounding_mode="floor") + + group_start_idx + + (beam_idx % group_size) + ) + + # Store scores, attentions and hidden_states when required + if return_dict_in_generate: + if output_scores: + scores += (processed_score,) + if output_logits: + raw_logits += (raw_logit_score,) + if output_attentions: + decoder_attentions += ( + (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) + ) + if self.config.is_encoder_decoder: + cross_attentions += (outputs.cross_attentions,) + + if output_hidden_states: + decoder_hidden_states += ( + (outputs.decoder_hidden_states,) + if self.config.is_encoder_decoder + else (outputs.hidden_states,) + ) + + input_ids = torch.cat([input_ids, current_tokens.unsqueeze(-1)], dim=-1) + + # This is needed to properly delete outputs.logits which may be very large for first iteration + # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration + # IMPORTANT: Note that this should appear BEFORE the call to _reorder_cache() to save the maximum memory + # (that way the memory peak does not include outputs.logits) + del outputs + + if model_kwargs.get("past_key_values", None) is not None: + model_kwargs["past_key_values"] = self._temporary_reorder_cache( + model_kwargs["past_key_values"], reordering_indices + ) + + # increase cur_len + cur_len = cur_len + 1 + + if beam_scorer.is_done or all(stopping_criteria(input_ids, scores)): + this_peer_finished = True + + final_beam_indices = sum(beam_indices, ()) if beam_indices is not None else None + sequence_outputs = beam_scorer.finalize( + input_ids, + beam_scores, + next_tokens, + next_indices, + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + max_length=stopping_criteria.max_length, + beam_indices=final_beam_indices, + decoder_prompt_len=decoder_prompt_len, + ) + + if return_dict_in_generate: + if not output_scores: + sequence_outputs["sequence_scores"] = None + + if self.config.is_encoder_decoder: + return GenerateBeamEncoderDecoderOutput( + sequences=sequence_outputs["sequences"], + sequences_scores=sequence_outputs["sequence_scores"], + scores=scores, + logits=raw_logits, + beam_indices=sequence_outputs["beam_indices"], + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return GenerateBeamDecoderOnlyOutput( + sequences=sequence_outputs["sequences"], + sequences_scores=sequence_outputs["sequence_scores"], + scores=scores, + logits=raw_logits, + beam_indices=sequence_outputs["beam_indices"], + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return sequence_outputs["sequences"] + + def _constrained_beam_search( + self, + input_ids: torch.LongTensor, + constrained_beam_scorer: ConstrainedBeamSearchScorer, + logits_processor: LogitsProcessorList, + stopping_criteria: StoppingCriteriaList, + generation_config: GenerationConfig, + synced_gpus: bool, + **model_kwargs, + ) -> Union[GenerateBeamOutput, torch.LongTensor]: + r""" + Generates sequences of token ids for models with a language modeling head using **constrained beam search + decoding** and can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models. + + Parameters: + input_ids (`torch.LongTensor` of shape `(batch_size*num_beams, sequence_length)`): + The sequence used as a prompt for the generation. + constrained_beam_scorer (`ConstrainedBeamSearchScorer`): + A derived instance of [`BeamScorer`] that defines how beam hypotheses are constructed, stored and + sorted during generation, while satisfying a list of positive constraints. For more information, the + documentation of [`ConstrainedBeamSearchScorer`] should be read. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + stopping_criteria (`StoppingCriteriaList`): + An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] + used to tell if the generation loop should stop. + generation_config ([`~generation.GenerationConfig`]): + The generation configuration to be used as parametrization of the decoding method. + synced_gpus (`bool`): + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). + model_kwargs: + Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is + an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.GenerateBeamDecoderOnlyOutput`], [`~generation.GenerateBeamEncoderDecoderOutput`] or + `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a + [`~generation.GenerateBeamDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.GenerateBeamEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + """ + # init values + pad_token_id = generation_config._pad_token_tensor + eos_token_id = generation_config._eos_token_tensor + output_attentions = generation_config.output_attentions + output_hidden_states = generation_config.output_hidden_states + output_scores = generation_config.output_scores + output_logits = generation_config.output_logits + return_dict_in_generate = generation_config.return_dict_in_generate + + batch_size = len(constrained_beam_scorer._beam_hyps) + num_beams = constrained_beam_scorer.num_beams + + batch_beam_size, cur_len = input_ids.shape + model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) + + if num_beams * batch_size != batch_beam_size: + raise ValueError( + f"Batch dimension of `input_ids` should be {num_beams * batch_size}, but is {batch_beam_size}." + ) + + # init attention / hidden states / scores tuples + scores = () if (return_dict_in_generate and output_scores) else None + raw_logits = () if (return_dict_in_generate and output_logits) else None + beam_indices = ( + tuple(() for _ in range(batch_beam_size)) if (return_dict_in_generate and output_scores) else None + ) + decoder_attentions = () if (return_dict_in_generate and output_attentions) else None + cross_attentions = () if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None + + # if model is an encoder-decoder, retrieve encoder attention weights and hidden states + if return_dict_in_generate and self.config.is_encoder_decoder: + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + # initialise score of first beam with 0 and the rest with -1e9. This makes sure that only tokens + # of the first beam are considered to avoid sampling the exact same tokens across all beams. + beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device) + beam_scores[:, 1:] = -1e9 + beam_scores = beam_scores.view((batch_size * num_beams,)) + + this_peer_finished = False + + decoder_prompt_len = input_ids.shape[-1] # record the prompt length of decoder + while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): + model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) + + # prepare variable output controls (note: some models won't accept all output controls) + model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) + model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) + + outputs = self(**model_inputs, return_dict=True) + + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + ) + if synced_gpus and this_peer_finished: + cur_len = cur_len + 1 + continue + + # Copy is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration + # (the clone itself is always small) + # .float() is needed to retain precision for later logits manipulations + next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device) + next_token_scores = nn.functional.log_softmax( + next_token_logits, dim=-1 + ) # (batch_size * num_beams, vocab_size) + + next_token_scores_processed = logits_processor(input_ids, next_token_scores) + + next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as( + next_token_scores_processed + ) + + scores_for_all_vocab = next_token_scores.clone() + + # Store scores, attentions and hidden_states when required + if return_dict_in_generate: + if output_scores: + scores += (next_token_scores,) + if output_logits: + raw_logits += (next_token_logits,) + if output_attentions: + decoder_attentions += ( + (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,) + ) + if self.config.is_encoder_decoder: + cross_attentions += (outputs.cross_attentions,) + + if output_hidden_states: + decoder_hidden_states += ( + (outputs.decoder_hidden_states,) + if self.config.is_encoder_decoder + else (outputs.hidden_states,) + ) + + # reshape for beam search + vocab_size = next_token_scores.shape[-1] + next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size) + + # Sample 1 + len(eos_token_id) next tokens for each beam so we have at least 1 non eos token per beam. + n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0 + next_token_scores, next_tokens = torch.topk( + next_token_scores, max(2, 1 + n_eos_tokens) * num_beams, dim=1, largest=True, sorted=True + ) + + next_indices = (next_tokens / vocab_size).long() + next_tokens = next_tokens % vocab_size + + # stateless + beam_outputs = constrained_beam_scorer.process( + input_ids, + next_token_scores, + next_tokens, + next_indices, + scores_for_all_vocab, + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + beam_indices=beam_indices, + decoder_prompt_len=decoder_prompt_len, + ) + beam_scores = beam_outputs["next_beam_scores"] + beam_next_tokens = beam_outputs["next_beam_tokens"] + beam_idx = beam_outputs["next_beam_indices"] + + input_ids = torch.cat([input_ids[beam_idx, :], beam_next_tokens.unsqueeze(-1)], dim=-1) + + # This is needed to properly delete outputs.logits which may be very large for first iteration + # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration + # IMPORTANT: Note that this should appear BEFORE the call to _reorder_cache() to save the maximum memory + # (that way the memory peak does not include outputs.logits) + del outputs + + if model_kwargs.get("past_key_values", None) is not None: + model_kwargs["past_key_values"] = self._temporary_reorder_cache( + model_kwargs["past_key_values"], beam_idx + ) + + if return_dict_in_generate and output_scores: + beam_indices = tuple((beam_indices[beam_idx[i]] + (beam_idx[i],) for i in range(len(beam_indices)))) + + # increase cur_len + cur_len = cur_len + 1 + + if constrained_beam_scorer.is_done or all(stopping_criteria(input_ids, scores)): + this_peer_finished = True + + sequence_outputs = constrained_beam_scorer.finalize( + input_ids, + beam_scores, + next_tokens, + next_indices, + pad_token_id=pad_token_id, + eos_token_id=eos_token_id, + max_length=stopping_criteria.max_length, + beam_indices=beam_indices, + decoder_prompt_len=decoder_prompt_len, + ) + + if return_dict_in_generate: + if not output_scores: + sequence_outputs["sequence_scores"] = None + if self.config.is_encoder_decoder: + return GenerateBeamEncoderDecoderOutput( + sequences=sequence_outputs["sequences"], + sequences_scores=sequence_outputs["sequence_scores"], + scores=scores, + logits=raw_logits, + beam_indices=sequence_outputs["beam_indices"], + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return GenerateBeamDecoderOnlyOutput( + sequences=sequence_outputs["sequences"], + sequences_scores=sequence_outputs["sequence_scores"], + scores=scores, + logits=raw_logits, + beam_indices=sequence_outputs["beam_indices"], + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return sequence_outputs["sequences"] + + def _assisted_decoding( + self, + input_ids: torch.LongTensor, + candidate_generator: CandidateGenerator, + logits_processor: LogitsProcessorList, + stopping_criteria: StoppingCriteriaList, + generation_config: GenerationConfig, + synced_gpus: bool, + streamer: Optional["BaseStreamer"], + **model_kwargs, + ) -> Union[GenerateNonBeamOutput, torch.LongTensor]: + r""" + Generates sequences of token ids for models with a language modeling head using **greedy decoding** or + **sample** (depending on `do_sample`), assisted by candidate sequences. Assisted generation is an example of a + candidate decoding strategy. Can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text + models. + + Parameters: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + The sequence used as a prompt for the generation. + candidate_generator (`CandidateGenerator`): + A derived instance of [`CandidateGenerator`] that defines how candidate sequences are generated. For + more information, the documentation of [`CandidateGenerator`] should be read. + logits_processor (`LogitsProcessorList`): + An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`] + used to modify the prediction scores of the language modeling head applied at each generation step. + stopping_criteria (`StoppingCriteriaList`): + An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`] + used to tell if the generation loop should stop. + generation_config ([`~generation.GenerationConfig`]): + The generation configuration to be used as parametrization of the decoding method. + synced_gpus (`bool`): + Whether to continue running the while loop until max_length (needed to avoid deadlocking with + `FullyShardedDataParallel` and DeepSpeed ZeRO Stage 3). + streamer (`BaseStreamer`, *optional*): + Streamer object that will be used to stream the generated sequences. Generated tokens are passed + through `streamer.put(token_ids)` and the streamer is responsible for any further processing. + model_kwargs: + Additional model specific keyword arguments will be forwarded to the `forward` function of the model. + If model is an encoder-decoder model the kwargs should include `encoder_outputs`. + + Return: + [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or + `torch.LongTensor`: A `torch.LongTensor` containing the generated tokens (default behaviour) or a + [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and + `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if + `model.config.is_encoder_decoder=True`. + """ + # init values + do_sample = generation_config.do_sample + output_attentions = generation_config.output_attentions + output_hidden_states = generation_config.output_hidden_states + output_scores = generation_config.output_scores + output_logits = generation_config.output_logits + return_dict_in_generate = generation_config.return_dict_in_generate + + # init attention / hidden states / scores tuples + scores = () if (return_dict_in_generate and output_scores) else None + raw_logits = () if (return_dict_in_generate and output_logits) else None + decoder_attentions = () if (return_dict_in_generate and output_attentions) else None + cross_attentions = () if (return_dict_in_generate and output_attentions) else None + decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None + + # if model is an encoder-decoder, retrieve encoder attention weights and hidden states + if return_dict_in_generate and self.config.is_encoder_decoder: + encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None + encoder_hidden_states = ( + model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None + ) + + # keep track of which sequences are already finished + batch_size = input_ids.shape[0] + unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) + model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs) + + this_peer_finished = False + is_first_iteration = True # to preserve the same API in the output as other generation methods + while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device): + cur_len = input_ids.shape[-1] + + # 1. Fetch candidate sequences from a `CandidateGenerator` and move to the correct device + candidate_input_ids, candidate_logits = candidate_generator.get_candidates(input_ids) + candidate_input_ids = candidate_input_ids.to(self.device) + if candidate_logits is not None: + candidate_logits = candidate_logits.to(self.device) + + candidate_length = candidate_input_ids.shape[1] - input_ids.shape[1] + is_done_candidate = stopping_criteria(candidate_input_ids, None) + + # 2. Use the original model to obtain the next token logits given the candidate sequence. We obtain + # `candidate_length + 1` relevant logits from this process: in the event that all candidates are correct, + # we use this forward pass to also pick the subsequent logits in the original model. + + # 2.1. Prepare the model inputs + candidate_kwargs = copy.copy(model_kwargs) + candidate_kwargs = _prepare_attention_mask( + candidate_kwargs, candidate_input_ids.shape[1], self.config.is_encoder_decoder + ) + candidate_kwargs = _prepare_token_type_ids(candidate_kwargs, candidate_input_ids.shape[1]) + if "cache_position" in candidate_kwargs: + candidate_kwargs["cache_position"] = torch.cat( + ( + candidate_kwargs["cache_position"], + torch.arange(cur_len, cur_len + candidate_length, device=input_ids.device, dtype=torch.long), + ), + dim=0, + ) + + model_inputs = self.prepare_inputs_for_generation(candidate_input_ids, **candidate_kwargs) + if "logits_to_keep" in model_inputs: + model_inputs["logits_to_keep"] = candidate_length + 1 + + # 2.2. Run a forward pass on the candidate sequence + # prepare variable output controls (note: some models won't accept all output controls) + model_inputs.update({"output_attentions": output_attentions} if output_attentions else {}) + model_inputs.update({"output_hidden_states": output_hidden_states} if output_hidden_states else {}) + + outputs = self(**model_inputs) + + # 2.3. Process the new logits + # .float() is needed to retain precision for later logits manipulations + new_logits = outputs.logits[:, -candidate_length - 1 :].to( + dtype=torch.float32, device=input_ids.device + ) # excludes the input prompt if present + next_token_logits = new_logits.clone() + if len(logits_processor) > 0: + for i in range(candidate_length + 1): + new_logits[:, i, :] = logits_processor(candidate_input_ids[:, : cur_len + i], new_logits[:, i, :]) + + # 3. Select the accepted tokens. There are two possible cases: + # Case 1: `do_sample=True` and we have logits for the candidates (originally from speculative decoding) + # 👉 Apply algorithm 1 from the speculative decoding paper (https://arxiv.org/pdf/2211.17192.pdf). + if do_sample and candidate_logits is not None: + valid_tokens, n_matches = _speculative_sampling( + candidate_input_ids, + candidate_logits, + candidate_length, + new_logits, + is_done_candidate, + ) + + # Case 2: all other cases (originally from assisted generation) 👉 Compare the tokens selected from the + # original model logits with the candidate tokens. We can keep the candidate tokens until the first + # mismatch, or until the max length is reached. + else: + if do_sample: + probs = new_logits.softmax(dim=-1) + selected_tokens = torch.multinomial(probs[0, :, :], num_samples=1).squeeze(1)[None, :] + else: + selected_tokens = new_logits.argmax(dim=-1) + + candidate_new_tokens = candidate_input_ids[:, cur_len:] + n_matches = ((~(candidate_new_tokens == selected_tokens[:, :-1])).cumsum(dim=-1) < 1).sum() + + # Ensure we don't generate beyond max_len or an EOS token + if is_done_candidate and n_matches == candidate_length: + n_matches -= 1 + valid_tokens = selected_tokens[:, : n_matches + 1] + + # 4. Update variables according to the number of matching assistant tokens. Remember: the token generated + # by the model after the last candidate match is also valid, as it is generated from a correct sequence. + # Because of this last token, assisted generation search reduces to a normal greedy search/sample if there + # is no match. + + # 4.1. Get the valid continuation, after the matching tokens + input_ids = torch.cat((input_ids, valid_tokens), dim=-1) + if streamer is not None: + streamer.put(valid_tokens.cpu()) + new_cur_len = input_ids.shape[-1] + + # 4.2. Discard past key values relative to unused assistant tokens + new_cache_size = new_cur_len - 1 + outputs.past_key_values = _crop_past_key_values(self, outputs.past_key_values, new_cache_size) + + # 5. Update the candidate generation strategy if needed + candidate_generator.update_candidate_strategy(input_ids, new_logits, n_matches) + + # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping + model_kwargs = self._update_model_kwargs_for_generation( + outputs, + model_kwargs, + is_encoder_decoder=self.config.is_encoder_decoder, + num_new_tokens=n_matches + 1, + ) + if synced_gpus and this_peer_finished: + continue + + # Store scores, attentions and hidden_states when required + # Assistant: modified to append one tuple element per token, as in the other generation methods. + if return_dict_in_generate: + newly_added_length = n_matches + 1 + if output_scores: + scores += tuple(new_logits[:, i, :] for i in range(newly_added_length)) + if output_logits: + raw_logits += tuple(next_token_logits[:, i, :] for i in range(newly_added_length)) + + newly_added_length = new_cur_len if is_first_iteration else newly_added_length + if output_attentions: + if self.config.is_encoder_decoder: + cross_attentions = _split_model_outputs( + cross_attentions, outputs.cross_attentions, cur_len, newly_added_length + ) + decoder_attentions = _split_model_outputs( + decoder_attentions, + outputs.decoder_attentions, + cur_len, + newly_added_length, + is_decoder_attention=True, + ) + # some (V)LLMs have hard requirement on SDPA and thus never return attn + elif outputs.attentions[0] is not None: + decoder_attentions = _split_model_outputs( + decoder_attentions, + outputs.attentions, + cur_len, + newly_added_length, + is_decoder_attention=True, + ) + if output_hidden_states: + if self.config.is_encoder_decoder: + decoder_hidden_states = _split_model_outputs( + decoder_hidden_states, outputs.decoder_hidden_states, cur_len, newly_added_length + ) + else: + decoder_hidden_states = _split_model_outputs( + decoder_hidden_states, outputs.hidden_states, cur_len, newly_added_length + ) + + unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores) + this_peer_finished = unfinished_sequences.max() == 0 + is_first_iteration = False + + if streamer is not None: + streamer.end() + + if ( + hasattr(candidate_generator, "assistant_model") + and candidate_generator.assistant_model.generation_config.num_assistant_tokens_schedule == "heuristic" + ): + candidate_generator.assistant_model.generation_config.num_assistant_tokens = ( + candidate_generator.num_assistant_tokens + ) + if return_dict_in_generate: + if self.config.is_encoder_decoder: + return GenerateEncoderDecoderOutput( + sequences=input_ids, + scores=scores, + logits=raw_logits, + encoder_attentions=encoder_attentions, + encoder_hidden_states=encoder_hidden_states, + decoder_attentions=decoder_attentions, + cross_attentions=cross_attentions, + decoder_hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return GenerateDecoderOnlyOutput( + sequences=input_ids, + scores=scores, + logits=raw_logits, + attentions=decoder_attentions, + hidden_states=decoder_hidden_states, + past_key_values=model_kwargs.get("past_key_values"), + ) + else: + return input_ids + + def _prefill_chunking(self, input_ids: torch.LongTensor, generation_config: GenerationConfig, **model_kwargs): + # Even if we are not compiling the forward, flex is always compiled when used. With chunk prefill, we may + # end up needing just a bit more graphs than the default (which is 8). Doing this avoids very cryptic warnings + torch._dynamo.config.cache_size_limit = 64 + + chunk_size = generation_config.prefill_chunk_size + # Only chunk up the token just before last, so that decoding is completely performed outside this function + # (here we simply prefill the cache) + input_chunks = torch.split(input_ids[:, :-1], chunk_size, dim=-1) + + if "past_key_values" not in model_kwargs: + raise ValueError("Cannot use prefill chunkink without a cache") + + model_forward = self.get_compiled_call(generation_config.compile_config) + attention_mask = model_kwargs.pop("attention_mask", None) + + past_length = 0 + for input_chunk in input_chunks: + current_length = past_length + input_chunk.shape[-1] + # Prepare inputs + if attention_mask is not None: + model_kwargs["attention_mask"] = attention_mask[:, :current_length] + model_kwargs["cache_position"] = torch.arange( + past_length, current_length, dtype=torch.long, device=input_chunk.device + ) + model_kwargs["position_ids"] = model_kwargs["cache_position"].unsqueeze(0) + model_inputs = self.prepare_inputs_for_generation(input_chunk, **model_kwargs) + + outputs = model_forward(**model_inputs, return_dict=True) + + model_kwargs["past_key_values"] = outputs.past_key_values + past_length = current_length + + model_kwargs["attention_mask"] = attention_mask + model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + 1 + _ = model_kwargs.pop("position_ids", None) + + return model_kwargs + + +def _speculative_sampling( + candidate_input_ids, + candidate_logits, + candidate_length, + new_logits, + is_done_candidate, +): + """ + Applies sampling as in the speculative decoding paper (https://arxiv.org/pdf/2211.17192.pdf, algorithm 1). Returns + the selected tokens, as well as the number of candidate matches. + + NOTE: Unless otherwise stated, the variable names match those in the paper. + """ + new_candidate_input_ids = candidate_input_ids[:, -candidate_length:] + # Gets the probabilities from the logits. q_i and p_i denote the assistant and model probabilities of the tokens + # selected by the assistant, respectively. + q = candidate_logits.softmax(dim=-1) + q_i = q[:, torch.arange(candidate_length), new_candidate_input_ids].squeeze(0, 1) + p = new_logits.softmax(dim=-1) + p_i = p[:, torch.arange(candidate_length), new_candidate_input_ids].squeeze(0, 1) + probability_ratio = p_i / q_i + + # When probability_ratio > 1 (i.e. q_i(x) < p_i(x), or "assistant probability of the candidate token is smaller + # than the model probability for the same token"), keep the token. Otherwise reject with p = 1 - probability_ratio + # (= keep with p = probability_ratio). Keep all the tokens until the first rejection + r_i = torch.rand_like(probability_ratio) + is_accepted = r_i <= probability_ratio + n_matches = ((~is_accepted).cumsum(dim=-1) < 1).sum() # this is `n` in algorithm 1 + + # Ensure we don't generate beyond max_len or an EOS token (not in algorithm 1, but needed for correct behavior) + if is_done_candidate and n_matches == candidate_length: + # Output length is assumed to be `n_matches + 1`. Since we won't generate another token with the target model + # due to acceptance on EOS we fix `n_matches` + n_matches -= 1 + valid_tokens = new_candidate_input_ids[:, : n_matches + 1] + else: + # Next token selection: if there is a rejection, adjust the distribution from the main model before sampling. + gamma = candidate_logits.shape[1] + p_n_plus_1 = p[:, n_matches, :] + if n_matches < gamma: + q_n_plus_1 = q[:, n_matches, :] + p_prime = torch.clamp((p_n_plus_1 - q_n_plus_1), min=0) + p_prime.div_(p_prime.sum()) + else: + p_prime = p_n_plus_1 + t = torch.multinomial(p_prime, num_samples=1).squeeze(1)[None, :] + + # The selected tokens include the matches (if any) plus the next sampled tokens + if n_matches > 0: + valid_tokens = torch.cat((new_candidate_input_ids[:, :n_matches], t), dim=-1) + else: + valid_tokens = t + + return valid_tokens, n_matches + + +def _split_model_outputs(outputs, new_outputs, cur_len, added_len, is_decoder_attention=False): + """ + Given the (decoder/cross attentions)/(decoder hidden states) for multiple generated tokens, splits it into a tuple + where each member corresponds to a single generated token. + """ + # Retrocompatibility: in our generation functions, the first iteration includes the attention/hidden states for the + # prompt. + if len(outputs) == 0: + new_tuple = () + for layer in new_outputs: + last_dim_size = cur_len if is_decoder_attention else layer.shape[-1] + new_tuple += (layer[..., :cur_len, :last_dim_size],) + outputs += (new_tuple,) + # The first iteration contains the prompt + 1 generated token, let's update the length variables accordingly + cur_len += 1 + added_len -= cur_len + + for i in range(added_len): + new_tuple = () + for layer in new_outputs: + last_dim_size = cur_len + i if is_decoder_attention else layer.shape[-1] + new_tuple += (layer[..., i : i + 1, :last_dim_size],) + outputs += (new_tuple,) + return outputs + + +def _ranking_fast( + context_hidden: torch.FloatTensor, + next_hidden: torch.FloatTensor, + next_top_k_probs: torch.FloatTensor, + cosine_matrix_mask: torch.LongTensor, + alpha: float, + beam_width: int, +) -> torch.FloatTensor: + """ + Reranks the top_k candidates based on a degeneration penalty (cosine similarity with previous tokens), as described + in the paper "A Contrastive Framework for Neural Text Generation". Returns the index of the best candidate for each + row in the batch. + """ + norm_context_hidden = context_hidden / context_hidden.norm(dim=2, keepdim=True) + norm_next_hidden = next_hidden / next_hidden.norm(dim=2, keepdim=True) + cosine_matrix = torch.matmul(norm_context_hidden, norm_next_hidden.transpose(1, 2)).squeeze(-1) # [B*K, S] + + # Penalize cosine_matrix based on the cosine_matrix_mask (ignore padding positions) + # Using a large negative value for masked positions + cosine_matrix_mask = cosine_matrix_mask.to(dtype=cosine_matrix.dtype) + cosine_matrix_mask = (1 - cosine_matrix_mask) * torch.finfo(cosine_matrix.dtype).min + cosine_matrix = cosine_matrix + cosine_matrix_mask + + degeneration_penalty, _ = torch.max(cosine_matrix, dim=-1) # [B*K] + next_top_k_probs = next_top_k_probs.view(-1) # [B*K] + contrastive_score = (1.0 - alpha) * next_top_k_probs - alpha * degeneration_penalty + contrastive_score = torch.stack(torch.split(contrastive_score, beam_width)) # [B, K] + _, selected_idx = contrastive_score.max(dim=-1) # [B] + return selected_idx + + +def _split(data, full_batch_size: int, split_size: int): + """ + Takes care of three cases: + 1. data is a tensor: e.g. last_hidden_state, pooler_output etc. split them on the batch_size dim + 2. data is a tuple: e.g. hidden_states, attentions etc. Keep the tuple as it is and split each tensor in it and + return a list of tuples + 3. data is a tuple of tuples, e.g. past_key_values. Keep the tuple as it is and split each tuple in it and + return a list of tuples of tuples + (see documentation of ModelOutput) + """ + if data is None: + return [None] * (full_batch_size // split_size) + if isinstance(data, torch.Tensor): + return [data[i : i + split_size] for i in range(0, full_batch_size, split_size)] + # New cache format + elif isinstance(data, DynamicCache) or ( + isinstance(data, EncoderDecoderCache) and isinstance(data.self_attention_cache, DynamicCache) + ): + return data.batch_split(full_batch_size, split_size) + elif isinstance(data, tuple): + # If the elements of the tuple are also tuples (e.g., past_key_values in our earlier example) + if isinstance(data[0], tuple): + return [ + tuple(tuple(tensor[i : i + split_size] for tensor in inner_tuple) for inner_tuple in data) + for i in range(0, full_batch_size, split_size) + ] + + else: + return [ + tuple(sub_tensor[i : i + split_size] for sub_tensor in data) + for i in range(0, full_batch_size, split_size) + ] + else: + raise TypeError(f"Unexpected attribute type: {type(data)}") + + +def _split_model_inputs( + model_input: Union[ModelOutput, Dict], split_size: int, full_batch_size: int, config: PretrainedConfig +) -> List[Union[ModelOutput, Dict]]: + """ + Split a ModelOutput object (or its subclasses) or Dict into a list of same-class objects based on a specified split + size. The input object is dict when it was prepared for forward pass and ModelOutput when it was returned from + previous forward pass. + """ + # Edge case: if model_input is None, return a list of Nones + # this happens with Whisper where encoder_outputs is None + if model_input is None: + return [model_input] * (full_batch_size // split_size) + # Infer the class from the object + model_output_cls = type(model_input) + if (full_batch_size % split_size) != 0: + raise ValueError("`full_batch_size` must be divisible by `split_size`") + + if split_size > full_batch_size: + raise ValueError("`split_size` must be smaller or equal to `full_batch_size`") + + # Helper function to split tensors or tuples of tensors + + # Find all the dataclass fields (e.g., last_hidden_state, pooler_output etc.) and split them + keys = ( + model_input.__dataclass_fields__.keys() if hasattr(model_input, "__dataclass_fields__") else model_input.keys() + ) + # We only keep keys that are in the model_input + keys = [k for k in keys if k in model_input] + # Here we can have four types of values: tensors, tuples of tensors and booleans, and encoder_outputs which is a + # ModelOutput object. + # bool should not be split but replicated for each split + bool_keys = [k for k in keys if isinstance(model_input[k], bool) or k == "cache_position"] + keys_to_ignore = ["cache_position", "encoder_outputs", "logits_to_keep"] + non_bool_keys = [k for k in keys if not isinstance(model_input[k], bool) and k not in keys_to_ignore] + + # we split the tensors and tuples of tensors + data_split_list = [ + {k: _split(model_input[k], full_batch_size, split_size)[i] for k in non_bool_keys} + for i in range(full_batch_size // split_size) + ] + # bool values are the same and replicated for each split + bool_data = {k: model_input[k] for k in bool_keys} + # encoder_outputs is a ModelOutput object and should be split by its own + if "encoder_outputs" in model_input: + encoder_outputs_split = _split_model_inputs( + model_input["encoder_outputs"], split_size, full_batch_size, config.get_text_config() + ) + data_split_list = [ + {**data_split, "encoder_outputs": encoder_outputs_split[i]} for i, data_split in enumerate(data_split_list) + ] + # logits_to_keep should be replicated for each split, similar to bool values + if "logits_to_keep" in model_input: + data_split_list = [ + {**data_split, "logits_to_keep": model_input["logits_to_keep"]} for data_split in data_split_list + ] + + # Convert each dictionary in the list to an object of the inferred class + split_model_inputs: List[Union[ModelOutput, Dict]] = [ + model_output_cls(**data_split, **bool_data) for data_split in data_split_list + ] + + return split_model_inputs + + +def stack_model_outputs(model_outputs: List[ModelOutput], config: PretrainedConfig) -> ModelOutput: + """ + Stack a list of ModelOutput objects (or its subclasses) along the batch_size dimension. The function infers the + specific ModelOutput subclass from the list provided. + """ + if not model_outputs: + raise ValueError("Input list is empty.") + + # Infer the class from the first object in the list + model_output_cls = type(model_outputs[0]) + + # Ensure all objects are of the same type + if not all(isinstance(obj, model_output_cls) for obj in model_outputs): + raise ValueError("All elements in the list should be of the same type.") + + # Helper function to concat tensors or tuples of tensors + def _concat(data): + """ + Reverse of `_split` function above. + """ + if any(data is None for data in data): + return None + if isinstance(data[0], torch.Tensor): + return torch.cat(data, dim=0) + # New cache format + elif isinstance(data[0], DynamicCache): + return DynamicCache.from_batch_splits(data) + elif isinstance(data[0], EncoderDecoderCache): + return EncoderDecoderCache.from_batch_splits(data) + elif isinstance(data[0], tuple): + # If the elements of the tuple are also tuples (e.g., past_key_values in our earlier example) + if isinstance(data[0][0], tuple): + return tuple( + tuple(torch.cat([attr[i][j] for attr in data], dim=0) for j in range(len(data[0][0]))) + for i in range(len(data[0])) + ) + else: + return tuple(torch.cat([attr[i] for attr in data], dim=0) for i in range(len(data[0]))) + elif isinstance(data[0], (int, float)): + # If the elements are integers or floats, return a tensor + return torch.tensor(data) + else: + raise TypeError(f"Unexpected attribute type: {type(data[0])}") + + # Use a dictionary comprehension to gather attributes from all objects and concatenate them + concatenated_data = { + k: _concat([getattr(model_output, k) for model_output in model_outputs]) + for k in model_output_cls.__dataclass_fields__.keys() + } + + # Return a new object of the inferred class with the concatenated attributes + return model_output_cls(**concatenated_data) + + +def _relative_top_filter( + scores: torch.FloatTensor, + baseline_scores: torch.FloatTensor, + relative_top: float = 0.1, + filter_value: float = -float("Inf"), + base_filter_value=-1e-3, + min_tokens_to_keep: int = 1, +) -> torch.FloatTensor: + """ + Reference: https://github.com/XiangLi1999/ContrastiveDecoding/blob/170e9142e92159c1237d731e240f5eb14aabf428/transformers/src/transformers/generation_logits_process.py#L235 + Apply filtering to only keep tokens with a probability above a certain threshold. The threshold is defined as `relative_top` * max probability in the distribution. + """ + scores_normalized = scores.log_softmax(dim=-1) + baseline_scores_normalized = baseline_scores.log_softmax(dim=-1) + sorted_logits, sorted_indices = torch.sort(scores_normalized, descending=True) + min_thresh = sorted_logits[..., min_tokens_to_keep - 1] + probs_max = torch.max(scores_normalized, dim=-1).values + probs_thresh = probs_max + np.log(relative_top) + probs_thresh = torch.min(min_thresh, probs_thresh) + probs_thresh = probs_thresh.unsqueeze(-1) + baseline_scores_normalized[scores_normalized < probs_thresh] = base_filter_value + scores_normalized[scores_normalized < probs_thresh] = filter_value + return scores_normalized, baseline_scores_normalized + + +def _dola_select_contrast( + candidate_premature_layers: List[int], + candidate_premature_logits: Dict[int, torch.FloatTensor], + final_logits: torch.FloatTensor, +) -> torch.FloatTensor: + if len(candidate_premature_layers) == 1: + base_logits = candidate_premature_logits[candidate_premature_layers[0]] + final_logits, base_logits = _relative_top_filter(final_logits, base_logits) + logits = final_logits - base_logits + return logits + + # 1. Stacking all premature_layers into a new dimension + stacked_premature_layers = torch.stack([candidate_premature_logits[i] for i in candidate_premature_layers], dim=0) + + # 2. Calculate the softmax values for mature_layer and all premature_layers + # shape: (batch_size, vocab_size) + softmax_mature_layer = F.softmax(final_logits, dim=-1) + # shape: (num_premature_layers, batch_size, vocab_size) + softmax_premature_layers = F.softmax(stacked_premature_layers, dim=-1) + + # 3. Calculate the average distribution + # shape: (num_premature_layers, batch_size, vocab_size) + avg_dist = 0.5 * (softmax_mature_layer[None, :, :] + softmax_premature_layers) + + # 4. Calculate log-softmax for the KL divergence + # shape: (batch_size, vocab_size) + log_softmax_mature_layer = F.log_softmax(final_logits, dim=-1) + # shape: (num_premature_layers, batch_size, vocab_size) + log_softmax_premature_layers = F.log_softmax(stacked_premature_layers, dim=-1) + + # 5. Calculate the KL divergences and then the JS divergences + # shape: (num_premature_layers, batch_size) + kl1 = F.kl_div(log_softmax_mature_layer[None, :, :], avg_dist, reduction="none").mean(-1) + # shape: (num_premature_layers, batch_size) + kl2 = F.kl_div(log_softmax_premature_layers, avg_dist, reduction="none").mean(-1) + js_divs = 0.5 * (kl1 + kl2) # shape: (num_premature_layers, batch_size) + + # 6. Reduce the batchmean + js_divs = js_divs.mean(-1) # shape: (num_premature_layers,) + premature_layer = candidate_premature_layers[int(js_divs.argmax().item())] + + base_logits = candidate_premature_logits[premature_layer] + final_logits, base_logits = _relative_top_filter(final_logits, base_logits) + logits = final_logits - base_logits + return logits diff --git a/lib/python3.12/site-packages/transformers/generation/watermarking.py b/lib/python3.12/site-packages/transformers/generation/watermarking.py new file mode 100644 index 0000000000000000000000000000000000000000..e5f900c3b7d693989568e40eee25d1a895d0059c --- /dev/null +++ b/lib/python3.12/site-packages/transformers/generation/watermarking.py @@ -0,0 +1,549 @@ +# coding=utf-8 +# Copyright 2024 The HuggingFace Inc. team and Google DeepMind. +# +# 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. + +import collections +from dataclasses import dataclass +from functools import lru_cache +from typing import Any, Dict, Optional, Tuple, Union + +import numpy as np +import torch +from torch import nn +from torch.nn import BCELoss + +from ..modeling_utils import PreTrainedModel +from ..utils import ModelOutput, is_torch_available, logging +from .configuration_utils import PretrainedConfig, WatermarkingConfig + + +if is_torch_available(): + import torch + + from .logits_process import SynthIDTextWatermarkLogitsProcessor, WatermarkLogitsProcessor + + +logger = logging.get_logger(__name__) + + +@dataclass +class WatermarkDetectorOutput: + """ + Outputs of a watermark detector. + + Args: + num_tokens_scored (np.array of shape (batch_size)): + Array containing the number of tokens scored for each element in the batch. + num_green_tokens (np.array of shape (batch_size)): + Array containing the number of green tokens for each element in the batch. + green_fraction (np.array of shape (batch_size)): + Array containing the fraction of green tokens for each element in the batch. + z_score (np.array of shape (batch_size)): + Array containing the z-score for each element in the batch. Z-score here shows + how many standard deviations away is the green token count in the input text + from the expected green token count for machine-generated text. + p_value (np.array of shape (batch_size)): + Array containing the p-value for each batch obtained from z-scores. + prediction (np.array of shape (batch_size)), *optional*: + Array containing boolean predictions whether a text is machine-generated for each element in the batch. + confidence (np.array of shape (batch_size)), *optional*: + Array containing confidence scores of a text being machine-generated for each element in the batch. + """ + + num_tokens_scored: Optional[np.array] = None + num_green_tokens: Optional[np.array] = None + green_fraction: Optional[np.array] = None + z_score: Optional[np.array] = None + p_value: Optional[np.array] = None + prediction: Optional[np.array] = None + confidence: Optional[np.array] = None + + +class WatermarkDetector: + r""" + Detector for detection of watermark generated text. The detector needs to be given the exact same settings that were + given during text generation to replicate the watermark greenlist generation and so detect the watermark. This includes + the correct device that was used during text generation, the correct watermarking arguments and the correct tokenizer vocab size. + The code was based on the [original repo](https://github.com/jwkirchenbauer/lm-watermarking/tree/main). + + See [the paper](https://arxiv.org/abs/2306.04634) for more information. + + Args: + model_config (`PretrainedConfig`): + The model config that will be used to get model specific arguments used when generating. + device (`str`): + The device which was used during watermarked text generation. + watermarking_config (Union[`WatermarkingConfig`, `Dict`]): + The exact same watermarking config and arguments used when generating text. + ignore_repeated_ngrams (`bool`, *optional*, defaults to `False`): + Whether to count every unique ngram only once or not. + max_cache_size (`int`, *optional*, defaults to 128): + The max size to be used for LRU caching of seeding/sampling algorithms called for every token. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, WatermarkDetector, WatermarkingConfig + + >>> model_id = "openai-community/gpt2" + >>> model = AutoModelForCausalLM.from_pretrained(model_id) + >>> tok = AutoTokenizer.from_pretrained(model_id) + >>> tok.pad_token_id = tok.eos_token_id + >>> tok.padding_side = "left" + + >>> inputs = tok(["This is the beginning of a long story", "Alice and Bob are"], padding=True, return_tensors="pt") + >>> input_len = inputs["input_ids"].shape[-1] + + >>> # first generate text with watermark and without + >>> watermarking_config = WatermarkingConfig(bias=2.5, seeding_scheme="selfhash") + >>> out_watermarked = model.generate(**inputs, watermarking_config=watermarking_config, do_sample=False, max_length=20) + >>> out = model.generate(**inputs, do_sample=False, max_length=20) + + >>> # now we can instantiate the detector and check the generated text + >>> detector = WatermarkDetector(model_config=model.config, device="cpu", watermarking_config=watermarking_config) + >>> detection_out_watermarked = detector(out_watermarked, return_dict=True) + >>> detection_out = detector(out, return_dict=True) + >>> detection_out_watermarked.prediction + array([ True, True]) + + >>> detection_out.prediction + array([False, False]) + ``` + """ + + def __init__( + self, + model_config: PretrainedConfig, + device: str, + watermarking_config: Union[WatermarkingConfig, Dict], + ignore_repeated_ngrams: bool = False, + max_cache_size: int = 128, + ): + if isinstance(watermarking_config, WatermarkingConfig): + watermarking_config = watermarking_config.to_dict() + + self.bos_token_id = ( + model_config.bos_token_id if not model_config.is_encoder_decoder else model_config.decoder_start_token_id + ) + self.greenlist_ratio = watermarking_config["greenlist_ratio"] + self.ignore_repeated_ngrams = ignore_repeated_ngrams + self.processor = WatermarkLogitsProcessor( + vocab_size=model_config.vocab_size, device=device, **watermarking_config + ) + + # Expensive re-seeding and sampling is cached. + self._get_ngram_score_cached = lru_cache(maxsize=max_cache_size)(self._get_ngram_score) + + def _get_ngram_score(self, prefix: torch.LongTensor, target: int): + greenlist_ids = self.processor._get_greenlist_ids(prefix) + return target in greenlist_ids + + def _score_ngrams_in_passage(self, input_ids: torch.LongTensor): + batch_size, seq_length = input_ids.shape + selfhash = int(self.processor.seeding_scheme == "selfhash") + n = self.processor.context_width + 1 - selfhash + indices = torch.arange(n).unsqueeze(0) + torch.arange(seq_length - n + 1).unsqueeze(1) + ngram_tensors = input_ids[:, indices] + + num_tokens_scored_batch = np.zeros(batch_size) + green_token_count_batch = np.zeros(batch_size) + for batch_idx in range(ngram_tensors.shape[0]): + frequencies_table = collections.Counter(ngram_tensors[batch_idx]) + ngram_to_watermark_lookup = {} + for ngram_example in frequencies_table.keys(): + prefix = ngram_example if selfhash else ngram_example[:-1] + target = ngram_example[-1] + ngram_to_watermark_lookup[ngram_example] = self._get_ngram_score_cached(prefix, target) + + if self.ignore_repeated_ngrams: + # counts a green/red hit once per unique ngram. + # num total tokens scored becomes the number unique ngrams. + num_tokens_scored_batch[batch_idx] = len(frequencies_table.keys()) + green_token_count_batch[batch_idx] = sum(ngram_to_watermark_lookup.values()) + else: + num_tokens_scored_batch[batch_idx] = sum(frequencies_table.values()) + green_token_count_batch[batch_idx] = sum( + freq * outcome + for freq, outcome in zip(frequencies_table.values(), ngram_to_watermark_lookup.values()) + ) + return num_tokens_scored_batch, green_token_count_batch + + def _compute_z_score(self, green_token_count: np.array, total_num_tokens: np.array) -> np.array: + expected_count = self.greenlist_ratio + numer = green_token_count - expected_count * total_num_tokens + denom = np.sqrt(total_num_tokens * expected_count * (1 - expected_count)) + z = numer / denom + return z + + def _compute_pval(self, x, loc=0, scale=1): + z = (x - loc) / scale + return 1 - (0.5 * (1 + np.sign(z) * (1 - np.exp(-2 * z**2 / np.pi)))) + + def __call__( + self, + input_ids: torch.LongTensor, + z_threshold: float = 3.0, + return_dict: bool = False, + ) -> Union[WatermarkDetectorOutput, np.array]: + """ + Args: + input_ids (`torch.LongTensor`): + The watermark generated text. It is advised to remove the prompt, which can affect the detection. + z_threshold (`Dict`, *optional*, defaults to `3.0`): + Changing this threshold will change the sensitivity of the detector. Higher z threshold gives less + sensitivity and vice versa for lower z threshold. + return_dict (`bool`, *optional*, defaults to `False`): + Whether to return `~generation.WatermarkDetectorOutput` or not. If not it will return boolean predictions, + ma + Return: + [`~generation.WatermarkDetectorOutput`] or `np.array`: A [`~generation.WatermarkDetectorOutput`] + if `return_dict=True` otherwise a `np.array`. + + """ + + # Let's assume that if one batch start with `bos`, all batched also do + if input_ids[0, 0] == self.bos_token_id: + input_ids = input_ids[:, 1:] + + if input_ids.shape[-1] - self.processor.context_width < 1: + raise ValueError( + f"Must have at least `1` token to score after the first " + f"min_prefix_len={self.processor.context_width} tokens required by the seeding scheme." + ) + + num_tokens_scored, green_token_count = self._score_ngrams_in_passage(input_ids) + z_score = self._compute_z_score(green_token_count, num_tokens_scored) + prediction = z_score > z_threshold + + if return_dict: + p_value = self._compute_pval(z_score) + confidence = 1 - p_value + + return WatermarkDetectorOutput( + num_tokens_scored=num_tokens_scored, + num_green_tokens=green_token_count, + green_fraction=green_token_count / num_tokens_scored, + z_score=z_score, + p_value=p_value, + prediction=prediction, + confidence=confidence, + ) + return prediction + + +class BayesianDetectorConfig(PretrainedConfig): + """ + This is the configuration class to store the configuration of a [`BayesianDetectorModel`]. It is used to + instantiate a Bayesian Detector model according to the specified arguments. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + watermarking_depth (`int`, *optional*): + The number of tournament layers. + base_rate (`float1`, *optional*, defaults to 0.5): + Prior probability P(w) that a text is watermarked. + """ + + def __init__(self, watermarking_depth: Optional[int] = None, base_rate: float = 0.5, **kwargs): + self.watermarking_depth = watermarking_depth + self.base_rate = base_rate + # These can be set later to store information about this detector. + self.model_name = None + self.watermarking_config = None + + super().__init__(**kwargs) + + def set_detector_information(self, model_name, watermarking_config): + self.model_name = model_name + self.watermarking_config = watermarking_config + + +@dataclass +class BayesianWatermarkDetectorModelOutput(ModelOutput): + """ + Base class for outputs of models predicting if the text is watermarked. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss. + posterior_probabilities (`torch.FloatTensor` of shape `(1,)`): + Multiple choice classification loss. + """ + + loss: Optional[torch.FloatTensor] = None + posterior_probabilities: Optional[torch.FloatTensor] = None + + +class BayesianDetectorWatermarkedLikelihood(nn.Module): + """Watermarked likelihood model for binary-valued g-values. + + This takes in g-values and returns p(g_values|watermarked). + """ + + def __init__(self, watermarking_depth: int): + """Initializes the model parameters.""" + super().__init__() + self.watermarking_depth = watermarking_depth + self.beta = torch.nn.Parameter(-2.5 + 0.001 * torch.randn(1, 1, watermarking_depth)) + self.delta = torch.nn.Parameter(0.001 * torch.randn(1, 1, self.watermarking_depth, watermarking_depth)) + + def _compute_latents(self, g_values: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Computes the unique token probability distribution given g-values. + + Args: + g_values (`torch.Tensor` of shape `(batch_size, seq_len, watermarking_depth)`): + PRF values. + + Returns: + p_one_unique_token and p_two_unique_tokens, both of shape + [batch_size, seq_len, watermarking_depth]. p_one_unique_token[i,t,l] + gives the probability of there being one unique token in a tournament + match on layer l, on timestep t, for batch item i. + p_one_unique_token[i,t,l] + p_two_unique_token[i,t,l] = 1. + """ + # Tile g-values to produce feature vectors for predicting the latents + # for each layer in the tournament; our model for the latents psi is a + # logistic regression model psi = sigmoid(delta * x + beta). + + # [batch_size, seq_len, watermarking_depth, watermarking_depth] + x = torch.repeat_interleave(torch.unsqueeze(g_values, dim=-2), self.watermarking_depth, axis=-2) + + # mask all elements above -1 diagonal for autoregressive factorization + x = torch.tril(x, diagonal=-1) + + # [batch_size, seq_len, watermarking_depth] + # (i, j, k, l) x (i, j, k, l) -> (i, j, k) einsum equivalent + logits = (self.delta[..., None, :] @ x.type(self.delta.dtype)[..., None]).squeeze() + self.beta + + p_two_unique_tokens = torch.sigmoid(logits) + p_one_unique_token = 1 - p_two_unique_tokens + return p_one_unique_token, p_two_unique_tokens + + def forward(self, g_values: torch.Tensor) -> torch.Tensor: + """Computes the likelihoods P(g_values|watermarked). + + Args: + g_values (`torch.Tensor` of shape `(batch_size, seq_len, watermarking_depth)`): + g-values (values 0 or 1) + + Returns: + p(g_values|watermarked) of shape [batch_size, seq_len, watermarking_depth]. + """ + p_one_unique_token, p_two_unique_tokens = self._compute_latents(g_values) + + # P(g_tl | watermarked) is equal to + # 0.5 * [ (g_tl+0.5) * p_two_unique_tokens + p_one_unique_token]. + return 0.5 * ((g_values + 0.5) * p_two_unique_tokens + p_one_unique_token) + + +class BayesianDetectorModel(PreTrainedModel): + r""" + Bayesian classifier for watermark detection. + + This detector uses Bayes' rule to compute a watermarking score, which is the sigmoid of the log of ratio of the + posterior probabilities P(watermarked|g_values) and P(unwatermarked|g_values). Please see the section on + BayesianScore in the paper for further details. + Paper URL: https://www.nature.com/articles/s41586-024-08025-4 + + Note that this detector only works with non-distortionary Tournament-based watermarking using the Bernoulli(0.5) + g-value distribution. + + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`BayesianDetectorConfig`]): Model configuration class with all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. + """ + + config_class = BayesianDetectorConfig + base_model_prefix = "model" + + def __init__(self, config): + super().__init__(config) + + self.watermarking_depth = config.watermarking_depth + self.base_rate = config.base_rate + self.likelihood_model_watermarked = BayesianDetectorWatermarkedLikelihood( + watermarking_depth=self.watermarking_depth + ) + self.prior = torch.nn.Parameter(torch.tensor([self.base_rate])) + + def _init_weights(self, module): + """Initialize the weights.""" + if isinstance(module, nn.Parameter): + module.weight.data.normal_(mean=0.0, std=0.02) + + def _compute_posterior( + self, + likelihoods_watermarked: torch.Tensor, + likelihoods_unwatermarked: torch.Tensor, + mask: torch.Tensor, + prior: float, + ) -> torch.Tensor: + """ + Compute posterior P(w|g) given likelihoods, mask and prior. + + Args: + likelihoods_watermarked (`torch.Tensor` of shape `(batch, length, depth)`): + Likelihoods P(g_values|watermarked) of g-values under watermarked model. + likelihoods_unwatermarked (`torch.Tensor` of shape `(batch, length, depth)`): + Likelihoods P(g_values|unwatermarked) of g-values under unwatermarked model. + mask (`torch.Tensor` of shape `(batch, length)`): + A binary array indicating which g-values should be used. g-values with mask value 0 are discarded. + prior (`float`): + the prior probability P(w) that the text is watermarked. + + Returns: + Posterior probability P(watermarked|g_values), shape [batch]. + """ + mask = torch.unsqueeze(mask, dim=-1) + prior = torch.clamp(prior, min=1e-5, max=1 - 1e-5) + log_likelihoods_watermarked = torch.log(torch.clamp(likelihoods_watermarked, min=1e-30, max=float("inf"))) + log_likelihoods_unwatermarked = torch.log(torch.clamp(likelihoods_unwatermarked, min=1e-30, max=float("inf"))) + log_odds = log_likelihoods_watermarked - log_likelihoods_unwatermarked + + # Sum relative surprisals (log odds) across all token positions and layers. + relative_surprisal_likelihood = torch.einsum("i...->i", log_odds * mask) + + # Compute the relative surprisal prior + relative_surprisal_prior = torch.log(prior) - torch.log(1 - prior) + + # Combine prior and likelihood. + # [batch_size] + relative_surprisal = relative_surprisal_prior + relative_surprisal_likelihood + + # Compute the posterior probability P(w|g) = sigmoid(relative_surprisal). + return torch.sigmoid(relative_surprisal) + + def forward( + self, + g_values: torch.Tensor, + mask: torch.Tensor, + labels: Optional[torch.Tensor] = None, + loss_batch_weight=1, + return_dict=False, + ) -> BayesianWatermarkDetectorModelOutput: + """ + Computes the watermarked posterior P(watermarked|g_values). + + Args: + g_values (`torch.Tensor` of shape `(batch_size, seq_len, watermarking_depth, ...)`): + g-values (with values 0 or 1) + mask: + A binary array shape [batch_size, seq_len] indicating which g-values should be used. g-values with mask + value 0 are discarded. + + Returns: + p(watermarked | g_values), of shape [batch_size]. + """ + + likelihoods_watermarked = self.likelihood_model_watermarked(g_values) + likelihoods_unwatermarked = 0.5 * torch.ones_like(g_values) + out = self._compute_posterior( + likelihoods_watermarked=likelihoods_watermarked, + likelihoods_unwatermarked=likelihoods_unwatermarked, + mask=mask, + prior=self.prior, + ) + + loss = None + if labels is not None: + loss_fct = BCELoss() + loss_unwweight = torch.sum(self.likelihood_model_watermarked.delta**2) + loss_weight = loss_unwweight * loss_batch_weight + loss = loss_fct(torch.clamp(out, 1e-5, 1 - 1e-5), labels) + loss_weight + + if not return_dict: + return (out,) if loss is None else (out, loss) + + return BayesianWatermarkDetectorModelOutput(loss=loss, posterior_probabilities=out) + + +class SynthIDTextWatermarkDetector: + r""" + SynthID text watermark detector class. + + This class has to be initialized with the trained bayesian detector module check script + in examples/synthid_text/detector_training.py for example in training/saving/loading this + detector module. The folder also showcases example use case of this detector. + + Parameters: + detector_module ([`BayesianDetectorModel`]): + Bayesian detector module object initialized with parameters. + Check https://github.com/huggingface/transformers-research-projects/tree/main/synthid_text for usage. + logits_processor (`SynthIDTextWatermarkLogitsProcessor`): + The logits processor used for watermarking. + tokenizer (`Any`): + The tokenizer used for the model. + + Examples: + ```python + >>> from transformers import ( + ... AutoTokenizer, BayesianDetectorModel, SynthIDTextWatermarkLogitsProcessor, SynthIDTextWatermarkDetector + ... ) + + >>> # Load the detector. See https://github.com/huggingface/transformers-research-projects/tree/main/synthid_text for training a detector. + >>> detector_model = BayesianDetectorModel.from_pretrained("joaogante/dummy_synthid_detector") + >>> logits_processor = SynthIDTextWatermarkLogitsProcessor( + ... **detector_model.config.watermarking_config, device="cpu" + ... ) + >>> tokenizer = AutoTokenizer.from_pretrained(detector_model.config.model_name) + >>> detector = SynthIDTextWatermarkDetector(detector_model, logits_processor, tokenizer) + + >>> # Test whether a certain string is watermarked + >>> test_input = tokenizer(["This is a test input"], return_tensors="pt") + >>> is_watermarked = detector(test_input.input_ids) + ``` + """ + + def __init__( + self, + detector_module: BayesianDetectorModel, + logits_processor: SynthIDTextWatermarkLogitsProcessor, + tokenizer: Any, + ): + self.detector_module = detector_module + self.logits_processor = logits_processor + self.tokenizer = tokenizer + + def __call__(self, tokenized_outputs: torch.Tensor): + # eos mask is computed, skip first ngram_len - 1 tokens + # eos_mask will be of shape [batch_size, output_len] + eos_token_mask = self.logits_processor.compute_eos_token_mask( + input_ids=tokenized_outputs, + eos_token_id=self.tokenizer.eos_token_id, + )[:, self.logits_processor.ngram_len - 1 :] + + # context repetition mask is computed + context_repetition_mask = self.logits_processor.compute_context_repetition_mask( + input_ids=tokenized_outputs, + ) + # context repitition mask shape [batch_size, output_len - (ngram_len - 1)] + + combined_mask = context_repetition_mask * eos_token_mask + + g_values = self.logits_processor.compute_g_values( + input_ids=tokenized_outputs, + ) + # g values shape [batch_size, output_len - (ngram_len - 1), depth] + return self.detector_module(g_values, combined_mask) diff --git a/lib/python3.12/site-packages/transformers/hf_argparser.py b/lib/python3.12/site-packages/transformers/hf_argparser.py new file mode 100644 index 0000000000000000000000000000000000000000..0c5610dbfdf9c764de2e9ca1bab171eeebf501bf --- /dev/null +++ b/lib/python3.12/site-packages/transformers/hf_argparser.py @@ -0,0 +1,443 @@ +# Copyright 2020 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. + +import dataclasses +import json +import os +import sys +import types +from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, ArgumentTypeError +from collections.abc import Iterable +from copy import copy +from enum import Enum +from inspect import isclass +from pathlib import Path +from typing import Any, Callable, Literal, NewType, Optional, Union, get_type_hints + +import yaml + + +DataClass = NewType("DataClass", Any) +DataClassType = NewType("DataClassType", Any) + + +# From https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse +def string_to_bool(v): + if isinstance(v, bool): + return v + if v.lower() in ("yes", "true", "t", "y", "1"): + return True + elif v.lower() in ("no", "false", "f", "n", "0"): + return False + else: + raise ArgumentTypeError( + f"Truthy value expected: got {v} but expected one of yes/no, true/false, t/f, y/n, 1/0 (case insensitive)." + ) + + +def make_choice_type_function(choices: list) -> Callable[[str], Any]: + """ + Creates a mapping function from each choices string representation to the actual value. Used to support multiple + value types for a single argument. + + Args: + choices (list): List of choices. + + Returns: + Callable[[str], Any]: Mapping function from string representation to actual value for each choice. + """ + str_to_choice = {str(choice): choice for choice in choices} + return lambda arg: str_to_choice.get(arg, arg) + + +def HfArg( + *, + aliases: Optional[Union[str, list[str]]] = None, + help: Optional[str] = None, + default: Any = dataclasses.MISSING, + default_factory: Callable[[], Any] = dataclasses.MISSING, + metadata: Optional[dict] = None, + **kwargs, +) -> dataclasses.Field: + """Argument helper enabling a concise syntax to create dataclass fields for parsing with `HfArgumentParser`. + + Example comparing the use of `HfArg` and `dataclasses.field`: + ``` + @dataclass + class Args: + regular_arg: str = dataclasses.field(default="Huggingface", metadata={"aliases": ["--example", "-e"], "help": "This syntax could be better!"}) + hf_arg: str = HfArg(default="Huggingface", aliases=["--example", "-e"], help="What a nice syntax!") + ``` + + Args: + aliases (Union[str, List[str]], optional): + Single string or list of strings of aliases to pass on to argparse, e.g. `aliases=["--example", "-e"]`. + Defaults to None. + help (str, optional): Help string to pass on to argparse that can be displayed with --help. Defaults to None. + default (Any, optional): + Default value for the argument. If not default or default_factory is specified, the argument is required. + Defaults to dataclasses.MISSING. + default_factory (Callable[[], Any], optional): + The default_factory is a 0-argument function called to initialize a field's value. It is useful to provide + default values for mutable types, e.g. lists: `default_factory=list`. Mutually exclusive with `default=`. + Defaults to dataclasses.MISSING. + metadata (dict, optional): Further metadata to pass on to `dataclasses.field`. Defaults to None. + + Returns: + Field: A `dataclasses.Field` with the desired properties. + """ + if metadata is None: + # Important, don't use as default param in function signature because dict is mutable and shared across function calls + metadata = {} + if aliases is not None: + metadata["aliases"] = aliases + if help is not None: + metadata["help"] = help + + return dataclasses.field(metadata=metadata, default=default, default_factory=default_factory, **kwargs) + + +class HfArgumentParser(ArgumentParser): + """ + This subclass of `argparse.ArgumentParser` uses type hints on dataclasses to generate arguments. + + The class is designed to play well with the native argparse. In particular, you can add more (non-dataclass backed) + arguments to the parser after initialization and you'll get the output back after parsing as an additional + namespace. Optional: To create sub argument groups use the `_argument_group_name` attribute in the dataclass. + + Args: + dataclass_types (`DataClassType` or `Iterable[DataClassType]`, *optional*): + Dataclass type, or list of dataclass types for which we will "fill" instances with the parsed args. + kwargs (`Dict[str, Any]`, *optional*): + Passed to `argparse.ArgumentParser()` in the regular way. + """ + + dataclass_types: Iterable[DataClassType] + + def __init__(self, dataclass_types: Optional[Union[DataClassType, Iterable[DataClassType]]] = None, **kwargs): + # Make sure dataclass_types is an iterable + if dataclass_types is None: + dataclass_types = [] + elif not isinstance(dataclass_types, Iterable): + dataclass_types = [dataclass_types] + + # To make the default appear when using --help + if "formatter_class" not in kwargs: + kwargs["formatter_class"] = ArgumentDefaultsHelpFormatter + super().__init__(**kwargs) + if dataclasses.is_dataclass(dataclass_types): + dataclass_types = [dataclass_types] + self.dataclass_types = list(dataclass_types) + for dtype in self.dataclass_types: + self._add_dataclass_arguments(dtype) + + @staticmethod + def _parse_dataclass_field(parser: ArgumentParser, field: dataclasses.Field): + # Long-option strings are conventionlly separated by hyphens rather + # than underscores, e.g., "--long-format" rather than "--long_format". + # Argparse converts hyphens to underscores so that the destination + # string is a valid attribute name. Hf_argparser should do the same. + long_options = [f"--{field.name}"] + if "_" in field.name: + long_options.append(f"--{field.name.replace('_', '-')}") + + kwargs = field.metadata.copy() + # field.metadata is not used at all by Data Classes, + # it is provided as a third-party extension mechanism. + if isinstance(field.type, str): + raise RuntimeError( + "Unresolved type detected, which should have been done with the help of " + "`typing.get_type_hints` method by default" + ) + + aliases = kwargs.pop("aliases", []) + if isinstance(aliases, str): + aliases = [aliases] + + origin_type = getattr(field.type, "__origin__", field.type) + if origin_type is Union or (hasattr(types, "UnionType") and isinstance(origin_type, types.UnionType)): + if str not in field.type.__args__ and ( + len(field.type.__args__) != 2 or type(None) not in field.type.__args__ + ): + raise ValueError( + "Only `Union[X, NoneType]` (i.e., `Optional[X]`) is allowed for `Union` because" + " the argument parser only supports one type per argument." + f" Problem encountered in field '{field.name}'." + ) + if type(None) not in field.type.__args__: + # filter `str` in Union + field.type = field.type.__args__[0] if field.type.__args__[1] is str else field.type.__args__[1] + origin_type = getattr(field.type, "__origin__", field.type) + elif bool not in field.type.__args__: + # filter `NoneType` in Union (except for `Union[bool, NoneType]`) + field.type = ( + field.type.__args__[0] if isinstance(None, field.type.__args__[1]) else field.type.__args__[1] + ) + origin_type = getattr(field.type, "__origin__", field.type) + + # A variable to store kwargs for a boolean field, if needed + # so that we can init a `no_*` complement argument (see below) + bool_kwargs = {} + if origin_type is Literal or (isinstance(field.type, type) and issubclass(field.type, Enum)): + if origin_type is Literal: + kwargs["choices"] = field.type.__args__ + else: + kwargs["choices"] = [x.value for x in field.type] + + kwargs["type"] = make_choice_type_function(kwargs["choices"]) + + if field.default is not dataclasses.MISSING: + kwargs["default"] = field.default + else: + kwargs["required"] = True + elif field.type is bool or field.type == Optional[bool]: + # Copy the correct kwargs to use to instantiate a `no_*` complement argument below. + # We do not initialize it here because the `no_*` alternative must be instantiated after the real argument + bool_kwargs = copy(kwargs) + + # Hack because type=bool in argparse does not behave as we want. + kwargs["type"] = string_to_bool + if field.type is bool or (field.default is not None and field.default is not dataclasses.MISSING): + # Default value is False if we have no default when of type bool. + default = False if field.default is dataclasses.MISSING else field.default + # This is the value that will get picked if we don't include --{field.name} in any way + kwargs["default"] = default + # This tells argparse we accept 0 or 1 value after --{field.name} + kwargs["nargs"] = "?" + # This is the value that will get picked if we do --{field.name} (without value) + kwargs["const"] = True + elif isclass(origin_type) and issubclass(origin_type, list): + kwargs["type"] = field.type.__args__[0] + kwargs["nargs"] = "+" + if field.default_factory is not dataclasses.MISSING: + kwargs["default"] = field.default_factory() + elif field.default is dataclasses.MISSING: + kwargs["required"] = True + else: + kwargs["type"] = field.type + if field.default is not dataclasses.MISSING: + kwargs["default"] = field.default + elif field.default_factory is not dataclasses.MISSING: + kwargs["default"] = field.default_factory() + else: + kwargs["required"] = True + parser.add_argument(*long_options, *aliases, **kwargs) + + # Add a complement `no_*` argument for a boolean field AFTER the initial field has already been added. + # Order is important for arguments with the same destination! + # We use a copy of earlier kwargs because the original kwargs have changed a lot before reaching down + # here and we do not need those changes/additional keys. + if field.default is True and (field.type is bool or field.type == Optional[bool]): + bool_kwargs["default"] = False + parser.add_argument( + f"--no_{field.name}", + f"--no-{field.name.replace('_', '-')}", + action="store_false", + dest=field.name, + **bool_kwargs, + ) + + def _add_dataclass_arguments(self, dtype: DataClassType): + if hasattr(dtype, "_argument_group_name"): + parser = self.add_argument_group(dtype._argument_group_name) + else: + parser = self + + try: + type_hints: dict[str, type] = get_type_hints(dtype) + except NameError: + raise RuntimeError( + f"Type resolution failed for {dtype}. Try declaring the class in global scope or " + "removing line of `from __future__ import annotations` which opts in Postponed " + "Evaluation of Annotations (PEP 563)" + ) + except TypeError as ex: + # Remove this block when we drop Python 3.9 support + if sys.version_info[:2] < (3, 10) and "unsupported operand type(s) for |" in str(ex): + python_version = ".".join(map(str, sys.version_info[:3])) + raise RuntimeError( + f"Type resolution failed for {dtype} on Python {python_version}. Try removing " + "line of `from __future__ import annotations` which opts in union types as " + "`X | Y` (PEP 604) via Postponed Evaluation of Annotations (PEP 563). To " + "support Python versions that lower than 3.10, you need to use " + "`typing.Union[X, Y]` instead of `X | Y` and `typing.Optional[X]` instead of " + "`X | None`." + ) from ex + raise + + for field in dataclasses.fields(dtype): + if not field.init: + continue + field.type = type_hints[field.name] + self._parse_dataclass_field(parser, field) + + def parse_args_into_dataclasses( + self, + args=None, + return_remaining_strings=False, + look_for_args_file=True, + args_filename=None, + args_file_flag=None, + ) -> tuple[DataClass, ...]: + """ + Parse command-line args into instances of the specified dataclass types. + + This relies on argparse's `ArgumentParser.parse_known_args`. See the doc at: + docs.python.org/3.7/library/argparse.html#argparse.ArgumentParser.parse_args + + Args: + args: + List of strings to parse. The default is taken from sys.argv. (same as argparse.ArgumentParser) + return_remaining_strings: + If true, also return a list of remaining argument strings. + look_for_args_file: + If true, will look for a ".args" file with the same base name as the entry point script for this + process, and will append its potential content to the command line args. + args_filename: + If not None, will uses this file instead of the ".args" file specified in the previous argument. + args_file_flag: + If not None, will look for a file in the command-line args specified with this flag. The flag can be + specified multiple times and precedence is determined by the order (last one wins). + + Returns: + Tuple consisting of: + + - the dataclass instances in the same order as they were passed to the initializer.abspath + - if applicable, an additional namespace for more (non-dataclass backed) arguments added to the parser + after initialization. + - The potential list of remaining argument strings. (same as argparse.ArgumentParser.parse_known_args) + """ + + if args_file_flag or args_filename or (look_for_args_file and len(sys.argv)): + args_files = [] + + if args_filename: + args_files.append(Path(args_filename)) + elif look_for_args_file and len(sys.argv): + args_files.append(Path(sys.argv[0]).with_suffix(".args")) + + # args files specified via command line flag should overwrite default args files so we add them last + if args_file_flag: + # Create special parser just to extract the args_file_flag values + args_file_parser = ArgumentParser() + args_file_parser.add_argument(args_file_flag, type=str, action="append") + + # Use only remaining args for further parsing (remove the args_file_flag) + cfg, args = args_file_parser.parse_known_args(args=args) + cmd_args_file_paths = vars(cfg).get(args_file_flag.lstrip("-"), None) + + if cmd_args_file_paths: + args_files.extend([Path(p) for p in cmd_args_file_paths]) + + file_args = [] + for args_file in args_files: + if args_file.exists(): + file_args += args_file.read_text().split() + + # in case of duplicate arguments the last one has precedence + # args specified via the command line should overwrite args from files, so we add them last + args = file_args + args if args is not None else file_args + sys.argv[1:] + namespace, remaining_args = self.parse_known_args(args=args) + outputs = [] + for dtype in self.dataclass_types: + keys = {f.name for f in dataclasses.fields(dtype) if f.init} + inputs = {k: v for k, v in vars(namespace).items() if k in keys} + for k in keys: + delattr(namespace, k) + obj = dtype(**inputs) + outputs.append(obj) + if len(namespace.__dict__) > 0: + # additional namespace. + outputs.append(namespace) + if return_remaining_strings: + return (*outputs, remaining_args) + else: + if remaining_args: + raise ValueError(f"Some specified arguments are not used by the HfArgumentParser: {remaining_args}") + + return (*outputs,) + + def parse_dict(self, args: dict[str, Any], allow_extra_keys: bool = False) -> tuple[DataClass, ...]: + """ + Alternative helper method that does not use `argparse` at all, instead uses a dict and populating the dataclass + types. + + Args: + args (`dict`): + dict containing config values + allow_extra_keys (`bool`, *optional*, defaults to `False`): + Defaults to False. If False, will raise an exception if the dict contains keys that are not parsed. + + Returns: + Tuple consisting of: + + - the dataclass instances in the same order as they were passed to the initializer. + """ + unused_keys = set(args.keys()) + outputs = [] + for dtype in self.dataclass_types: + keys = {f.name for f in dataclasses.fields(dtype) if f.init} + inputs = {k: v for k, v in args.items() if k in keys} + unused_keys.difference_update(inputs.keys()) + obj = dtype(**inputs) + outputs.append(obj) + if not allow_extra_keys and unused_keys: + raise ValueError(f"Some keys are not used by the HfArgumentParser: {sorted(unused_keys)}") + return tuple(outputs) + + def parse_json_file( + self, json_file: Union[str, os.PathLike], allow_extra_keys: bool = False + ) -> tuple[DataClass, ...]: + """ + Alternative helper method that does not use `argparse` at all, instead loading a json file and populating the + dataclass types. + + Args: + json_file (`str` or `os.PathLike`): + File name of the json file to parse + allow_extra_keys (`bool`, *optional*, defaults to `False`): + Defaults to False. If False, will raise an exception if the json file contains keys that are not + parsed. + + Returns: + Tuple consisting of: + + - the dataclass instances in the same order as they were passed to the initializer. + """ + with open(Path(json_file), encoding="utf-8") as open_json_file: + data = json.loads(open_json_file.read()) + outputs = self.parse_dict(data, allow_extra_keys=allow_extra_keys) + return tuple(outputs) + + def parse_yaml_file( + self, yaml_file: Union[str, os.PathLike], allow_extra_keys: bool = False + ) -> tuple[DataClass, ...]: + """ + Alternative helper method that does not use `argparse` at all, instead loading a yaml file and populating the + dataclass types. + + Args: + yaml_file (`str` or `os.PathLike`): + File name of the yaml file to parse + allow_extra_keys (`bool`, *optional*, defaults to `False`): + Defaults to False. If False, will raise an exception if the json file contains keys that are not + parsed. + + Returns: + Tuple consisting of: + + - the dataclass instances in the same order as they were passed to the initializer. + """ + outputs = self.parse_dict(yaml.safe_load(Path(yaml_file).read_text()), allow_extra_keys=allow_extra_keys) + return tuple(outputs) diff --git a/lib/python3.12/site-packages/transformers/hyperparameter_search.py b/lib/python3.12/site-packages/transformers/hyperparameter_search.py new file mode 100644 index 0000000000000000000000000000000000000000..e8558ceed32f4640c727a192d3ab00ed9104986f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/hyperparameter_search.py @@ -0,0 +1,141 @@ +# Copyright 2023-present the HuggingFace Inc. team. +# +# 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 typing import Optional + +from .integrations import ( + is_optuna_available, + is_ray_tune_available, + is_sigopt_available, + is_wandb_available, + run_hp_search_optuna, + run_hp_search_ray, + run_hp_search_sigopt, + run_hp_search_wandb, +) +from .trainer_utils import ( + HPSearchBackend, + default_hp_space_optuna, + default_hp_space_ray, + default_hp_space_sigopt, + default_hp_space_wandb, +) +from .utils import logging + + +logger = logging.get_logger(__name__) + + +class HyperParamSearchBackendBase: + name: str + pip_package: Optional[str] = None + + @staticmethod + def is_available(): + raise NotImplementedError + + def run(self, trainer, n_trials: int, direction: str, **kwargs): + raise NotImplementedError + + def default_hp_space(self, trial): + raise NotImplementedError + + def ensure_available(self): + if not self.is_available(): + raise RuntimeError( + f"You picked the {self.name} backend, but it is not installed. Run {self.pip_install()}." + ) + + @classmethod + def pip_install(cls): + return f"`pip install {cls.pip_package or cls.name}`" + + +class OptunaBackend(HyperParamSearchBackendBase): + name = "optuna" + + @staticmethod + def is_available(): + return is_optuna_available() + + def run(self, trainer, n_trials: int, direction: str, **kwargs): + return run_hp_search_optuna(trainer, n_trials, direction, **kwargs) + + def default_hp_space(self, trial): + return default_hp_space_optuna(trial) + + +class RayTuneBackend(HyperParamSearchBackendBase): + name = "ray" + pip_package = "'ray[tune]'" + + @staticmethod + def is_available(): + return is_ray_tune_available() + + def run(self, trainer, n_trials: int, direction: str, **kwargs): + return run_hp_search_ray(trainer, n_trials, direction, **kwargs) + + def default_hp_space(self, trial): + return default_hp_space_ray(trial) + + +class SigOptBackend(HyperParamSearchBackendBase): + name = "sigopt" + + @staticmethod + def is_available(): + return is_sigopt_available() + + def run(self, trainer, n_trials: int, direction: str, **kwargs): + return run_hp_search_sigopt(trainer, n_trials, direction, **kwargs) + + def default_hp_space(self, trial): + return default_hp_space_sigopt(trial) + + +class WandbBackend(HyperParamSearchBackendBase): + name = "wandb" + + @staticmethod + def is_available(): + return is_wandb_available() + + def run(self, trainer, n_trials: int, direction: str, **kwargs): + return run_hp_search_wandb(trainer, n_trials, direction, **kwargs) + + def default_hp_space(self, trial): + return default_hp_space_wandb(trial) + + +ALL_HYPERPARAMETER_SEARCH_BACKENDS = { + HPSearchBackend(backend.name): backend for backend in [OptunaBackend, RayTuneBackend, SigOptBackend, WandbBackend] +} + + +def default_hp_search_backend() -> str: + available_backends = [backend for backend in ALL_HYPERPARAMETER_SEARCH_BACKENDS.values() if backend.is_available()] + if len(available_backends) > 0: + name = available_backends[0].name + if len(available_backends) > 1: + logger.info( + f"{len(available_backends)} hyperparameter search backends available. Using {name} as the default." + ) + return name + raise RuntimeError( + "No hyperparameter search backend available.\n" + + "\n".join( + f" - To install {backend.name} run {backend.pip_install()}" + for backend in ALL_HYPERPARAMETER_SEARCH_BACKENDS.values() + ) + ) diff --git a/lib/python3.12/site-packages/transformers/image_processing_base.py b/lib/python3.12/site-packages/transformers/image_processing_base.py new file mode 100644 index 0000000000000000000000000000000000000000..5398abe028222b14a372e4e8d3a05a2e1dbf9883 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/image_processing_base.py @@ -0,0 +1,558 @@ +# Copyright 2020 The HuggingFace Inc. team. +# +# 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. + + +import copy +import json +import os +import warnings +from io import BytesIO +from typing import Any, Optional, TypeVar, Union + +import numpy as np +import requests + +from .dynamic_module_utils import custom_object_save +from .feature_extraction_utils import BatchFeature as BaseBatchFeature +from .utils import ( + IMAGE_PROCESSOR_NAME, + PushToHubMixin, + add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, + cached_file, + copy_func, + download_url, + is_offline_mode, + is_remote_url, + is_vision_available, + logging, +) + + +if is_vision_available(): + from PIL import Image + + +ImageProcessorType = TypeVar("ImageProcessorType", bound="ImageProcessingMixin") + + +logger = logging.get_logger(__name__) + + +# TODO: Move BatchFeature to be imported by both image_processing_utils and image_processing_utils +# We override the class string here, but logic is the same. +class BatchFeature(BaseBatchFeature): + r""" + Holds the output of the image processor specific `__call__` methods. + + This class is derived from a python dictionary and can be used as a dictionary. + + Args: + data (`dict`): + Dictionary of lists/arrays/tensors returned by the __call__ method ('pixel_values', etc.). + tensor_type (`Union[None, str, TensorType]`, *optional*): + You can give a tensor_type here to convert the lists of integers in PyTorch/TensorFlow/Numpy Tensors at + initialization. + """ + + +# TODO: (Amy) - factor out the common parts of this and the feature extractor +class ImageProcessingMixin(PushToHubMixin): + """ + This is an image processor mixin used to provide saving/loading functionality for sequential and image feature + extractors. + """ + + _auto_class = None + + def __init__(self, **kwargs): + """Set elements of `kwargs` as attributes.""" + # This key was saved while we still used `XXXFeatureExtractor` for image processing. Now we use + # `XXXImageProcessor`, this attribute and its value are misleading. + kwargs.pop("feature_extractor_type", None) + # Pop "processor_class" as it should be saved as private attribute + self._processor_class = kwargs.pop("processor_class", None) + # Additional attributes without default values + 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 _set_processor_class(self, processor_class: str): + """Sets processor class as an attribute.""" + self._processor_class = processor_class + + @classmethod + def from_pretrained( + cls: type[ImageProcessorType], + pretrained_model_name_or_path: Union[str, os.PathLike], + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + **kwargs, + ) -> ImageProcessorType: + r""" + Instantiate a type of [`~image_processing_utils.ImageProcessingMixin`] from an image processor. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + This can be either: + + - a string, the *model id* of a pretrained image_processor hosted inside a model repo on + huggingface.co. + - a path to a *directory* containing a image processor file saved using the + [`~image_processing_utils.ImageProcessingMixin.save_pretrained`] method, e.g., + `./my_model_directory/`. + - a path or url to a saved image processor JSON *file*, e.g., + `./my_model_directory/preprocessor_config.json`. + cache_dir (`str` or `os.PathLike`, *optional*): + Path to a directory in which a downloaded pretrained model image processor should be cached if the + standard cache should not be used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the image processor files and override the cached versions if + they exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`. + + + + return_unused_kwargs (`bool`, *optional*, defaults to `False`): + If `False`, then this function returns just the final image processor object. If `True`, then this + functions returns a `Tuple(image_processor, unused_kwargs)` where *unused_kwargs* is a dictionary + consisting of the key/value pairs whose keys are not image processor attributes: i.e., the part of + `kwargs` which has not been used to update `image_processor` and is otherwise ignored. + 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. + kwargs (`Dict[str, Any]`, *optional*): + The values in kwargs of any keys which are image processor attributes will be used to override the + loaded values. Behavior concerning key/value pairs whose keys are *not* image processor attributes is + controlled by the `return_unused_kwargs` keyword parameter. + + Returns: + A image processor of type [`~image_processing_utils.ImageProcessingMixin`]. + + Examples: + + ```python + # We can't instantiate directly the base class *ImageProcessingMixin* so let's show the examples on a + # derived class: *CLIPImageProcessor* + image_processor = CLIPImageProcessor.from_pretrained( + "openai/clip-vit-base-patch32" + ) # Download image_processing_config from huggingface.co and cache. + image_processor = CLIPImageProcessor.from_pretrained( + "./test/saved_model/" + ) # E.g. image processor (or model) was saved using *save_pretrained('./test/saved_model/')* + image_processor = CLIPImageProcessor.from_pretrained("./test/saved_model/preprocessor_config.json") + image_processor = CLIPImageProcessor.from_pretrained( + "openai/clip-vit-base-patch32", do_normalize=False, foo=False + ) + assert image_processor.do_normalize is False + image_processor, unused_kwargs = CLIPImageProcessor.from_pretrained( + "openai/clip-vit-base-patch32", do_normalize=False, foo=False, return_unused_kwargs=True + ) + assert image_processor.do_normalize is False + assert unused_kwargs == {"foo": False} + ```""" + kwargs["cache_dir"] = cache_dir + kwargs["force_download"] = force_download + kwargs["local_files_only"] = local_files_only + kwargs["revision"] = revision + + use_auth_token = kwargs.pop("use_auth_token", None) + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if token is not None: + kwargs["token"] = token + + image_processor_dict, kwargs = cls.get_image_processor_dict(pretrained_model_name_or_path, **kwargs) + + return cls.from_dict(image_processor_dict, **kwargs) + + def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs): + """ + Save an image processor object to the directory `save_directory`, so that it can be re-loaded using the + [`~image_processing_utils.ImageProcessingMixin.from_pretrained`] class method. + + Args: + save_directory (`str` or `os.PathLike`): + Directory where the image processor JSON file will be saved (will be created if it does not exist). + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + use_auth_token = kwargs.pop("use_auth_token", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if kwargs.get("token", None) is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + kwargs["token"] = use_auth_token + + if os.path.isfile(save_directory): + raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file") + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = self._create_repo(repo_id, **kwargs) + files_timestamps = self._get_files_timestamps(save_directory) + + # If we have a custom config, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + custom_object_save(self, save_directory, config=self) + + # If we save using the predefined names, we can load using `from_pretrained` + output_image_processor_file = os.path.join(save_directory, IMAGE_PROCESSOR_NAME) + + self.to_json_file(output_image_processor_file) + logger.info(f"Image processor saved in {output_image_processor_file}") + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=kwargs.get("token"), + ) + + return [output_image_processor_file] + + @classmethod + def get_image_processor_dict( + cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs + ) -> tuple[dict[str, Any], dict[str, Any]]: + """ + From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a + image processor of type [`~image_processor_utils.ImageProcessingMixin`] using `from_dict`. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`): + The identifier of the pre-trained checkpoint from which we want the dictionary of parameters. + 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. + image_processor_filename (`str`, *optional*, defaults to `"config.json"`): + The name of the file in the model directory to use for the image processor config. + + Returns: + `Tuple[Dict, Dict]`: The dictionary(ies) that will be used to instantiate the image processor object. + """ + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + resume_download = kwargs.pop("resume_download", None) + proxies = kwargs.pop("proxies", None) + token = kwargs.pop("token", None) + use_auth_token = kwargs.pop("use_auth_token", None) + local_files_only = kwargs.pop("local_files_only", False) + revision = kwargs.pop("revision", None) + subfolder = kwargs.pop("subfolder", "") + image_processor_filename = kwargs.pop("image_processor_filename", IMAGE_PROCESSOR_NAME) + + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + user_agent = {"file_type": "image processor", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + is_local = os.path.isdir(pretrained_model_name_or_path) + if os.path.isdir(pretrained_model_name_or_path): + image_processor_file = os.path.join(pretrained_model_name_or_path, image_processor_filename) + if os.path.isfile(pretrained_model_name_or_path): + resolved_image_processor_file = pretrained_model_name_or_path + is_local = True + elif is_remote_url(pretrained_model_name_or_path): + image_processor_file = pretrained_model_name_or_path + resolved_image_processor_file = download_url(pretrained_model_name_or_path) + else: + image_processor_file = image_processor_filename + try: + # Load from local folder or from cache or download from model Hub and cache + resolved_image_processor_file = cached_file( + pretrained_model_name_or_path, + image_processor_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + ) + except OSError: + # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted to + # the original exception. + raise + except Exception: + # For any other exception, we throw a generic error. + raise OSError( + f"Can't load image processor for '{pretrained_model_name_or_path}'. If you were trying to load" + " it from 'https://huggingface.co/models', make sure you don't have a local directory with the" + f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a" + f" directory containing a {image_processor_filename} file" + ) + + try: + # Load image_processor dict + with open(resolved_image_processor_file, encoding="utf-8") as reader: + text = reader.read() + image_processor_dict = json.loads(text) + + except json.JSONDecodeError: + raise OSError( + f"It looks like the config file at '{resolved_image_processor_file}' is not a valid JSON file." + ) + + if is_local: + logger.info(f"loading configuration file {resolved_image_processor_file}") + else: + logger.info( + f"loading configuration file {image_processor_file} from cache at {resolved_image_processor_file}" + ) + if "auto_map" in image_processor_dict: + image_processor_dict["auto_map"] = add_model_info_to_auto_map( + image_processor_dict["auto_map"], pretrained_model_name_or_path + ) + if "custom_pipelines" in image_processor_dict: + image_processor_dict["custom_pipelines"] = add_model_info_to_custom_pipelines( + image_processor_dict["custom_pipelines"], pretrained_model_name_or_path + ) + + return image_processor_dict, kwargs + + @classmethod + def from_dict(cls, image_processor_dict: dict[str, Any], **kwargs): + """ + Instantiates a type of [`~image_processing_utils.ImageProcessingMixin`] from a Python dictionary of parameters. + + Args: + image_processor_dict (`Dict[str, Any]`): + Dictionary that will be used to instantiate the image processor object. Such a dictionary can be + retrieved from a pretrained checkpoint by leveraging the + [`~image_processing_utils.ImageProcessingMixin.to_dict`] method. + kwargs (`Dict[str, Any]`): + Additional parameters from which to initialize the image processor object. + + Returns: + [`~image_processing_utils.ImageProcessingMixin`]: The image processor object instantiated from those + parameters. + """ + image_processor_dict = image_processor_dict.copy() + return_unused_kwargs = kwargs.pop("return_unused_kwargs", False) + + # The `size` parameter is a dict and was previously an int or tuple in feature extractors. + # We set `size` here directly to the `image_processor_dict` so that it is converted to the appropriate + # dict within the image processor and isn't overwritten if `size` is passed in as a kwarg. + if "size" in kwargs and "size" in image_processor_dict: + image_processor_dict["size"] = kwargs.pop("size") + if "crop_size" in kwargs and "crop_size" in image_processor_dict: + image_processor_dict["crop_size"] = kwargs.pop("crop_size") + + image_processor = cls(**image_processor_dict) + + # Update image_processor with kwargs if needed + to_remove = [] + for key, value in kwargs.items(): + if hasattr(image_processor, key): + setattr(image_processor, key, value) + to_remove.append(key) + for key in to_remove: + kwargs.pop(key, None) + + logger.info(f"Image processor {image_processor}") + if return_unused_kwargs: + return image_processor, kwargs + else: + return image_processor + + def to_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. + + Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this image processor instance. + """ + output = copy.deepcopy(self.__dict__) + output["image_processor_type"] = self.__class__.__name__ + + return output + + @classmethod + def from_json_file(cls, json_file: Union[str, os.PathLike]): + """ + Instantiates a image processor of type [`~image_processing_utils.ImageProcessingMixin`] from the path to a JSON + file of parameters. + + Args: + json_file (`str` or `os.PathLike`): + Path to the JSON file containing the parameters. + + Returns: + A image processor of type [`~image_processing_utils.ImageProcessingMixin`]: The image_processor object + instantiated from that JSON file. + """ + with open(json_file, encoding="utf-8") as reader: + text = reader.read() + image_processor_dict = json.loads(text) + return cls(**image_processor_dict) + + def to_json_string(self) -> str: + """ + Serializes this instance to a JSON string. + + Returns: + `str`: String containing all the attributes that make up this feature_extractor instance in JSON format. + """ + dictionary = self.to_dict() + + for key, value in dictionary.items(): + if isinstance(value, np.ndarray): + dictionary[key] = value.tolist() + + # make sure private name "_processor_class" is correctly + # saved as "processor_class" + _processor_class = dictionary.pop("_processor_class", None) + if _processor_class is not None: + dictionary["processor_class"] = _processor_class + + return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" + + def to_json_file(self, json_file_path: Union[str, os.PathLike]): + """ + Save this instance to a JSON file. + + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file in which this image_processor instance's parameters will be saved. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + writer.write(self.to_json_string()) + + def __repr__(self): + return f"{self.__class__.__name__} {self.to_json_string()}" + + @classmethod + def register_for_auto_class(cls, auto_class="AutoImageProcessor"): + """ + Register this class with a given auto class. This should only be used for custom image processors as the ones + in the library are already mapped with `AutoImageProcessor `. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"AutoImageProcessor "`): + The auto class to register this new image processor with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + def fetch_images(self, image_url_or_urls: Union[str, list[str]]): + """ + Convert a single or a list of urls into the corresponding `PIL.Image` objects. + + If a single url is passed, the return value will be a single object. If a list is passed a list of objects is + returned. + """ + headers = { + "User-Agent": ( + "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0" + " Safari/537.36" + ) + } + if isinstance(image_url_or_urls, list): + return [self.fetch_images(x) for x in image_url_or_urls] + elif isinstance(image_url_or_urls, str): + response = requests.get(image_url_or_urls, stream=True, headers=headers) + response.raise_for_status() + return Image.open(BytesIO(response.content)) + else: + raise TypeError(f"only a single or a list of entries is supported but got type={type(image_url_or_urls)}") + + +ImageProcessingMixin.push_to_hub = copy_func(ImageProcessingMixin.push_to_hub) +if ImageProcessingMixin.push_to_hub.__doc__ is not None: + ImageProcessingMixin.push_to_hub.__doc__ = ImageProcessingMixin.push_to_hub.__doc__.format( + object="image processor", object_class="AutoImageProcessor", object_files="image processor file" + ) diff --git a/lib/python3.12/site-packages/transformers/image_processing_utils.py b/lib/python3.12/site-packages/transformers/image_processing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ec0f817728d327fe95b140f3f30f33bac7042ad9 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/image_processing_utils.py @@ -0,0 +1,308 @@ +# Copyright 2022 The HuggingFace Inc. team. +# +# 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. + +import math +from collections.abc import Iterable +from typing import Optional, Union + +import numpy as np + +from .image_processing_base import BatchFeature, ImageProcessingMixin +from .image_transforms import center_crop, normalize, rescale +from .image_utils import ChannelDimension, get_image_size +from .utils import logging + + +logger = logging.get_logger(__name__) + + +INIT_SERVICE_KWARGS = [ + "processor_class", + "image_processor_type", +] + + +class BaseImageProcessor(ImageProcessingMixin): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def __call__(self, images, **kwargs) -> BatchFeature: + """Preprocess an image or a batch of images.""" + return self.preprocess(images, **kwargs) + + def preprocess(self, images, **kwargs) -> BatchFeature: + raise NotImplementedError("Each image processor must implement its own preprocess method") + + def rescale( + self, + image: np.ndarray, + scale: float, + data_format: Optional[Union[str, ChannelDimension]] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, + **kwargs, + ) -> np.ndarray: + """ + Rescale an image by a scale factor. image = image * scale. + + Args: + image (`np.ndarray`): + Image to rescale. + scale (`float`): + The scaling factor to rescale pixel values by. + data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the output image. If unset, the channel dimension format of the input + image is used. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the input image. If unset, the channel dimension format is inferred + from the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + + Returns: + `np.ndarray`: The rescaled image. + """ + return rescale(image, scale=scale, data_format=data_format, input_data_format=input_data_format, **kwargs) + + def normalize( + self, + image: np.ndarray, + mean: Union[float, Iterable[float]], + std: Union[float, Iterable[float]], + data_format: Optional[Union[str, ChannelDimension]] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, + **kwargs, + ) -> np.ndarray: + """ + Normalize an image. image = (image - image_mean) / image_std. + + Args: + image (`np.ndarray`): + Image to normalize. + mean (`float` or `Iterable[float]`): + Image mean to use for normalization. + std (`float` or `Iterable[float]`): + Image standard deviation to use for normalization. + data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the output image. If unset, the channel dimension format of the input + image is used. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the input image. If unset, the channel dimension format is inferred + from the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + + Returns: + `np.ndarray`: The normalized image. + """ + return normalize( + image, mean=mean, std=std, data_format=data_format, input_data_format=input_data_format, **kwargs + ) + + def center_crop( + self, + image: np.ndarray, + size: dict[str, int], + data_format: Optional[Union[str, ChannelDimension]] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, + **kwargs, + ) -> np.ndarray: + """ + Center crop an image to `(size["height"], size["width"])`. If the input size is smaller than `crop_size` along + any edge, the image is padded with 0's and then center cropped. + + Args: + image (`np.ndarray`): + Image to center crop. + size (`Dict[str, int]`): + Size of the output image. + data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the output image. If unset, the channel dimension format of the input + image is used. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the input image. If unset, the channel dimension format is inferred + from the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + """ + size = get_size_dict(size) + if "height" not in size or "width" not in size: + raise ValueError(f"The size dictionary must have keys 'height' and 'width'. Got {size.keys()}") + return center_crop( + image, + size=(size["height"], size["width"]), + data_format=data_format, + input_data_format=input_data_format, + **kwargs, + ) + + def to_dict(self): + encoder_dict = super().to_dict() + encoder_dict.pop("_valid_processor_keys", None) + return encoder_dict + + +VALID_SIZE_DICT_KEYS = ( + {"height", "width"}, + {"shortest_edge"}, + {"shortest_edge", "longest_edge"}, + {"longest_edge"}, + {"max_height", "max_width"}, +) + + +def is_valid_size_dict(size_dict): + if not isinstance(size_dict, dict): + return False + + size_dict_keys = set(size_dict.keys()) + for allowed_keys in VALID_SIZE_DICT_KEYS: + if size_dict_keys == allowed_keys: + return True + return False + + +def convert_to_size_dict( + size, max_size: Optional[int] = None, default_to_square: bool = True, height_width_order: bool = True +): + # By default, if size is an int we assume it represents a tuple of (size, size). + if isinstance(size, int) and default_to_square: + if max_size is not None: + raise ValueError("Cannot specify both size as an int, with default_to_square=True and max_size") + return {"height": size, "width": size} + # In other configs, if size is an int and default_to_square is False, size represents the length of + # the shortest edge after resizing. + elif isinstance(size, int) and not default_to_square: + size_dict = {"shortest_edge": size} + if max_size is not None: + size_dict["longest_edge"] = max_size + return size_dict + # Otherwise, if size is a tuple it's either (height, width) or (width, height) + elif isinstance(size, (tuple, list)) and height_width_order: + return {"height": size[0], "width": size[1]} + elif isinstance(size, (tuple, list)) and not height_width_order: + return {"height": size[1], "width": size[0]} + elif size is None and max_size is not None: + if default_to_square: + raise ValueError("Cannot specify both default_to_square=True and max_size") + return {"longest_edge": max_size} + + raise ValueError(f"Could not convert size input to size dict: {size}") + + +def get_size_dict( + size: Union[int, Iterable[int], dict[str, int]] = None, + max_size: Optional[int] = None, + height_width_order: bool = True, + default_to_square: bool = True, + param_name="size", +) -> dict: + """ + Converts the old size parameter in the config into the new dict expected in the config. This is to ensure backwards + compatibility with the old image processor configs and removes ambiguity over whether the tuple is in (height, + width) or (width, height) format. + + - If `size` is tuple, it is converted to `{"height": size[0], "width": size[1]}` or `{"height": size[1], "width": + size[0]}` if `height_width_order` is `False`. + - If `size` is an int, and `default_to_square` is `True`, it is converted to `{"height": size, "width": size}`. + - If `size` is an int and `default_to_square` is False, it is converted to `{"shortest_edge": size}`. If `max_size` + is set, it is added to the dict as `{"longest_edge": max_size}`. + + Args: + size (`Union[int, Iterable[int], Dict[str, int]]`, *optional*): + The `size` parameter to be cast into a size dictionary. + max_size (`Optional[int]`, *optional*): + The `max_size` parameter to be cast into a size dictionary. + height_width_order (`bool`, *optional*, defaults to `True`): + If `size` is a tuple, whether it's in (height, width) or (width, height) order. + default_to_square (`bool`, *optional*, defaults to `True`): + If `size` is an int, whether to default to a square image or not. + """ + if not isinstance(size, dict): + size_dict = convert_to_size_dict(size, max_size, default_to_square, height_width_order) + logger.info( + f"{param_name} should be a dictionary on of the following set of keys: {VALID_SIZE_DICT_KEYS}, got {size}." + f" Converted to {size_dict}.", + ) + else: + size_dict = size + + if not is_valid_size_dict(size_dict): + raise ValueError( + f"{param_name} must have one of the following set of keys: {VALID_SIZE_DICT_KEYS}, got {size_dict.keys()}" + ) + return size_dict + + +def select_best_resolution(original_size: tuple, possible_resolutions: list) -> tuple: + """ + Selects the best resolution from a list of possible resolutions based on the original size. + + This is done by calculating the effective and wasted resolution for each possible resolution. + + The best fit resolution is the one that maximizes the effective resolution and minimizes the wasted resolution. + + Args: + original_size (tuple): + The original size of the image in the format (height, width). + possible_resolutions (list): + A list of possible resolutions in the format [(height1, width1), (height2, width2), ...]. + + Returns: + tuple: The best fit resolution in the format (height, width). + """ + original_height, original_width = original_size + best_fit = None + max_effective_resolution = 0 + min_wasted_resolution = float("inf") + + for height, width in possible_resolutions: + scale = min(width / original_width, height / original_height) + downscaled_width, downscaled_height = int(original_width * scale), int(original_height * scale) + effective_resolution = min(downscaled_width * downscaled_height, original_width * original_height) + wasted_resolution = (width * height) - effective_resolution + + if effective_resolution > max_effective_resolution or ( + effective_resolution == max_effective_resolution and wasted_resolution < min_wasted_resolution + ): + max_effective_resolution = effective_resolution + min_wasted_resolution = wasted_resolution + best_fit = (height, width) + + return best_fit + + +def get_patch_output_size(image, target_resolution, input_data_format): + """ + Given an image and a target resolution, calculate the output size of the image after cropping to the target + """ + original_height, original_width = get_image_size(image, channel_dim=input_data_format) + target_height, target_width = target_resolution + + scale_w = target_width / original_width + scale_h = target_height / original_height + + if scale_w < scale_h: + new_width = target_width + new_height = min(math.ceil(original_height * scale_w), target_height) + else: + new_height = target_height + new_width = min(math.ceil(original_width * scale_h), target_width) + + return new_height, new_width diff --git a/lib/python3.12/site-packages/transformers/image_processing_utils_fast.py b/lib/python3.12/site-packages/transformers/image_processing_utils_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..b671a1119111353f6569e4ba0f379c8e113de70b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/image_processing_utils_fast.py @@ -0,0 +1,792 @@ +# Copyright 2024 The HuggingFace Inc. team. +# +# 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 collections.abc import Iterable +from functools import lru_cache, partial +from typing import Any, Optional, TypedDict, Union + +import numpy as np + +from .image_processing_utils import ( + BaseImageProcessor, + BatchFeature, + get_size_dict, +) +from .image_transforms import ( + convert_to_rgb, + get_resize_output_image_size, + get_size_with_aspect_ratio, + group_images_by_shape, + reorder_images, +) +from .image_utils import ( + ChannelDimension, + ImageInput, + ImageType, + SizeDict, + get_image_size, + get_image_size_for_max_height_width, + get_image_type, + infer_channel_dimension_format, + make_flat_list_of_images, + validate_kwargs, + validate_preprocess_arguments, +) +from .processing_utils import Unpack +from .utils import ( + TensorType, + add_start_docstrings, + is_torch_available, + is_torchvision_available, + is_torchvision_v2_available, + is_vision_available, + logging, +) + + +if is_vision_available(): + from .image_utils import PILImageResampling + +if is_torch_available(): + import torch + +if is_torchvision_available(): + from .image_utils import pil_torch_interpolation_mapping + + if is_torchvision_v2_available(): + from torchvision.transforms.v2 import functional as F + else: + from torchvision.transforms import functional as F + +logger = logging.get_logger(__name__) + + +@lru_cache(maxsize=10) +def validate_fast_preprocess_arguments( + do_rescale: Optional[bool] = None, + rescale_factor: Optional[float] = None, + do_normalize: Optional[bool] = None, + image_mean: Optional[Union[float, list[float]]] = None, + image_std: Optional[Union[float, list[float]]] = None, + do_pad: Optional[bool] = None, + size_divisibility: Optional[int] = None, + do_center_crop: Optional[bool] = None, + crop_size: Optional[SizeDict] = None, + do_resize: Optional[bool] = None, + size: Optional[SizeDict] = None, + resample: Optional["PILImageResampling"] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + data_format: Optional[ChannelDimension] = ChannelDimension.FIRST, +): + """ + Checks validity of typically used arguments in an `ImageProcessorFast` `preprocess` method. + Raises `ValueError` if arguments incompatibility is caught. + """ + validate_preprocess_arguments( + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + do_pad=do_pad, + size_divisibility=size_divisibility, + do_center_crop=do_center_crop, + crop_size=crop_size, + do_resize=do_resize, + size=size, + resample=resample, + ) + # Extra checks for ImageProcessorFast + if return_tensors is not None and return_tensors != "pt": + raise ValueError("Only returning PyTorch tensors is currently supported.") + + if data_format != ChannelDimension.FIRST: + raise ValueError("Only channel first data format is currently supported.") + + +def safe_squeeze(tensor: "torch.Tensor", axis: Optional[int] = None) -> "torch.Tensor": + """ + Squeezes a tensor, but only if the axis specified has dim 1. + """ + if axis is None: + return tensor.squeeze() + + try: + return tensor.squeeze(axis=axis) + except ValueError: + return tensor + + +def max_across_indices(values: Iterable[Any]) -> list[Any]: + """ + Return the maximum value across all indices of an iterable of values. + """ + return [max(values_i) for values_i in zip(*values)] + + +def get_max_height_width(images: list["torch.Tensor"]) -> tuple[int]: + """ + Get the maximum height and width across all images in a batch. + """ + + _, max_height, max_width = max_across_indices([img.shape for img in images]) + + return (max_height, max_width) + + +def divide_to_patches( + image: Union[np.array, "torch.Tensor"], patch_size: int +) -> list[Union[np.array, "torch.Tensor"]]: + """ + Divides an image into patches of a specified size. + + Args: + image (`Union[np.array, "torch.Tensor"]`): + The input image. + patch_size (`int`): + The size of each patch. + Returns: + list: A list of Union[np.array, "torch.Tensor"] representing the patches. + """ + patches = [] + height, width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + for i in range(0, height, patch_size): + for j in range(0, width, patch_size): + patch = image[:, i : i + patch_size, j : j + patch_size] + patches.append(patch) + + return patches + + +class DefaultFastImageProcessorKwargs(TypedDict, total=False): + do_resize: Optional[bool] + size: Optional[dict[str, int]] + default_to_square: Optional[bool] + resample: Optional[Union["PILImageResampling", "F.InterpolationMode"]] + do_center_crop: Optional[bool] + crop_size: Optional[dict[str, int]] + do_rescale: Optional[bool] + rescale_factor: Optional[Union[int, float]] + do_normalize: Optional[bool] + image_mean: Optional[Union[float, list[float]]] + image_std: Optional[Union[float, list[float]]] + do_convert_rgb: Optional[bool] + return_tensors: Optional[Union[str, TensorType]] + data_format: Optional[ChannelDimension] + input_data_format: Optional[Union[str, ChannelDimension]] + device: Optional["torch.device"] + + +BASE_IMAGE_PROCESSOR_FAST_DOCSTRING = r""" + + Args: + do_resize (`bool`, *optional*, defaults to `self.do_resize`): + Whether to resize the image's (height, width) dimensions to the specified `size`. Can be overridden by the + `do_resize` parameter in the `preprocess` method. + size (`dict`, *optional*, defaults to `self.size`): + Size of the output image after resizing. Can be overridden by the `size` parameter in the `preprocess` + method. + default_to_square (`bool`, *optional*, defaults to `self.default_to_square`): + Whether to default to a square image when resizing, if size is an int. + resample (`PILImageResampling`, *optional*, defaults to `self.resample`): + Resampling filter to use if resizing the image. Only has an effect if `do_resize` is set to `True`. Can be + overridden by the `resample` parameter in the `preprocess` method. + do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`): + Whether to center crop the image to the specified `crop_size`. Can be overridden by `do_center_crop` in the + `preprocess` method. + crop_size (`Dict[str, int]` *optional*, defaults to `self.crop_size`): + Size of the output image after applying `center_crop`. Can be overridden by `crop_size` in the `preprocess` + method. + do_rescale (`bool`, *optional*, defaults to `self.do_rescale`): + Whether to rescale the image by the specified scale `rescale_factor`. Can be overridden by the + `do_rescale` parameter in the `preprocess` method. + rescale_factor (`int` or `float`, *optional*, defaults to `self.rescale_factor`): + Scale factor to use if rescaling the image. Only has an effect if `do_rescale` is set to `True`. Can be + overridden by the `rescale_factor` parameter in the `preprocess` method. + do_normalize (`bool`, *optional*, defaults to `self.do_normalize`): + Whether to normalize the image. Can be overridden by the `do_normalize` parameter in the `preprocess` + method. Can be overridden by the `do_normalize` parameter in the `preprocess` method. + image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`): + Mean to use if normalizing the image. This is a float or list of floats the length of the number of + channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method. Can be + overridden by the `image_mean` parameter in the `preprocess` method. + image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`): + Standard deviation to use if normalizing the image. This is a float or list of floats the length of the + number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method. + Can be overridden by the `image_std` parameter in the `preprocess` method. + do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`): + Whether to convert the image to RGB. + return_tensors (`str` or `TensorType`, *optional*, defaults to `self.return_tensors`): + Returns stacked tensors if set to `pt, otherwise returns a list of tensors. + data_format (`ChannelDimension` or `str`, *optional*, defaults to `self.data_format`): + Only `ChannelDimension.FIRST` is supported. Added for compatibility with slow processors. + input_data_format (`ChannelDimension` or `str`, *optional*, defaults to `self.input_data_format`): + The channel dimension format for the input image. If unset, the channel dimension format is inferred + from the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. + device (`torch.device`, *optional*, defaults to `self.device`): + The device to process the images on. If unset, the device is inferred from the input images.""" + +BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS = r""" + Preprocess an image or batch of images. + + Args: + images (`ImageInput`): + Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If + passing in images with pixel values between 0 and 1, set `do_rescale=False`. + do_resize (`bool`, *optional*, defaults to `self.do_resize`): + Whether to resize the image. + size (`Dict[str, int]`, *optional*, defaults to `self.size`): + Describes the maximum input dimensions to the model. + resample (`PILImageResampling` or `InterpolationMode`, *optional*, defaults to `self.resample`): + Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only + has an effect if `do_resize` is set to `True`. + do_center_crop (`bool`, *optional*, defaults to `self.do_center_crop`): + Whether to center crop the image. + crop_size (`Dict[str, int]`, *optional*, defaults to `self.crop_size`): + Size of the output image after applying `center_crop`. + do_rescale (`bool`, *optional*, defaults to `self.do_rescale`): + Whether to rescale the image. + rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`): + Rescale factor to rescale the image by if `do_rescale` is set to `True`. + do_normalize (`bool`, *optional*, defaults to `self.do_normalize`): + Whether to normalize the image. + image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`): + Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`. + image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`): + Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to + `True`. + do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`): + Whether to convert the image to RGB. + return_tensors (`str` or `TensorType`, *optional*, defaults to `self.return_tensors`): + Returns stacked tensors if set to `pt, otherwise returns a list of tensors. + data_format (`ChannelDimension` or `str`, *optional*, defaults to `self.data_format`): + Only `ChannelDimension.FIRST` is supported. Added for compatibility with slow processors. + input_data_format (`ChannelDimension` or `str`, *optional*, defaults to `self.input_data_format`): + The channel dimension format for the input image. If unset, the channel dimension format is inferred + from the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. + device (`torch.device`, *optional*, defaults to `self.device`): + The device to process the images on. If unset, the device is inferred from the input images.""" + + +@add_start_docstrings( + "Constructs a fast base image processor.", + BASE_IMAGE_PROCESSOR_FAST_DOCSTRING, +) +class BaseImageProcessorFast(BaseImageProcessor): + resample = None + image_mean = None + image_std = None + size = None + default_to_square = True + crop_size = None + do_resize = None + do_center_crop = None + do_rescale = None + rescale_factor = 1 / 255 + do_normalize = None + do_convert_rgb = None + return_tensors = None + data_format = ChannelDimension.FIRST + input_data_format = None + device = None + model_input_names = ["pixel_values"] + valid_kwargs = DefaultFastImageProcessorKwargs + unused_kwargs = None + + def __init__( + self, + **kwargs: Unpack[DefaultFastImageProcessorKwargs], + ) -> None: + super().__init__(**kwargs) + kwargs = self.filter_out_unused_kwargs(kwargs) + size = kwargs.pop("size", self.size) + self.size = ( + get_size_dict(size=size, default_to_square=kwargs.pop("default_to_square", self.default_to_square)) + if size is not None + else None + ) + crop_size = kwargs.pop("crop_size", self.crop_size) + self.crop_size = get_size_dict(crop_size, param_name="crop_size") if crop_size is not None else None + for key in self.valid_kwargs.__annotations__.keys(): + kwarg = kwargs.pop(key, None) + if kwarg is not None: + setattr(self, key, kwarg) + else: + setattr(self, key, getattr(self, key, None)) + + def resize( + self, + image: "torch.Tensor", + size: SizeDict, + interpolation: "F.InterpolationMode" = None, + antialias: bool = True, + **kwargs, + ) -> "torch.Tensor": + """ + Resize an image to `(size["height"], size["width"])`. + + Args: + image (`torch.Tensor`): + Image to resize. + size (`SizeDict`): + Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image. + resample (`InterpolationMode`, *optional*, defaults to `InterpolationMode.BILINEAR`): + `InterpolationMode` filter to use when resizing the image e.g. `InterpolationMode.BICUBIC`. + + Returns: + `torch.Tensor`: The resized image. + """ + interpolation = interpolation if interpolation is not None else F.InterpolationMode.BILINEAR + if size.shortest_edge and size.longest_edge: + # Resize the image so that the shortest edge or the longest edge is of the given size + # while maintaining the aspect ratio of the original image. + new_size = get_size_with_aspect_ratio( + image.size()[-2:], + size.shortest_edge, + size.longest_edge, + ) + elif size.shortest_edge: + new_size = get_resize_output_image_size( + image, + size=size.shortest_edge, + default_to_square=False, + input_data_format=ChannelDimension.FIRST, + ) + elif size.max_height and size.max_width: + new_size = get_image_size_for_max_height_width(image.size()[-2:], size.max_height, size.max_width) + elif size.height and size.width: + new_size = (size.height, size.width) + else: + raise ValueError( + "Size must contain 'height' and 'width' keys, or 'max_height' and 'max_width', or 'shortest_edge' key. Got" + f" {size}." + ) + return F.resize(image, new_size, interpolation=interpolation, antialias=antialias) + + def rescale( + self, + image: "torch.Tensor", + scale: float, + **kwargs, + ) -> "torch.Tensor": + """ + Rescale an image by a scale factor. image = image * scale. + + Args: + image (`torch.Tensor`): + Image to rescale. + scale (`float`): + The scaling factor to rescale pixel values by. + + Returns: + `torch.Tensor`: The rescaled image. + """ + return image * scale + + def normalize( + self, + image: "torch.Tensor", + mean: Union[float, Iterable[float]], + std: Union[float, Iterable[float]], + **kwargs, + ) -> "torch.Tensor": + """ + Normalize an image. image = (image - image_mean) / image_std. + + Args: + image (`torch.Tensor`): + Image to normalize. + mean (`torch.Tensor`, `float` or `Iterable[float]`): + Image mean to use for normalization. + std (`torch.Tensor`, `float` or `Iterable[float]`): + Image standard deviation to use for normalization. + + Returns: + `torch.Tensor`: The normalized image. + """ + return F.normalize(image, mean, std) + + @lru_cache(maxsize=10) + def _fuse_mean_std_and_rescale_factor( + self, + do_normalize: Optional[bool] = None, + image_mean: Optional[Union[float, list[float]]] = None, + image_std: Optional[Union[float, list[float]]] = None, + do_rescale: Optional[bool] = None, + rescale_factor: Optional[float] = None, + device: Optional["torch.device"] = None, + ) -> tuple: + if do_rescale and do_normalize: + # Fused rescale and normalize + image_mean = torch.tensor(image_mean, device=device) * (1.0 / rescale_factor) + image_std = torch.tensor(image_std, device=device) * (1.0 / rescale_factor) + do_rescale = False + return image_mean, image_std, do_rescale + + def rescale_and_normalize( + self, + images: "torch.Tensor", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: Union[float, list[float]], + image_std: Union[float, list[float]], + ) -> "torch.Tensor": + """ + Rescale and normalize images. + """ + image_mean, image_std, do_rescale = self._fuse_mean_std_and_rescale_factor( + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + do_rescale=do_rescale, + rescale_factor=rescale_factor, + device=images.device, + ) + # if/elif as we use fused rescale and normalize if both are set to True + if do_normalize: + images = self.normalize(images.to(dtype=torch.float32), image_mean, image_std) + elif do_rescale: + images = self.rescale(images, rescale_factor) + + return images + + def center_crop( + self, + image: "torch.Tensor", + size: dict[str, int], + **kwargs, + ) -> "torch.Tensor": + """ + Center crop an image to `(size["height"], size["width"])`. If the input size is smaller than `crop_size` along + any edge, the image is padded with 0's and then center cropped. + + Args: + image (`"torch.Tensor"`): + Image to center crop. + size (`Dict[str, int]`): + Size of the output image. + + Returns: + `torch.Tensor`: The center cropped image. + """ + if size.height is None or size.width is None: + raise ValueError(f"The size dictionary must have keys 'height' and 'width'. Got {size.keys()}") + return F.center_crop(image, (size["height"], size["width"])) + + def convert_to_rgb( + self, + image: ImageInput, + ) -> ImageInput: + """ + Converts an image to RGB format. Only converts if the image is of type PIL.Image.Image, otherwise returns the image + as is. + Args: + image (ImageInput): + The image to convert. + + Returns: + ImageInput: The converted image. + """ + return convert_to_rgb(image) + + def filter_out_unused_kwargs(self, kwargs: dict): + """ + Filter out the unused kwargs from the kwargs dictionary. + """ + if self.unused_kwargs is None: + return kwargs + + for kwarg_name in self.unused_kwargs: + if kwarg_name in kwargs: + logger.warning_once(f"This processor does not use the `{kwarg_name}` parameter. It will be ignored.") + kwargs.pop(kwarg_name) + return kwargs + + def _prepare_images_structure( + self, + images: ImageInput, + ) -> ImageInput: + """ + Prepare the images structure for processing. + + Args: + images (`ImageInput`): + The input images to process. + + Returns: + `ImageInput`: The images with a valid nesting. + """ + return make_flat_list_of_images(images) + + def _process_image( + self, + image: ImageInput, + do_convert_rgb: Optional[bool] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, + device: Optional["torch.device"] = None, + ) -> "torch.Tensor": + image_type = get_image_type(image) + if image_type not in [ImageType.PIL, ImageType.TORCH, ImageType.NUMPY]: + raise ValueError(f"Unsupported input image type {image_type}") + + if do_convert_rgb: + image = self.convert_to_rgb(image) + + if image_type == ImageType.PIL: + image = F.pil_to_tensor(image) + elif image_type == ImageType.NUMPY: + # not using F.to_tensor as it doesn't handle (C, H, W) numpy arrays + image = torch.from_numpy(image).contiguous() + + # Infer the channel dimension format if not provided + if input_data_format is None: + input_data_format = infer_channel_dimension_format(image) + + if input_data_format == ChannelDimension.LAST: + # We force the channel dimension to be first for torch tensors as this is what torchvision expects. + image = image.permute(2, 0, 1).contiguous() + + # Now that we have torch tensors, we can move them to the right device + if device is not None: + image = image.to(device) + + return image + + def _prepare_input_images( + self, + images: ImageInput, + do_convert_rgb: Optional[bool] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, + device: Optional["torch.device"] = None, + ) -> list["torch.Tensor"]: + """ + Prepare the input images for processing. + """ + images = self._prepare_images_structure(images) + process_image_fn = partial( + self._process_image, + do_convert_rgb=do_convert_rgb, + input_data_format=input_data_format, + device=device, + ) + # todo: yoni - check if we can parallelize this efficiently + processed_images = [] + for image in images: + processed_images.append(process_image_fn(image)) + + return processed_images + + def _further_process_kwargs( + self, + size: Optional[SizeDict] = None, + crop_size: Optional[SizeDict] = None, + default_to_square: Optional[bool] = None, + image_mean: Optional[Union[float, list[float]]] = None, + image_std: Optional[Union[float, list[float]]] = None, + data_format: Optional[ChannelDimension] = None, + **kwargs, + ) -> dict: + """ + Update kwargs that need further processing before being validated + Can be overridden by subclasses to customize the processing of kwargs. + """ + if kwargs is None: + kwargs = {} + if size is not None: + size = SizeDict(**get_size_dict(size=size, default_to_square=default_to_square)) + if crop_size is not None: + crop_size = SizeDict(**get_size_dict(crop_size, param_name="crop_size")) + if isinstance(image_mean, list): + image_mean = tuple(image_mean) + if isinstance(image_std, list): + image_std = tuple(image_std) + if data_format is None: + data_format = ChannelDimension.FIRST + + kwargs["size"] = size + kwargs["crop_size"] = crop_size + kwargs["default_to_square"] = default_to_square + kwargs["image_mean"] = image_mean + kwargs["image_std"] = image_std + kwargs["data_format"] = data_format + + return kwargs + + def _validate_preprocess_kwargs( + self, + do_rescale: Optional[bool] = None, + rescale_factor: Optional[float] = None, + do_normalize: Optional[bool] = None, + image_mean: Optional[Union[float, tuple[float]]] = None, + image_std: Optional[Union[float, tuple[float]]] = None, + do_resize: Optional[bool] = None, + size: Optional[SizeDict] = None, + do_center_crop: Optional[bool] = None, + crop_size: Optional[SizeDict] = None, + resample: Optional[Union["PILImageResampling", "F.InterpolationMode"]] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + data_format: Optional[ChannelDimension] = None, + **kwargs, + ): + """ + validate the kwargs for the preprocess method. + """ + validate_fast_preprocess_arguments( + do_rescale=do_rescale, + rescale_factor=rescale_factor, + do_normalize=do_normalize, + image_mean=image_mean, + image_std=image_std, + do_resize=do_resize, + size=size, + do_center_crop=do_center_crop, + crop_size=crop_size, + resample=resample, + return_tensors=return_tensors, + data_format=data_format, + ) + + @add_start_docstrings(BASE_IMAGE_PROCESSOR_FAST_DOCSTRING_PREPROCESS) + def preprocess(self, images: ImageInput, **kwargs: Unpack[DefaultFastImageProcessorKwargs]) -> BatchFeature: + validate_kwargs(captured_kwargs=kwargs.keys(), valid_processor_keys=self.valid_kwargs.__annotations__.keys()) + # Set default kwargs from self. This ensures that if a kwarg is not provided + # by the user, it gets its default value from the instance, or is set to None. + for kwarg_name in self.valid_kwargs.__annotations__: + kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None)) + + # Extract parameters that are only used for preparing the input images + do_convert_rgb = kwargs.pop("do_convert_rgb") + input_data_format = kwargs.pop("input_data_format") + device = kwargs.pop("device") + # Prepare input images + images = self._prepare_input_images( + images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device + ) + + # Update kwargs that need further processing before being validated + kwargs = self._further_process_kwargs(**kwargs) + + # Validate kwargs + self._validate_preprocess_kwargs(**kwargs) + + # torch resize uses interpolation instead of resample + resample = kwargs.pop("resample") + kwargs["interpolation"] = ( + pil_torch_interpolation_mapping[resample] if isinstance(resample, (PILImageResampling, int)) else resample + ) + + # Pop kwargs that are not needed in _preprocess + kwargs.pop("default_to_square") + kwargs.pop("data_format") + + return self._preprocess(images=images, **kwargs) + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + size: SizeDict, + interpolation: Optional["F.InterpolationMode"], + do_center_crop: bool, + crop_size: SizeDict, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: Optional[Union[float, list[float]]], + image_std: Optional[Union[float, list[float]]], + return_tensors: Optional[Union[str, TensorType]], + **kwargs, + ) -> BatchFeature: + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images) + resized_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(image=stacked_images, size=size, interpolation=interpolation) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + # Group images by size for further processing + # Needed in case do_resize is False, or resize returns images with different sizes + grouped_images, grouped_images_index = group_images_by_shape(resized_images) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + if do_center_crop: + stacked_images = self.center_crop(stacked_images, crop_size) + # Fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images + + return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + def to_dict(self): + encoder_dict = super().to_dict() + encoder_dict.pop("_valid_processor_keys", None) + return encoder_dict + + +class SemanticSegmentationMixin: + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple] = None): + """ + Converts the output of [`MobileNetV2ForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch. + + Args: + outputs ([`MobileNetV2ForSemanticSegmentation`]): + Raw outputs of the model. + target_sizes (`List[Tuple]` of length `batch_size`, *optional*): + List of tuples corresponding to the requested final size (height, width) of each prediction. If unset, + predictions will not be resized. + + Returns: + semantic_segmentation: `List[torch.Tensor]` of length `batch_size`, where each item is a semantic + segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is + specified). Each entry of each `torch.Tensor` correspond to a semantic class id. + """ + logits = outputs.logits + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if len(logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + # if is_torch_tensor(target_sizes): + # target_sizes = target_sizes.numpy() + + semantic_segmentation = [] + + for idx in range(len(logits)): + resized_logits = torch.nn.functional.interpolate( + logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = logits.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation diff --git a/lib/python3.12/site-packages/transformers/image_transforms.py b/lib/python3.12/site-packages/transformers/image_transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..5b0ba3f9122f17c4355c94f471af5720bde2d814 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/image_transforms.py @@ -0,0 +1,881 @@ +# Copyright 2022 The HuggingFace Inc. team. +# +# 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 collections.abc import Collection, Iterable +from math import ceil +from typing import Optional, Union + +import numpy as np + +from .image_utils import ( + ChannelDimension, + ImageInput, + get_channel_dimension_axis, + get_image_size, + infer_channel_dimension_format, +) +from .utils import ExplicitEnum, TensorType, is_jax_tensor, is_tf_tensor, is_torch_tensor +from .utils.import_utils import ( + is_flax_available, + is_tf_available, + is_torch_available, + is_vision_available, + requires_backends, +) + + +if is_vision_available(): + import PIL + + from .image_utils import PILImageResampling + +if is_torch_available(): + import torch + +if is_tf_available(): + import tensorflow as tf + +if is_flax_available(): + import jax.numpy as jnp + + +def to_channel_dimension_format( + image: np.ndarray, + channel_dim: Union[ChannelDimension, str], + input_channel_dim: Optional[Union[ChannelDimension, str]] = None, +) -> np.ndarray: + """ + Converts `image` to the channel dimension format specified by `channel_dim`. + + Args: + image (`numpy.ndarray`): + The image to have its channel dimension set. + channel_dim (`ChannelDimension`): + The channel dimension format to use. + input_channel_dim (`ChannelDimension`, *optional*): + The channel dimension format of the input image. If not provided, it will be inferred from the input image. + + Returns: + `np.ndarray`: The image with the channel dimension set to `channel_dim`. + """ + if not isinstance(image, np.ndarray): + raise TypeError(f"Input image must be of type np.ndarray, got {type(image)}") + + if input_channel_dim is None: + input_channel_dim = infer_channel_dimension_format(image) + + target_channel_dim = ChannelDimension(channel_dim) + if input_channel_dim == target_channel_dim: + return image + + if target_channel_dim == ChannelDimension.FIRST: + image = image.transpose((2, 0, 1)) + elif target_channel_dim == ChannelDimension.LAST: + image = image.transpose((1, 2, 0)) + else: + raise ValueError(f"Unsupported channel dimension format: {channel_dim}") + + return image + + +def rescale( + image: np.ndarray, + scale: float, + data_format: Optional[ChannelDimension] = None, + dtype: np.dtype = np.float32, + input_data_format: Optional[Union[str, ChannelDimension]] = None, +) -> np.ndarray: + """ + Rescales `image` by `scale`. + + Args: + image (`np.ndarray`): + The image to rescale. + scale (`float`): + The scale to use for rescaling the image. + data_format (`ChannelDimension`, *optional*): + The channel dimension format of the image. If not provided, it will be the same as the input image. + dtype (`np.dtype`, *optional*, defaults to `np.float32`): + The dtype of the output image. Defaults to `np.float32`. Used for backwards compatibility with feature + extractors. + input_data_format (`ChannelDimension`, *optional*): + The channel dimension format of the input image. If not provided, it will be inferred from the input image. + + Returns: + `np.ndarray`: The rescaled image. + """ + if not isinstance(image, np.ndarray): + raise TypeError(f"Input image must be of type np.ndarray, got {type(image)}") + + rescaled_image = image.astype(np.float64) * scale # Numpy type promotion has changed, so always upcast first + if data_format is not None: + rescaled_image = to_channel_dimension_format(rescaled_image, data_format, input_data_format) + + rescaled_image = rescaled_image.astype(dtype) # Finally downcast to the desired dtype at the end + + return rescaled_image + + +def _rescale_for_pil_conversion(image): + """ + Detects whether or not the image needs to be rescaled before being converted to a PIL image. + + The assumption is that if the image is of type `np.float` and all values are between 0 and 1, it needs to be + rescaled. + """ + if image.dtype == np.uint8: + do_rescale = False + elif np.allclose(image, image.astype(int)): + if np.all(0 <= image) and np.all(image <= 255): + do_rescale = False + else: + raise ValueError( + "The image to be converted to a PIL image contains values outside the range [0, 255], " + f"got [{image.min()}, {image.max()}] which cannot be converted to uint8." + ) + elif np.all(0 <= image) and np.all(image <= 1): + do_rescale = True + else: + raise ValueError( + "The image to be converted to a PIL image contains values outside the range [0, 1], " + f"got [{image.min()}, {image.max()}] which cannot be converted to uint8." + ) + return do_rescale + + +def to_pil_image( + image: Union[np.ndarray, "PIL.Image.Image", "torch.Tensor", "tf.Tensor", "jnp.ndarray"], + do_rescale: Optional[bool] = None, + image_mode: Optional[str] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, +) -> "PIL.Image.Image": + """ + Converts `image` to a PIL Image. Optionally rescales it and puts the channel dimension back as the last axis if + needed. + + Args: + image (`PIL.Image.Image` or `numpy.ndarray` or `torch.Tensor` or `tf.Tensor`): + The image to convert to the `PIL.Image` format. + do_rescale (`bool`, *optional*): + Whether or not to apply the scaling factor (to make pixel values integers between 0 and 255). Will default + to `True` if the image type is a floating type and casting to `int` would result in a loss of precision, + and `False` otherwise. + image_mode (`str`, *optional*): + The mode to use for the PIL image. If unset, will use the default mode for the input image type. + input_data_format (`ChannelDimension`, *optional*): + The channel dimension format of the input image. If unset, will use the inferred format from the input. + + Returns: + `PIL.Image.Image`: The converted image. + """ + requires_backends(to_pil_image, ["vision"]) + + if isinstance(image, PIL.Image.Image): + return image + + # Convert all tensors to numpy arrays before converting to PIL image + if is_torch_tensor(image) or is_tf_tensor(image): + image = image.numpy() + elif is_jax_tensor(image): + image = np.array(image) + elif not isinstance(image, np.ndarray): + raise ValueError(f"Input image type not supported: {type(image)}") + + # If the channel has been moved to first dim, we put it back at the end. + image = to_channel_dimension_format(image, ChannelDimension.LAST, input_data_format) + + # If there is a single channel, we squeeze it, as otherwise PIL can't handle it. + image = np.squeeze(image, axis=-1) if image.shape[-1] == 1 else image + + # PIL.Image can only store uint8 values so we rescale the image to be between 0 and 255 if needed. + do_rescale = _rescale_for_pil_conversion(image) if do_rescale is None else do_rescale + + if do_rescale: + image = rescale(image, 255) + + image = image.astype(np.uint8) + return PIL.Image.fromarray(image, mode=image_mode) + + +def get_size_with_aspect_ratio(image_size, size, max_size=None) -> tuple[int, int]: + """ + Computes the output image size given the input image size and the desired output size. + + Args: + image_size (`Tuple[int, int]`): + The input image size. + size (`int`): + The desired output size. + max_size (`int`, *optional*): + The maximum allowed output size. + """ + height, width = image_size + raw_size = None + if max_size is not None: + min_original_size = float(min((height, width))) + max_original_size = float(max((height, width))) + if max_original_size / min_original_size * size > max_size: + raw_size = max_size * min_original_size / max_original_size + size = int(round(raw_size)) + + if (height <= width and height == size) or (width <= height and width == size): + oh, ow = height, width + elif width < height: + ow = size + if max_size is not None and raw_size is not None: + oh = int(raw_size * height / width) + else: + oh = int(size * height / width) + else: + oh = size + if max_size is not None and raw_size is not None: + ow = int(raw_size * width / height) + else: + ow = int(size * width / height) + + return (oh, ow) + + +# Logic adapted from torchvision resizing logic: https://github.com/pytorch/vision/blob/511924c1ced4ce0461197e5caa64ce5b9e558aab/torchvision/transforms/functional.py#L366 +def get_resize_output_image_size( + input_image: np.ndarray, + size: Union[int, tuple[int, int], list[int], tuple[int]], + default_to_square: bool = True, + max_size: Optional[int] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, +) -> tuple: + """ + Find the target (height, width) dimension of the output image after resizing given the input image and the desired + size. + + Args: + input_image (`np.ndarray`): + The image to resize. + size (`int` or `Tuple[int, int]` or List[int] or `Tuple[int]`): + The size to use for resizing the image. If `size` is a sequence like (h, w), output size will be matched to + this. + + If `size` is an int and `default_to_square` is `True`, then image will be resized to (size, size). If + `size` is an int and `default_to_square` is `False`, then smaller edge of the image will be matched to this + number. i.e, if height > width, then image will be rescaled to (size * height / width, size). + default_to_square (`bool`, *optional*, defaults to `True`): + How to convert `size` when it is a single int. If set to `True`, the `size` will be converted to a square + (`size`,`size`). If set to `False`, will replicate + [`torchvision.transforms.Resize`](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.Resize) + with support for resizing only the smallest edge and providing an optional `max_size`. + max_size (`int`, *optional*): + The maximum allowed for the longer edge of the resized image: if the longer edge of the image is greater + than `max_size` after being resized according to `size`, then the image is resized again so that the longer + edge is equal to `max_size`. As a result, `size` might be overruled, i.e the smaller edge may be shorter + than `size`. Only used if `default_to_square` is `False`. + input_data_format (`ChannelDimension`, *optional*): + The channel dimension format of the input image. If unset, will use the inferred format from the input. + + Returns: + `tuple`: The target (height, width) dimension of the output image after resizing. + """ + if isinstance(size, (tuple, list)): + if len(size) == 2: + return tuple(size) + elif len(size) == 1: + # Perform same logic as if size was an int + size = size[0] + else: + raise ValueError("size must have 1 or 2 elements if it is a list or tuple") + + if default_to_square: + return (size, size) + + height, width = get_image_size(input_image, input_data_format) + short, long = (width, height) if width <= height else (height, width) + requested_new_short = size + + new_short, new_long = requested_new_short, int(requested_new_short * long / short) + + if max_size is not None: + if max_size <= requested_new_short: + raise ValueError( + f"max_size = {max_size} must be strictly greater than the requested " + f"size for the smaller edge size = {size}" + ) + if new_long > max_size: + new_short, new_long = int(max_size * new_short / new_long), max_size + + return (new_long, new_short) if width <= height else (new_short, new_long) + + +def resize( + image: np.ndarray, + size: tuple[int, int], + resample: "PILImageResampling" = None, + reducing_gap: Optional[int] = None, + data_format: Optional[ChannelDimension] = None, + return_numpy: bool = True, + input_data_format: Optional[Union[str, ChannelDimension]] = None, +) -> np.ndarray: + """ + Resizes `image` to `(height, width)` specified by `size` using the PIL library. + + Args: + image (`np.ndarray`): + The image to resize. + size (`Tuple[int, int]`): + The size to use for resizing the image. + resample (`int`, *optional*, defaults to `PILImageResampling.BILINEAR`): + The filter to user for resampling. + reducing_gap (`int`, *optional*): + Apply optimization by resizing the image in two steps. The bigger `reducing_gap`, the closer the result to + the fair resampling. See corresponding Pillow documentation for more details. + data_format (`ChannelDimension`, *optional*): + The channel dimension format of the output image. If unset, will use the inferred format from the input. + return_numpy (`bool`, *optional*, defaults to `True`): + Whether or not to return the resized image as a numpy array. If False a `PIL.Image.Image` object is + returned. + input_data_format (`ChannelDimension`, *optional*): + The channel dimension format of the input image. If unset, will use the inferred format from the input. + + Returns: + `np.ndarray`: The resized image. + """ + requires_backends(resize, ["vision"]) + + resample = resample if resample is not None else PILImageResampling.BILINEAR + + if not len(size) == 2: + raise ValueError("size must have 2 elements") + + # For all transformations, we want to keep the same data format as the input image unless otherwise specified. + # The resized image from PIL will always have channels last, so find the input format first. + if input_data_format is None: + input_data_format = infer_channel_dimension_format(image) + data_format = input_data_format if data_format is None else data_format + + # To maintain backwards compatibility with the resizing done in previous image feature extractors, we use + # the pillow library to resize the image and then convert back to numpy + do_rescale = False + if not isinstance(image, PIL.Image.Image): + do_rescale = _rescale_for_pil_conversion(image) + image = to_pil_image(image, do_rescale=do_rescale, input_data_format=input_data_format) + height, width = size + # PIL images are in the format (width, height) + resized_image = image.resize((width, height), resample=resample, reducing_gap=reducing_gap) + + if return_numpy: + resized_image = np.array(resized_image) + # If the input image channel dimension was of size 1, then it is dropped when converting to a PIL image + # so we need to add it back if necessary. + resized_image = np.expand_dims(resized_image, axis=-1) if resized_image.ndim == 2 else resized_image + # The image is always in channels last format after converting from a PIL image + resized_image = to_channel_dimension_format( + resized_image, data_format, input_channel_dim=ChannelDimension.LAST + ) + # If an image was rescaled to be in the range [0, 255] before converting to a PIL image, then we need to + # rescale it back to the original range. + resized_image = rescale(resized_image, 1 / 255) if do_rescale else resized_image + return resized_image + + +def normalize( + image: np.ndarray, + mean: Union[float, Collection[float]], + std: Union[float, Collection[float]], + data_format: Optional[ChannelDimension] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, +) -> np.ndarray: + """ + Normalizes `image` using the mean and standard deviation specified by `mean` and `std`. + + image = (image - mean) / std + + Args: + image (`np.ndarray`): + The image to normalize. + mean (`float` or `Collection[float]`): + The mean to use for normalization. + std (`float` or `Collection[float]`): + The standard deviation to use for normalization. + data_format (`ChannelDimension`, *optional*): + The channel dimension format of the output image. If unset, will use the inferred format from the input. + input_data_format (`ChannelDimension`, *optional*): + The channel dimension format of the input image. If unset, will use the inferred format from the input. + """ + if not isinstance(image, np.ndarray): + raise ValueError("image must be a numpy array") + + if input_data_format is None: + input_data_format = infer_channel_dimension_format(image) + + channel_axis = get_channel_dimension_axis(image, input_data_format=input_data_format) + num_channels = image.shape[channel_axis] + + # We cast to float32 to avoid errors that can occur when subtracting uint8 values. + # We preserve the original dtype if it is a float type to prevent upcasting float16. + if not np.issubdtype(image.dtype, np.floating): + image = image.astype(np.float32) + + if isinstance(mean, Collection): + if len(mean) != num_channels: + raise ValueError(f"mean must have {num_channels} elements if it is an iterable, got {len(mean)}") + else: + mean = [mean] * num_channels + mean = np.array(mean, dtype=image.dtype) + + if isinstance(std, Collection): + if len(std) != num_channels: + raise ValueError(f"std must have {num_channels} elements if it is an iterable, got {len(std)}") + else: + std = [std] * num_channels + std = np.array(std, dtype=image.dtype) + + if input_data_format == ChannelDimension.LAST: + image = (image - mean) / std + else: + image = ((image.T - mean) / std).T + + image = to_channel_dimension_format(image, data_format, input_data_format) if data_format is not None else image + return image + + +def center_crop( + image: np.ndarray, + size: tuple[int, int], + data_format: Optional[Union[str, ChannelDimension]] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, +) -> np.ndarray: + """ + Crops the `image` to the specified `size` using a center crop. Note that if the image is too small to be cropped to + the size given, it will be padded (so the returned result will always be of size `size`). + + Args: + image (`np.ndarray`): + The image to crop. + size (`Tuple[int, int]`): + The target size for the cropped image. + data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the output image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + If unset, will use the inferred format of the input image. + input_data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + If unset, will use the inferred format of the input image. + Returns: + `np.ndarray`: The cropped image. + """ + requires_backends(center_crop, ["vision"]) + + if not isinstance(image, np.ndarray): + raise TypeError(f"Input image must be of type np.ndarray, got {type(image)}") + + if not isinstance(size, Iterable) or len(size) != 2: + raise ValueError("size must have 2 elements representing the height and width of the output image") + + if input_data_format is None: + input_data_format = infer_channel_dimension_format(image) + output_data_format = data_format if data_format is not None else input_data_format + + # We perform the crop in (C, H, W) format and then convert to the output format + image = to_channel_dimension_format(image, ChannelDimension.FIRST, input_data_format) + + orig_height, orig_width = get_image_size(image, ChannelDimension.FIRST) + crop_height, crop_width = size + crop_height, crop_width = int(crop_height), int(crop_width) + + # In case size is odd, (image_shape[0] + size[0]) // 2 won't give the proper result. + top = (orig_height - crop_height) // 2 + bottom = top + crop_height + # In case size is odd, (image_shape[1] + size[1]) // 2 won't give the proper result. + left = (orig_width - crop_width) // 2 + right = left + crop_width + + # Check if cropped area is within image boundaries + if top >= 0 and bottom <= orig_height and left >= 0 and right <= orig_width: + image = image[..., top:bottom, left:right] + image = to_channel_dimension_format(image, output_data_format, ChannelDimension.FIRST) + return image + + # Otherwise, we may need to pad if the image is too small. Oh joy... + new_height = max(crop_height, orig_height) + new_width = max(crop_width, orig_width) + new_shape = image.shape[:-2] + (new_height, new_width) + new_image = np.zeros_like(image, shape=new_shape) + + # If the image is too small, pad it with zeros + top_pad = ceil((new_height - orig_height) / 2) + bottom_pad = top_pad + orig_height + left_pad = ceil((new_width - orig_width) / 2) + right_pad = left_pad + orig_width + new_image[..., top_pad:bottom_pad, left_pad:right_pad] = image + + top += top_pad + bottom += top_pad + left += left_pad + right += left_pad + + new_image = new_image[..., max(0, top) : min(new_height, bottom), max(0, left) : min(new_width, right)] + new_image = to_channel_dimension_format(new_image, output_data_format, ChannelDimension.FIRST) + + return new_image + + +def _center_to_corners_format_torch(bboxes_center: "torch.Tensor") -> "torch.Tensor": + center_x, center_y, width, height = bboxes_center.unbind(-1) + bbox_corners = torch.stack( + # top left x, top left y, bottom right x, bottom right y + [(center_x - 0.5 * width), (center_y - 0.5 * height), (center_x + 0.5 * width), (center_y + 0.5 * height)], + dim=-1, + ) + return bbox_corners + + +def _center_to_corners_format_numpy(bboxes_center: np.ndarray) -> np.ndarray: + center_x, center_y, width, height = bboxes_center.T + bboxes_corners = np.stack( + # top left x, top left y, bottom right x, bottom right y + [center_x - 0.5 * width, center_y - 0.5 * height, center_x + 0.5 * width, center_y + 0.5 * height], + axis=-1, + ) + return bboxes_corners + + +def _center_to_corners_format_tf(bboxes_center: "tf.Tensor") -> "tf.Tensor": + center_x, center_y, width, height = tf.unstack(bboxes_center, axis=-1) + bboxes_corners = tf.stack( + # top left x, top left y, bottom right x, bottom right y + [center_x - 0.5 * width, center_y - 0.5 * height, center_x + 0.5 * width, center_y + 0.5 * height], + axis=-1, + ) + return bboxes_corners + + +# 2 functions below inspired by https://github.com/facebookresearch/detr/blob/master/util/box_ops.py +def center_to_corners_format(bboxes_center: TensorType) -> TensorType: + """ + Converts bounding boxes from center format to corners format. + + center format: contains the coordinate for the center of the box and its width, height dimensions + (center_x, center_y, width, height) + corners format: contains the coordinates for the top-left and bottom-right corners of the box + (top_left_x, top_left_y, bottom_right_x, bottom_right_y) + """ + # Function is used during model forward pass, so we use the input framework if possible, without + # converting to numpy + if is_torch_tensor(bboxes_center): + return _center_to_corners_format_torch(bboxes_center) + elif isinstance(bboxes_center, np.ndarray): + return _center_to_corners_format_numpy(bboxes_center) + elif is_tf_tensor(bboxes_center): + return _center_to_corners_format_tf(bboxes_center) + + raise ValueError(f"Unsupported input type {type(bboxes_center)}") + + +def _corners_to_center_format_torch(bboxes_corners: "torch.Tensor") -> "torch.Tensor": + top_left_x, top_left_y, bottom_right_x, bottom_right_y = bboxes_corners.unbind(-1) + b = [ + (top_left_x + bottom_right_x) / 2, # center x + (top_left_y + bottom_right_y) / 2, # center y + (bottom_right_x - top_left_x), # width + (bottom_right_y - top_left_y), # height + ] + return torch.stack(b, dim=-1) + + +def _corners_to_center_format_numpy(bboxes_corners: np.ndarray) -> np.ndarray: + top_left_x, top_left_y, bottom_right_x, bottom_right_y = bboxes_corners.T + bboxes_center = np.stack( + [ + (top_left_x + bottom_right_x) / 2, # center x + (top_left_y + bottom_right_y) / 2, # center y + (bottom_right_x - top_left_x), # width + (bottom_right_y - top_left_y), # height + ], + axis=-1, + ) + return bboxes_center + + +def _corners_to_center_format_tf(bboxes_corners: "tf.Tensor") -> "tf.Tensor": + top_left_x, top_left_y, bottom_right_x, bottom_right_y = tf.unstack(bboxes_corners, axis=-1) + bboxes_center = tf.stack( + [ + (top_left_x + bottom_right_x) / 2, # center x + (top_left_y + bottom_right_y) / 2, # center y + (bottom_right_x - top_left_x), # width + (bottom_right_y - top_left_y), # height + ], + axis=-1, + ) + return bboxes_center + + +def corners_to_center_format(bboxes_corners: TensorType) -> TensorType: + """ + Converts bounding boxes from corners format to center format. + + corners format: contains the coordinates for the top-left and bottom-right corners of the box + (top_left_x, top_left_y, bottom_right_x, bottom_right_y) + center format: contains the coordinate for the center of the box and its the width, height dimensions + (center_x, center_y, width, height) + """ + # Inverse function accepts different input types so implemented here too + if is_torch_tensor(bboxes_corners): + return _corners_to_center_format_torch(bboxes_corners) + elif isinstance(bboxes_corners, np.ndarray): + return _corners_to_center_format_numpy(bboxes_corners) + elif is_tf_tensor(bboxes_corners): + return _corners_to_center_format_tf(bboxes_corners) + + raise ValueError(f"Unsupported input type {type(bboxes_corners)}") + + +# 2 functions below copied from https://github.com/cocodataset/panopticapi/blob/master/panopticapi/utils.py +# Copyright (c) 2018, Alexander Kirillov +# All rights reserved. +def rgb_to_id(color): + """ + Converts RGB color to unique ID. + """ + if isinstance(color, np.ndarray) and len(color.shape) == 3: + if color.dtype == np.uint8: + color = color.astype(np.int32) + return color[:, :, 0] + 256 * color[:, :, 1] + 256 * 256 * color[:, :, 2] + return int(color[0] + 256 * color[1] + 256 * 256 * color[2]) + + +def id_to_rgb(id_map): + """ + Converts unique ID to RGB color. + """ + if isinstance(id_map, np.ndarray): + id_map_copy = id_map.copy() + rgb_shape = tuple(list(id_map.shape) + [3]) + rgb_map = np.zeros(rgb_shape, dtype=np.uint8) + for i in range(3): + rgb_map[..., i] = id_map_copy % 256 + id_map_copy //= 256 + return rgb_map + color = [] + for _ in range(3): + color.append(id_map % 256) + id_map //= 256 + return color + + +class PaddingMode(ExplicitEnum): + """ + Enum class for the different padding modes to use when padding images. + """ + + CONSTANT = "constant" + REFLECT = "reflect" + REPLICATE = "replicate" + SYMMETRIC = "symmetric" + + +def pad( + image: np.ndarray, + padding: Union[int, tuple[int, int], Iterable[tuple[int, int]]], + mode: PaddingMode = PaddingMode.CONSTANT, + constant_values: Union[float, Iterable[float]] = 0.0, + data_format: Optional[Union[str, ChannelDimension]] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, +) -> np.ndarray: + """ + Pads the `image` with the specified (height, width) `padding` and `mode`. + + Args: + image (`np.ndarray`): + The image to pad. + padding (`int` or `Tuple[int, int]` or `Iterable[Tuple[int, int]]`): + Padding to apply to the edges of the height, width axes. Can be one of three formats: + - `((before_height, after_height), (before_width, after_width))` unique pad widths for each axis. + - `((before, after),)` yields same before and after pad for height and width. + - `(pad,)` or int is a shortcut for before = after = pad width for all axes. + mode (`PaddingMode`): + The padding mode to use. Can be one of: + - `"constant"`: pads with a constant value. + - `"reflect"`: pads with the reflection of the vector mirrored on the first and last values of the + vector along each axis. + - `"replicate"`: pads with the replication of the last value on the edge of the array along each axis. + - `"symmetric"`: pads with the reflection of the vector mirrored along the edge of the array. + constant_values (`float` or `Iterable[float]`, *optional*): + The value to use for the padding if `mode` is `"constant"`. + data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the output image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + If unset, will use same as the input image. + input_data_format (`str` or `ChannelDimension`, *optional*): + The channel dimension format for the input image. Can be one of: + - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format. + If unset, will use the inferred format of the input image. + + Returns: + `np.ndarray`: The padded image. + + """ + if input_data_format is None: + input_data_format = infer_channel_dimension_format(image) + + def _expand_for_data_format(values): + """ + Convert values to be in the format expected by np.pad based on the data format. + """ + if isinstance(values, (int, float)): + values = ((values, values), (values, values)) + elif isinstance(values, tuple) and len(values) == 1: + values = ((values[0], values[0]), (values[0], values[0])) + elif isinstance(values, tuple) and len(values) == 2 and isinstance(values[0], int): + values = (values, values) + elif isinstance(values, tuple) and len(values) == 2 and isinstance(values[0], tuple): + values = values + else: + raise ValueError(f"Unsupported format: {values}") + + # add 0 for channel dimension + values = ((0, 0), *values) if input_data_format == ChannelDimension.FIRST else (*values, (0, 0)) + + # Add additional padding if there's a batch dimension + values = (0, *values) if image.ndim == 4 else values + return values + + padding = _expand_for_data_format(padding) + + if mode == PaddingMode.CONSTANT: + constant_values = _expand_for_data_format(constant_values) + image = np.pad(image, padding, mode="constant", constant_values=constant_values) + elif mode == PaddingMode.REFLECT: + image = np.pad(image, padding, mode="reflect") + elif mode == PaddingMode.REPLICATE: + image = np.pad(image, padding, mode="edge") + elif mode == PaddingMode.SYMMETRIC: + image = np.pad(image, padding, mode="symmetric") + else: + raise ValueError(f"Invalid padding mode: {mode}") + + image = to_channel_dimension_format(image, data_format, input_data_format) if data_format is not None else image + return image + + +# TODO (Amy): Accept 1/3/4 channel numpy array as input and return np.array as default +def convert_to_rgb(image: ImageInput) -> ImageInput: + """ + Converts an image to RGB format. Only converts if the image is of type PIL.Image.Image, otherwise returns the image + as is. + Args: + image (Image): + The image to convert. + """ + requires_backends(convert_to_rgb, ["vision"]) + + if not isinstance(image, PIL.Image.Image): + return image + + if image.mode == "RGB": + return image + + image = image.convert("RGB") + return image + + +def flip_channel_order( + image: np.ndarray, + data_format: Optional[ChannelDimension] = None, + input_data_format: Optional[Union[str, ChannelDimension]] = None, +) -> np.ndarray: + """ + Flips the channel order of the image. + + If the image is in RGB format, it will be converted to BGR and vice versa. + + Args: + image (`np.ndarray`): + The image to flip. + data_format (`ChannelDimension`, *optional*): + The channel dimension format for the output image. Can be one of: + - `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `ChannelDimension.LAST`: image in (height, width, num_channels) format. + If unset, will use same as the input image. + input_data_format (`ChannelDimension`, *optional*): + The channel dimension format for the input image. Can be one of: + - `ChannelDimension.FIRST`: image in (num_channels, height, width) format. + - `ChannelDimension.LAST`: image in (height, width, num_channels) format. + If unset, will use the inferred format of the input image. + """ + input_data_format = infer_channel_dimension_format(image) if input_data_format is None else input_data_format + + if input_data_format == ChannelDimension.LAST: + image = image[..., ::-1] + elif input_data_format == ChannelDimension.FIRST: + image = image[::-1, ...] + else: + raise ValueError(f"Unsupported channel dimension: {input_data_format}") + + if data_format is not None: + image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) + return image + + +def _cast_tensor_to_float(x): + if x.is_floating_point(): + return x + return x.float() + + +def group_images_by_shape( + images: list["torch.Tensor"], +) -> tuple[dict[tuple[int, int], list["torch.Tensor"]], dict[int, tuple[tuple[int, int], int]]]: + """ + Groups images by shape. + Returns a dictionary with the shape as key and a list of images with that shape as value, + and a dictionary with the index of the image in the original list as key and the shape and index in the grouped list as value. + """ + grouped_images = {} + grouped_images_index = {} + for i, image in enumerate(images): + shape = image.shape[1:] + if shape not in grouped_images: + grouped_images[shape] = [] + grouped_images[shape].append(image) + grouped_images_index[i] = (shape, len(grouped_images[shape]) - 1) + # stack images with the same shape + grouped_images = {shape: torch.stack(images, dim=0) for shape, images in grouped_images.items()} + return grouped_images, grouped_images_index + + +def reorder_images( + processed_images: dict[tuple[int, int], "torch.Tensor"], grouped_images_index: dict[int, tuple[int, int]] +) -> list["torch.Tensor"]: + """ + Reconstructs a list of images in the original order. + """ + return [ + processed_images[grouped_images_index[i][0]][grouped_images_index[i][1]] + for i in range(len(grouped_images_index)) + ] + + +class NumpyToTensor: + """ + Convert a numpy array to a PyTorch tensor. + """ + + def __call__(self, image: np.ndarray): + # Same as in PyTorch, we assume incoming numpy images are in HWC format + # c.f. https://github.com/pytorch/vision/blob/61d97f41bc209e1407dcfbd685d2ee2da9c1cdad/torchvision/transforms/functional.py#L154 + return torch.from_numpy(image.transpose(2, 0, 1)).contiguous() diff --git a/lib/python3.12/site-packages/transformers/image_utils.py b/lib/python3.12/site-packages/transformers/image_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f07ac1ae7d91dea63514f74ee86a46051383d0ef --- /dev/null +++ b/lib/python3.12/site-packages/transformers/image_utils.py @@ -0,0 +1,1345 @@ +# Copyright 2021 The HuggingFace Inc. team. +# +# 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. + +import base64 +import os +from collections.abc import Iterable +from contextlib import redirect_stdout +from dataclasses import dataclass +from io import BytesIO +from typing import TYPE_CHECKING, Callable, Optional, Union + +import numpy as np +import requests +from packaging import version + +from .utils import ( + ExplicitEnum, + is_av_available, + is_cv2_available, + is_decord_available, + is_jax_tensor, + is_numpy_array, + is_tf_tensor, + is_torch_available, + is_torch_tensor, + is_torchvision_available, + is_vision_available, + is_yt_dlp_available, + logging, + requires_backends, + to_numpy, +) +from .utils.constants import ( # noqa: F401 + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + OPENAI_CLIP_MEAN, + OPENAI_CLIP_STD, +) + + +if is_vision_available(): + import PIL.Image + import PIL.ImageOps + + if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"): + PILImageResampling = PIL.Image.Resampling + else: + PILImageResampling = PIL.Image + + if is_torchvision_available(): + from torchvision import io as torchvision_io + from torchvision.transforms import InterpolationMode + + pil_torch_interpolation_mapping = { + PILImageResampling.NEAREST: InterpolationMode.NEAREST, + PILImageResampling.BOX: InterpolationMode.BOX, + PILImageResampling.BILINEAR: InterpolationMode.BILINEAR, + PILImageResampling.HAMMING: InterpolationMode.HAMMING, + PILImageResampling.BICUBIC: InterpolationMode.BICUBIC, + PILImageResampling.LANCZOS: InterpolationMode.LANCZOS, + } + + +if TYPE_CHECKING: + if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) + + +ImageInput = Union[ + "PIL.Image.Image", np.ndarray, "torch.Tensor", list["PIL.Image.Image"], list[np.ndarray], list["torch.Tensor"] +] # noqa + + +VideoInput = Union[ + list["PIL.Image.Image"], + "np.ndarray", + "torch.Tensor", + list["np.ndarray"], + list["torch.Tensor"], + list[list["PIL.Image.Image"]], + list[list["np.ndarray"]], + list[list["torch.Tensor"]], +] # noqa + + +class ChannelDimension(ExplicitEnum): + FIRST = "channels_first" + LAST = "channels_last" + + +class AnnotationFormat(ExplicitEnum): + COCO_DETECTION = "coco_detection" + COCO_PANOPTIC = "coco_panoptic" + + +class AnnotionFormat(ExplicitEnum): + COCO_DETECTION = AnnotationFormat.COCO_DETECTION.value + COCO_PANOPTIC = AnnotationFormat.COCO_PANOPTIC.value + + +@dataclass +class VideoMetadata: + total_num_frames: int + fps: float + duration: float + video_backend: str + + +AnnotationType = dict[str, Union[int, str, list[dict]]] + + +def is_pil_image(img): + return is_vision_available() and isinstance(img, PIL.Image.Image) + + +class ImageType(ExplicitEnum): + PIL = "pillow" + TORCH = "torch" + NUMPY = "numpy" + TENSORFLOW = "tensorflow" + JAX = "jax" + + +def get_image_type(image): + if is_pil_image(image): + return ImageType.PIL + if is_torch_tensor(image): + return ImageType.TORCH + if is_numpy_array(image): + return ImageType.NUMPY + if is_tf_tensor(image): + return ImageType.TENSORFLOW + if is_jax_tensor(image): + return ImageType.JAX + raise ValueError(f"Unrecognised image type {type(image)}") + + +def is_valid_image(img): + return is_pil_image(img) or is_numpy_array(img) or is_torch_tensor(img) or is_tf_tensor(img) or is_jax_tensor(img) + + +def is_valid_list_of_images(images: list): + return images and all(is_valid_image(image) for image in images) + + +def valid_images(imgs): + # If we have an list of images, make sure every image is valid + if isinstance(imgs, (list, tuple)): + for img in imgs: + if not valid_images(img): + return False + # If not a list of tuple, we have been given a single image or batched tensor of images + elif not is_valid_image(imgs): + return False + return True + + +def is_batched(img): + if isinstance(img, (list, tuple)): + return is_valid_image(img[0]) + return False + + +def is_scaled_image(image: np.ndarray) -> bool: + """ + Checks to see whether the pixel values have already been rescaled to [0, 1]. + """ + if image.dtype == np.uint8: + return False + + # It's possible the image has pixel values in [0, 255] but is of floating type + return np.min(image) >= 0 and np.max(image) <= 1 + + +def make_list_of_images(images, expected_ndims: int = 3) -> list[ImageInput]: + """ + Ensure that the output is a list of images. If the input is a single image, it is converted to a list of length 1. + If the input is a batch of images, it is converted to a list of images. + + Args: + images (`ImageInput`): + Image of images to turn into a list of images. + expected_ndims (`int`, *optional*, defaults to 3): + Expected number of dimensions for a single input image. If the input image has a different number of + dimensions, an error is raised. + """ + if is_batched(images): + return images + + # Either the input is a single image, in which case we create a list of length 1 + if is_pil_image(images): + # PIL images are never batched + return [images] + + if is_valid_image(images): + if images.ndim == expected_ndims + 1: + # Batch of images + images = list(images) + elif images.ndim == expected_ndims: + # Single image + images = [images] + else: + raise ValueError( + f"Invalid image shape. Expected either {expected_ndims + 1} or {expected_ndims} dimensions, but got" + f" {images.ndim} dimensions." + ) + return images + raise ValueError( + "Invalid image type. Expected either PIL.Image.Image, numpy.ndarray, torch.Tensor, tf.Tensor or " + f"jax.ndarray, but got {type(images)}." + ) + + +def make_flat_list_of_images( + images: Union[list[ImageInput], ImageInput], +) -> ImageInput: + """ + Ensure that the output is a flat list of images. If the input is a single image, it is converted to a list of length 1. + If the input is a nested list of images, it is converted to a flat list of images. + Args: + images (`Union[List[ImageInput], ImageInput]`): + The input image. + Returns: + list: A list of images or a 4d array of images. + """ + # If the input is a nested list of images, we flatten it + if ( + isinstance(images, (list, tuple)) + and all(isinstance(images_i, (list, tuple)) for images_i in images) + and all(is_valid_list_of_images(images_i) for images_i in images) + ): + return [img for img_list in images for img in img_list] + + if isinstance(images, (list, tuple)) and is_valid_list_of_images(images): + if is_pil_image(images[0]) or images[0].ndim == 3: + return images + if images[0].ndim == 4: + return [img for img_list in images for img in img_list] + + if is_valid_image(images): + if is_pil_image(images) or images.ndim == 3: + return [images] + if images.ndim == 4: + return list(images) + + raise ValueError(f"Could not make a flat list of images from {images}") + + +def make_nested_list_of_images( + images: Union[list[ImageInput], ImageInput], +) -> ImageInput: + """ + Ensure that the output is a nested list of images. + Args: + images (`Union[List[ImageInput], ImageInput]`): + The input image. + Returns: + list: A list of list of images or a list of 4d array of images. + """ + # If it's a list of batches, it's already in the right format + if ( + isinstance(images, (list, tuple)) + and all(isinstance(images_i, (list, tuple)) for images_i in images) + and all(is_valid_list_of_images(images_i) for images_i in images) + ): + return images + + # If it's a list of images, it's a single batch, so convert it to a list of lists + if isinstance(images, (list, tuple)) and is_valid_list_of_images(images): + if is_pil_image(images[0]) or images[0].ndim == 3: + return [images] + if images[0].ndim == 4: + return [list(image) for image in images] + + # If it's a single image, convert it to a list of lists + if is_valid_image(images): + if is_pil_image(images) or images.ndim == 3: + return [[images]] + if images.ndim == 4: + return [list(images)] + + raise ValueError("Invalid input type. Must be a single image, a list of images, or a list of batches of images.") + + +def make_batched_videos(videos) -> VideoInput: + """ + Ensure that the input is a list of videos. + Args: + videos (`VideoInput`): + Video or videos to turn into a list of videos. + Returns: + list: A list of videos. + """ + if isinstance(videos, (list, tuple)) and isinstance(videos[0], (list, tuple)) and is_valid_image(videos[0][0]): + # case 1: nested batch of videos so we flatten it + if not is_pil_image(videos[0][0]) and videos[0][0].ndim == 4: + videos = [[video for batch_list in batched_videos for video in batch_list] for batched_videos in videos] + # case 2: list of videos represented as list of video frames + return videos + + elif isinstance(videos, (list, tuple)) and is_valid_image(videos[0]): + if is_pil_image(videos[0]) or videos[0].ndim == 3: + return [videos] + elif videos[0].ndim == 4: + return [list(video) for video in videos] + + elif is_valid_image(videos): + if is_pil_image(videos) or videos.ndim == 3: + return [[videos]] + elif videos.ndim == 4: + return [list(videos)] + + raise ValueError(f"Could not make batched video from {videos}") + + +def to_numpy_array(img) -> np.ndarray: + if not is_valid_image(img): + raise ValueError(f"Invalid image type: {type(img)}") + + if is_vision_available() and isinstance(img, PIL.Image.Image): + return np.array(img) + return to_numpy(img) + + +def infer_channel_dimension_format( + image: np.ndarray, num_channels: Optional[Union[int, tuple[int, ...]]] = None +) -> ChannelDimension: + """ + Infers the channel dimension format of `image`. + + Args: + image (`np.ndarray`): + The image to infer the channel dimension of. + num_channels (`int` or `Tuple[int, ...]`, *optional*, defaults to `(1, 3)`): + The number of channels of the image. + + Returns: + The channel dimension of the image. + """ + num_channels = num_channels if num_channels is not None else (1, 3) + num_channels = (num_channels,) if isinstance(num_channels, int) else num_channels + + if image.ndim == 3: + first_dim, last_dim = 0, 2 + elif image.ndim == 4: + first_dim, last_dim = 1, 3 + else: + raise ValueError(f"Unsupported number of image dimensions: {image.ndim}") + + if image.shape[first_dim] in num_channels and image.shape[last_dim] in num_channels: + logger.warning( + f"The channel dimension is ambiguous. Got image shape {image.shape}. Assuming channels are the first dimension." + ) + return ChannelDimension.FIRST + elif image.shape[first_dim] in num_channels: + return ChannelDimension.FIRST + elif image.shape[last_dim] in num_channels: + return ChannelDimension.LAST + raise ValueError("Unable to infer channel dimension format") + + +def get_channel_dimension_axis( + image: np.ndarray, input_data_format: Optional[Union[ChannelDimension, str]] = None +) -> int: + """ + Returns the channel dimension axis of the image. + + Args: + image (`np.ndarray`): + The image to get the channel dimension axis of. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format of the image. If `None`, will infer the channel dimension from the image. + + Returns: + The channel dimension axis of the image. + """ + if input_data_format is None: + input_data_format = infer_channel_dimension_format(image) + if input_data_format == ChannelDimension.FIRST: + return image.ndim - 3 + elif input_data_format == ChannelDimension.LAST: + return image.ndim - 1 + raise ValueError(f"Unsupported data format: {input_data_format}") + + +def get_image_size(image: np.ndarray, channel_dim: ChannelDimension = None) -> tuple[int, int]: + """ + Returns the (height, width) dimensions of the image. + + Args: + image (`np.ndarray`): + The image to get the dimensions of. + channel_dim (`ChannelDimension`, *optional*): + Which dimension the channel dimension is in. If `None`, will infer the channel dimension from the image. + + Returns: + A tuple of the image's height and width. + """ + if channel_dim is None: + channel_dim = infer_channel_dimension_format(image) + + if channel_dim == ChannelDimension.FIRST: + return image.shape[-2], image.shape[-1] + elif channel_dim == ChannelDimension.LAST: + return image.shape[-3], image.shape[-2] + else: + raise ValueError(f"Unsupported data format: {channel_dim}") + + +def get_image_size_for_max_height_width( + image_size: tuple[int, int], + max_height: int, + max_width: int, +) -> tuple[int, int]: + """ + Computes the output image size given the input image and the maximum allowed height and width. Keep aspect ratio. + Important, even if image_height < max_height and image_width < max_width, the image will be resized + to at least one of the edges be equal to max_height or max_width. + + For example: + - input_size: (100, 200), max_height: 50, max_width: 50 -> output_size: (25, 50) + - input_size: (100, 200), max_height: 200, max_width: 500 -> output_size: (200, 400) + + Args: + image_size (`Tuple[int, int]`): + The image to resize. + max_height (`int`): + The maximum allowed height. + max_width (`int`): + The maximum allowed width. + """ + height, width = image_size + height_scale = max_height / height + width_scale = max_width / width + min_scale = min(height_scale, width_scale) + new_height = int(height * min_scale) + new_width = int(width * min_scale) + return new_height, new_width + + +def is_valid_annotation_coco_detection(annotation: dict[str, Union[list, tuple]]) -> bool: + if ( + isinstance(annotation, dict) + and "image_id" in annotation + and "annotations" in annotation + and isinstance(annotation["annotations"], (list, tuple)) + and ( + # an image can have no annotations + len(annotation["annotations"]) == 0 or isinstance(annotation["annotations"][0], dict) + ) + ): + return True + return False + + +def is_valid_annotation_coco_panoptic(annotation: dict[str, Union[list, tuple]]) -> bool: + if ( + isinstance(annotation, dict) + and "image_id" in annotation + and "segments_info" in annotation + and "file_name" in annotation + and isinstance(annotation["segments_info"], (list, tuple)) + and ( + # an image can have no segments + len(annotation["segments_info"]) == 0 or isinstance(annotation["segments_info"][0], dict) + ) + ): + return True + return False + + +def valid_coco_detection_annotations(annotations: Iterable[dict[str, Union[list, tuple]]]) -> bool: + return all(is_valid_annotation_coco_detection(ann) for ann in annotations) + + +def valid_coco_panoptic_annotations(annotations: Iterable[dict[str, Union[list, tuple]]]) -> bool: + return all(is_valid_annotation_coco_panoptic(ann) for ann in annotations) + + +def load_image(image: Union[str, "PIL.Image.Image"], timeout: Optional[float] = None) -> "PIL.Image.Image": + """ + Loads `image` to a PIL Image. + + Args: + image (`str` or `PIL.Image.Image`): + The image to convert to the PIL Image format. + timeout (`float`, *optional*): + The timeout value in seconds for the URL request. + + Returns: + `PIL.Image.Image`: A PIL Image. + """ + requires_backends(load_image, ["vision"]) + if isinstance(image, str): + if image.startswith("http://") or image.startswith("https://"): + # We need to actually check for a real protocol, otherwise it's impossible to use a local file + # like http_huggingface_co.png + image = PIL.Image.open(BytesIO(requests.get(image, timeout=timeout).content)) + elif os.path.isfile(image): + image = PIL.Image.open(image) + else: + if image.startswith("data:image/"): + image = image.split(",")[1] + + # Try to load as base64 + try: + b64 = base64.decodebytes(image.encode()) + image = PIL.Image.open(BytesIO(b64)) + except Exception as e: + raise ValueError( + f"Incorrect image source. Must be a valid URL starting with `http://` or `https://`, a valid path to an image file, or a base64 encoded string. Got {image}. Failed with {e}" + ) + elif isinstance(image, PIL.Image.Image): + image = image + else: + raise TypeError( + "Incorrect format used for image. Should be an url linking to an image, a base64 string, a local path, or a PIL image." + ) + image = PIL.ImageOps.exif_transpose(image) + image = image.convert("RGB") + return image + + +def default_sample_indices_fn(metadata: VideoMetadata, num_frames=None, fps=None, **kwargs): + """ + A default sampling function that replicates the logic used in get_uniform_frame_indices, + while optionally handling `fps` if `num_frames` is not provided. + + Args: + metadata (`VideoMetadata`): + `VideoMetadata` object containing metadata about the video, such as "total_num_frames" or "fps". + num_frames (`int`, *optional*): + Number of frames to sample uniformly. + fps (`int`, *optional*): + Desired frames per second. Takes priority over num_frames if both are provided. + + Returns: + `np.ndarray`: Array of frame indices to sample. + """ + total_num_frames = metadata.total_num_frames + video_fps = metadata.fps + + # If num_frames is not given but fps is, calculate num_frames from fps + if num_frames is None and fps is not None: + num_frames = int(total_num_frames / video_fps * fps) + if num_frames > total_num_frames: + raise ValueError( + f"When loading the video with fps={fps}, we computed num_frames={num_frames} " + f"which exceeds total_num_frames={total_num_frames}. Check fps or video metadata." + ) + + if num_frames is not None: + indices = np.arange(0, total_num_frames, total_num_frames / num_frames, dtype=int) + else: + indices = np.arange(0, total_num_frames, dtype=int) + return indices + + +def read_video_opencv( + video_path: str, + sample_indices_fn: Callable, + **kwargs, +): + """ + Decode a video using the OpenCV backend. + + Args: + video_path (`str`): + Path to the video file. + sample_indices_fn (`Callable`): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniform sampling with fps is performed. + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + Tuple[`np.array`, `VideoMetadata`]: A tuple containing: + - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]). + - `VideoMetadata` object. + """ + # Lazy import cv2 + requires_backends(read_video_opencv, ["cv2"]) + import cv2 + + video = cv2.VideoCapture(video_path) + total_num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) + video_fps = video.get(cv2.CAP_PROP_FPS) + duration = total_num_frames / video_fps if video_fps else 0 + metadata = VideoMetadata( + total_num_frames=int(total_num_frames), fps=float(video_fps), duration=float(duration), video_backend="opencv" + ) + indices = sample_indices_fn(metadata=metadata, **kwargs) + + index = 0 + frames = [] + while video.isOpened(): + success, frame = video.read() + if not success: + break + if index in indices: + height, width, channel = frame.shape + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frames.append(frame[0:height, 0:width, 0:channel]) + if success: + index += 1 + if index >= total_num_frames: + break + + video.release() + metadata.frames_indices = indices + return np.stack(frames), metadata + + +def read_video_decord( + video_path: str, + sample_indices_fn: Optional[Callable] = None, + **kwargs, +): + """ + Decode a video using the Decord backend. + + Args: + video_path (`str`): + Path to the video file. + sample_indices_fn (`Callable`, *optional*): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniform sampling with fps is performed. + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + Tuple[`np.array`, `VideoMetadata`]: A tuple containing: + - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]). + - `VideoMetadata` object. + """ + # Lazy import from decord + requires_backends(read_video_decord, ["decord"]) + from decord import VideoReader, cpu + + vr = VideoReader(uri=video_path, ctx=cpu(0)) # decord has problems with gpu + video_fps = vr.get_avg_fps() + total_num_frames = len(vr) + duration = total_num_frames / video_fps if video_fps else 0 + metadata = VideoMetadata( + total_num_frames=int(total_num_frames), fps=float(video_fps), duration=float(duration), video_backend="decord" + ) + + indices = sample_indices_fn(metadata=metadata, **kwargs) + + frames = vr.get_batch(indices).asnumpy() + metadata.frames_indices = indices + return frames, metadata + + +def read_video_pyav( + video_path: str, + sample_indices_fn: Callable, + **kwargs, +): + """ + Decode the video with PyAV decoder. + + Args: + video_path (`str`): + Path to the video file. + sample_indices_fn (`Callable`, *optional*): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniform sampling with fps is performed. + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + Tuple[`np.array`, `VideoMetadata`]: A tuple containing: + - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]). + - `VideoMetadata` object. + """ + # Lazy import av + requires_backends(read_video_pyav, ["av"]) + import av + + container = av.open(video_path) + total_num_frames = container.streams.video[0].frames + video_fps = container.streams.video[0].average_rate # should we better use `av_guess_frame_rate`? + duration = total_num_frames / video_fps if video_fps else 0 + metadata = VideoMetadata( + total_num_frames=int(total_num_frames), fps=float(video_fps), duration=float(duration), video_backend="pyav" + ) + indices = sample_indices_fn(metadata=metadata, **kwargs) + + frames = [] + container.seek(0) + end_index = indices[-1] + for i, frame in enumerate(container.decode(video=0)): + if i > end_index: + break + if i >= 0 and i in indices: + frames.append(frame) + + video = np.stack([x.to_ndarray(format="rgb24") for x in frames]) + metadata.frames_indices = indices + return video, metadata + + +def read_video_torchvision( + video_path: str, + sample_indices_fn: Callable, + **kwargs, +): + """ + Decode the video with torchvision decoder. + + Args: + video_path (`str`): + Path to the video file. + sample_indices_fn (`Callable`, *optional*): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniform sampling with fps is performed. + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + Tuple[`np.array`, `VideoMetadata`]: A tuple containing: + - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]). + - `VideoMetadata` object. + """ + video, _, info = torchvision_io.read_video( + video_path, + start_pts=0.0, + end_pts=None, + pts_unit="sec", + output_format="THWC", + ) + video_fps = info["video_fps"] + total_num_frames = video.size(0) + duration = total_num_frames / video_fps if video_fps else 0 + metadata = VideoMetadata( + total_num_frames=int(total_num_frames), + fps=float(video_fps), + duration=float(duration), + video_backend="torchvision", + ) + + indices = sample_indices_fn(metadata=metadata, **kwargs) + + video = video[indices].contiguous().numpy() + metadata.frames_indices = indices + return video, metadata + + +VIDEO_DECODERS = { + "decord": read_video_decord, + "opencv": read_video_opencv, + "pyav": read_video_pyav, + "torchvision": read_video_torchvision, +} + + +def load_video( + video: Union[str, "VideoInput"], + num_frames: Optional[int] = None, + fps: Optional[int] = None, + backend: str = "opencv", + sample_indices_fn: Optional[Callable] = None, + **kwargs, +) -> np.array: + """ + Loads `video` to a numpy array. + + Args: + video (`str` or `VideoInput`): + The video to convert to the numpy array format. Can be a link to video or local path. + num_frames (`int`, *optional*): + Number of frames to sample uniformly. If not passed, the whole video is loaded. + fps (`int`, *optional*): + Number of frames to sample per second. Should be passed only when `num_frames=None`. + If not specified and `num_frames==None`, all frames are sampled. + backend (`str`, *optional*, defaults to `"opencv"`): + The backend to use when loading the video. Can be any of ["decord", "pyav", "opencv", "torchvision"]. Defaults to "opencv". + sample_indices_fn (`Callable`, *optional*): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniformt sampling with fps is performed, otherwise `sample_indices_fn` has priority over other args. + The function expects at input the all args along with all kwargs passed to `load_video` and should output valid + indices at which the video should be sampled. For example: + + Example: + def sample_indices_fn(metadata, **kwargs): + return np.linspace(0, metadata.total_num_frames - 1, num_frames, dtype=int) + + Returns: + Tuple[`np.array`, Dict]: A tuple containing: + - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]). + - Metadata dictionary. + """ + + # If `sample_indices_fn` is given, we can accept any args as those might be needed by custom `sample_indices_fn` + if fps is not None and num_frames is not None and sample_indices_fn is None: + raise ValueError( + "`num_frames`, `fps`, and `sample_indices_fn` are mutually exclusive arguments, please use only one!" + ) + + # If user didn't pass a sampling function, create one on the fly with default logic + if sample_indices_fn is None: + + def sample_indices_fn_func(metadata, **fn_kwargs): + return default_sample_indices_fn(metadata, num_frames=num_frames, fps=fps, **fn_kwargs) + + sample_indices_fn = sample_indices_fn_func + + if video.startswith("https://www.youtube.com") or video.startswith("http://www.youtube.com"): + if not is_yt_dlp_available(): + raise ImportError("To load a video from YouTube url you have to install `yt_dlp` first.") + # Lazy import from yt_dlp + requires_backends(load_video, ["yt_dlp"]) + from yt_dlp import YoutubeDL + + buffer = BytesIO() + with redirect_stdout(buffer), YoutubeDL() as f: + f.download([video]) + bytes_obj = buffer.getvalue() + file_obj = BytesIO(bytes_obj) + elif video.startswith("http://") or video.startswith("https://"): + file_obj = BytesIO(requests.get(video).content) + elif os.path.isfile(video): + file_obj = video + elif is_valid_image(video) or (isinstance(video, (list, tuple)) and is_valid_image(video[0])): + file_obj = None + else: + raise TypeError("Incorrect format used for video. Should be an url linking to an video or a local path.") + + # can also load with decord, but not cv2/torchvision + # both will fail in case of url links + video_is_url = video.startswith("http://") or video.startswith("https://") + if video_is_url and backend in ["opencv", "torchvision"]: + raise ValueError( + "If you are trying to load a video from URL, you can decode the video only with `pyav` or `decord` as backend" + ) + + if file_obj is None: + return video + + if ( + (not is_decord_available() and backend == "decord") + or (not is_av_available() and backend == "pyav") + or (not is_cv2_available() and backend == "opencv") + or (not is_torchvision_available() and backend == "torchvision") + ): + raise ImportError( + f"You chose backend={backend} for loading the video but the required library is not found in your environment " + f"Make sure to install {backend} before loading the video." + ) + + video_decoder = VIDEO_DECODERS[backend] + video, metadata = video_decoder(file_obj, sample_indices_fn, **kwargs) + return video, metadata + + +def load_images( + images: Union[list, tuple, str, "PIL.Image.Image"], timeout: Optional[float] = None +) -> Union["PIL.Image.Image", list["PIL.Image.Image"], list[list["PIL.Image.Image"]]]: + """Loads images, handling different levels of nesting. + + Args: + images: A single image, a list of images, or a list of lists of images to load. + timeout: Timeout for loading images. + + Returns: + A single image, a list of images, a list of lists of images. + """ + if isinstance(images, (list, tuple)): + if len(images) and isinstance(images[0], (list, tuple)): + return [[load_image(image, timeout=timeout) for image in image_group] for image_group in images] + else: + return [load_image(image, timeout=timeout) for image in images] + else: + return load_image(images, timeout=timeout) + + +def validate_preprocess_arguments( + do_rescale: Optional[bool] = None, + rescale_factor: Optional[float] = None, + do_normalize: Optional[bool] = None, + image_mean: Optional[Union[float, list[float]]] = None, + image_std: Optional[Union[float, list[float]]] = None, + do_pad: Optional[bool] = None, + size_divisibility: Optional[int] = None, + do_center_crop: Optional[bool] = None, + crop_size: Optional[dict[str, int]] = None, + do_resize: Optional[bool] = None, + size: Optional[dict[str, int]] = None, + resample: Optional["PILImageResampling"] = None, +): + """ + Checks validity of typically used arguments in an `ImageProcessor` `preprocess` method. + Raises `ValueError` if arguments incompatibility is caught. + Many incompatibilities are model-specific. `do_pad` sometimes needs `size_divisor`, + sometimes `size_divisibility`, and sometimes `size`. New models and processors added should follow + existing arguments when possible. + + """ + if do_rescale and rescale_factor is None: + raise ValueError("`rescale_factor` must be specified if `do_rescale` is `True`.") + + if do_pad and size_divisibility is None: + # Here, size_divisor might be passed as the value of size + raise ValueError( + "Depending on the model, `size_divisibility`, `size_divisor`, `pad_size` or `size` must be specified if `do_pad` is `True`." + ) + + if do_normalize and (image_mean is None or image_std is None): + raise ValueError("`image_mean` and `image_std` must both be specified if `do_normalize` is `True`.") + + if do_center_crop and crop_size is None: + raise ValueError("`crop_size` must be specified if `do_center_crop` is `True`.") + + if do_resize and (size is None or resample is None): + raise ValueError("`size` and `resample` must be specified if `do_resize` is `True`.") + + +# In the future we can add a TF implementation here when we have TF models. +class ImageFeatureExtractionMixin: + """ + Mixin that contain utilities for preparing image features. + """ + + def _ensure_format_supported(self, image): + if not isinstance(image, (PIL.Image.Image, np.ndarray)) and not is_torch_tensor(image): + raise ValueError( + f"Got type {type(image)} which is not supported, only `PIL.Image.Image`, `np.array` and " + "`torch.Tensor` are." + ) + + def to_pil_image(self, image, rescale=None): + """ + Converts `image` to a PIL Image. Optionally rescales it and puts the channel dimension back as the last axis if + needed. + + Args: + image (`PIL.Image.Image` or `numpy.ndarray` or `torch.Tensor`): + The image to convert to the PIL Image format. + rescale (`bool`, *optional*): + Whether or not to apply the scaling factor (to make pixel values integers between 0 and 255). Will + default to `True` if the image type is a floating type, `False` otherwise. + """ + self._ensure_format_supported(image) + + if is_torch_tensor(image): + image = image.numpy() + + if isinstance(image, np.ndarray): + if rescale is None: + # rescale default to the array being of floating type. + rescale = isinstance(image.flat[0], np.floating) + # If the channel as been moved to first dim, we put it back at the end. + if image.ndim == 3 and image.shape[0] in [1, 3]: + image = image.transpose(1, 2, 0) + if rescale: + image = image * 255 + image = image.astype(np.uint8) + return PIL.Image.fromarray(image) + return image + + def convert_rgb(self, image): + """ + Converts `PIL.Image.Image` to RGB format. + + Args: + image (`PIL.Image.Image`): + The image to convert. + """ + self._ensure_format_supported(image) + if not isinstance(image, PIL.Image.Image): + return image + + return image.convert("RGB") + + def rescale(self, image: np.ndarray, scale: Union[float, int]) -> np.ndarray: + """ + Rescale a numpy image by scale amount + """ + self._ensure_format_supported(image) + return image * scale + + def to_numpy_array(self, image, rescale=None, channel_first=True): + """ + Converts `image` to a numpy array. Optionally rescales it and puts the channel dimension as the first + dimension. + + Args: + image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`): + The image to convert to a NumPy array. + rescale (`bool`, *optional*): + Whether or not to apply the scaling factor (to make pixel values floats between 0. and 1.). Will + default to `True` if the image is a PIL Image or an array/tensor of integers, `False` otherwise. + channel_first (`bool`, *optional*, defaults to `True`): + Whether or not to permute the dimensions of the image to put the channel dimension first. + """ + self._ensure_format_supported(image) + + if isinstance(image, PIL.Image.Image): + image = np.array(image) + + if is_torch_tensor(image): + image = image.numpy() + + rescale = isinstance(image.flat[0], np.integer) if rescale is None else rescale + + if rescale: + image = self.rescale(image.astype(np.float32), 1 / 255.0) + + if channel_first and image.ndim == 3: + image = image.transpose(2, 0, 1) + + return image + + def expand_dims(self, image): + """ + Expands 2-dimensional `image` to 3 dimensions. + + Args: + image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`): + The image to expand. + """ + self._ensure_format_supported(image) + + # Do nothing if PIL image + if isinstance(image, PIL.Image.Image): + return image + + if is_torch_tensor(image): + image = image.unsqueeze(0) + else: + image = np.expand_dims(image, axis=0) + return image + + def normalize(self, image, mean, std, rescale=False): + """ + Normalizes `image` with `mean` and `std`. Note that this will trigger a conversion of `image` to a NumPy array + if it's a PIL Image. + + Args: + image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`): + The image to normalize. + mean (`List[float]` or `np.ndarray` or `torch.Tensor`): + The mean (per channel) to use for normalization. + std (`List[float]` or `np.ndarray` or `torch.Tensor`): + The standard deviation (per channel) to use for normalization. + rescale (`bool`, *optional*, defaults to `False`): + Whether or not to rescale the image to be between 0 and 1. If a PIL image is provided, scaling will + happen automatically. + """ + self._ensure_format_supported(image) + + if isinstance(image, PIL.Image.Image): + image = self.to_numpy_array(image, rescale=True) + # If the input image is a PIL image, it automatically gets rescaled. If it's another + # type it may need rescaling. + elif rescale: + if isinstance(image, np.ndarray): + image = self.rescale(image.astype(np.float32), 1 / 255.0) + elif is_torch_tensor(image): + image = self.rescale(image.float(), 1 / 255.0) + + if isinstance(image, np.ndarray): + if not isinstance(mean, np.ndarray): + mean = np.array(mean).astype(image.dtype) + if not isinstance(std, np.ndarray): + std = np.array(std).astype(image.dtype) + elif is_torch_tensor(image): + import torch + + if not isinstance(mean, torch.Tensor): + if isinstance(mean, np.ndarray): + mean = torch.from_numpy(mean) + else: + mean = torch.tensor(mean) + if not isinstance(std, torch.Tensor): + if isinstance(std, np.ndarray): + std = torch.from_numpy(std) + else: + std = torch.tensor(std) + + if image.ndim == 3 and image.shape[0] in [1, 3]: + return (image - mean[:, None, None]) / std[:, None, None] + else: + return (image - mean) / std + + def resize(self, image, size, resample=None, default_to_square=True, max_size=None): + """ + Resizes `image`. Enforces conversion of input to PIL.Image. + + Args: + image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`): + The image to resize. + size (`int` or `Tuple[int, int]`): + The size to use for resizing the image. If `size` is a sequence like (h, w), output size will be + matched to this. + + If `size` is an int and `default_to_square` is `True`, then image will be resized to (size, size). If + `size` is an int and `default_to_square` is `False`, then smaller edge of the image will be matched to + this number. i.e, if height > width, then image will be rescaled to (size * height / width, size). + resample (`int`, *optional*, defaults to `PILImageResampling.BILINEAR`): + The filter to user for resampling. + default_to_square (`bool`, *optional*, defaults to `True`): + How to convert `size` when it is a single int. If set to `True`, the `size` will be converted to a + square (`size`,`size`). If set to `False`, will replicate + [`torchvision.transforms.Resize`](https://pytorch.org/vision/stable/transforms.html#torchvision.transforms.Resize) + with support for resizing only the smallest edge and providing an optional `max_size`. + max_size (`int`, *optional*, defaults to `None`): + The maximum allowed for the longer edge of the resized image: if the longer edge of the image is + greater than `max_size` after being resized according to `size`, then the image is resized again so + that the longer edge is equal to `max_size`. As a result, `size` might be overruled, i.e the smaller + edge may be shorter than `size`. Only used if `default_to_square` is `False`. + + Returns: + image: A resized `PIL.Image.Image`. + """ + resample = resample if resample is not None else PILImageResampling.BILINEAR + + self._ensure_format_supported(image) + + if not isinstance(image, PIL.Image.Image): + image = self.to_pil_image(image) + + if isinstance(size, list): + size = tuple(size) + + if isinstance(size, int) or len(size) == 1: + if default_to_square: + size = (size, size) if isinstance(size, int) else (size[0], size[0]) + else: + width, height = image.size + # specified size only for the smallest edge + short, long = (width, height) if width <= height else (height, width) + requested_new_short = size if isinstance(size, int) else size[0] + + if short == requested_new_short: + return image + + new_short, new_long = requested_new_short, int(requested_new_short * long / short) + + if max_size is not None: + if max_size <= requested_new_short: + raise ValueError( + f"max_size = {max_size} must be strictly greater than the requested " + f"size for the smaller edge size = {size}" + ) + if new_long > max_size: + new_short, new_long = int(max_size * new_short / new_long), max_size + + size = (new_short, new_long) if width <= height else (new_long, new_short) + + return image.resize(size, resample=resample) + + def center_crop(self, image, size): + """ + Crops `image` to the given size using a center crop. Note that if the image is too small to be cropped to the + size given, it will be padded (so the returned result has the size asked). + + Args: + image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor` of shape (n_channels, height, width) or (height, width, n_channels)): + The image to resize. + size (`int` or `Tuple[int, int]`): + The size to which crop the image. + + Returns: + new_image: A center cropped `PIL.Image.Image` or `np.ndarray` or `torch.Tensor` of shape: (n_channels, + height, width). + """ + self._ensure_format_supported(image) + + if not isinstance(size, tuple): + size = (size, size) + + # PIL Image.size is (width, height) but NumPy array and torch Tensors have (height, width) + if is_torch_tensor(image) or isinstance(image, np.ndarray): + if image.ndim == 2: + image = self.expand_dims(image) + image_shape = image.shape[1:] if image.shape[0] in [1, 3] else image.shape[:2] + else: + image_shape = (image.size[1], image.size[0]) + + top = (image_shape[0] - size[0]) // 2 + bottom = top + size[0] # In case size is odd, (image_shape[0] + size[0]) // 2 won't give the proper result. + left = (image_shape[1] - size[1]) // 2 + right = left + size[1] # In case size is odd, (image_shape[1] + size[1]) // 2 won't give the proper result. + + # For PIL Images we have a method to crop directly. + if isinstance(image, PIL.Image.Image): + return image.crop((left, top, right, bottom)) + + # Check if image is in (n_channels, height, width) or (height, width, n_channels) format + channel_first = True if image.shape[0] in [1, 3] else False + + # Transpose (height, width, n_channels) format images + if not channel_first: + if isinstance(image, np.ndarray): + image = image.transpose(2, 0, 1) + if is_torch_tensor(image): + image = image.permute(2, 0, 1) + + # Check if cropped area is within image boundaries + if top >= 0 and bottom <= image_shape[0] and left >= 0 and right <= image_shape[1]: + return image[..., top:bottom, left:right] + + # Otherwise, we may need to pad if the image is too small. Oh joy... + new_shape = image.shape[:-2] + (max(size[0], image_shape[0]), max(size[1], image_shape[1])) + if isinstance(image, np.ndarray): + new_image = np.zeros_like(image, shape=new_shape) + elif is_torch_tensor(image): + new_image = image.new_zeros(new_shape) + + top_pad = (new_shape[-2] - image_shape[0]) // 2 + bottom_pad = top_pad + image_shape[0] + left_pad = (new_shape[-1] - image_shape[1]) // 2 + right_pad = left_pad + image_shape[1] + new_image[..., top_pad:bottom_pad, left_pad:right_pad] = image + + top += top_pad + bottom += top_pad + left += left_pad + right += left_pad + + new_image = new_image[ + ..., max(0, top) : min(new_image.shape[-2], bottom), max(0, left) : min(new_image.shape[-1], right) + ] + + return new_image + + def flip_channel_order(self, image): + """ + Flips the channel order of `image` from RGB to BGR, or vice versa. Note that this will trigger a conversion of + `image` to a NumPy array if it's a PIL Image. + + Args: + image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`): + The image whose color channels to flip. If `np.ndarray` or `torch.Tensor`, the channel dimension should + be first. + """ + self._ensure_format_supported(image) + + if isinstance(image, PIL.Image.Image): + image = self.to_numpy_array(image) + + return image[::-1, :, :] + + def rotate(self, image, angle, resample=None, expand=0, center=None, translate=None, fillcolor=None): + """ + Returns a rotated copy of `image`. This method returns a copy of `image`, rotated the given number of degrees + counter clockwise around its centre. + + Args: + image (`PIL.Image.Image` or `np.ndarray` or `torch.Tensor`): + The image to rotate. If `np.ndarray` or `torch.Tensor`, will be converted to `PIL.Image.Image` before + rotating. + + Returns: + image: A rotated `PIL.Image.Image`. + """ + resample = resample if resample is not None else PIL.Image.NEAREST + + self._ensure_format_supported(image) + + if not isinstance(image, PIL.Image.Image): + image = self.to_pil_image(image) + + return image.rotate( + angle, resample=resample, expand=expand, center=center, translate=translate, fillcolor=fillcolor + ) + + +def validate_annotations( + annotation_format: AnnotationFormat, + supported_annotation_formats: tuple[AnnotationFormat, ...], + annotations: list[dict], +) -> None: + if annotation_format not in supported_annotation_formats: + raise ValueError(f"Unsupported annotation format: {format} must be one of {supported_annotation_formats}") + + if annotation_format is AnnotationFormat.COCO_DETECTION: + if not valid_coco_detection_annotations(annotations): + raise ValueError( + "Invalid COCO detection annotations. Annotations must a dict (single image) or list of dicts " + "(batch of images) with the following keys: `image_id` and `annotations`, with the latter " + "being a list of annotations in the COCO format." + ) + + if annotation_format is AnnotationFormat.COCO_PANOPTIC: + if not valid_coco_panoptic_annotations(annotations): + raise ValueError( + "Invalid COCO panoptic annotations. Annotations must a dict (single image) or list of dicts " + "(batch of images) with the following keys: `image_id`, `file_name` and `segments_info`, with " + "the latter being a list of annotations in the COCO format." + ) + + +def validate_kwargs(valid_processor_keys: list[str], captured_kwargs: list[str]): + unused_keys = set(captured_kwargs).difference(set(valid_processor_keys)) + if unused_keys: + unused_key_str = ", ".join(unused_keys) + # TODO raise a warning here instead of simply logging? + logger.warning(f"Unused or unrecognized kwargs: {unused_key_str}.") + + +@dataclass(frozen=True) +class SizeDict: + """ + Hashable dictionary to store image size information. + """ + + height: Optional[int] = None + width: Optional[int] = None + longest_edge: Optional[int] = None + shortest_edge: Optional[int] = None + max_height: Optional[int] = None + max_width: Optional[int] = None + + def __getitem__(self, key): + if hasattr(self, key): + return getattr(self, key) + raise KeyError(f"Key {key} not found in SizeDict.") diff --git a/lib/python3.12/site-packages/transformers/integrations/__init__.py b/lib/python3.12/site-packages/transformers/integrations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8d03c5cf790ee8bb3eac6448e8927445132d2758 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/__init__.py @@ -0,0 +1,289 @@ +# Copyright 2023 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. +from typing import TYPE_CHECKING + +from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_torch_greater_or_equal + + +_import_structure = { + "aqlm": ["replace_with_aqlm_linear"], + "awq": [ + "fuse_awq_modules", + "post_init_awq_exllama_modules", + "post_init_awq_ipex_modules", + "replace_quantization_scales", + "replace_with_awq_linear", + ], + "bitnet": [ + "BitLinear", + "pack_weights", + "replace_with_bitnet_linear", + "unpack_weights", + ], + "bitsandbytes": [ + "dequantize_and_replace", + "get_keys_to_not_convert", + "replace_8bit_linear", + "replace_with_bnb_linear", + "set_module_8bit_tensor_to_device", + "set_module_quantized_tensor_to_device", + "validate_bnb_backend_availability", + ], + "deepspeed": [ + "HfDeepSpeedConfig", + "HfTrainerDeepSpeedConfig", + "deepspeed_config", + "deepspeed_init", + "deepspeed_load_checkpoint", + "deepspeed_optim_sched", + "is_deepspeed_available", + "is_deepspeed_zero3_enabled", + "set_hf_deepspeed_config", + "unset_hf_deepspeed_config", + ], + "eetq": ["replace_with_eetq_linear"], + "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"], + "finegrained_fp8": ["FP8Linear", "replace_with_fp8_linear"], + "fsdp": ["is_fsdp_managed_module"], + "ggml": [ + "GGUF_CONFIG_MAPPING", + "GGUF_TOKENIZER_MAPPING", + "_gguf_parse_value", + "load_dequant_gguf_tensor", + "load_gguf", + ], + "higgs": [ + "HiggsLinear", + "dequantize_higgs", + "quantize_with_higgs", + "replace_with_higgs_linear", + ], + "hqq": ["prepare_for_hqq_linear"], + "hub_kernels": [ + "LayerRepository", + "register_kernel_mapping", + "replace_kernel_forward_from_hub", + "use_kernel_forward_from_hub", + ], + "integration_utils": [ + "INTEGRATION_TO_CALLBACK", + "AzureMLCallback", + "ClearMLCallback", + "CodeCarbonCallback", + "CometCallback", + "DagsHubCallback", + "DVCLiveCallback", + "FlyteCallback", + "MLflowCallback", + "NeptuneCallback", + "NeptuneMissingConfiguration", + "SwanLabCallback", + "TensorBoardCallback", + "WandbCallback", + "get_available_reporting_integrations", + "get_reporting_integration_callbacks", + "hp_params", + "is_azureml_available", + "is_clearml_available", + "is_codecarbon_available", + "is_comet_available", + "is_dagshub_available", + "is_dvclive_available", + "is_flyte_deck_standard_available", + "is_flytekit_available", + "is_mlflow_available", + "is_neptune_available", + "is_optuna_available", + "is_ray_available", + "is_ray_tune_available", + "is_sigopt_available", + "is_swanlab_available", + "is_tensorboard_available", + "is_wandb_available", + "rewrite_logs", + "run_hp_search_optuna", + "run_hp_search_ray", + "run_hp_search_sigopt", + "run_hp_search_wandb", + ], + "peft": ["PeftAdapterMixin"], + "quanto": ["replace_with_quanto_layers"], + "spqr": ["replace_with_spqr_linear"], + "vptq": ["replace_with_vptq_linear"], +} + +try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["executorch"] = [ + "TorchExportableModuleWithStaticCache", + "convert_and_export_with_cache", + ] + +try: + if not is_torch_greater_or_equal("2.3"): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["tensor_parallel"] = [ + "shard_and_distribute_module", + "SUPPORTED_TP_STYLES", + "translate_to_torch_parallel_style", + ] +try: + if not is_torch_greater_or_equal("2.5"): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["flex_attention"] = [ + "make_flex_block_causal_mask", + ] + +if TYPE_CHECKING: + from .aqlm import replace_with_aqlm_linear + from .awq import ( + fuse_awq_modules, + post_init_awq_exllama_modules, + post_init_awq_ipex_modules, + replace_quantization_scales, + replace_with_awq_linear, + ) + from .bitnet import ( + BitLinear, + pack_weights, + replace_with_bitnet_linear, + unpack_weights, + ) + from .bitsandbytes import ( + dequantize_and_replace, + get_keys_to_not_convert, + replace_8bit_linear, + replace_with_bnb_linear, + set_module_8bit_tensor_to_device, + set_module_quantized_tensor_to_device, + validate_bnb_backend_availability, + ) + from .deepspeed import ( + HfDeepSpeedConfig, + HfTrainerDeepSpeedConfig, + deepspeed_config, + deepspeed_init, + deepspeed_load_checkpoint, + deepspeed_optim_sched, + is_deepspeed_available, + is_deepspeed_zero3_enabled, + set_hf_deepspeed_config, + unset_hf_deepspeed_config, + ) + from .eetq import replace_with_eetq_linear + from .fbgemm_fp8 import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts, replace_with_fbgemm_fp8_linear + from .finegrained_fp8 import FP8Linear, replace_with_fp8_linear + from .fsdp import is_fsdp_managed_module + from .ggml import ( + GGUF_CONFIG_MAPPING, + GGUF_TOKENIZER_MAPPING, + _gguf_parse_value, + load_dequant_gguf_tensor, + load_gguf, + ) + from .higgs import HiggsLinear, dequantize_higgs, quantize_with_higgs, replace_with_higgs_linear + from .hqq import prepare_for_hqq_linear + from .hub_kernels import ( + LayerRepository, + register_kernel_mapping, + replace_kernel_forward_from_hub, + use_kernel_forward_from_hub, + ) + from .integration_utils import ( + INTEGRATION_TO_CALLBACK, + AzureMLCallback, + ClearMLCallback, + CodeCarbonCallback, + CometCallback, + DagsHubCallback, + DVCLiveCallback, + FlyteCallback, + MLflowCallback, + NeptuneCallback, + NeptuneMissingConfiguration, + SwanLabCallback, + TensorBoardCallback, + WandbCallback, + get_available_reporting_integrations, + get_reporting_integration_callbacks, + hp_params, + is_azureml_available, + is_clearml_available, + is_codecarbon_available, + is_comet_available, + is_dagshub_available, + is_dvclive_available, + is_flyte_deck_standard_available, + is_flytekit_available, + is_mlflow_available, + is_neptune_available, + is_optuna_available, + is_ray_available, + is_ray_tune_available, + is_sigopt_available, + is_swanlab_available, + is_tensorboard_available, + is_wandb_available, + rewrite_logs, + run_hp_search_optuna, + run_hp_search_ray, + run_hp_search_sigopt, + run_hp_search_wandb, + ) + from .peft import PeftAdapterMixin + from .quanto import replace_with_quanto_layers + from .spqr import replace_with_spqr_linear + from .vptq import replace_with_vptq_linear + + try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .executorch import TorchExportableModuleWithStaticCache, convert_and_export_with_cache + + try: + if not is_torch_greater_or_equal("2.3"): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .tensor_parallel import ( + SUPPORTED_TP_STYLES, + shard_and_distribute_module, + translate_to_torch_parallel_style, + ) + + try: + if not is_torch_greater_or_equal("2.5"): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .flex_attention import make_flex_block_causal_mask +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d9a44bd1c8d625b4ac44130ee878c50200194e9 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/accelerate.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/accelerate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32c5b6829c89657f01c345238d08396949d9deeb Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/accelerate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/aqlm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/aqlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..38e07fee9376397fa637f6b597122498aa4cfdce Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/aqlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/awq.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/awq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f5b503c33e0858cddbcb57913794c7c74a69877e Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/awq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/bitnet.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/bitnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60475677133351837ce6169dd1df52e1c645d24b Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/bitnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/bitsandbytes.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/bitsandbytes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b107dc7c9156438933f35080c0b5632c924970c Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/bitsandbytes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/compressed_tensors.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/compressed_tensors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db35603fe7d3348a5bb5ef761d50c9762221ea22 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/compressed_tensors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/deepspeed.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/deepspeed.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ead2b5744368a70ea75102838b3aee329ba4b8f Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/deepspeed.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/eetq.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/eetq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..562a7a815a4c991bc3efc786df30b26461ff51ef Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/eetq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/executorch.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/executorch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f597ef3dc8c2462e68bd816fa8cdf3c532ca115 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/executorch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/fbgemm_fp8.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/fbgemm_fp8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..574bd3dd2b480fe8dde073a9cb7c075b94e23384 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/fbgemm_fp8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/finegrained_fp8.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/finegrained_fp8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c97e0467935eca1d6c1753395ceb3cee55e3532 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/finegrained_fp8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/flash_attention.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/flash_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c3394eaf2083d12ad40bd470c155e2c837726c4b Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/flash_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/flex_attention.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/flex_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c34fdc93a1a95cea605323c631a5fe7469f4124 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/flex_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/fsdp.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/fsdp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a27126a77d4b027f19a538cfb232e408df0e47ad Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/fsdp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/ggml.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/ggml.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a2cab1df183bb2fcaa712fd6b450f34952081b8 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/ggml.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/higgs.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/higgs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9fa2511f6290f059ca55db5047ae3707f0f847d0 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/higgs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/hqq.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/hqq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e93814e24fff288e440fb04400fa418dc771f84 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/hqq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/hub_kernels.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/hub_kernels.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f0ede4fee035be1f02a44bc81f81838afe90f0ea Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/hub_kernels.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/mistral.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/mistral.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e1b3465a078678a3d67f1e93d5c0beac983c849 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/mistral.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/npu_flash_attention.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/npu_flash_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ae1ee4186240dc6d5846f26c84ec3303cd71ffdc Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/npu_flash_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/peft.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/peft.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c351d12c01c9e2dea9459053c57041149720f0f Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/peft.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/quanto.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/quanto.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8dc7b28d6592636bd9ac4c382342c1a9a5d34c9d Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/quanto.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/sdpa_attention.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/sdpa_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e98bacf00055bb6104207ad815347c06d2b22fd2 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/sdpa_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/spqr.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/spqr.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea9173e388565fb5263a1f4b30af22d6c94de56b Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/spqr.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/tensor_parallel.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/tensor_parallel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..198a4842c6a98d1abaf7829b2ad8c9826cca5248 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/tensor_parallel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/tiktoken.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/tiktoken.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f9ff29ddb8eb805e462c98ba8192c707db2125c7 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/tiktoken.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/tpu.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/tpu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c236ea42153e04583da19adff5af3d7a4950cc6 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/tpu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/__pycache__/vptq.cpython-312.pyc b/lib/python3.12/site-packages/transformers/integrations/__pycache__/vptq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b7d4dc20caa2b5f5787b9ae9575350184d3713ff Binary files /dev/null and b/lib/python3.12/site-packages/transformers/integrations/__pycache__/vptq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/integrations/accelerate.py b/lib/python3.12/site-packages/transformers/integrations/accelerate.py new file mode 100644 index 0000000000000000000000000000000000000000..83efac9661af0423901f26783c76194aa7c205d6 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/accelerate.py @@ -0,0 +1,196 @@ +# Copyright 2025 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. +""" +Since, https://github.com/huggingface/transformers/pull/36963, loading is always performed with models on meta +device. But since the `init_empty_weights` and `find_tied_parameters` functions are from accelerate, and accelerate is +somewhat still a soft dependency, we copy the functions here to be used natively in Transformers. + +The `init_empty_weights` and `init_on_device` functions were copied from `accelerate.big_modeling.py`, and the +`find_tied_parameters` was copied from `accelerate.utils.modeling.py` +""" + +from contextlib import contextmanager + +from ..utils import is_torch_available, logging + + +if is_torch_available(): + import torch + import torch.nn as nn + + +logger = logging.get_logger(__name__) + + +@contextmanager +def init_empty_weights(include_buffers: bool = False): + """ + A context manager under which models are initialized with all parameters on the meta device, therefore creating an + empty model. Useful when just initializing the model would blow the available RAM. + + Args: + include_buffers (`bool`, *optional*): + Whether or not to also put all buffers on the meta device while initializing. + + Example: + + ```python + import torch.nn as nn + from accelerate import init_empty_weights + + # Initialize a model with 100 billions parameters in no time and without using any RAM. + with init_empty_weights(): + tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)]) + ``` + + + + Any model created under this context manager has no weights. As such you can't do something like + `model.to(some_device)` with it. To load weights inside your empty model, see [`load_checkpoint_and_dispatch`]. + Make sure to overwrite the default device_map param for [`load_checkpoint_and_dispatch`], otherwise dispatch is not + called. + + + """ + with init_on_device(torch.device("meta"), include_buffers=include_buffers) as f: + yield f + + +@contextmanager +def init_on_device(device: "torch.device", include_buffers: bool = False): + """ + A context manager under which models are initialized with all parameters on the specified device. + + Args: + device (`torch.device`): + Device to initialize all parameters on. + include_buffers (`bool`, *optional*): + Whether or not to also put all buffers on the meta device while initializing. + + Example: + + ```python + import torch.nn as nn + from accelerate import init_on_device + + with init_on_device(device=torch.device("cuda")): + tst = nn.Linear(100, 100) # on `cuda` device + ``` + """ + if include_buffers: + with device: + yield + return + + old_register_parameter = nn.Module.register_parameter + if include_buffers: + old_register_buffer = nn.Module.register_buffer + + def register_empty_parameter(module, name, param): + old_register_parameter(module, name, param) + if param is not None: + param_cls = type(module._parameters[name]) + kwargs = module._parameters[name].__dict__ + kwargs["requires_grad"] = param.requires_grad + module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) + + def register_empty_buffer(module, name, buffer, persistent=True): + old_register_buffer(module, name, buffer, persistent=persistent) + if buffer is not None: + module._buffers[name] = module._buffers[name].to(device) + + # Patch tensor creation + if include_buffers: + tensor_constructors_to_patch = { + torch_function_name: getattr(torch, torch_function_name) + for torch_function_name in ["empty", "zeros", "ones", "full"] + } + else: + tensor_constructors_to_patch = {} + + def patch_tensor_constructor(fn): + def wrapper(*args, **kwargs): + kwargs["device"] = device + return fn(*args, **kwargs) + + return wrapper + + try: + nn.Module.register_parameter = register_empty_parameter + if include_buffers: + nn.Module.register_buffer = register_empty_buffer + for torch_function_name in tensor_constructors_to_patch.keys(): + setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name))) + yield + finally: + nn.Module.register_parameter = old_register_parameter + if include_buffers: + nn.Module.register_buffer = old_register_buffer + for torch_function_name, old_torch_function in tensor_constructors_to_patch.items(): + setattr(torch, torch_function_name, old_torch_function) + + +def find_tied_parameters(model: "nn.Module", **kwargs): + """ + Find the tied parameters in a given model. + + + + The signature accepts keyword arguments, but they are for the recursive part of this function and you should ignore + them. + + + + Args: + model (`torch.nn.Module`): The model to inspect. + + Returns: + List[List[str]]: A list of lists of parameter names being all tied together. + + Example: + + ```py + >>> from collections import OrderedDict + >>> import torch.nn as nn + + >>> model = nn.Sequential(OrderedDict([("linear1", nn.Linear(4, 4)), ("linear2", nn.Linear(4, 4))])) + >>> model.linear2.weight = model.linear1.weight + >>> find_tied_parameters(model) + [['linear1.weight', 'linear2.weight']] + ``` + """ + + # get ALL model parameters and thier names + all_named_parameters = dict(model.named_parameters(remove_duplicate=False)) + + # get ONLY unique named parameters, + # if parameter is tied and have multiple names, it will be included only once + no_duplicate_named_parameters = dict(model.named_parameters(remove_duplicate=True)) + + # the difference of the two sets will give us the tied parameters + tied_param_names = set(all_named_parameters.keys()) - set(no_duplicate_named_parameters.keys()) + + # 'tied_param_names' contains the names of parameters that are tied in the model, but we do not know + # which names refer to the same parameter. To identify this, we need to group them together. + tied_param_groups = {} + for tied_param_name in tied_param_names: + tied_param = all_named_parameters[tied_param_name] + for param_name, param in no_duplicate_named_parameters.items(): + # compare if parameters are the same, if so, group thier names together + if param is tied_param: + if param_name not in tied_param_groups: + tied_param_groups[param_name] = [] + tied_param_groups[param_name].append(tied_param_name) + + return [sorted([weight] + list(set(tied))) for weight, tied in tied_param_groups.items()] diff --git a/lib/python3.12/site-packages/transformers/integrations/aqlm.py b/lib/python3.12/site-packages/transformers/integrations/aqlm.py new file mode 100644 index 0000000000000000000000000000000000000000..0626da7aced5bc259295ac51c58931858aac8dc8 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/aqlm.py @@ -0,0 +1,100 @@ +# 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. +"AQLM (Additive Quantization of Language Model) integration file" + +from ..utils import ACCELERATE_MIN_VERSION, is_accelerate_available, is_aqlm_available, is_torch_available + + +if is_torch_available(): + import torch.nn as nn + + +def replace_with_aqlm_linear( + model, + quantization_config=None, + linear_weights_not_to_quantize=None, + current_key_name=None, + has_been_replaced=False, +): + """ + Public method that recursively replaces the Linear layers of the given model with AQLM quantized layers. + `accelerate` is needed to use this method. Returns the converted model and a boolean that indicates if the + conversion has been successfull or not. + + Args: + model (`torch.nn.Module`): + The model to convert, can be any `torch.nn.Module` instance. + quantization_config (`AqlmConfig`): + The quantization config object that contains the quantization parameters. + linear_weights_not_to_quantize (`list[str]`, *optional*): + A list of nn.Linear weights to not convert. If a parameter path is in the list (e.g. `lm_head.weight`), the corresponding module will not be + converted. + current_key_name (`list`, *optional*): + A list that contains the current key name. This is used for recursion and should not be passed by the user. + has_been_replaced (`bool`, *optional*): + A boolean that indicates if the conversion has been successful or not. This is used for recursion and + should not be passed by the user. + """ + if not is_aqlm_available(): + raise ValueError("AQLM is not available. Please install it with `pip install aqlm[cpu,gpu]`") + + if not is_accelerate_available(): + raise ValueError( + f"AQLM requires Accelerate to be installed: `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`" + ) + + if linear_weights_not_to_quantize is None: + linear_weights_not_to_quantize = [] + + from accelerate import init_empty_weights + from aqlm import QuantizedLinear + + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + + if isinstance(module, nn.Linear): + # Check if the current key is not in the `linear_weights_not_to_quantize` + if ".".join(current_key_name) + ".weight" not in linear_weights_not_to_quantize: + with init_empty_weights(): + in_features = module.in_features + out_features = module.out_features + + model._modules[name] = QuantizedLinear( + in_features, + out_features, + bias=module.bias is not None, + in_group_size=quantization_config.in_group_size, + out_group_size=quantization_config.out_group_size, + num_codebooks=quantization_config.num_codebooks, + nbits_per_codebook=quantization_config.nbits_per_codebook, + ) + has_been_replaced = True + + # Store the module class in case we need to transpose the weight later + model._modules[name].source_cls = type(module) + # Force requires grad to False to avoid unexpected errors + model._modules[name].requires_grad_(False) + if len(list(module.children())) > 0: + _, has_been_replaced = replace_with_aqlm_linear( + module, + quantization_config=quantization_config, + linear_weights_not_to_quantize=linear_weights_not_to_quantize, + current_key_name=current_key_name, + has_been_replaced=has_been_replaced, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced diff --git a/lib/python3.12/site-packages/transformers/integrations/awq.py b/lib/python3.12/site-packages/transformers/integrations/awq.py new file mode 100644 index 0000000000000000000000000000000000000000..c860ea1f53744bba6a4d63bc73ee51a187fc92f8 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/awq.py @@ -0,0 +1,498 @@ +# Copyright 2023 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. +"AWQ (Activation aware Weight Quantization) integration file" + +import importlib + +from packaging import version + +from ..activations import ACT2FN +from ..modeling_utils import PreTrainedModel +from ..utils import is_auto_awq_available, is_ipex_available, is_torch_available, logging +from ..utils.quantization_config import ( + AwqBackendPackingMethod, + AwqConfig, + AWQLinearVersion, + ExllamaVersion, +) + + +if is_torch_available(): + import torch + import torch.nn as nn + +logger = logging.get_logger(__name__) + +AWQ_FUSED_MAPPINGS = { + "mistral": { + "attention": ["q_proj", "k_proj", "v_proj", "o_proj"], + "mlp": ["gate_proj", "up_proj", "down_proj"], + "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"], + "use_alibi": False, + }, + "mixtral": { + "attention": ["q_proj", "k_proj", "v_proj", "o_proj"], + "mlp": ["w1", "w3", "w2"], + "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"], + "use_alibi": False, + "rope_theta": 1000000.0, + }, + "llama": { + "attention": ["q_proj", "k_proj", "v_proj", "o_proj"], + "mlp": ["gate_proj", "up_proj", "down_proj"], + "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"], + "use_alibi": False, + }, + "llava": { + "attention": ["q_proj", "k_proj", "v_proj", "o_proj"], + "mlp": ["gate_proj", "up_proj", "down_proj"], + "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"], + "use_alibi": False, + }, +} + +AWQ_SCALES_MAPPINGS = { + "starcoder2": {"act": "act", "layer_before_act": "c_fc"}, + "RefinedWebModel": {"act": "act", "layer_before_act": "dense_h_to_4h"}, + "falcon": {"act": "act", "layer_before_act": "dense_h_to_4h"}, + "mpt": {"act": "act", "layer_before_act": "up_proj"}, + "gptj": {"act": "act", "layer_before_act": "fc_in"}, + "gpt_neox": {"act": "act", "layer_before_act": "dense_h_to_4h"}, + "gpt_bigcode": {"act": "act", "layer_before_act": "c_fc"}, + "bloom": {"act": "gelu_impl", "layer_before_act": "dense_h_to_4h"}, +} + + +def replace_quantization_scales(model, model_type): + from awq.modules.act import ScaledActivation + + if model_type not in AWQ_SCALES_MAPPINGS: + return model + for name, module in model.named_children(): + act_name = AWQ_SCALES_MAPPINGS[model_type]["act"] + layer_before_act_name = AWQ_SCALES_MAPPINGS[model_type]["layer_before_act"] + if name == act_name and hasattr(model, layer_before_act_name): + layer_before_act = getattr(model, AWQ_SCALES_MAPPINGS[model_type]["layer_before_act"]) + size = layer_before_act.out_features + scale_like = torch.ones(size) + model._modules[name] = ScaledActivation(module, scale_like) + _ = replace_quantization_scales(module, model_type) + return model + + +def replace_with_awq_linear( + model, + modules_to_not_convert=None, + quantization_config=None, + current_key_name=None, + has_been_replaced=False, +) -> bool: + """ + Public method that recursively replaces the Linear layers of the given model with AWQ quantized layers. + `accelerate` is needed to use this method. Returns the converted model and a boolean that indicates if the + conversion has been successfull or not. + + During the module replacement, we also infer the backend to use through the `quantization_config` object. + + Args: + model (`torch.nn.Module`): + The model to convert, can be any `torch.nn.Module` instance. + quantization_config (`AwqConfig`): + The quantization config object that contains the quantization parameters. + modules_to_not_convert (`list`, *optional*): + A list of modules to not convert. If a module name is in the list (e.g. `lm_head`), it will not be + converted. + current_key_name (`list`, *optional*): + A list that contains the current key name. This is used for recursion and should not be passed by the user. + has_been_replaced (`bool`, *optional*): + A boolean that indicates if the conversion has been successful or not. This is used for recursion and + should not be passed by the user. + """ + if modules_to_not_convert is None: + modules_to_not_convert = [] + + backend = quantization_config.backend + + if not is_auto_awq_available(): + raise ValueError( + "AWQ (either `autoawq` or `llmawq`) is not available. Please install it with `pip install autoawq` or check out the installation guide in https://github.com/mit-han-lab/llm-awq" + ) + + if backend == AwqBackendPackingMethod.AUTOAWQ: + if quantization_config.version == AWQLinearVersion.GEMM: + from awq.modules.linear.gemm import WQLinear_GEMM + + target_cls = WQLinear_GEMM + elif quantization_config.version == AWQLinearVersion.GEMV: + from awq.modules.linear.gemv import WQLinear_GEMV + + target_cls = WQLinear_GEMV + elif quantization_config.version == AWQLinearVersion.EXLLAMA: + if quantization_config.exllama_config["version"] == ExllamaVersion.ONE: + from awq.modules.linear.exllama import WQLinear_Exllama + + target_cls = WQLinear_Exllama + elif quantization_config.exllama_config["version"] == ExllamaVersion.TWO: + from awq.modules.linear.exllamav2 import WQLinear_ExllamaV2 + + target_cls = WQLinear_ExllamaV2 + else: + raise ValueError(f"Unrecognized Exllama version: {quantization_config.exllama_config['version']}") + elif quantization_config.version == AWQLinearVersion.IPEX: + from awq.modules.linear.gemm_ipex import WQLinear_IPEX + + target_cls = WQLinear_IPEX + else: + raise ValueError(f"Unrecognized AWQ version: {quantization_config.version}") + else: + from awq.quantize.qmodule import WQLinear + + target_cls = WQLinear + + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + + if isinstance(module, nn.Linear) and name not in modules_to_not_convert: + # Check if the current key is not in the `modules_to_not_convert` + if not any(key in ".".join(current_key_name) for key in modules_to_not_convert): + in_features = module.in_features + out_features = module.out_features + + model._modules[name] = target_cls( + w_bit=quantization_config.bits, + group_size=quantization_config.group_size, + in_features=in_features, + out_features=out_features, + bias=module.bias is not None, + dev=module.weight.device, + ) + has_been_replaced = True + + # Force requires grad to False to avoid unexpected errors + model._modules[name].requires_grad_(False) + if len(list(module.children())) > 0: + _, has_been_replaced = replace_with_awq_linear( + module, + modules_to_not_convert=modules_to_not_convert, + current_key_name=current_key_name, + quantization_config=quantization_config, + has_been_replaced=has_been_replaced, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced + + +def get_modules_to_fuse(model, quantization_config): + """ + Returns the fusing mapping given the quantization config and the model + + Args: + model (`~PreTrainedModel`): + The model to fuse - note this model should have been converted into AWQ format beforehand. + quantization_config (`~transformers.quantization_config.AWQConfig`): + The quantization configuration to use. + """ + if not isinstance(model, PreTrainedModel): + raise TypeError(f"The model should be an instance of `PreTrainedModel`, got {model.__class__.__name__}") + + # Always default to `quantization_config.modules_to_fuse` + if quantization_config.modules_to_fuse is not None: + current_fused_mapping = quantization_config.modules_to_fuse + current_fused_mapping["max_seq_len"] = quantization_config.fuse_max_seq_len + elif model.config.model_type in AWQ_FUSED_MAPPINGS: + current_fused_mapping = AWQ_FUSED_MAPPINGS[model.config.model_type] + + # Properly deal with the case where we have a multi-modal model as well (e.g. Llava) + config = model.config.get_text_config(decoder=True) + + # Handle hidden_size, num_attention_heads, num_key_value_heads on our own. + hidden_size = config.hidden_size + num_attention_heads = config.num_attention_heads + num_key_value_heads = getattr(config, "num_key_value_heads", num_attention_heads) + + # Fill `current_fused_mapping` with the expected values + current_fused_mapping["hidden_size"] = hidden_size + current_fused_mapping["num_attention_heads"] = num_attention_heads + current_fused_mapping["num_key_value_heads"] = num_key_value_heads + current_fused_mapping["max_seq_len"] = quantization_config.fuse_max_seq_len + else: + raise ValueError( + "Fusing mapping not found either on the quantization config or the supported `AWQ_FUSED_MAPPINGS`. Please pass a `fused_mapping` argument" + " in the `quantization_config` or raise an issue on transformers https://github.com/huggingface/transformers to add its support." + ) + return current_fused_mapping + + +def fuse_awq_modules(model, quantization_config): + """ + Optionally fuse some modules in the model to speedup inference. + + Args: + model (`~PreTrainedModel`): + The model to fuse - note this model should have been converted into AWQ format beforehand. + quantization_config (`Union[AwqConfig, dict]`): + The quantization configuration to use. + """ + # We need to convert it from dict in order to get an AwqConfig object + # otherwise the fields `backend` etc. will not be available + # https://github.com/huggingface/transformers/pull/27411#discussion_r1414044495 + if isinstance(quantization_config, dict): + quantization_config = AwqConfig.from_dict(quantization_config) + backend = quantization_config.backend + + modules_to_fuse = get_modules_to_fuse(model, quantization_config) + modules_to_not_convert = getattr(quantization_config, "modules_to_not_convert", None) + + if backend == AwqBackendPackingMethod.AUTOAWQ: + from awq.modules.fused.attn import QuantAttentionFused + from awq.modules.fused.mlp import QuantFusedMLP + from awq.modules.fused.norm import FasterTransformerRMSNorm + else: + raise ValueError("Fusing is only supported for the AutoAWQ backend") + + fused_attention_modules = [] + + for name, module in model.named_modules(): + if modules_to_not_convert is not None: + if any(module_name_to_not_convert in name for module_name_to_not_convert in modules_to_not_convert): + continue + + # Replace layer norms + _fuse_awq_layernorm(modules_to_fuse["layernorm"], module, FasterTransformerRMSNorm) + + # Replace MLP layers if awq version is not ipex. + if quantization_config.version != "ipex": + _fuse_awq_mlp(model, name, modules_to_fuse["mlp"], module, QuantFusedMLP) + else: + logger.info("The IPEX version AWQ does not support fuse mlp for now.") + + # Replace attention layers + attention_has_been_fused = _fuse_awq_attention_layers( + model, module, modules_to_fuse, name, QuantAttentionFused + ) + + if attention_has_been_fused: + fused_attention_modules.append(name.split(".")[0]) + + # For AWQ fused + Llama we need to set `config._attn_implementation` = "custom" to avoid unexpected behavior and pass + # `None` attention mask to the fused attention modules as now the attention mask is dropped by our models and dealt + # by the `AttentionMaskConverter` module. + if len(fused_attention_modules) > 0: + for module_name, module in model.named_modules(): + if any( + module_name in fused_attention_modules for fused_attention_parent_module in fused_attention_modules + ): + if hasattr(module, "config") and hasattr(module.config, "_attn_implementation"): + module.config._attn_implementation = "custom" + return model + + +def _fuse_awq_layernorm(fuse_module_names, module, target_cls): + """ + Fuse the LayerNorm layers into a target class using autoawq + + Args: + fuse_module_names (`List[str]`): + The list of module names to fuse + module (`nn.Module`): + The pytorch parent module that has layernorm modules to fuse + target_cls (`~autoawq.FasterTransformerRMSNorm`): + The `FasterTransformerRMSNorm` class as it only supports that class + for now. + """ + for module_name in fuse_module_names: + if hasattr(module, module_name): + old_module = getattr(module, module_name) + module._modules[module_name] = target_cls( + old_module.weight, + old_module.variance_epsilon, + ).to(old_module.weight.device) + del old_module + + +def _fuse_awq_mlp(model, current_module_name, fuse_module_names, module, target_cls): + """ + Fuse the MLP layers into a target class using autoawq + + Args: + model (`~PreTrainedModel`): + The input pretrained model + current_module_name (`str`): + The current submodule name + fuse_module_names (`List[str]`): + The list of module names to fuse. For the MLP layers it has to be an array + of length 3 that consists of the 3 MLP layers in the order (gate (dense layer post-attention) / up / down layers) + module (`nn.Module`): + The pytorch parent module that has layernorm modules to fuse + target_cls (`~autoawq.QuantFusedMLP`): + The `QuantFusedMLP` class as it only supports that class + for now. + """ + if len(fuse_module_names) == 0: + return + + if hasattr(module, fuse_module_names[0]): + gate_proj = getattr(module, fuse_module_names[0]) + up_proj = getattr(module, fuse_module_names[1]) + down_proj = getattr(module, fuse_module_names[2]) + + previous_device = gate_proj.qweight.device + + # Deal also with the case model has `text_config` attribute + config = model.config.get_text_config(decoder=True) + hidden_act = config.hidden_act + activation_fn = ACT2FN[hidden_act] + new_module = target_cls(gate_proj, down_proj, up_proj, activation_fn) + + parent_name, child_name = current_module_name.rsplit(".", 1) + parent = model.get_submodule(parent_name) + setattr(parent, child_name, new_module.to(previous_device)) + + del gate_proj, up_proj, down_proj + + +def _fuse_awq_attention_layers(model, module, modules_to_fuse, current_module_name, target_cls): + """ + Fuse the Attention layers into a target class using autoawq + + Args: + model (`~PreTrainedModel`): + The input pretrained model + module (`nn.Module`): + The pytorch parent module that has layernorm modules to fuse + modules_to_fuse (`List[str]`): + The module fusing mapping. The dictionary has to contain a field `attention` with attention module names + in the correct order: q, k, v, o layer + current_module_name (`str`): + The current submodule name + target_cls (`~autoawq.QuantAttentionFused`): + The `QuantAttentionFused` class as it only supports that class + for now. + """ + from awq.modules.linear import WQLinear_GEMM, WQLinear_GEMV + + module_has_been_fused = False + + if len(modules_to_fuse["attention"]) == 0: + return module_has_been_fused + + if hasattr(module, modules_to_fuse["attention"][0]): + # First, we pack the QKV layers together + q_proj = getattr(module, modules_to_fuse["attention"][0]) + + if isinstance(q_proj, WQLinear_GEMV): + linear_target_cls = WQLinear_GEMV + cat_dim = 0 + elif isinstance(q_proj, WQLinear_GEMM): + linear_target_cls = WQLinear_GEMM + cat_dim = 1 + elif is_ipex_available() and version.parse(importlib.metadata.version("autoawq")) > version.parse("0.2.6"): + from awq.modules.linear import WQLinear_IPEX + + if isinstance(q_proj, WQLinear_IPEX): + linear_target_cls = WQLinear_IPEX + cat_dim = 1 + else: + raise ValueError("Unsupported q_proj type: {type(q_proj)}") + + previous_device = q_proj.qweight.device + + k_proj = getattr(module, modules_to_fuse["attention"][1]) + v_proj = getattr(module, modules_to_fuse["attention"][2]) + o_proj = getattr(module, modules_to_fuse["attention"][3]) + + bias = torch.cat([q_proj.bias, k_proj.bias, v_proj.bias], dim=0) if q_proj.bias is not None else None + + qkv_layer = linear_target_cls( + q_proj.w_bit, + q_proj.group_size, + q_proj.in_features, + q_proj.out_features + k_proj.out_features + v_proj.out_features, + q_proj.bias is not None, + next(iter(module.state_dict().values())).device, + ) + + qkv_layer.qweight = torch.cat([q_proj.qweight, k_proj.qweight, v_proj.qweight], dim=cat_dim) + qkv_layer.qzeros = torch.cat([q_proj.qzeros, k_proj.qzeros, v_proj.qzeros], dim=cat_dim) + qkv_layer.scales = torch.cat([q_proj.scales, k_proj.scales, v_proj.scales], dim=cat_dim) + + if isinstance(qkv_layer, WQLinear_GEMV): + qkv_layer.split_k_iters = q_proj.split_k_iters + + qkv_layer.bias = bias + + fused_attention_layer = target_cls( + modules_to_fuse["hidden_size"], + modules_to_fuse["num_attention_heads"], + modules_to_fuse["num_key_value_heads"], + qkv_layer, + o_proj, + previous_device, + modules_to_fuse["max_seq_len"], + use_alibi=modules_to_fuse["use_alibi"], + # The default value in autoawq is set to 10000.0 + rope_theta=modules_to_fuse.get("rope_theta", 10000.0), + ) + + fused_attention_layer.is_hf_transformers = True + + parent_name, child_name = current_module_name.rsplit(".", 1) + parent = model.get_submodule(parent_name) + setattr(parent, child_name, fused_attention_layer.to(previous_device)) + + del q_proj, k_proj, v_proj, o_proj + module_has_been_fused = True + + return module_has_been_fused + + +def post_init_awq_exllama_modules(model, exllama_config): + """ + Runs post init for Exllama layers which performs: + - Weights unpacking, reordering and repacking + - Devices scratch space allocation + """ + + if exllama_config["version"] == ExllamaVersion.ONE: + from awq.modules.linear.exllama import exllama_post_init + + model = exllama_post_init(model) + elif exllama_config["version"] == ExllamaVersion.TWO: + from awq.modules.linear.exllamav2 import exllamav2_post_init + + model = exllamav2_post_init( + model, + max_input_len=exllama_config["max_input_len"], + max_batch_size=exllama_config["max_batch_size"], + ) + else: + raise ValueError(f"Unrecognized Exllama version: {exllama_config['version']}") + + return model + + +def post_init_awq_ipex_modules(model): + """ + Runs post init for IPEX layers which performs: + - Weights packing, reordering and repacking + """ + + from awq.modules.linear.gemm_ipex import ipex_post_init + + model = ipex_post_init(model) + + return model diff --git a/lib/python3.12/site-packages/transformers/integrations/bitnet.py b/lib/python3.12/site-packages/transformers/integrations/bitnet.py new file mode 100644 index 0000000000000000000000000000000000000000..0b50f9738afb69e680371b5f868bc4508f2a60e7 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/bitnet.py @@ -0,0 +1,288 @@ +from ..utils import is_accelerate_available, is_torch_available, logging + + +if is_accelerate_available(): + from accelerate import init_empty_weights + +if is_torch_available(): + import torch + import torch.nn as nn + import torch.nn.functional as F + +logger = logging.get_logger(__name__) + + +# the weights are ternary so can be represented with 2 bits, and they are packed in uint8 tensors, hence the number of values per item is 4 +VALUES_PER_ITEM = 4 + + +def pack_weights(quantized_weights: torch.Tensor) -> torch.Tensor: + """ + Packs a tensor of quantized weights into a compact format using 2 bits per value. + + Parameters: + ----------- + quantized_weights : torch.Tensor + A tensor containing ternary quantized weights with values in {-1, 0, 1}. These values are adjusted to + {0, 1, 2} before being packed. + + Returns: + -------- + torch.Tensor + A packed tensor where each element stores 4 quantized values (each using 2 bits) in an 8-bit format. + """ + + original_shape = quantized_weights.shape + + row_dim = (original_shape[0] + VALUES_PER_ITEM - 1) // VALUES_PER_ITEM + + if len(original_shape) == 1: + packed_tensor_shape = (row_dim,) + else: + packed_tensor_shape = (row_dim, *original_shape[1:]) + + quantized_weights += 1 + packed = torch.zeros(packed_tensor_shape, device=quantized_weights.device, dtype=torch.uint8) + unpacked = quantized_weights.to(torch.uint8) + + it = min(VALUES_PER_ITEM, (original_shape[0] // row_dim) + 1) + for i in range(it): + start = i * row_dim + end = min(start + row_dim, original_shape[0]) + packed[: (end - start)] |= unpacked[start:end] << 2 * i + + return packed + + +@torch.compile +def unpack_weights(packed: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """ + Unpacks a tensor of quantized weights that were stored in a packed format using 2 bits per value. + + Parameters: + ----------- + packed : torch.Tensor + A tensor containing packed weights where each element represents 4 quantized values (using 2 bits per value). + dtype : torch.dtype + The dtype of the returned Tensor + Returns: + -------- + torch.Tensor + A tensor of unpacked weights, where each value is converted from its packed 2-bit representation. + + Example: + -------- + packed = torch.tensor([[0b10100001, 0b00011000], + [0b10010000, 0b00001010]], dtype=torch.uint8) + + # Unpack the values + unpacked = unpack_weights(packed) + + # Resulting unpacked tensor + print(unpacked) + # Output: tensor([[ 0, -1], + [-1, 1], + [-1, 1], + [-1, 1], + [ 1, 0], + [ 0, -1], + [ 1, -1], + [ 1, -1]]) + + Explanation of the example: + --------------------------- + Let's take the first value for example 0b10100001, we we will only focus on the first column, + because every element is unpacked across the first dimension + - First 2 bits: `01` → 0 at [0][0] + - Second 2 bits: `00` → -1 at [0][2] + - Third 2 bits: `10` → 1 at [0][4] + - Fourth 2 bits: `10` → 1 at [0][6] + the second value of the same row (0b10010000) will give the values for [0][1], [0][3], [0][5], [0][7] + + We subtract 1 because during the packing process, it's easier to work with values like 0, 1, and 2. To make this possible, + we add 1 to the original ternary weights (which are typically -1, 0, and 1) when packing them. When unpacking, we reverse + this by subtracting 1 to restore the original ternary values. + """ + packed_shape = packed.shape + + if len(packed_shape) == 1: + original_row_dim = packed_shape[0] * VALUES_PER_ITEM + unpacked_shape = (original_row_dim,) + else: + original_row_dim = packed_shape[0] * VALUES_PER_ITEM + unpacked_shape = (original_row_dim, *packed_shape[1:]) + + unpacked = torch.zeros(unpacked_shape, device=packed.device, dtype=torch.uint8) + + for i in range(VALUES_PER_ITEM): + start = i * packed_shape[0] + end = start + packed_shape[0] + mask = 3 << (2 * i) + unpacked[start:end] = (packed & mask) >> (2 * i) + + return unpacked.to(dtype) - 1 + + +class BitLinear(nn.Module): + def __init__(self, in_features: int, out_features: int, bias: bool, device=None, dtype=None): + super().__init__() + self.dtype = dtype + self.in_features = in_features + self.out_features = out_features + self.register_buffer( + "weight", + torch.zeros( + (out_features // VALUES_PER_ITEM, in_features), + dtype=torch.uint8, + device=device, + ), + ) + self.register_buffer( + "weight_scale", + torch.ones( + (1), + dtype=dtype, + device=device, + ), + ) + if bias: + self.register_buffer("bias", torch.zeros((out_features), dtype=dtype, device=device)) + else: + self.bias = None + + @torch.compile + def activation_quant(self, input, num_bits=8): + """ + Activation function : Performs symmetric, per-token quantization on the input activations. + Parameters: + ----------- + x : torch.Tensor + Input activations to be quantized. + num_bits : int, optional (default=8) + Number of bits to use for quantization, determining the quantization range. + + Returns: + -------- + result : torch.Tensor + Quantized activation tensor, with values mapped to an `int8` range. + scale : torch.Tensor + The per-channel scaling factors used to quantize the tensor. + """ + Qn = -(2 ** (num_bits - 1)) + Qp = 2 ** (num_bits - 1) - 1 + scale = Qp / input.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-5) + result = (input * scale).round().clamp(Qn, Qp) + return result.to(torch.int8), scale + + @torch.compile + def post_quant_process(self, input, input_scale, weight_scale): + out = input / (input_scale * weight_scale) + return out + + def forward(self, input): + w = self.weight + w_quant = unpack_weights(w, dtype=self.dtype) + input_quant, input_scale = self.activation_quant(input) + y = F.linear(input_quant.to(self.dtype), w_quant) + y = self.post_quant_process(y, self.weight_scale, input_scale) + if self.bias is not None: + y += self.bias.view(1, -1).expand_as(y) + return y + + +def _replace_with_bitnet_linear( + model, + modules_to_not_convert=None, + current_key_name=None, + quantization_config=None, + has_been_replaced=False, + pre_quantized=False, +): + """ + Private method that wraps the recursion for module replacement. + + Returns the converted model and a boolean that indicates if the conversion has been successfull or not. + """ + + if current_key_name is None: + current_key_name = [] + + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + + # Check if the current key is not in the `modules_to_not_convert` + if not any(key in ".".join(current_key_name) for key in modules_to_not_convert): + with init_empty_weights(): + if isinstance(module, nn.Linear) and name not in modules_to_not_convert: + in_features = module.in_features + out_features = module.out_features + model._modules[name] = BitLinear( + in_features=in_features, + out_features=out_features, + bias=module.bias is not None, + device=module.weight.device, + dtype=module.weight.dtype, + ) + has_been_replaced = True + model._modules[name].requires_grad_(False) + + if len(list(module.children())) > 0: + _, has_been_replaced = _replace_with_bitnet_linear( + module, + modules_to_not_convert=modules_to_not_convert, + current_key_name=current_key_name, + quantization_config=quantization_config, + has_been_replaced=has_been_replaced, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced + + +def replace_with_bitnet_linear( + model, + modules_to_not_convert=None, + current_key_name=None, + quantization_config=None, + pre_quantized=False, +): + """ + A helper function to replace all `torch.nn.Linear` modules by `BitLinear158` modules`. + + The function will be run recursively and replace all `torch.nn.Linear` modules except for the `lm_head` that should + be kept as a `torch.nn.Linear` module. The replacement is done under `init_empty_weights` context manager so no + CPU/GPU memory is required to run this function. Each weight will be quantized along the channel. + + Parameters: + model (`torch.nn.Module`): + Input model or `torch.nn.Module` as the function is run recursively. + modules_to_not_convert (`List[`str`]`, *optional*, defaults to `["lm_head"]`): + Names of the modules to not convert in `EetqLinear`. In practice we keep the `lm_head` in full precision + for numerical stability reasons. + current_key_name (`List[`str`]`, *optional*): + An array to track the current key of the recursion. This is used to check whether the current key (part of + it) is not in the list of modules to not convert (for instances modules that are offloaded to `cpu` or + `disk`). + """ + modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert + if quantization_config and quantization_config.modules_to_not_convert is not None: + modules_to_not_convert.extend(quantization_config.modules_to_not_convert) + modules_to_not_convert = list(set(modules_to_not_convert)) + model, has_been_replaced = _replace_with_bitnet_linear( + model, + modules_to_not_convert, + current_key_name, + quantization_config, + pre_quantized=pre_quantized, + ) + + if not has_been_replaced: + logger.warning( + "You are loading your model using bitnet but no linear modules were found in your model." + " Please double check your model architecture, or submit an issue on github if you think this is" + " a bug." + ) + + return model diff --git a/lib/python3.12/site-packages/transformers/integrations/bitsandbytes.py b/lib/python3.12/site-packages/transformers/integrations/bitsandbytes.py new file mode 100644 index 0000000000000000000000000000000000000000..3973dc58c062bc9b1755a24c708e39e1ad4016ce --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/bitsandbytes.py @@ -0,0 +1,561 @@ +import importlib.metadata +import inspect +import warnings +from copy import deepcopy +from inspect import signature + +from packaging import version + +from ..utils import ( + get_available_devices, + is_accelerate_available, + is_bitsandbytes_available, + is_bitsandbytes_multi_backend_available, + is_ipex_available, + is_torch_available, + logging, +) + + +if is_bitsandbytes_available(): + import bitsandbytes as bnb + import torch + import torch.nn as nn + + from ..pytorch_utils import Conv1D + +if is_accelerate_available(): + import accelerate + from accelerate import init_empty_weights + from accelerate.hooks import add_hook_to_module, remove_hook_from_module + from accelerate.utils import find_tied_parameters + +logger = logging.get_logger(__name__) + + +def set_module_quantized_tensor_to_device(module, tensor_name, device, value=None, quantized_stats=None): + """ + A helper function to set a given tensor (parameter of buffer) of a module on a specific device (note that doing + `param.to(device)` creates a new tensor not linked to the parameter, which is why we need this function). The + function is adapted from `set_module_tensor_to_device` function from accelerate that is adapted to support the + class `Int8Params` from `bitsandbytes`. + + Args: + module (`torch.nn.Module`): + The module in which the tensor we want to move lives. + tensor_name (`str`): + The full name of the parameter/buffer. + device (`int`, `str` or `torch.device`): + The device on which to set the tensor. + value (`torch.Tensor`, *optional*): + The value of the tensor (useful when going from the meta device to any other device). + quantized_stats (`dict[str, Any]`, *optional*): + Dict with items for either 4-bit or 8-bit serialization + """ + # Recurse if needed + if "." in tensor_name: + splits = tensor_name.split(".") + for split in splits[:-1]: + new_module = getattr(module, split) + if new_module is None: + raise ValueError(f"{module} has no attribute {split}.") + module = new_module + tensor_name = splits[-1] + + if tensor_name not in module._parameters and tensor_name not in module._buffers: + raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.") + is_buffer = tensor_name in module._buffers + old_value = getattr(module, tensor_name) + + if old_value.device == torch.device("meta") and device not in ["meta", torch.device("meta")] and value is None: + raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {device}.") + + prequantized_loading = quantized_stats is not None + if is_buffer or not is_bitsandbytes_available(): + is_8bit = False + is_4bit = False + else: + is_4bit = hasattr(bnb.nn, "Params4bit") and isinstance(module._parameters[tensor_name], bnb.nn.Params4bit) + is_8bit = isinstance(module._parameters[tensor_name], bnb.nn.Int8Params) + + if is_8bit or is_4bit: + param = module._parameters[tensor_name] + if param.device.type != "cuda": + if value is None: + new_value = old_value.to(device) + elif isinstance(value, torch.Tensor): + new_value = value.to("cpu") + else: + new_value = torch.tensor(value, device="cpu") + + # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization. + # Since weights are saved in the correct "orientation", we skip transposing when loading. + if issubclass(module.source_cls, Conv1D) and not prequantized_loading: + new_value = new_value.T + + kwargs = old_value.__dict__ + + if prequantized_loading != (new_value.dtype in (torch.int8, torch.uint8)): + raise ValueError( + f"Value dtype `{new_value.dtype}` is not compatible with parameter quantization status." + ) + + if is_8bit: + is_8bit_serializable = version.parse(importlib.metadata.version("bitsandbytes")) > version.parse( + "0.37.2" + ) + if new_value.dtype in (torch.int8, torch.uint8) and not is_8bit_serializable: + raise ValueError( + "Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. " + "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`." + ) + new_value = bnb.nn.Int8Params(new_value, requires_grad=False, **kwargs).to(device) + if prequantized_loading: + setattr(new_value, "SCB", quantized_stats["SCB"].to(device)) + elif is_4bit: + if prequantized_loading: + is_4bit_serializable = version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse( + "0.41.3" + ) + if new_value.dtype in (torch.int8, torch.uint8) and not is_4bit_serializable: + raise ValueError( + "Detected 4-bit weights but the version of bitsandbytes is not compatible with 4-bit serialization. " + "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`." + ) + new_value = bnb.nn.Params4bit.from_prequantized( + data=new_value, + quantized_stats=quantized_stats, + requires_grad=False, + device=device, + **kwargs, + ) + else: + new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(device) + module._parameters[tensor_name] = new_value + + else: + if value is None: + new_value = old_value.to(device) + elif isinstance(value, torch.Tensor): + new_value = value.to(device) + else: + new_value = torch.tensor(value, device=device) + + if is_buffer: + module._buffers[tensor_name] = new_value + else: + new_value = nn.Parameter(new_value, requires_grad=old_value.requires_grad) + module._parameters[tensor_name] = new_value + + +def _replace_with_bnb_linear( + model, + modules_to_not_convert=None, + current_key_name=None, + quantization_config=None, + has_been_replaced=False, +): + """ + Private method that wraps the recursion for module replacement. + + Returns the converted model and a boolean that indicates if the conversion has been successfull or not. + """ + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + + if (isinstance(module, nn.Linear) or isinstance(module, Conv1D)) and name not in modules_to_not_convert: + # Check if the current key is not in the `modules_to_not_convert` + current_key_name_str = ".".join(current_key_name) + if not any( + (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert + ): + with init_empty_weights(): + if isinstance(module, Conv1D): + in_features, out_features = module.weight.shape + else: + in_features = module.in_features + out_features = module.out_features + + if quantization_config.quantization_method() == "llm_int8": + model._modules[name] = bnb.nn.Linear8bitLt( + in_features, + out_features, + module.bias is not None, + has_fp16_weights=quantization_config.llm_int8_has_fp16_weight, + threshold=quantization_config.llm_int8_threshold, + ) + has_been_replaced = True + else: + if ( + quantization_config.llm_int8_skip_modules is not None + and name in quantization_config.llm_int8_skip_modules + ): + pass + else: + extra_kwargs = ( + {"quant_storage": quantization_config.bnb_4bit_quant_storage} + if "quant_storage" in list(signature(bnb.nn.Linear4bit).parameters) + else {} + ) + model._modules[name] = bnb.nn.Linear4bit( + in_features, + out_features, + module.bias is not None, + quantization_config.bnb_4bit_compute_dtype, + compress_statistics=quantization_config.bnb_4bit_use_double_quant, + quant_type=quantization_config.bnb_4bit_quant_type, + **extra_kwargs, + ) + has_been_replaced = True + # Store the module class in case we need to transpose the weight later + model._modules[name].source_cls = type(module) + # Force requires grad to False to avoid unexpected errors + model._modules[name].requires_grad_(False) + if len(list(module.children())) > 0: + _, has_been_replaced = _replace_with_bnb_linear( + module, + modules_to_not_convert, + current_key_name, + quantization_config, + has_been_replaced=has_been_replaced, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced + + +def replace_with_bnb_linear(model, modules_to_not_convert=None, current_key_name=None, quantization_config=None): + """ + A helper function to replace all `torch.nn.Linear` modules by `bnb.nn.Linear8bit` modules from the `bitsandbytes` + library. This will enable running your models using mixed int8 precision as described by the paper `LLM.int8(): + 8-bit Matrix Multiplication for Transformers at Scale`. Make sure `bitsandbytes` compiled with the correct CUDA + version of your hardware is installed before running this function. `pip install -i https://test.pypi.org/simple/ + bitsandbytes` + + The function will be run recursively and replace all `torch.nn.Linear` modules except for the `lm_head` that should + be kept as a `torch.nn.Linear` module. The replacement is done under `init_empty_weights` context manager so no + CPU/GPU memory is required to run this function. Int8 mixed-precision matrix decomposition works by separating a + matrix multiplication into two streams: (1) and systematic feature outlier stream matrix multiplied in fp16 + (0.01%), (2) a regular stream of int8 matrix multiplication (99.9%). With this method, int8 inference with no + predictive degradation is possible for very large models (>=176B parameters). + + Parameters: + model (`torch.nn.Module`): + Input model or `torch.nn.Module` as the function is run recursively. + modules_to_not_convert (`List[`str`]`, *optional*, defaults to `["lm_head"]`): + Names of the modules to not convert in `Linear8bitLt`. In practice we keep the `lm_head` in full precision + for numerical stability reasons. + current_key_name (`List[`str`]`, *optional*): + An array to track the current key of the recursion. This is used to check whether the current key (part of + it) is not in the list of modules to not convert (for instances modules that are offloaded to `cpu` or + `disk`). + quantization_config ('transformers.utils.quantization_config.BitsAndBytesConfig'): + To configure and manage settings related to quantization, a technique used to compress neural network models + by reducing the precision of the weights and activations, thus making models more efficient in terms of both + storage and computation. + """ + modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert + model, has_been_replaced = _replace_with_bnb_linear( + model, modules_to_not_convert, current_key_name, quantization_config + ) + + if not has_been_replaced: + logger.warning( + "You are loading your model in 8bit or 4bit but no linear modules were found in your model." + " Please double check your model architecture, or submit an issue on github if you think this is" + " a bug." + ) + + return model + + +# For backward compatibility +def replace_8bit_linear(*args, **kwargs): + warnings.warn( + "`replace_8bit_linear` will be deprecated in a future version, please use `replace_with_bnb_linear` instead", + FutureWarning, + ) + return replace_with_bnb_linear(*args, **kwargs) + + +# For backward compatiblity +def set_module_8bit_tensor_to_device(*args, **kwargs): + warnings.warn( + "`set_module_8bit_tensor_to_device` will be deprecated in a future version, please use `set_module_quantized_tensor_to_device` instead", + FutureWarning, + ) + return set_module_quantized_tensor_to_device(*args, **kwargs) + + +def get_keys_to_not_convert(model): + r""" + An utility function to get the key of the module to keep in full precision if any For example for CausalLM modules + we may want to keep the lm_head in full precision for numerical stability reasons. For other architectures, we want + to keep the tied weights of the model. The function will return a list of the keys of the modules to not convert in + int8. + + Parameters: + model (`torch.nn.Module`): + Input model + """ + # Create a copy of the model and tie the weights, then + # check if it contains tied weights + tied_model = deepcopy(model) # this has 0 cost since it is done inside `init_empty_weights` context manager` + tied_model.tie_weights() + + tied_params = find_tied_parameters(tied_model) + # For compatibility with Accelerate < 0.18 + if isinstance(tied_params, dict): + tied_keys = sum(list(tied_params.values()), []) + list(tied_params.keys()) + else: + tied_keys = sum(tied_params, []) + has_tied_params = len(tied_keys) > 0 + + # If there is not tied weights, we want to keep the lm_head（output_embedding) in full precision + if not has_tied_params: + output_emb = model.get_output_embeddings() + if output_emb is not None: + list_last_module = [name for name, module in model.named_modules() if id(module) == id(output_emb)] + return list_last_module + + # otherwise, no tied weights, no output embedding defined, simply keep the last module in full precision + list_modules = list(model.named_parameters()) + list_last_module = [list_modules[-1][0]] + # add last module together with tied weights + intersection = set(list_last_module) - set(tied_keys) + list_untouched = list(set(tied_keys)) + list(intersection) + + # remove ".weight" from the keys + names_to_remove = [".weight", ".bias"] + filtered_module_names = [] + for name in list_untouched: + for name_to_remove in names_to_remove: + if name_to_remove in name: + name = name.replace(name_to_remove, "") + filtered_module_names.append(name) + + return filtered_module_names + + +# Copied from PEFT: https://github.com/huggingface/peft/blob/47b3712898539569c02ec5b3ed4a6c36811331a1/src/peft/utils/integrations.py#L41 +def dequantize_bnb_weight(weight: "torch.nn.Parameter", dtype: "torch.dtype", state=None): + """ + Helper function to dequantize 4bit or 8bit bnb weights. + + If the weight is not a bnb quantized weight, it will be returned as is. + """ + if not isinstance(weight, torch.nn.Parameter): + raise TypeError(f"Input weight should be of type nn.Parameter, got {type(weight)} instead") + + cls_name = weight.__class__.__name__ + if cls_name not in ("Params4bit", "Int8Params"): + return weight + + if cls_name == "Params4bit": + output_tensor = bnb.functional.dequantize_4bit(weight.data, weight.quant_state) + logger.warning_once( + f"The model is going to be dequantized in {output_tensor.dtype} - if you want to upcast it to another dtype, make sure to pass the desired dtype when quantizing the model through `bnb_4bit_quant_type` argument of `BitsAndBytesConfig`" + ) + return output_tensor.to(dtype) + + if state.SCB is None: + state.SCB = weight.SCB + + if hasattr(bnb.functional, "int8_vectorwise_dequant"): + # Use bitsandbytes API if available (requires v0.45.0+) + dequantized = bnb.functional.int8_vectorwise_dequant(weight.data, state.SCB) + else: + # Multiply by (scale/127) to dequantize. + dequantized = weight.data * state.SCB.view(-1, 1) * 7.874015718698502e-3 + + return dequantized.to(dtype) + + +def _create_accelerate_new_hook(old_hook): + r""" + Creates a new hook based on the old hook. Use it only if you know what you are doing ! + This method is a copy of: https://github.com/huggingface/peft/blob/748f7968f3a31ec06a1c2b0328993319ad9a150a/src/peft/utils/other.py#L245 + with some changes + """ + old_hook_cls = getattr(accelerate.hooks, old_hook.__class__.__name__) + old_hook_attr = old_hook.__dict__ + filtered_old_hook_attr = {} + old_hook_init_signature = inspect.signature(old_hook_cls.__init__) + for k in old_hook_attr.keys(): + if k in old_hook_init_signature.parameters: + filtered_old_hook_attr[k] = old_hook_attr[k] + new_hook = old_hook_cls(**filtered_old_hook_attr) + return new_hook + + +def _dequantize_and_replace( + model, + dtype, + modules_to_not_convert=None, + current_key_name=None, + quantization_config=None, + has_been_replaced=False, +): + """ + Converts a quantized model into its dequantized original version. The newly converted model will have + some performance drop compared to the original model before quantization - use it only for specific usecases + such as QLoRA adapters merging. + + Returns the converted model and a boolean that indicates if the conversion has been successfull or not. + """ + quant_method = quantization_config.quantization_method() + + target_cls = bnb.nn.Linear8bitLt if quant_method == "llm_int8" else bnb.nn.Linear4bit + + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + + if isinstance(module, target_cls) and name not in modules_to_not_convert: + # Check if the current key is not in the `modules_to_not_convert` + current_key_name_str = ".".join(current_key_name) + + if not any( + (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert + ): + bias = getattr(module, "bias", None) + + device = module.weight.device + with init_empty_weights(): + new_module = torch.nn.Linear(module.in_features, module.out_features, bias=bias is not None) + + if quant_method == "llm_int8": + state = module.state + else: + state = None + + new_module.weight = torch.nn.Parameter(dequantize_bnb_weight(module.weight, dtype, state)) + + if bias is not None: + new_module.bias = bias + + # Create a new hook and attach it in case we use accelerate + if hasattr(module, "_hf_hook"): + old_hook = module._hf_hook + new_hook = _create_accelerate_new_hook(old_hook) + + remove_hook_from_module(module) + add_hook_to_module(new_module, new_hook) + + new_module.to(device) + model._modules[name] = new_module + has_been_replaced = True + if len(list(module.children())) > 0: + _, has_been_replaced = _dequantize_and_replace( + module, + dtype, + modules_to_not_convert, + current_key_name, + quantization_config, + has_been_replaced=has_been_replaced, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced + + +def dequantize_and_replace( + model, + modules_to_not_convert=None, + quantization_config=None, +): + model, has_been_replaced = _dequantize_and_replace( + model, + model.dtype, + modules_to_not_convert=modules_to_not_convert, + quantization_config=quantization_config, + ) + + if not has_been_replaced: + logger.warning( + "For some reason the model has not been properly dequantized. You might see unexpected behavior." + ) + + return model + + +def _validate_bnb_multi_backend_availability(raise_exception): + import bitsandbytes as bnb + + bnb_supported_devices = getattr(bnb, "supported_torch_devices", set()) + available_devices = set(get_available_devices()) + + if available_devices == {"cpu"} and not is_ipex_available(): + from importlib.util import find_spec + + if find_spec("intel_extension_for_pytorch"): + logger.warning( + "You have Intel IPEX installed but if you're intending to use it for CPU, it might not have the right version. Be sure to double check that your PyTorch and IPEX installs are compatible." + ) + + available_devices = frozenset( + [device for device in available_devices if device != "cpu"] + ) # Only Intel CPU is supported by BNB at the moment + + if not available_devices.intersection(bnb_supported_devices): + if raise_exception: + bnb_supported_devices_with_info = set( # noqa: C401 + '"cpu" (needs an Intel CPU and intel_extension_for_pytorch installed and compatible with the PyTorch version)' + if device == "cpu" + else device + for device in bnb_supported_devices + ) + err_msg = ( + f"None of the available devices `available_devices = {available_devices or None}` are supported by the bitsandbytes version you have installed: `bnb_supported_devices = {bnb_supported_devices_with_info}`. " + "Please check the docs to see if the backend you intend to use is available and how to install it: https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend" + ) + + logger.error(err_msg) + raise RuntimeError(err_msg) + + logger.warning("No supported devices found for bitsandbytes multi-backend.") + return False + + logger.debug("Multi-backend validation successful.") + return True + + +def _validate_bnb_cuda_backend_availability(raise_exception): + if not is_torch_available(): + return False + + import torch + + if not torch.cuda.is_available(): + log_msg = ( + "CUDA is required but not available for bitsandbytes. Please consider installing the multi-platform enabled version of bitsandbytes, which is currently a work in progress. " + "Please check currently supported platforms and installation instructions at https://huggingface.co/docs/bitsandbytes/main/en/installation#multi-backend" + ) + if raise_exception: + logger.error(log_msg) + raise RuntimeError(log_msg) + + logger.warning(log_msg) + return False + + logger.debug("CUDA backend validation successful.") + return True + + +def validate_bnb_backend_availability(raise_exception=False): + """ + Validates if the available devices are supported by bitsandbytes, optionally raising an exception if not. + """ + if not is_bitsandbytes_available(): + if importlib.util.find_spec("bitsandbytes") and version.parse( + importlib.metadata.version("bitsandbytes") + ) < version.parse("0.43.1"): + return _validate_bnb_cuda_backend_availability(raise_exception) + return False + + if is_bitsandbytes_multi_backend_available(): + return _validate_bnb_multi_backend_availability(raise_exception) + return _validate_bnb_cuda_backend_availability(raise_exception) diff --git a/lib/python3.12/site-packages/transformers/integrations/compressed_tensors.py b/lib/python3.12/site-packages/transformers/integrations/compressed_tensors.py new file mode 100644 index 0000000000000000000000000000000000000000..752227914d98b40e919ab34e26e88868a9334c61 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/compressed_tensors.py @@ -0,0 +1,54 @@ +# Copyright 2025 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. +from transformers.utils import is_torch_available + + +if is_torch_available(): + import torch + import torch.nn as nn + +from transformers.models.llama4.modeling_llama4 import Llama4TextMLP + + +def skip(*args, **kwargs): + pass + + +class CompressedExpertsLinear(nn.Module): + """ + A module that implements a compressed version of a list of expert modules. + This is specifically designed to work with Llama4TextExperts in MoE layers. + """ + + def __init__(self, config): + # Skip random weight initialization for experts. Otherwise, + # the init of this module would take over minutes. For a model + # with tens of layers of experts, it would easily take over 20 minutes. + nn.init.kaiming_uniform_ = skip + nn.init.uniform_ = skip + nn.init.normal_ = skip + super().__init__() + self.num_experts = config.num_local_experts + self.expert_modules = nn.ModuleList([Llama4TextMLP(config) for _ in range(self.num_experts)]) + + def forward( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + hidden_states = hidden_states.reshape(self.num_experts, -1, hidden_states.shape[-1]) + expert_routed_out_list = [] + for expert_idx in range(self.num_experts): + expert_routed_out_list.append(self.expert_modules[expert_idx](hidden_states[expert_idx])) + routed_out = torch.cat(expert_routed_out_list, dim=0) + return routed_out diff --git a/lib/python3.12/site-packages/transformers/integrations/deepspeed.py b/lib/python3.12/site-packages/transformers/integrations/deepspeed.py new file mode 100644 index 0000000000000000000000000000000000000000..1700301db51fd73bd9023f68170066d88c6971eb --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/deepspeed.py @@ -0,0 +1,498 @@ +# Copyright 2020 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. +""" +Integration with Deepspeed +""" + +import copy +import importlib.metadata as importlib_metadata +import importlib.util +import weakref +from functools import partialmethod + +from ..dependency_versions_check import dep_version_check +from ..utils import is_accelerate_available, is_torch_available, logging + + +if is_torch_available(): + import torch + from torch import nn + + +logger = logging.get_logger(__name__) + + +def is_deepspeed_available(): + package_exists = importlib.util.find_spec("deepspeed") is not None + + # Check we're not importing a "deepspeed" directory somewhere but the actual library by trying to grab the version + # AND checking it has an author field in the metadata that is HuggingFace. + if package_exists: + try: + _ = importlib_metadata.metadata("deepspeed") + return True + except importlib_metadata.PackageNotFoundError: + return False + + +if is_accelerate_available() and is_deepspeed_available(): + from accelerate.utils.deepspeed import HfDeepSpeedConfig as DeepSpeedConfig +else: + # Inherits from a dummy `object` if accelerate is not available, so that python succeeds to import this file. + # Deepspeed glue code will never inherit this dummy object as it checks if accelerate is available. + from builtins import object as DeepSpeedConfig + + +class HfDeepSpeedConfig(DeepSpeedConfig): + """ + This object contains a DeepSpeed configuration dictionary and can be quickly queried for things like zero stage. + + A `weakref` of this object is stored in the module's globals to be able to access the config from areas where + things like the Trainer object is not available (e.g. `from_pretrained` and `_get_resized_embeddings`). Therefore + it's important that this object remains alive while the program is still running. + + [`Trainer`] uses the `HfTrainerDeepSpeedConfig` subclass instead. That subclass has logic to sync the configuration + with values of [`TrainingArguments`] by replacing special placeholder values: `"auto"`. Without this special logic + the DeepSpeed configuration is not modified in any way. + + Args: + config_file_or_dict (`Union[str, Dict]`): path to DeepSpeed config file or dict. + + """ + + def __init__(self, config_file_or_dict): + # set global weakref object + set_hf_deepspeed_config(self) + dep_version_check("accelerate") + dep_version_check("deepspeed") + super().__init__(config_file_or_dict) + + +class HfTrainerDeepSpeedConfig(HfDeepSpeedConfig): + """ + The `HfTrainerDeepSpeedConfig` object is meant to be created during `TrainingArguments` object creation and has the + same lifespan as the latter. + """ + + def __init__(self, config_file_or_dict): + super().__init__(config_file_or_dict) + self._dtype = None + self.mismatches = [] + + def dtype(self): + if self._dtype is None: + raise ValueError("trainer_config_process() wasn't called yet to tell dtype") + return self._dtype + + def is_auto(self, ds_key_long): + val = self.get_value(ds_key_long) + if val is None: + return False + else: + return val == "auto" + + def fill_match(self, ds_key_long, hf_val, hf_key=None, must_match=True): + """ + A utility method that massages the config file and can optionally verify that the values match. + + 1. Replace "auto" values with `TrainingArguments` value. + + 2. If it wasn't "auto" and `must_match` is true, then check that DS config matches Trainer + config values and if mismatched add the entry to `self.mismatched` - will assert during + `trainer_config_finalize` for one or more mismatches. + + """ + config, ds_key = self.find_config_node(ds_key_long) + if config is None: + return + + if config.get(ds_key) == "auto": + config[ds_key] = hf_val + return + + if not must_match: + return + + ds_val = config.get(ds_key) + if ds_val is not None and ds_val != hf_val: + self.mismatches.append(f"- ds {ds_key_long}={ds_val} vs hf {hf_key}={hf_val}") + + fill_only = partialmethod(fill_match, must_match=False) + + def trainer_config_process(self, args, auto_find_batch_size=False): + """ + Adjust the config with `TrainingArguments` values. This stage is run during `TrainingArguments` object + creation. + """ + # DeepSpeed does: + # train_batch_size = world_size * train_micro_batch_size_per_gpu * gradient_accumulation_steps + train_batch_size = args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps + self.fill_match( + "train_micro_batch_size_per_gpu", + args.per_device_train_batch_size, + "per_device_train_batch_size", + not auto_find_batch_size, + ) + self.fill_match( + "gradient_accumulation_steps", + args.gradient_accumulation_steps, + "gradient_accumulation_steps", + ) + self.fill_match( + "train_batch_size", + train_batch_size, + "train_batch_size (calculated)", + not auto_find_batch_size, + ) + self.fill_match("gradient_clipping", args.max_grad_norm, "max_grad_norm") + + self.fill_match("optimizer.params.lr", args.learning_rate, "learning_rate") + self.fill_match( + "optimizer.params.betas", + [args.adam_beta1, args.adam_beta2], + "adam_beta1+adam_beta2", + ) + self.fill_match("optimizer.params.eps", args.adam_epsilon, "adam_epsilon") + self.fill_match("optimizer.params.weight_decay", args.weight_decay, "weight_decay") + + self.fill_only("scheduler.params.warmup_min_lr", 0) # not a trainer arg + self.fill_match("scheduler.params.warmup_max_lr", args.learning_rate, "learning_rate") + # total_num_steps - will get set in trainer_config_finalize + + # fp16 + if args.fp16 or args.fp16_full_eval: + fp16_backend = "apex" if args.fp16_backend == "apex" else "amp" + else: + fp16_backend = None + + if args.save_on_each_node: + # deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True + self.config["checkpoint"] = self.config.get("checkpoint", {}) + self.config["checkpoint"]["use_node_local_storage"] = args.save_on_each_node + + # amp: similar to the pytorch native amp - it has a bunch of optional params but we won't set + # any here unless the user did the work + self.fill_match( + "fp16.enabled", + ((args.fp16 or args.fp16_full_eval) and fp16_backend == "amp"), + "fp16|fp16_full_eval+fp16_backend(amp)", + ) + + # apex: delegates amp work to apex (which needs to be available), but it cannot be used with any + # ZeRO features + self.fill_match("amp.enabled", fp16_backend == "apex", "fp16+fp16_backend(apex)") + self.fill_match("amp.opt_level", args.fp16_opt_level, "fp16_opt_level") + + self.fill_match("bf16.enabled", (args.bf16 or args.bf16_full_eval), "bf16|bf16_full_eval") + + # deepspeed's default mode is fp16 unless there is a config that says differently + if self.is_true("bf16.enabled"): + self._dtype = torch.bfloat16 + elif self.is_false("fp16.enabled"): + self._dtype = torch.float32 + else: + self._dtype = torch.float16 + + def trainer_config_finalize(self, args, model, num_training_steps): + """ + This stage is run after we have the model and know num_training_steps. + + Now we can complete the configuration process. + """ + # zero + + # deal with config keys that use `auto` value and rely on model's hidden_size + hidden_size_based_keys = [ + "zero_optimization.reduce_bucket_size", + "zero_optimization.stage3_prefetch_bucket_size", + "zero_optimization.stage3_param_persistence_threshold", + ] + hidden_size_auto_keys = [x for x in hidden_size_based_keys if self.is_auto(x)] + + if len(hidden_size_auto_keys) > 0: + if hasattr(model.config, "hidden_size"): + hidden_size = model.config.hidden_size + elif hasattr(model.config, "hidden_sizes"): + # if there are many hidden sizes pick the largest one + hidden_size = max(model.config.hidden_sizes) + elif hasattr(model.config, "text_config") and hasattr(model.config.text_config, "hidden_size"): + hidden_size = model.config.text_config.hidden_size + elif hasattr(model.config, "text_config") and hasattr(model.config.text_config, "hidden_sizes"): + # if there are many hidden sizes pick the largest one + hidden_size = max(model.config.text_config.hidden_sizes) + else: + raise ValueError( + "The model's config file has neither `hidden_size` nor `hidden_sizes` entry, " + "therefore it's not possible to automatically fill out the following `auto` entries " + f"in the DeepSpeed config file: {hidden_size_auto_keys}. You can fix that by replacing " + "`auto` values for these keys with an integer value of your choice." + ) + + self.fill_only("zero_optimization.reduce_bucket_size", hidden_size * hidden_size) + if self.is_zero3(): + # automatically assign the optimal config values based on model config + self.fill_only( + "zero_optimization.stage3_prefetch_bucket_size", + int(0.9 * hidden_size * hidden_size), + ) + self.fill_only( + "zero_optimization.stage3_param_persistence_threshold", + 10 * hidden_size, + ) + + # scheduler + self.fill_match( + "scheduler.params.total_num_steps", + num_training_steps, + "num_training_steps (calculated)", + ) + self.fill_match( + "scheduler.params.warmup_num_steps", + args.get_warmup_steps(num_training_steps), + "warmup_steps", + ) + + if len(self.mismatches) > 0: + mismatches = "\n".join(self.mismatches) + raise ValueError( + "Please correct the following DeepSpeed config values that mismatch TrainingArguments" + f" values:\n{mismatches}\nThe easiest method is to set these DeepSpeed config values to 'auto'." + ) + + +# keep the config object global to be able to access it anywhere during TrainingArguments life-cycle +_hf_deepspeed_config_weak_ref = None + + +def set_hf_deepspeed_config(hf_deepspeed_config_obj): + # this is a special weakref global object to allow us to get to Deepspeed config from APIs + # that don't have an easy way to get to the Deepspeed config outside of the Trainer domain. + global _hf_deepspeed_config_weak_ref + # will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed) + _hf_deepspeed_config_weak_ref = weakref.ref(hf_deepspeed_config_obj) + + +def unset_hf_deepspeed_config(): + # useful for unit tests to ensure the global state doesn't leak - call from `tearDown` method + global _hf_deepspeed_config_weak_ref + _hf_deepspeed_config_weak_ref = None + + +def is_deepspeed_zero3_enabled(): + if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None: + return _hf_deepspeed_config_weak_ref().is_zero3() + else: + return False + + +def deepspeed_config(): + if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None: + return _hf_deepspeed_config_weak_ref().config + else: + return None + + +def _load_state_dict_into_zero3_model(model_to_load, state_dict): + """ + Loads state dict into a model specifically for Zero3, since DeepSpeed does not support the `transformers` + tensor parallelism API. + + Nearly identical code to PyTorch's `_load_from_state_dict` + """ + # copy state_dict so `_load_state_dict_into_zero3_model` can modify it + metadata = getattr(state_dict, "_metadata", None) + state_dict = state_dict.copy() + if metadata is not None: + state_dict._metadata = metadata + + error_msgs = [] + + # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants + # so we need to apply the function recursively. + def load(module: nn.Module, state_dict, prefix="", assign_to_params_buffers=False): + local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) + local_metadata["assign_to_params_buffers"] = assign_to_params_buffers + + args = (state_dict, prefix, local_metadata, True, [], [], error_msgs) + # Parameters of module and children will start with prefix. We can exit early if there are none in this + # state_dict + if is_deepspeed_zero3_enabled() and len([key for key in state_dict if key.startswith(prefix)]) > 0: + import deepspeed + + # In sharded models, each shard has only part of the full state_dict, so only gather + # parameters that are in the current state_dict. + named_parameters = dict(module.named_parameters(prefix=prefix[:-1], recurse=False)) + params_to_gather = [named_parameters[k] for k in state_dict.keys() if k in named_parameters] + if len(params_to_gather) > 0: + # because zero3 puts placeholders in model params, this context + # manager gathers (unpartitions) the params of the current layer, then loads from + # the state dict and then re-partitions them again + with deepspeed.zero.GatheredParameters(params_to_gather, modifier_rank=0): + if torch.distributed.get_rank() == 0: + module._load_from_state_dict(*args) + + for name, child in module._modules.items(): + if child is not None: + load(child, state_dict, prefix + name + ".", assign_to_params_buffers) + + load(model_to_load, state_dict, assign_to_params_buffers=False) + + return error_msgs + + +def deepspeed_optim_sched(trainer, hf_deepspeed_config, args, num_training_steps, model_parameters): + """ + A convenience wrapper that deals with optimizer and lr scheduler configuration. + """ + from accelerate.utils import DummyOptim, DummyScheduler + + config = hf_deepspeed_config.config + + # Mixing and matching DS schedulers and optimizers is supported unless Offload is enabled in which case it's: + # 1. DS scheduler + DS optimizer: Yes + # 2. HF scheduler + HF optimizer: Mostly* + # 3. DS scheduler + HF optimizer: Mostly* + # 4. HF scheduler + DS optimizer: Yes + # + # Mostly*: All non-native DeepSpeed optimizers that have both CPU and GPU implementation should work (except LAMB) + + optimizer = None + if "optimizer" in config: + if args.adafactor: + raise ValueError( + "--adafactor was passed, but also found `optimizer` configured in the DeepSpeed config. " + "Only one optimizer can be configured." + ) + optimizer = DummyOptim(params=model_parameters) + else: + if hf_deepspeed_config.is_offload(): + logger.info( + "Detected ZeRO Offload and non-DeepSpeed optimizers: This combination should work as long as the" + " custom optimizer has both CPU and GPU implementation (except LAMB)" + ) + + # ds supports Adam, OneBitAdam, and Lamb optimizers and can import other optimizers from torch. + # But trainer uses AdamW by default. + optimizer = trainer.create_optimizer() + # To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer` + config["zero_allow_untested_optimizer"] = True + + lr_scheduler = None + if "scheduler" in config: + lr_scheduler = DummyScheduler(optimizer) + else: + if isinstance(optimizer, DummyOptim): + + def _lr_scheduler_callable(optimizer): + # create a shallow copy first, so later modifications do not affect original trainer + trainer_copy = copy.copy(trainer) + # at the time _lr_scheduler_callable is called, trainer.lr_scheduler has been set + # update it to None so that we can re-create a new scheduler + trainer_copy.lr_scheduler = None + lr_scheduler = trainer_copy.create_scheduler( + num_training_steps=num_training_steps, optimizer=optimizer + ) + return lr_scheduler + + lr_scheduler = DummyScheduler(optimizer, lr_scheduler_callable=_lr_scheduler_callable) + else: + lr_scheduler = trainer.create_scheduler(num_training_steps=num_training_steps, optimizer=optimizer) + + return optimizer, lr_scheduler + + +def deepspeed_init(trainer, num_training_steps, inference=False): + """ + Init DeepSpeed, after updating the DeepSpeed configuration with any relevant Trainer's args. + + If `resume_from_checkpoint` was passed then an attempt to resume from a previously saved checkpoint will be made. + + Args: + trainer: Trainer object + num_training_steps: per single gpu + resume_from_checkpoint: path to a checkpoint if to resume from after normal DeepSpeedEngine load + inference: launch in inference mode (no optimizer and no lr scheduler) + auto_find_batch_size: whether to ignore the `train_micro_batch_size_per_gpu` argument as it's being + set automatically by the auto batch size finder + + Returns: optimizer, lr_scheduler + + We may use `deepspeed_init` more than once during the life of Trainer, when we do - it's a temp hack based on: + https://github.com/deepspeedai/DeepSpeed/issues/1394#issuecomment-937405374 until Deepspeed fixes a bug where it + can't resume from a checkpoint after it did some stepping https://github.com/deepspeedai/DeepSpeed/issues/1612 + + """ + from deepspeed.utils import logger as ds_logger + + model = trainer.model + args = trainer.args + + hf_deepspeed_config = trainer.accelerator.state.deepspeed_plugin.hf_ds_config + + # resume config update - some bits like `model` and `num_training_steps` only become available during train + hf_deepspeed_config.trainer_config_finalize(args, model, num_training_steps) + + # set the Deepspeed log level consistent with the Trainer + ds_logger.setLevel(args.get_process_log_level()) + + if inference: + # only Z3 makes sense for the inference + if not hf_deepspeed_config.is_zero3(): + raise ValueError("ZeRO inference only makes sense with ZeRO Stage 3 - please adjust your config") + + # in case the training config is re-used for inference + hf_deepspeed_config.del_config_sub_tree("optimizer") + hf_deepspeed_config.del_config_sub_tree("lr_scheduler") + optimizer, lr_scheduler = None, None + model_parameters = None + else: + trainer.optimizer = None # important for when deepspeed_init is used as re-init + tp_size = hf_deepspeed_config.config.get("tensor_parallel", {}).get("autotp_size", 0) + if tp_size > 1: + import deepspeed + + model = deepspeed.tp_model_init(model=model, tp_size=tp_size, dtype=hf_deepspeed_config.dtype()) + model_parameters = list(filter(lambda p: p.requires_grad, model.parameters())) + optimizer, lr_scheduler = deepspeed_optim_sched( + trainer, hf_deepspeed_config, args, num_training_steps, model_parameters + ) + + # keep for quick debug: + # from pprint import pprint; pprint(config) + + return optimizer, lr_scheduler + + +def deepspeed_load_checkpoint(deepspeed_engine, checkpoint_path, load_module_strict=True): + # it's possible that the user is trying to resume from model_path, which doesn't necessarily + # contain a deepspeed checkpoint. e.g. examples just check if the dir exists and assume it's + # a resume from a checkpoint and not just a local pretrained weight. So we check here if the + # path contains what looks like a deepspeed checkpoint + import glob + + deepspeed_checkpoint_dirs = sorted(glob.glob(f"{checkpoint_path}/global_step*")) + + if len(deepspeed_checkpoint_dirs) > 0: + logger.info(f"Attempting to resume from {checkpoint_path}") + # this magically updates self.optimizer and self.lr_scheduler + load_path, _ = deepspeed_engine.load_checkpoint( + checkpoint_path, + load_module_strict=load_module_strict, + load_optimizer_states=True, + load_lr_scheduler_states=True, + ) + if load_path is None: + raise ValueError(f"[deepspeed] failed to resume from checkpoint {checkpoint_path}") + else: + raise ValueError(f"Can't find a valid checkpoint at {checkpoint_path}") diff --git a/lib/python3.12/site-packages/transformers/integrations/eetq.py b/lib/python3.12/site-packages/transformers/integrations/eetq.py new file mode 100644 index 0000000000000000000000000000000000000000..97698cf1aa37c6279d4f6932ef496ed708e324a7 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/eetq.py @@ -0,0 +1,121 @@ +# coding=utf-8 +# Copyright 2024 NetEase, Inc. and the HuggingFace Inc. 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. +from ..utils import is_accelerate_available, is_eetq_available, logging + + +if is_eetq_available(): + import eetq + import torch.nn as nn + +if is_accelerate_available(): + from accelerate import init_empty_weights + +logger = logging.get_logger(__name__) + + +def _replace_with_eetq_linear( + model, + modules_to_not_convert=None, + current_key_name=None, + quantization_config=None, + has_been_replaced=False, + pre_quantized=False, +): + """ + Private method that wraps the recursion for module replacement. + + Returns the converted model and a boolean that indicates if the conversion has been successfull or not. + """ + if current_key_name is None: + current_key_name = [] + + for name, module in model.named_children(): + current_key_name.append(name) + + if (isinstance(module, nn.Linear)) and name not in modules_to_not_convert: + # Check if the current key is not in the `modules_to_not_convert` + current_key_name_str = ".".join(current_key_name) + if not any( + (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert + ): + with init_empty_weights(): + in_features = module.in_features + out_features = module.out_features + model._modules[name] = eetq.EetqLinear( + in_features, out_features, module.bias is not None, module.weight.device + ) + if pre_quantized: + model._modules[name].register_scale(module.weight.device) + has_been_replaced = True + + # Force requires grad to False to avoid unexpected errors + model._modules[name].requires_grad_(False) + if len(list(module.children())) > 0: + _, has_been_replaced = _replace_with_eetq_linear( + module, + modules_to_not_convert, + current_key_name, + quantization_config, + has_been_replaced=has_been_replaced, + pre_quantized=pre_quantized, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced + + +def replace_with_eetq_linear( + model, modules_to_not_convert=None, current_key_name=None, quantization_config=None, pre_quantized=False +): + """ + A helper function to replace all `torch.nn.Linear` modules by `eetq.EetqLinear` modules from the `eetq` + library. This will enable running your models using high performance int8 weight-only gemm kerner from + FasterTransformer and TensorRT-LLM. Make sure `eetq` compiled with the correct CUDA + version of your hardware is installed before running this function. EETQ shall be installed via the source + 'https://github.com/NetEase-FuXi/EETQ' + + The function will be run recursively and replace all `torch.nn.Linear` modules except for the `lm_head` that should + be kept as a `torch.nn.Linear` module. The replacement is done under `init_empty_weights` context manager so no + CPU/GPU memory is required to run this function. Each weight will be quantized along the channel. + + Parameters: + model (`torch.nn.Module`): + Input model or `torch.nn.Module` as the function is run recursively. + modules_to_not_convert (`List[`str`]`, *optional*, defaults to `["lm_head"]`): + Names of the modules to not convert in `EetqLinear`. In practice we keep the `lm_head` in full precision + for numerical stability reasons. + current_key_name (`List[`str`]`, *optional*): + An array to track the current key of the recursion. This is used to check whether the current key (part of + it) is not in the list of modules to not convert (for instances modules that are offloaded to `cpu` or + `disk`). + """ + + modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert + + if quantization_config.modules_to_not_convert is not None: + modules_to_not_convert.extend(quantization_config.modules_to_not_convert) + modules_to_not_convert = list(set(modules_to_not_convert)) + model, has_been_replaced = _replace_with_eetq_linear( + model, modules_to_not_convert, current_key_name, quantization_config, pre_quantized=pre_quantized + ) + + if not has_been_replaced: + logger.warning( + "You are loading your model using eetq but no linear modules were found in your model." + " Please double check your model architecture, or submit an issue on github if you think this is" + " a bug." + ) + + return model diff --git a/lib/python3.12/site-packages/transformers/integrations/executorch.py b/lib/python3.12/site-packages/transformers/integrations/executorch.py new file mode 100644 index 0000000000000000000000000000000000000000..591c556e59f07e39e758f04e29c74871930034d9 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/executorch.py @@ -0,0 +1,408 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# 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 typing import Optional + +import torch + +from transformers.generation.configuration_utils import GenerationConfig + +from ..utils.import_utils import is_torch_available + + +if is_torch_available(): + from transformers import PreTrainedModel, StaticCache + from transformers.pytorch_utils import is_torch_greater_or_equal, is_torch_greater_or_equal_than_2_3 + + +class TorchExportableModuleWithStaticCache(torch.nn.Module): + """ + A wrapper module designed to make a `PreTrainedModel` exportable with `torch.export`, + specifically for use with static caching. This module ensures that the exported model + is compatible with further lowering and execution in `ExecuTorch`. + + Note: + This class is specifically designed to support export process using `torch.export` + in a way that ensures the model can be further lowered and run efficiently in `ExecuTorch`. + """ + + def __init__(self, model: PreTrainedModel): + """ + Initializes the wrapper module with the pretrained model. + + Args: + model (`PreTrainedModel`): The pretrained model to wrap. The model must have caching + enabled and use a 'static' caching implementation. + + Raises: + AssertionError: If the pretrained model does not have caching enabled or if it does + not use a 'static' caching implementation in `model.generation_config`. + """ + super().__init__() + + # Sanity checks + if model.generation_config is None: + raise AssertionError( + "The model must have a generation config to be exported with static caching. " + "Please set `generation_config`." + ) + + if not model.generation_config.use_cache: + raise AssertionError( + "The model must have caching enabled to be exported with static caching. " + "Please set `generation_config.use_cache=True`." + ) + + if model.generation_config.cache_implementation != "static": + raise AssertionError( + "The model must use a 'static' caching implementation to be exported with static caching. " + "Please set `generation_config.cache_implementation='static'`." + ) + + self.model = model + self.static_cache = StaticCache( + config=self.model.config, + max_batch_size=self.model.generation_config.cache_config.batch_size, + max_cache_len=self.model.generation_config.cache_config.max_cache_len, + device=self.model.generation_config.cache_config.device, + dtype=self.model.dtype, + ) + for i in range(len(self.static_cache.key_cache)): + self.register_buffer(f"key_cache_{i}", self.static_cache.key_cache[i], persistent=False) + self.register_buffer(f"value_cache_{i}", self.static_cache.value_cache[i], persistent=False) + + self.is_causal = any("CausalLM" in arch for arch in self.model.config.architectures) + if self.is_causal: + causal_mask = torch.tril( + torch.ones( + self.static_cache.max_cache_len, + self.static_cache.max_cache_len, + dtype=torch.bool, + ) + ) + self.register_buffer("mask", causal_mask, persistent=False) + + def forward(self, input_ids: torch.Tensor, cache_position: torch.Tensor): + """ + Forward pass of the module, which is compatible with the ExecuTorch runtime. + + Args: + input_ids (`torch.Tensor`): Tensor representing current input token id to the module. + cache_position (`torch.Tensor`): Tensor representing current input position in the cache. + + Returns: + torch.Tensor: Logits output from the model. + + This forward adapter serves two primary purposes: + + 1. **Making the Model `torch.export`-Compatible**: + The adapter hides unsupported objects, such as the `Cache`, from the graph inputs and outputs, + enabling the model to be exportable using `torch.export` without encountering issues. + + 2. **Ensuring Compatibility with `ExecuTorch` runtime**: + The adapter matches the model's forward signature with that in `executorch/extension/llm/runner`, + ensuring that the exported model can be executed in `ExecuTorch` out-of-the-box. + """ + _, seqlen = input_ids.shape + attn_mask = self.mask[cache_position, :seqlen] if self.is_causal else None + position_ids = cache_position.unsqueeze(0) + past_key_values = self.static_cache + + outs = self.model( + input_ids=input_ids, + attention_mask=attn_mask, + position_ids=position_ids, + cache_position=cache_position, + past_key_values=past_key_values, + use_cache=True, + ) + return outs.logits + + @staticmethod + def generate( + exported_program: torch.export.ExportedProgram, prompt_token_ids: torch.Tensor, max_new_tokens: int + ) -> torch.Tensor: + """ + Generate a sequence of tokens using an exported program. + + This util function is designed to test exported models by simulating the generation process. + It processes the input prompt tokens sequentially (no parallel prefill). + This generate function is not intended to replace the original `generate` method, and the support + for leveraging the original `generate` is potentially planed! + + Args: + exported_program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`. + prompt_token_ids (`torch.Tensor`): Tensor representing the input prompt token IDs. + max_new_tokens (`int`): Maximum number of new tokens to generate. Note that the total generation + length is limited by both `max_new_tokens` and the model's cache size. + + Returns: + torch.Tensor: A tensor containing the generated sequence of token IDs, including the original prompt tokens. + """ + prompt_token_len = prompt_token_ids.shape[-1] + max_generation_length = prompt_token_len + max_new_tokens + for buffer_name, buffer in exported_program.named_buffers(): + if buffer_name.startswith("key_cache"): + max_cache_len = buffer.shape[2] + max_generation_length = min(max_generation_length, max_cache_len) + break + + response_tokens = [] + for input_pos in range(min(max_generation_length, prompt_token_len)): + result = exported_program.module().forward( + input_ids=prompt_token_ids[:, input_pos : input_pos + 1], + cache_position=torch.tensor([input_pos], dtype=torch.long), + ) + response_tokens.append(prompt_token_ids[0][input_pos].item()) + + current_token = torch.argmax(result[:, -1, :], dim=-1).item() + response_tokens.append(current_token) + + while len(response_tokens) < max_generation_length: + result = exported_program.module().forward( + input_ids=torch.tensor([[current_token]], dtype=torch.long), + cache_position=torch.tensor([len(response_tokens)], dtype=torch.long), + ) + current_token = torch.argmax(result[:, -1, :], dim=-1).item() + response_tokens.append(current_token) + + return torch.tensor([response_tokens], dtype=torch.long) + + +def convert_and_export_with_cache( + model: PreTrainedModel, + example_input_ids: Optional[torch.Tensor] = None, + example_cache_position: Optional[torch.Tensor] = None, +): + """ + Convert a `PreTrainedModel` into an exportable module and export it using `torch.export`, + ensuring the exported model is compatible with `ExecuTorch`. + + Args: + model (`PreTrainedModel`): The pretrained model to be exported. + example_input_ids (`torch.Tensor`): Example input token id used by `torch.export`. + example_cache_position (`torch.Tensor`): Example current cache position used by `torch.export`. + + Returns: + Exported program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`. + """ + if not is_torch_greater_or_equal_than_2_3: + raise ImportError("torch >= 2.3 is required.") + + import torch.export._trace + + with torch.no_grad(): + # TODO: The default inputs only work for text models. We need to add support for vision/audio models. + example_input_ids = ( + example_input_ids if example_input_ids is not None else torch.tensor([[1]], dtype=torch.long) + ) + example_cache_position = ( + example_cache_position if example_cache_position is not None else torch.tensor([0], dtype=torch.long) + ) + + if is_torch_greater_or_equal("2.5.0"): + exported_program = torch.export.export( + TorchExportableModuleWithStaticCache(model), + args=(example_input_ids,), + kwargs={"cache_position": example_cache_position}, + strict=True, + ) + else: + # We have to keep this path for BC. + # + # Due to issue https://github.com/pytorch/pytorch/issues/128394, we need to switch to use an internal + # export API and pre_dispatch=False. Switch to use the public API once the issue is included in 2.5 release. + exported_program = torch.export._trace._export( + TorchExportableModuleWithStaticCache(model), + args=(example_input_ids,), + kwargs={"cache_position": example_cache_position}, + pre_dispatch=False, + strict=True, + ) + return exported_program + + +class Seq2SeqLMEncoderExportableModule(torch.nn.Module): + """ + A wrapper module designed to make a Seq2Seq LM encoder exportable with `torch.export`. + This module ensures that the exported encoder model is compatible with ExecuTorch. + """ + + def __init__(self, encoder_model): + super().__init__() + self.encoder = encoder_model + + def forward(self, input_ids): + return self.encoder(input_ids=input_ids).last_hidden_state + + +class Seq2SeqLMDecoderExportableModuleWithStaticCache(torch.nn.Module): + """ + A wrapper module designed to make a Seq2Seq LM decoder exportable with `torch.export`, + specifically for use with static caching. This module ensures the exported decoder + is compatible with ExecuTorch. + """ + + def __init__(self, model, max_static_cache_length, batch_size): + super().__init__() + + # Get the decoder component + self.decoder = model.get_decoder() + self.lm_head = model.lm_head + self.config = model.config + + # Initialize static cache + self.static_cache = StaticCache( + config=self.config, + max_batch_size=batch_size, + max_cache_len=max_static_cache_length, + device="cpu", + dtype=torch.float32, + ) + + # Register cache buffers to make them exportable + for i in range(len(self.static_cache.key_cache)): + self.register_buffer(f"key_cache_{i}", self.static_cache.key_cache[i], persistent=False) + self.register_buffer(f"value_cache_{i}", self.static_cache.value_cache[i], persistent=False) + + def forward(self, decoder_input_ids, encoder_hidden_states, cache_position): + # Get outputs from decoder + outputs = self.decoder( + input_ids=decoder_input_ids, + encoder_hidden_states=encoder_hidden_states, + past_key_values=self.static_cache, + use_cache=True, + cache_position=cache_position, + ) + + # Apply language model head + lm_logits = self.lm_head(outputs[0]) + + return lm_logits + + +class Seq2SeqLMExportableModule(torch.nn.Module): + def __init__( + self, model, batch_size=1, max_hidden_seq_length=4096, cache_implementation="static", max_cache_length=1024 + ): + super().__init__() + + self.full_model = model + self.encoder = model.get_encoder() + self.config = model.config + self.max_hidden_seq_length = max_hidden_seq_length + self.generation_config = GenerationConfig( + use_cache=True, + max_length=max_cache_length, + cache_implementation=cache_implementation, + cache_config={ + "batch_size": batch_size, + "max_cache_len": max_cache_length, + }, + ) + self.exported_encoder = None + self.exported_decoder = None + + def _export_encoder(self, encoder_input_ids): + wrapped_encoder = Seq2SeqLMEncoderExportableModule(self.encoder).to("cpu").eval() + + # Define dynamic sequence length for encoder + seq_len_dim = torch.export.Dim("encoder_seq_length", max=self.max_hidden_seq_length) + + # Export the encoder + with torch.no_grad(): + exported_encoder = torch.export.export( + wrapped_encoder, (encoder_input_ids,), dynamic_shapes={"input_ids": {1: seq_len_dim}}, strict=True + ) + + return exported_encoder + + def _export_decoder(self, decoder_input_ids, encoder_hidden_states, cache_position): + wrapped_decoder = ( + Seq2SeqLMDecoderExportableModuleWithStaticCache( + model=self.full_model, + max_static_cache_length=self.generation_config.cache_config.max_cache_len, + batch_size=self.generation_config.cache_config.batch_size, + ) + .to("cpu") + .eval() + ) + + # Define dynamic dimension for encoder output sequence length + encoder_seq_len_dim = torch.export.Dim("encoder_hidden_seq_length", max=self.max_hidden_seq_length) + + # Export the decoder + with torch.no_grad(): + exported_decoder = torch.export.export( + wrapped_decoder, + (decoder_input_ids, encoder_hidden_states, cache_position), + dynamic_shapes={ + "decoder_input_ids": None, + "encoder_hidden_states": {1: encoder_seq_len_dim}, + "cache_position": None, + }, + strict=True, + ) + + return exported_decoder + + def export(self, encoder_input_ids=None, decoder_input_ids=None, encoder_hidden_states=None, cache_position=None): + example_encoder_input_ids = ( + encoder_input_ids if encoder_input_ids is not None else torch.ones((1, 10), dtype=torch.long) + ) + example_decoder_input_ids = ( + decoder_input_ids if decoder_input_ids is not None else torch.tensor([[0]], dtype=torch.long) + ) # Start token + example_cache_position = cache_position if cache_position is not None else torch.tensor([0], dtype=torch.long) + example_encoder_hidden_states = ( + encoder_hidden_states + if encoder_hidden_states is not None + else torch.zeros( + (self.generation_config.cache_config.batch_size, 10, self.config.d_model), dtype=torch.float32 + ) + ) + self.exported_encoder = self._export_encoder(example_encoder_input_ids) + self.exported_decoder = self._export_decoder( + example_decoder_input_ids, example_encoder_hidden_states, example_cache_position + ) + + # Return self to allow chaining + return self + + def generate(self, prompt_token_ids, max_new_tokens): + with torch.no_grad(): + # Run encoder + encoder_output = self.exported_encoder.module()(prompt_token_ids) + + # Initialize with start token (0 for T5) + decoder_input_ids = torch.tensor([[0]], dtype=torch.long) + generated_ids = [0] + + # Generate tokens one by one + for i in range(max_new_tokens - 1): + # Run decoder for next token prediction + logits = self.exported_decoder.module()( + decoder_input_ids, encoder_output, torch.tensor([i], dtype=torch.long) + ) + + # Get next token + next_token = torch.argmax(logits[:, -1, :], dim=-1).item() + generated_ids.append(next_token) + + # Update input for next iteration + decoder_input_ids = torch.tensor([[next_token]], dtype=torch.long) + + # Check if EOS token + if next_token == self.config.eos_token_id: + break + + return generated_ids diff --git a/lib/python3.12/site-packages/transformers/integrations/fbgemm_fp8.py b/lib/python3.12/site-packages/transformers/integrations/fbgemm_fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..5cca37f51510eeb21b94133a152225c27a080d7f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/fbgemm_fp8.py @@ -0,0 +1,287 @@ +# 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. + +from ..activations import ACT2FN +from ..utils import is_accelerate_available, is_fbgemm_gpu_available, is_torch_available, logging + + +if is_torch_available(): + import torch + from torch import nn + +if is_accelerate_available(): + from accelerate import init_empty_weights + +if is_fbgemm_gpu_available(): + import fbgemm_gpu.experimental.gen_ai # noqa: F401 + +logger = logging.get_logger(__name__) + + +class FbgemmFp8Linear(torch.nn.Linear): + def __init__(self, in_features, out_features, bias, weight_dtype=torch.float32): + super().__init__(in_features, out_features, bias) + self.in_features = in_features + self.out_features = out_features + + self.weight = torch.nn.Parameter(torch.zeros((out_features, in_features), dtype=torch.float8_e4m3fn)) + self.weight_scale = torch.nn.Parameter(torch.zeros((out_features, 1), dtype=weight_dtype)) + self.register_buffer("input_scale_ub", torch.zeros([1], dtype=torch.float), persistent=False) + + if bias: + self.bias = torch.nn.Parameter(torch.zeros((self.out_features), dtype=weight_dtype)) + else: + self.bias = None + + def forward(self, x): + # quantize_fp8_per_row will squash the leading dimensions, so save the desired shape here + output_shape = (*x.shape[:-1], -1) + # x_quantized and x_scale are not necessarily on the same device as x, this is an issue. + # https://github.com/pytorch/FBGEMM/blob/e08af8539c391437f447173863df0f3f6f6f1855/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cu#L1237C3-L1237C45 + x_quantized, x_scale = torch.ops.fbgemm.quantize_fp8_per_row( + x.view(-1, x.shape[-1]), scale_ub=self.input_scale_ub + ) + # moving x_quantized, x_scale here creates glibberish output ... However, if we move the output, it works + # x_quantized, x_scale = x_quantized.to(x.device), x_scale.to(x.device) + + # The computation still happens on the device where self.weight is even if x_quantized is not on the same device as self.weight + weight_scale_float32 = self.weight_scale.to(torch.float32) + output = torch.ops.fbgemm.f8f8bf16_rowwise( + x_quantized, self.weight, x_scale, weight_scale_float32, use_fast_accum=True + ) + output = output + self.bias if self.bias is not None else output + # Hacky for now, we have the output to the device of x + output = output.to(x.device) + output = output.reshape(output_shape) + del x_quantized, x_scale + return output + + +class FbgemmFp8Llama4TextExperts(nn.Module): + def __init__(self, config, dtype=torch.float32): + super().__init__() + self.num_experts = config.num_local_experts + self.intermediate_size = config.intermediate_size + self.hidden_size = config.hidden_size + self.expert_dim = self.intermediate_size + self.act_fn = ACT2FN[config.hidden_act] + # Register FP8 buffers for gate_up_proj + self.gate_up_proj = torch.nn.Parameter( + torch.zeros((self.num_experts, self.hidden_size, 2 * self.expert_dim), dtype=torch.float8_e4m3fn) + ) + self.gate_up_proj_scale = torch.nn.Parameter( + torch.zeros((self.num_experts, 1, self.expert_dim * 2), dtype=torch.float32) + ) + # Register FP8 buffers for down_proj + self.down_proj = torch.nn.Parameter( + torch.zeros((self.num_experts, self.expert_dim, self.hidden_size), dtype=torch.float8_e4m3fn) + ) + self.down_proj_scale = torch.nn.Parameter( + torch.zeros((self.num_experts, self.hidden_size, 1), dtype=torch.float32) + ) + # Register input scale upper bound + self.register_buffer("input_scale_ub", torch.zeros([1], dtype=torch.float), persistent=False) + + def forward(self, hidden_states): + """ + Args: + hidden_states (torch.Tensor): (batch_size * token_num, hidden_size) + Returns: + torch.Tensor: (batch_size * token_num, hidden_size) + """ + # Reshape hidden states for expert computation + hidden_states = hidden_states.view(self.num_experts, -1, self.hidden_size) + num_tokens = None + + # Pre-allocate tensor for all expert outputs with same shape as hidden_states + next_states = torch.empty_like(hidden_states) + + for i in range(self.num_experts): + # Extract expert's hidden states + expert_hidden = hidden_states[i] + expert_hidden_reshaped = expert_hidden.reshape(-1, self.hidden_size) + # Quantize for this expert + expert_quantized, expert_scale = torch.ops.fbgemm.quantize_fp8_per_row( + expert_hidden_reshaped, num_tokens, self.input_scale_ub + ) + sharded_expert_dim = self.gate_up_proj.shape[-1] // 2 + gate_up_proj_scale_float32 = self.gate_up_proj_scale.to(torch.float32) + + gate = torch.ops.fbgemm.f8f8bf16_rowwise( + expert_quantized, + self.gate_up_proj[i].transpose(0, 1)[:sharded_expert_dim].contiguous(), + expert_scale, + gate_up_proj_scale_float32[i][0][:sharded_expert_dim].view(-1, 1).contiguous(), + use_fast_accum=True, + ) + + up = torch.ops.fbgemm.f8f8bf16_rowwise( + expert_quantized, + self.gate_up_proj[i].transpose(0, 1)[sharded_expert_dim:].contiguous(), + expert_scale, + gate_up_proj_scale_float32[i][0][sharded_expert_dim:].view(-1, 1).contiguous(), + use_fast_accum=True, + ) + + activated = up * self.act_fn(gate) + + activated_quantized, activated_scale = torch.ops.fbgemm.quantize_fp8_per_row( + activated, num_tokens, self.input_scale_ub + ) + + down_proj_scale_float32 = self.down_proj_scale.to(torch.float32) + expert_output = torch.ops.fbgemm.f8f8bf16_rowwise( + activated_quantized, + self.down_proj[i].transpose(0, 1).contiguous(), + activated_scale, + down_proj_scale_float32[i].view(-1, 1).contiguous(), + use_fast_accum=True, + ) + + next_states[i] = expert_output + next_states = next_states.to(hidden_states.device) + return next_states.view(-1, self.hidden_size) + + +def _replace_with_fbgemm_fp8_linear( + model, + modules_to_not_convert=None, + current_key_name=None, + quantization_config=None, + has_been_replaced=False, + pre_quantized=False, + config=None, + tp_plan=None, +): + """ + Private method that wraps the recursion for module replacement. + + Returns the converted model and a boolean that indicates if the conversion has been successfull or not. + """ + + import re + + if current_key_name is None: + current_key_name = [] + + for name, module in model.named_children(): + current_key_name.append(name) + + if (isinstance(module, nn.Linear)) and name not in modules_to_not_convert: + # Check if the current key is not in the `modules_to_not_convert` + current_key_name_str = ".".join(current_key_name) + if not any( + (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert + ): + with init_empty_weights(include_buffers=True): + in_features = module.in_features + out_features = module.out_features + model._modules[name] = FbgemmFp8Linear( + in_features, + out_features, + module.bias is not None, + ) + has_been_replaced = True + + # Force requires grad to False to avoid unexpected errors + model._modules[name].requires_grad_(False) + # set non persistant buffer outside of init_empty_weights + model._modules[name].input_scale_ub = torch.tensor( + [quantization_config.activation_scale_ub], + dtype=torch.float, + ) + if module.__class__.__name__ == "Llama4TextExperts" and name not in modules_to_not_convert: + current_key_name_str = ".".join(current_key_name) + if not any( + (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert + ): + with init_empty_weights(include_buffers=True): + tp_plan[re.sub(r"\d+", "*", current_key_name_str + ".gate_up_proj_scale")] = tp_plan[ + re.sub(r"\d+", "*", current_key_name_str + ".gate_up_proj") + ] + tp_plan[re.sub(r"\d+", "*", current_key_name_str + ".down_proj_scale")] = None + model._modules[name] = FbgemmFp8Llama4TextExperts( + config.text_config, + ) + model._modules[name].input_scale_ub = torch.tensor( + [quantization_config.activation_scale_ub], dtype=torch.float + ) + + if len(list(module.children())) > 0: + _, has_been_replaced = _replace_with_fbgemm_fp8_linear( + module, + modules_to_not_convert, + current_key_name, + quantization_config, + has_been_replaced=has_been_replaced, + pre_quantized=pre_quantized, + config=config, + tp_plan=tp_plan, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced + + +def replace_with_fbgemm_fp8_linear( + model, + modules_to_not_convert=None, + current_key_name=None, + quantization_config=None, + pre_quantized=False, + config=None, + tp_plan=None, +): + """ + A helper function to replace all `torch.nn.Linear` modules by `FbgemmFp8Linear` modules. + This will enable running your models using high performance fp8 kernel from FBGEMM library. + + The function will be run recursively and replace all `torch.nn.Linear` modules except for the `lm_head` that should + be kept as a `torch.nn.Linear` module. The replacement is done under `init_empty_weights` context manager so no + CPU/GPU memory is required to run this function. Each weight will be quantized along the channel. + + Parameters: + model (`torch.nn.Module`): + Input model or `torch.nn.Module` as the function is run recursively. + modules_to_not_convert (`List[`str`]`, *optional*, defaults to `["lm_head"]`): + Names of the modules to not convert in `FP8Linear`. In practice we keep the `lm_head` in full precision + for numerical stability reasons. + current_key_name (`List[`str`]`, *optional*): + An array to track the current key of the recursion. This is used to check whether the current key (part of + it) is not in the list of modules to not convert (for instances modules that are offloaded to `cpu` or + `disk`). + """ + + modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert + + if quantization_config.modules_to_not_convert is not None: + modules_to_not_convert.extend(quantization_config.modules_to_not_convert) + modules_to_not_convert = list(set(modules_to_not_convert)) + model, has_been_replaced = _replace_with_fbgemm_fp8_linear( + model, + modules_to_not_convert, + current_key_name, + quantization_config, + pre_quantized=pre_quantized, + config=config, + tp_plan=tp_plan, + ) + if not has_been_replaced: + logger.warning( + "You are loading your model using FP8 quantization but no linear modules were found in your model." + " Please double check your model architecture, or submit an issue on github if you think this is" + " a bug." + ) + + return model diff --git a/lib/python3.12/site-packages/transformers/integrations/finegrained_fp8.py b/lib/python3.12/site-packages/transformers/integrations/finegrained_fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..d4e472a990dcfc1226d03a12bdeb7f3bd993dc7f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/finegrained_fp8.py @@ -0,0 +1,426 @@ +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. 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. + +from typing import List, Optional, Tuple + +from ..utils import is_accelerate_available, is_torch_available, logging + + +if is_torch_available(): + import torch + import torch.nn as nn + import triton + import triton.language as tl + from torch.nn import functional as F + +if is_accelerate_available(): + from accelerate import init_empty_weights + + +logger = logging.get_logger(__name__) + + +# Copied from https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py +@triton.jit +def act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + x = tl.load(x_ptr + offs).to(tl.float32) + s = tl.max(tl.abs(x)) / 448.0 + y = x / s + y = y.to(y_ptr.dtype.element_ty) + tl.store(y_ptr + offs, y) + tl.store(s_ptr + pid, s) + + +def act_quant(x: torch.Tensor, block_size: int = 128) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.is_contiguous() + assert x.shape[-1] % block_size == 0 + y = torch.empty_like(x, dtype=torch.float8_e4m3fn) + s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32) + + def grid(meta): + return (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),) + + act_quant_kernel[grid](x, y, s, BLOCK_SIZE=block_size) + return y, s + + +# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py +@triton.jit +def _w8a8_block_fp8_matmul( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N, + K, + # Block size for block-wise quantization + group_n, + group_k, + # Stride for inputs and output + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_As_m, + stride_As_k, + stride_Bs_k, + stride_Bs_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Triton-accelerated function used to perform linear operations (dot + product) on input tensors `A` and `B` with block-wise quantization, and + store the result in output tensor `C`. + """ + + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +def w8a8_block_fp8_matmul_triton( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """This function performs matrix multiplication with block-wise + quantization. + It takes two input tensors `A` and `B` with scales `As` and `Bs`. + The output is returned in the specified `output_dtype`. + Args: + A: The input tensor, e.g., activation. + B: The input tensor, e.g., weight. + As: The per-token-group quantization scale for `A`. + Bs: The per-block quantization scale for `B`. + block_size: The block size for per-block quantization. It should + be 2-dim, e.g., [128, 128]. + output_dytpe: The dtype of the returned tensor. + Returns: + torch.Tensor: The result of matmul. + """ + assert len(block_size) == 2 + block_n, block_k = block_size[0], block_size[1] + + assert A.shape[-1] == B.shape[-1] + assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() + assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] + M = A.numel() // A.shape[-1] + + assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 + N, K = B.shape + assert triton.cdiv(N, block_n) == Bs.shape[0] + assert triton.cdiv(K, block_k) == Bs.shape[1] + + C_shape = A.shape[:-1] + (N,) + C = A.new_empty(C_shape, dtype=output_dtype) + + BLOCK_SIZE_M = 128 + if M < BLOCK_SIZE_M: + BLOCK_SIZE_M = triton.next_power_of_2(M) + BLOCK_SIZE_M = max(BLOCK_SIZE_M, 16) + BLOCK_SIZE_K = block_k + assert block_k % BLOCK_SIZE_K == 0 + BLOCK_SIZE_N = block_n + + def grid(META): + return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),) + + _w8a8_block_fp8_matmul[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + GROUP_SIZE_M=8, + ) + + return C + + +# Python version of the above triton function, it's much slower than the triton version, for testing +@torch.compile +def w8a8_block_fp8_matmul_compile( + input_q: torch.Tensor, # [batch, seq_len, hidden_dim] + weight_q: torch.Tensor, # [out_features, hidden_dim] + input_scale: torch.Tensor, # [batch * seq_len, num_input_groups] + weight_scale: torch.Tensor, # [num_weight_blocks_m, num_weight_blocks_n] + block_size: Optional[Tuple[int, int]] = None, # (M=128, N=128) for weights for example + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Performs blocked matrix multiplication with FP8 quantized matrices. + + Args: + input_q: Quantized input tensor with 1x128 block quantization + weight_q: Quantized weight tensor with 128x128 block quantization + input_scale: Scaling factors for input blocks + weight_scale: Scaling factors for weight blocks + block_size: Tuple of (M, N) for weight block dimensions + output_dtype: Desired output dtype + """ + batch_size, seq_len, hidden_dim = input_q.shape if input_q.ndim == 3 else (1, input_q.shape[0], input_q.shape[1]) + out_features = weight_q.shape[0] + + # Reshape input for batched matmul + input_reshaped = input_q.view(-1, hidden_dim) # [batch*seq_len, hidden_dim] + input_scale_reshaped = input_scale.view(input_scale.shape[0], -1) # [batch*seq_len, 1] + # Calculate number of blocks + num_weight_blocks_m = out_features // block_size[0] + num_weight_blocks_n = hidden_dim // block_size[1] + + output = torch.zeros((batch_size * seq_len, out_features), dtype=torch.float32, device=input_q.device) + + for i in range(num_weight_blocks_m): + m_start = i * block_size[0] + m_end = m_start + block_size[0] + + for j in range(num_weight_blocks_n): + n_start = j * block_size[1] + n_end = n_start + block_size[1] + + # Extract current blocks + input_block = input_reshaped[:, n_start:n_end] + weight_block = weight_q[m_start:m_end, n_start:n_end] + + # Get corresponding scales + curr_input_scale = input_scale_reshaped[:, j : j + 1] # [batch*seq_len, 1] + curr_weight_scale = weight_scale[i, j] # scalar + + block_result = ( + torch._scaled_mm( + input_block, + weight_block.t(), + scale_a=torch.tensor(1, dtype=torch.float32, device=input_q.device), + scale_b=curr_weight_scale, + out_dtype=output_dtype, + ) + * curr_input_scale + ) + + output[:, m_start:m_end] += block_result + + output = output.view(batch_size, seq_len, out_features) + + return output.to(output_dtype) + + +class FP8Linear(nn.Linear): + dtype = torch.float8_e4m3fn + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = False, + dtype=None, + block_size: Optional[Tuple[int, int]] = None, + device=None, + activation_scheme="dynamic", + ): + super().__init__(in_features, out_features) + self.in_features = in_features + self.out_features = out_features + + self.weight = torch.nn.Parameter(torch.empty(out_features, in_features, dtype=FP8Linear.dtype, device=device)) + + if self.weight.element_size() == 1: + scale_out_features = (out_features + block_size[0] - 1) // block_size[0] + scale_in_features = (in_features + block_size[1] - 1) // block_size[1] + self.weight_scale_inv = nn.Parameter( + torch.empty(scale_out_features, scale_in_features, dtype=torch.float32, device=device) + ) + else: + self.register_parameter("weight_scale_inv", None) + + self.block_size = block_size + + self.activation_scheme = activation_scheme + + if bias: + self.bias = nn.Parameter(torch.empty(self.out_features)) + else: + self.register_parameter("bias", None) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + if self.weight.element_size() > 1: + return F.linear(input, self.weight, self.bias) + else: + # Context manager used to switch among the available cuda devices + # with torch.cuda.device(input.device): + qinput, scale = act_quant(input, self.block_size[1]) + # Blocks the CPU until all CUDA operations on the specified device are complete. It is used to ensure that the results of the + # preceding operations are ready before proceeding + # torch.cuda.synchronize(device=self.weight.device) + with torch.cuda.device(input.device): + output = w8a8_block_fp8_matmul_triton( + qinput, + self.weight, + scale, + self.weight_scale_inv, + self.block_size, + output_dtype=input.dtype, + ) + torch.cuda.synchronize() + if self.bias is not None: + output = output + self.bias + return output.to(dtype=input.dtype) + + +def _replace_with_fp8_linear( + model, + tp_plan=None, + modules_to_not_convert=None, + current_key_name=None, + quantization_config=None, + has_been_replaced=False, +): + """Replace Linear layers with FP8Linear.""" + if current_key_name is None: + current_key_name = [] + + for name, module in model.named_children(): + current_key_name.append(name) + + if isinstance(module, nn.Linear) and name not in (modules_to_not_convert or []): + current_key_name_str = ".".join(current_key_name) + if not any(key in current_key_name_str for key in (modules_to_not_convert or [])): + with init_empty_weights(): + model._modules[name] = FP8Linear( + in_features=module.in_features, + out_features=module.out_features, + bias=module.bias is not None, + device=module.weight.device, + dtype=module.weight.dtype, + activation_scheme=quantization_config.activation_scheme, + block_size=quantization_config.weight_block_size, + ) + has_been_replaced = True + # when changing a layer the TP PLAN for that layer should be updated. TODO + + if len(list(module.children())) > 0: + _, has_been_replaced = _replace_with_fp8_linear( + module, + tp_plan, + modules_to_not_convert, + current_key_name, + quantization_config, + has_been_replaced=has_been_replaced, + ) + + current_key_name.pop(-1) + + return model, has_been_replaced + + +def replace_with_fp8_linear( + model, + modules_to_not_convert=None, + quantization_config=None, +): + """Helper function to replace model layers with FP8 versions.""" + modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert + + if quantization_config.modules_to_not_convert is not None: + modules_to_not_convert.extend(quantization_config.modules_to_not_convert) + modules_to_not_convert = list(set(modules_to_not_convert)) + model, has_been_replaced = _replace_with_fp8_linear( + model, + tp_plan=model._tp_plan, + modules_to_not_convert=modules_to_not_convert, + quantization_config=quantization_config, + ) + + if not has_been_replaced: + logger.warning( + "You are loading your model using fp8 but no linear modules were found in your model." + " Please double check your model architecture." + ) + + return model diff --git a/lib/python3.12/site-packages/transformers/integrations/flash_attention.py b/lib/python3.12/site-packages/transformers/integrations/flash_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..a78166ed040b620c6e24a7e2c64c06c96f2d89d9 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/flash_attention.py @@ -0,0 +1,65 @@ +from typing import Optional, Tuple + +import torch + +from ..modeling_flash_attention_utils import _flash_attention_forward, flash_attn_supports_top_left_mask + + +_use_top_left_mask = flash_attn_supports_top_left_mask() + + +def flash_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + dropout: float = 0.0, + scaling: Optional[float] = None, + sliding_window: Optional[int] = None, + softcap: Optional[float] = None, + **kwargs, +) -> Tuple[torch.Tensor, None]: + # This is before the transpose + seq_len = query.shape[2] + + # FA2 uses non-transposed inputs + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in the correct dtype just to be sure everything works as expected. + # This might slowdown training & inference so it is recommended to not cast the LayerNorms + # in fp32. (usually our RMSNorm modules handle it correctly) + target_dtype = None + if query.dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(module.config, "_pre_quantization_dtype"): + target_dtype = module.config._pre_quantization_dtype + else: + target_dtype = next(layer for layer in module.modules() if isinstance(layer, torch.nn.Linear)).weight.dtype + + # FA2 always relies on the value set in the module, so remove it if present in kwargs to avoid passing it twice + kwargs.pop("is_causal", None) + + attn_output = _flash_attention_forward( + query, + key, + value, + attention_mask, + query_length=seq_len, + is_causal=module.is_causal, + dropout=dropout, + softmax_scale=scaling, + sliding_window=sliding_window, + softcap=softcap, + use_top_left_mask=_use_top_left_mask, + target_dtype=target_dtype, + **kwargs, + ) + + return attn_output, None diff --git a/lib/python3.12/site-packages/transformers/integrations/flex_attention.py b/lib/python3.12/site-packages/transformers/integrations/flex_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..ec4ebcb22d284196126a5af04eaab2213d1bb4ca --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/flex_attention.py @@ -0,0 +1,253 @@ +""" +Partially inspired by torchtune's flex attention implementation + +Citation: +@software{torchtune, + title = {torchtune: PyTorch's finetuning library}, + author = {torchtune maintainers and contributors}, + url = {https//github.com/pytorch/torchtune}, + license = {BSD-3-Clause}, + month = apr, + year = {2024} +} +""" +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# 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 typing import Optional, Tuple, Union + +import torch + +from ..utils import is_torch_flex_attn_available +from ..utils.import_utils import _torch_version + + +if is_torch_flex_attn_available(): + from torch.nn.attention.flex_attention import BlockMask, flex_attention + from torch.nn.attention.flex_attention import ( + create_block_mask as create_block_causal_mask_flex, + ) + + +class WrappedFlexAttention: + """ + We are doing a singleton class so that flex attention is compiled once when it's first called. + """ + + _instance = None + _is_flex_compiled = False + _compiled_flex_attention = None + + def __new__(cls, *args, **kwargs): + if cls._instance is None: + # Create a new instance if one doesn't already exist + cls._instance = super().__new__(cls) + return cls._instance + + @torch.compiler.disable(recursive=False) + def __init__(self, training): + """ + Initialize or update the singleton instance. + """ + if not self._is_flex_compiled: + # In PyTorch 2.6.0, there's a known issue with flex attention compilation which may + # cause errors. The suggested fix is to compile with "max-autotune-no-cudagraphs" + # see https://github.com/pytorch/pytorch/issues/146260 for training + if _torch_version == "2.6.0" and training: + self._compiled_flex_attention = torch.compile( + flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs" + ) + else: + self._compiled_flex_attention = torch.compile(flex_attention, dynamic=False) + self._is_flex_compiled = True + + def __call__(self): + return self._compiled_flex_attention + + +Offset = Union[torch.Tensor, int] + + +def make_flex_block_causal_mask( + attention_mask_2d: torch.Tensor, + attention_chunk_size: Optional[int] = None, + query_length=None, + key_length=None, + offsets: Optional[Tuple[Offset, Offset]] = None, +) -> "BlockMask": + """ + Create a block causal document mask for a batch of sequences, both packed and unpacked. + Create Block causal logic and passing it into :func:`torch.nn.attention.flex_attention.create_block_mask`. + The resultant BlockMask is a compressed representation of the full block causal + mask. BlockMask is essential for performant computation of flex attention. + See: https://pytorch.org/blog/flexattention/ + + Args: + attention_mask_2d (torch.Tensor): Attention mask for packed and padded sequences + of shape (batch_size, total_seq_len). e.g. + + For unpacked sequence: + [[1, 1, 1, 1, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0]] + + For packed sequence: + [[1, 1, 1, 2, 2, 2, 0], + [1, 1, 2, 2, 2, 3, 3]] + + Returns: + BlockMask + """ + batch_size, total_seq_len = attention_mask_2d.shape + if not key_length: + key_length = total_seq_len + if not query_length: + query_length = total_seq_len + attention_mask_2d = torch.nn.functional.pad(attention_mask_2d, value=0, pad=(0, key_length)) + device = attention_mask_2d.device + document_ids = attention_mask_2d.clone() + + if attention_chunk_size is not None: + # we create an arange, then we just // by chunk size to get [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3] + document_ids = (document_ids.fill_(1).cumsum(-1) - 1) // (attention_chunk_size) + + # Instead of passing a tensor mask, flex attention requires a mask_mod function + # that determines which elements of QK^T should be included in the attention + # computation prior to the softmax. For sample packing, we need both the + # logic for both causal mask and document mask. See PyTorch's official + # blog post for more details: https://pytorch.org/blog/flexattention/#mask-mods + def causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx): + """ + Defines the logic of a block causal mask by combining both a standard causal mask + and a block diagonal document mask. + + See :func:`~torchtune.modules.attention_utils.create_block_causal_mask` + for an illustration. + """ + causal_mask = q_idx >= kv_idx # not valid when decoding + document_mask = document_ids[batch_idx, q_idx] == document_ids[batch_idx, kv_idx] + padding_mask = attention_mask_2d[batch_idx, q_idx] > 0 + final_mask = causal_mask & padding_mask & document_mask + return final_mask + + if offsets is not None: + q_offset = offsets[0] + kv_offset = offsets[1] + + def mask_mod(batch_idx, head_idx, q_idx, kv_idx): + offset_q = q_idx + q_offset + offset_kv = kv_idx + kv_offset + return causal_mask_mod(batch_idx, head_idx, offset_q, offset_kv) + else: + mask_mod = causal_mask_mod + return create_block_causal_mask_flex( + mask_mod=mask_mod, + B=batch_size, + H=None, # attention head + Q_LEN=query_length, + KV_LEN=key_length, + device=device, + _compile=True, + ) + + +@torch.compiler.disable(recursive=False) +def compile_friendly_flex_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + training=False, + **kwargs, +) -> torch.Tensor: + # First call initialise singleton wrapper object, second call invokes the object method to return compiled flex attention + flex_attention_compiled = WrappedFlexAttention(training)() + return flex_attention_compiled( + query, + key, + value, + **kwargs, + ) + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def flex_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Union[torch.Tensor, "BlockMask"], + scaling: Optional[float] = None, + softcap: Optional[float] = None, + head_mask: Optional[torch.Tensor] = None, + **kwargs, +) -> Tuple[torch.Tensor, torch.Tensor]: + block_mask = None + causal_mask = None + if isinstance(attention_mask, BlockMask): + block_mask = attention_mask + else: + causal_mask = attention_mask + + if causal_mask is not None: + causal_mask = causal_mask[:, :, :, : key.shape[-2]] + + def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): + if softcap is not None: + score = softcap * torch.tanh(score / softcap) + if causal_mask is not None: + score = score + causal_mask[batch_idx][0][q_idx][kv_idx] + if head_mask is not None: + score = score + head_mask[batch_idx][head_idx][0][0] + return score + + enable_gqa = True + num_local_query_heads = query.shape[1] + + # When running TP this helps: + if not ((num_local_query_heads & (num_local_query_heads - 1)) == 0): + key = repeat_kv(key, query.shape[1] // key.shape[1]) + value = repeat_kv(value, query.shape[1] // value.shape[1]) + enable_gqa = False + + kernel_options = kwargs.get("kernel_options", None) + attn_output, attention_weights = compile_friendly_flex_attention( + query, + key, + value, + score_mod=score_mod, + block_mask=block_mask, + enable_gqa=enable_gqa, + scale=scaling, + kernel_options=kernel_options, + # Last time checked on PyTorch == 2.5.1: Flex Attention always computes the lse regardless. + # For simplification, we thus always return it as no additional computations are introduced. + return_lse=True, + training=module.training, + ) + # lse is returned in float32 + attention_weights = attention_weights.to(value.dtype) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attention_weights diff --git a/lib/python3.12/site-packages/transformers/integrations/fsdp.py b/lib/python3.12/site-packages/transformers/integrations/fsdp.py new file mode 100644 index 0000000000000000000000000000000000000000..332231c04f433f4ebd1364347c5cbce62dd498e6 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/fsdp.py @@ -0,0 +1,38 @@ +# 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. +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ..utils import is_torch_available + + +if TYPE_CHECKING: + from torch import nn + + +def is_fsdp_managed_module(module: nn.Module) -> bool: + if not is_torch_available(): + return False + + import torch + + if not torch.distributed.is_available(): + return False + + import torch.distributed.fsdp + + return isinstance(module, torch.distributed.fsdp.FullyShardedDataParallel) or getattr( + module, "_is_fsdp_managed_module", False + ) diff --git a/lib/python3.12/site-packages/transformers/integrations/ggml.py b/lib/python3.12/site-packages/transformers/integrations/ggml.py new file mode 100644 index 0000000000000000000000000000000000000000..db7d93cad0427acd2c12683edd177bf28c1ceb07 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/ggml.py @@ -0,0 +1,692 @@ +# coding=utf-8 +# Copyright 2024 The ggml.ai team and The HuggingFace Inc. team. and pygguf author (github.com/99991) +# https://github.com/99991/pygguf +# +# 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. +""" +Integration with GGML / The file is copied and adapted from https://github.com/99991/pygguf +with extra methods beings exposed +""" + +from array import array + +import numpy as np +from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors +from tokenizers.models import BPE, Unigram + +from .. import AddedToken +from ..convert_slow_tokenizer import GemmaConverter, GPT2Converter, LlamaConverter, Qwen2Converter, T5Converter +from ..utils import logging +from ..utils.logging import tqdm + + +logger = logging.get_logger(__name__) + + +GGUF_CONFIG_MAPPING = { + "general": { + "architecture": "model_type", + "name": "_model_name_or_path", + }, + "llama": { + "context_length": "max_position_embeddings", + "block_count": "num_hidden_layers", + "feed_forward_length": "intermediate_size", + "embedding_length": "hidden_size", + # NOTE: rope.dimension_count==head_dim only suitable for llama/mistral + "rope.dimension_count": "head_dim", + "rope.freq_base": "rope_theta", + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_rms_epsilon": "rms_norm_eps", + "vocab_size": "vocab_size", + }, + "mistral": { + "context_length": "max_position_embeddings", + "block_count": "num_hidden_layers", + "feed_forward_length": "intermediate_size", + "embedding_length": "hidden_size", + # NOTE: rope.dimension_count==head_dim only suitable for llama/mistral + "rope.dimension_count": "head_dim", + "rope.freq_base": "rope_theta", + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_rms_epsilon": "rms_norm_eps", + "vocab_size": "vocab_size", + }, + "qwen2": { + "context_length": "max_position_embeddings", + "block_count": "num_hidden_layers", + "feed_forward_length": "intermediate_size", + "embedding_length": "hidden_size", + "rope.dimension_count": None, + "rope.freq_base": "rope_theta", + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_rms_epsilon": "rms_norm_eps", + "vocab_size": "vocab_size", + }, + "qwen2moe": { + "context_length": "max_position_embeddings", + "block_count": "num_hidden_layers", + "feed_forward_length": "intermediate_size", + "embedding_length": "hidden_size", + "rope.dimension_count": None, + "rope.freq_base": "rope_theta", + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_rms_epsilon": "rms_norm_eps", + "vocab_size": "vocab_size", + "expert_count": "num_experts", + "expert_used_count": "num_experts_per_tok", + }, + "falcon": { + "context_length": "max_position_embeddings", + "block_count": "num_hidden_layers", + "feed_forward_length": "intermediate_size", + "embedding_length": "hidden_size", + "rope.dimension_count": None, + "rope.freq_base": "rope_theta", + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_rms_epsilon": "rms_norm_eps", + "vocab_size": "vocab_size", + }, + "tokenizer": { + "ggml.bos_token_id": "bos_token_id", + "ggml.eos_token_id": "eos_token_id", + "ggml.unknown_token_id": "unk_token_id", + "ggml.padding_token_id": "pad_token_id", + }, + "phi3": { + "context_length": "max_position_embeddings", + "block_count": "num_hidden_layers", + "feed_forward_length": "intermediate_size", + "embedding_length": "hidden_size", + "rope.dimension_count": None, + "rope.freq_base": "rope_theta", + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_rms_epsilon": "rms_norm_eps", + "vocab_size": "vocab_size", + }, + "bloom": { + "block_count": "n_layer", + "embedding_length": "hidden_size", + "attention.head_count": "n_head", + "vocab_size": "vocab_size", + "attention.layer_norm_epsilon": "layer_norm_epsilon", + }, + "t5": { + "context_length": "n_positions", + "block_count": "num_layers", + "feed_forward_length": "d_ff", + "embedding_length": "d_model", + "attention.key_length": "d_kv", + "attention.head_count": "num_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_epsilon": "layer_norm_epsilon", + "attention.relative_buckets_count": "relative_attention_num_buckets", + "decoder_start_token_id": "decoder_start_token_id", + "vocab_size": "vocab_size", + }, + "stablelm": { + "context_length": "max_position_embeddings", + "block_count": "num_hidden_layers", + "feed_forward_length": "intermediate_size", + "embedding_length": "hidden_size", + "rope.dimension_count": None, + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_epsilon": "layer_norm_eps", + "vocab_size": "vocab_size", + }, + "gpt2": { + "block_count": "n_layer", + "context_length": "n_ctx", + "embedding_length": "n_embd", + "feed_forward_length": "feed_forward_length", + "attention.head_count": "n_head", + "attention.layer_norm_epsilon": "layer_norm_epsilon", + }, + "starcoder2": { + "block_count": "num_hidden_layers", + "context_length": "max_position_embeddings", + "embedding_length": "hidden_size", + "feed_forward_length": "intermediate_size", + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_epsilon": "norm_epsilon", + }, + "mamba": { + "vocab_size": "vocab_size", + "context_length": "max_position_embeddings", + "embedding_length": "hidden_size", + "attention.layer_norm_rms_epsilon": "layer_norm_epsilon", + "block_count": "num_hidden_layers", + "ssm.conv_kernel": "conv_kernel", + "ssm.state_size": "state_size", + "ssm.time_step_rank": "time_step_rank", + "ssm.inner_size": "intermediate_size", + }, + "nemotron": { + "context_length": "max_position_embeddings", + "block_count": "num_hidden_layers", + "feed_forward_length": "intermediate_size", + "embedding_length": "hidden_size", + "rope.dimension_count": None, + "rope.freq_base": "rope_theta", + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_rms_epsilon": "norm_eps", + "vocab_size": "vocab_size", + }, + "gemma2": { + "context_length": "max_position_embeddings", + "block_count": "num_hidden_layers", + "feed_forward_length": "intermediate_size", + "embedding_length": "hidden_size", + "rope.dimension_count": None, + "rope.freq_base": "rope_theta", + # NOTE: Gemma2 has key_length==value_length==head_dim + # See: https://github.com/ggerganov/llama.cpp/blob/2e2f8f093cd4fb6bbb87ba84f6b9684fa082f3fa/convert_hf_to_gguf.py#L3293-L3294 + "attention.key_length": "head_dim", + "attention.head_count": "num_attention_heads", + "attention.head_count_kv": "num_key_value_heads", + "attention.layer_norm_rms_epsilon": "rms_norm_eps", + "vocab_size": "vocab_size", + }, +} + +GGUF_TOKENIZER_MAPPING = { + "tokenizer": { + "ggml.model": "tokenizer_type", + "ggml.tokens": "tokens", + "ggml.scores": "scores", + "ggml.token_type": "token_type", + "ggml.merges": "merges", + "ggml.bos_token_id": "bos_token_id", + "ggml.eos_token_id": "eos_token_id", + "ggml.unknown_token_id": "unk_token_id", + "ggml.padding_token_id": "pad_token_id", + "ggml.add_space_prefix": "add_prefix_space", + }, + "tokenizer_config": { + "chat_template": "chat_template", + "ggml.model": "model_type", + "ggml.bos_token_id": "bos_token_id", + "ggml.eos_token_id": "eos_token_id", + "ggml.unknown_token_id": "unk_token_id", + "ggml.padding_token_id": "pad_token_id", + }, +} + + +def _gguf_parse_value(_value, data_type): + if not isinstance(data_type, list): + data_type = [data_type] + if len(data_type) == 1: + data_type = data_type[0] + array_data_type = None + else: + if data_type[0] != 9: + raise ValueError("Received multiple types, therefore expected the first type to indicate an array.") + data_type, array_data_type = data_type + + if data_type in [0, 1, 2, 3, 4, 5, 10, 11]: + _value = int(_value[0]) + elif data_type in [6, 12]: + _value = float(_value[0]) + elif data_type in [7]: + _value = bool(_value[0]) + elif data_type in [8]: + _value = array("B", list(_value)).tobytes().decode() + elif data_type in [9]: + _value = _gguf_parse_value(_value, array_data_type) + return _value + + +class GGUFTokenizerSkeleton: + def __init__(self, dict_): + for k, v in dict_.items(): + setattr(self, k, v) + + if not hasattr(self, "merges"): + if not hasattr(self, "tokens") or not hasattr(self, "scores"): + raise ValueError( + "tokens and scores need to be passed for a LLaMa tokenizer without merges to be instantiated." + ) + tokens = self.tokens + scores = self.scores + vocab = {t: scores[i] for i, t in enumerate(tokens)} + + logger.warning("Merges were not in checkpoint, building merges on the fly.") + merges = [] + for merge, piece_score in tqdm(vocab.items()): + local = [] + for index in range(1, len(merge)): + piece_l, piece_r = merge[:index], merge[index:] + if piece_l in tokens and piece_r in tokens: + local.append((piece_l, piece_r, piece_score)) + local = sorted(local, key=lambda x: (vocab[x[0]], vocab[x[1]]), reverse=True) + merges.extend(local) + merges = sorted(merges, key=lambda val: val[2], reverse=True) + merges = [(val[0], val[1]) for val in merges] + self.merges = merges + else: + self.merges = [tuple(merge.split(" ")) for merge in self.merges] + if not hasattr(self, "scores"): + self.scores = [None for _ in range(len(self.tokens))] + + if not hasattr(self, "added_tokens"): + self.added_tokens = [] + + if not hasattr(self, "unk_token_id"): + self.unk_token_id = None + + # Llama2 uses the field `unknown_token_id` + if hasattr(self, "unknown_token_id") and self.unk_token_id is None: + self.unk_token_id = self.unknown_token_id + + +class GGUFLlamaConverter(LlamaConverter): + def __init__(self, tokenizer_dict): + self.proto = GGUFTokenizerSkeleton(tokenizer_dict) + self.original_tokenizer = self.proto + self.additional_kwargs = {} + self.is_llama_3_tokenizer = getattr(self.proto, "tokenizer_type", "llama") != "llama" + + def vocab(self, proto): + return list(zip(proto.tokens, proto.scores)) + + def merges(self, proto): + return proto.merges + + def tokenizer(self, proto): + vocab_scores = self.vocab(self.proto) + merges = self.merges(self.proto) + bpe_vocab = {word: i for i, (word, _score) in enumerate(vocab_scores)} + + unk_token = proto.tokens[proto.unk_token_id] if proto.unk_token_id is not None else None + bos_token = proto.tokens[proto.bos_token_id] if getattr(proto, "bos_token_id", None) is not None else None + eos_token = proto.tokens[proto.bos_token_id] if getattr(proto, "eos_token_id", None) is not None else None + + tokenizer = Tokenizer( + BPE( + bpe_vocab, + merges, + unk_token=unk_token, + fuse_unk=True, + byte_fallback=True, + ) + ) + + special_tokens = [] + + if not hasattr(self.proto, "token_type"): + if unk_token is not None: + special_tokens.append(AddedToken(unk_token, normalized=False, special=True)) + + if bos_token is not None: + special_tokens.append(AddedToken(bos_token, normalized=False, special=True)) + + if eos_token is not None: + special_tokens.append(AddedToken(eos_token, normalized=False, special=True)) + else: + # 3 stands for special tokens + special_tokens_idx = np.where(np.array(self.proto.token_type) == 3)[0] + + for idx in special_tokens_idx: + special_tokens.append(AddedToken(self.proto.tokens[idx], normalized=False, special=True)) + + if len(special_tokens) != 0: + tokenizer.add_special_tokens(special_tokens) + + if len(self.proto.added_tokens) != 0: + tokenizer.add_tokens( + [AddedToken(added_token, normalized=False, special=False) for added_token in self.proto.added_tokens] + ) + + self.additional_kwargs["unk_token"] = unk_token + self.additional_kwargs["eos_token"] = bos_token + self.additional_kwargs["bos_token"] = eos_token + + if self.is_llama_3_tokenizer: + self.additional_kwargs["add_prefix_space"] = None + self.additional_kwargs["clean_up_tokenization_spaces"] = True + + self.additional_kwargs["legacy"] = False + self.original_tokenizer.legacy = False + + return tokenizer + + def decoder(self, replacement, add_prefix_space): + sequence = [ + decoders.ByteFallback(), + decoders.Fuse(), + decoders.Replace("▁", " "), + ] + + if self.is_llama_3_tokenizer: + sequence += [decoders.ByteLevel(add_prefix_space=False, trim_offsets=False, use_regex=True)] + + if add_prefix_space: + sequence += [decoders.Strip(content=" ", left=1)] + return decoders.Sequence(sequence) + + def converted(self): + # Copied partly from converted method in SpmConverter class + tokenizer = self.tokenizer(self.proto) + + # Tokenizer assemble + normalizer = self.normalizer(self.proto) + if normalizer is not None: + tokenizer.normalizer = normalizer + + replacement = "▁" + add_prefix_space = True + if hasattr(self.original_tokenizer, "add_prefix_space"): + add_prefix_space = self.original_tokenizer.add_prefix_space + + pre_tokenizer = self.pre_tokenizer(replacement, add_prefix_space) + if pre_tokenizer is not None: + tokenizer.pre_tokenizer = pre_tokenizer + + tokenizer.decoder = self.decoder(replacement, add_prefix_space) + post_processor = self.post_processor() + if post_processor: + tokenizer.post_processor = post_processor + + # HACK: patch the llama-3 tokenizer to use the correspinding pre-tokenizer + # and normalizer + if self.is_llama_3_tokenizer: + tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel( + add_prefix_space=False, trim_offsets=False, use_regex=True + ) + # This is tricky as the additional kwargs are passed after legacy is force-set in LlamaTokenizer's + # init. + tokenizer.normalizer = normalizers.Sequence([]) + + return tokenizer + + +class GGUFQwen2Converter(Qwen2Converter): + def __init__(self, tokenizer_dict): + self.original_tokenizer = GGUFTokenizerSkeleton(tokenizer_dict) + self.additional_kwargs = {} + + def converted(self) -> Tokenizer: + vocab = {word: i for i, word in enumerate(self.original_tokenizer.tokens)} + merges = self.original_tokenizer.merges + tokenizer = super().converted(vocab, merges) + + tokenizer.add_special_tokens( + [ + AddedToken("<|endoftext|>", normalized=False, special=True), + AddedToken("<|im_start|>", normalized=False, special=True), + AddedToken("<|im_end|>", normalized=False, special=True), + ] + ) + return tokenizer + + +class GGUFPhi3Converter(LlamaConverter): + def __init__(self, tokenizer_dict): + self.proto = GGUFTokenizerSkeleton(tokenizer_dict) + self.original_tokenizer = self.proto + self.additional_kwargs = {} + + def vocab(self, proto): + return list(zip(proto.tokens, proto.scores)) + + def merges(self, proto): + return proto.merges + + def tokenizer(self, proto): + vocab_scores = self.vocab(self.proto) + merges = self.merges(self.proto) + bpe_vocab = {word: i for i, (word, _score) in enumerate(vocab_scores)} + + tokenizer = Tokenizer(BPE(bpe_vocab, merges)) + # add the special tokens from phi3 tokenizer config + tokenizer.add_special_tokens( + [ + AddedToken("", rstrip=True, lstrip=False, normalized=False, special=True), + AddedToken("<|endoftext|>", normalized=False, special=True), + AddedToken("<|assistant|>", rstrip=True, normalized=False, special=True), + AddedToken("<|placeholder1|>", rstrip=True, normalized=False, special=True), + AddedToken("<|placeholder2|>", rstrip=True, normalized=False, special=True), + AddedToken("<|placeholder3|>", rstrip=True, normalized=False, special=True), + AddedToken("<|placeholder4|>", rstrip=True, normalized=False, special=True), + AddedToken("<|system|>", rstrip=True, normalized=False, special=True), + AddedToken("<|end|>", rstrip=True, normalized=False, special=True), + AddedToken("<|placeholder5|>", rstrip=True, normalized=False, special=True), + AddedToken("<|placeholder6|>", rstrip=True, normalized=False, special=True), + AddedToken("<|user|>", rstrip=True, normalized=False, special=True), + ] + ) + + self.additional_kwargs["unk_token"] = ( + proto.tokens[proto.unk_token_id] if proto.unk_token_id is not None else None + ) + self.additional_kwargs["eos_token"] = ( + proto.tokens[proto.eos_token_id] if proto.eos_token_id is not None else None + ) + self.additional_kwargs["bos_token"] = ( + proto.tokens[proto.bos_token_id] if proto.bos_token_id is not None else None + ) + self.additional_kwargs["pad_token"] = ( + proto.tokens[proto.pad_token_id] if proto.pad_token_id is not None else None + ) + + return tokenizer + + def decoder(self, replacement, add_prefix_space): + sequence = [ + decoders.ByteFallback(), + decoders.Fuse(), + decoders.Replace(replacement, " "), + ] + + if add_prefix_space: + sequence += [decoders.Strip(content=" ", left=1)] + return decoders.Sequence(sequence) + + def converted(self) -> Tokenizer: + tokenizer = self.tokenizer(self.proto) + + replacement = "▁" + add_prefix_space = True + if hasattr(self.original_tokenizer, "add_prefix_space"): + add_prefix_space = self.original_tokenizer.add_prefix_space + + tokenizer.decoder = self.decoder(replacement, add_prefix_space) + + return tokenizer + + +class GGUFGPTConverter(GPT2Converter): + def __init__(self, tokenizer_dict): + self.original_tokenizer = GGUFTokenizerSkeleton(tokenizer_dict) + self.additional_kwargs = {} + + def converted(self) -> Tokenizer: + vocab = {word: i for i, word in enumerate(self.original_tokenizer.tokens)} + merges = self.original_tokenizer.merges + tokenizer = super().converted(vocab, merges) + return tokenizer + + +class GGUFT5Converter(T5Converter): + def __init__(self, tokenizer_dict): + # set dummy data to avoid unnecessary merges calculation + tokenizer_dict["merges"] = ["dummy text"] + + self.proto = GGUFTokenizerSkeleton(tokenizer_dict) + self.token2id = {k: v for v, k in enumerate(self.proto.tokens)} + self.original_tokenizer = self.proto + self.additional_kwargs = {} + + def vocab(self, proto): + return list(zip(proto.tokens, proto.scores)) + + def normalizer(self, proto): + if getattr(self.original_tokenizer, "legacy", True): + sequence = [] + if getattr(self.original_tokenizer, "add_prefix_space", True): + sequence += [normalizers.Prepend(prepend="▁")] + sequence += [normalizers.Replace(pattern=" ", content="▁")] + return normalizers.Sequence(sequence) + return None # non-legacy, no normalizer + + def post_processor(self): + return processors.TemplateProcessing( + single=["$A", ""], + pair=["$A", "", "$B", ""], + special_tokens=[ + ("", self.token2id[""]), + ], + ) + + def converted(self) -> Tokenizer: + vocab_scores = self.vocab(self.proto) + tokenizer = Tokenizer( + Unigram( + vocab_scores, + unk_id=self.proto.unk_token_id, + byte_fallback=False, + ) + ) + + # Tokenizer assemble + normalizer = self.normalizer(self.proto) + if normalizer is not None: + tokenizer.normalizer = normalizer + + replacement = "▁" + add_prefix_space = True + if hasattr(self.original_tokenizer, "add_prefix_space"): + add_prefix_space = self.original_tokenizer.add_prefix_space + + pre_tokenizer = self.pre_tokenizer(replacement, add_prefix_space) + if pre_tokenizer is not None: + tokenizer.pre_tokenizer = pre_tokenizer + + tokenizer.decoder = self.decoder(replacement, add_prefix_space) + post_processor = self.post_processor() + if post_processor: + tokenizer.post_processor = post_processor + + return tokenizer + + +class GGUFGemmaConverter(GemmaConverter): + def __init__(self, tokenizer_dict): + # set dummy data to avoid unnecessary merges calculation + tokenizer_dict["merges"] = ["dummy text"] + + self.proto = GGUFTokenizerSkeleton(tokenizer_dict) + self.original_tokenizer = self.proto + self.additional_kwargs = {} + + def vocab(self, proto): + original_vocab = list(zip(proto.tokens, proto.scores)) + updated_vocab = [] + + for token, score in original_vocab: + if token == "<0x09>": + updated_vocab.append(("\t", score)) + elif " " in token and len(token.strip()) == 0: + underscores = "▁" * len(token) + updated_vocab.append((underscores, score)) + else: + updated_vocab.append((token, score)) + + return updated_vocab + + def normalizer(self, proto): + return normalizers.Replace(" ", "▁") + + def decoder(self, replacement, add_prefix_space): + sequence = [ + decoders.Replace("▁", " "), + decoders.ByteFallback(), + decoders.Fuse(), + ] + + if add_prefix_space: + sequence += [decoders.Strip(content=" ", left=1)] + return decoders.Sequence(sequence) + + def converted(self) -> Tokenizer: + vocab_scores = self.vocab(self.proto) + tokenizer = Tokenizer( + Unigram( + vocab_scores, + unk_id=self.proto.unk_token_id, + byte_fallback=self.handle_byte_fallback, + ) + ) + + normalizer = self.normalizer(self.proto) + if normalizer is not None: + tokenizer.normalizer = normalizer + + replacement = "▁" + add_prefix_space = True + if hasattr(self.original_tokenizer, "add_prefix_space"): + add_prefix_space = self.original_tokenizer.add_prefix_space + + tokenizer.decoder = self.decoder(replacement, add_prefix_space) + pre_tokenizer = self.pre_tokenizer(replacement, add_prefix_space) + if pre_tokenizer is not None: + tokenizer.pre_tokenizer = pre_tokenizer + + return tokenizer + + +GGUF_TO_FAST_CONVERTERS = { + "llama": GGUFLlamaConverter, + "qwen2": GGUFQwen2Converter, + "qwen2_moe": GGUFQwen2Converter, + "phi3": GGUFPhi3Converter, + "bloom": GGUFGPTConverter, + "falcon": GGUFGPTConverter, + "stablelm": GGUFGPTConverter, + "gpt2": GGUFGPTConverter, + "starcoder2": GGUFGPTConverter, + "t5": GGUFT5Converter, + "mamba": GGUFGPTConverter, + "nemotron": GGUFGPTConverter, + "gemma2": GGUFGemmaConverter, +} + + +def convert_gguf_tokenizer(architecture, tokenizer_dict) -> Tokenizer: + """ + Utilities to convert a slow tokenizer instance in a fast tokenizer instance. + + Args: + architecture (`str`): The model architecture derived from gguf file. + transformer_tokenizer ([`~tokenization_utils_base.PreTrainedTokenizer`]): + Instance of a slow tokenizer to convert in the backend tokenizer for + [`~tokenization_utils_base.PreTrainedTokenizerFast`]. + + Return: + A instance of [`~tokenizers.Tokenizer`] to be used as the backend tokenizer of a + [`~tokenization_utils_base.PreTrainedTokenizerFast`] + """ + tokenizer_class_name = architecture + converter = GGUF_TO_FAST_CONVERTERS[tokenizer_class_name](tokenizer_dict) + fast_tokenizer = converter.converted() + return fast_tokenizer, converter.additional_kwargs diff --git a/lib/python3.12/site-packages/transformers/integrations/higgs.py b/lib/python3.12/site-packages/transformers/integrations/higgs.py new file mode 100644 index 0000000000000000000000000000000000000000..dd31764dfe0ce731077a3fca9c6e35579c65b50b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/higgs.py @@ -0,0 +1,655 @@ +# 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. +"HIGGS through FLUTE (Flexible Lookup Table Engine for LUT-quantized LLMs) integration file" + +from math import sqrt + +from ..utils import ( + is_flute_available, + is_hadamard_available, + is_torch_available, +) + + +if is_torch_available(): + import torch + from torch import nn + + +if is_flute_available(): + from flute.integrations.higgs import prepare_data_transposed + from flute.tune import TuneMetaData, qgemm_v2 + +if is_hadamard_available(): + from fast_hadamard_transform import hadamard_transform + + +def pad_to_block(tensor, dims, had_block_size, value=0): + pad_dims = [0 for _ in range(2 * len(tensor.shape))] + for dim in dims: + size = tensor.shape[dim] + next_multiple_of_1024 = ((size - 1) // had_block_size + 1) * had_block_size + delta = next_multiple_of_1024 - size + pad_dims[-2 * dim - 1] = delta + + return nn.functional.pad(tensor, pad_dims, "constant", value) + + +def get_higgs_grid(p: int, n: int): + if (p, n) == (2, 256): + return torch.tensor( + [ + [-2.501467704772949, 0.17954708635807037], + [-0.6761789321899414, 1.2728623151779175], + [-1.8025816679000854, 0.7613157629966736], + [-0.538287878036499, -2.6028504371643066], + [0.8415029644966125, -0.8600977659225464], + [0.7023013234138489, 3.3138747215270996], + [0.5699077844619751, 2.5782253742218018], + [3.292393207550049, -0.6016128063201904], + [0.5561617016792297, -1.7723814249038696], + [-2.1012380123138428, 0.020958125591278076], + [0.46085724234580994, 0.8428705334663391], + [1.4548040628433228, -0.6156039237976074], + [3.210029363632202, 0.3546904921531677], + [0.8893890976905823, -0.5967988967895508], + [0.8618854284286499, -3.2061192989349365], + [1.1360996961593628, -0.23852407932281494], + [1.6646337509155273, -0.9265465140342712], + [1.4767773151397705, 1.2476022243499756], + [-1.0511897802352905, 1.94503915309906], + [-1.56318998336792, -0.3264186680316925], + [-0.1829211413860321, 0.2922491431236267], + [-0.8950616717338562, -1.3887052536010742], + [-0.08206957578659058, -1.329533576965332], + [-0.487422913312912, 1.4817842245101929], + [-1.6769757270812988, -2.8269758224487305], + [-1.5057679414749146, 1.8905963897705078], + [1.8335362672805786, 1.0515104532241821], + [0.3273945450782776, 1.0491033792495728], + [-3.295924186706543, -0.7021600008010864], + [-1.8428784608840942, -1.2315762042999268], + [-0.8575026392936707, -1.7005949020385742], + [-1.120667815208435, 0.6467998027801514], + [-0.1588846743106842, -1.804071068763733], + [-0.8539647459983826, 0.5645008683204651], + [-1.4192019701004028, -0.6175029873847961], + [1.0799058675765991, 1.7871345281600952], + [1.171311855316162, 0.7511613965034485], + [2.162078380584717, 0.8044339418411255], + [1.3969420194625854, -1.243762493133545], + [-0.23818807303905487, 0.053944624960422516], + [2.304199457168579, -1.2667627334594727], + [1.4225027561187744, 0.568610668182373], + [0.376836895942688, -0.7134661674499512], + [2.0404467582702637, 0.4087389409542084], + [0.7639489769935608, -1.1367933750152588], + [0.3622530400753021, -1.4827953577041626], + [0.4100743532180786, 0.36108437180519104], + [-1.5867475271224976, -1.618212342262268], + [-2.2769672870635986, -1.2132309675216675], + [0.9184022545814514, -0.34428009390830994], + [-0.3902314603328705, 0.21785245835781097], + [3.120687484741211, 1.3077973127365112], + [1.587440848350525, -1.6506884098052979], + [-1.718808889389038, -0.038405973464250565], + [-0.6888407468795776, -0.8402308821678162], + [-0.7981445789337158, -1.1117373704910278], + [-2.4124443531036377, 1.3419722318649292], + [-0.6611530184745789, 0.9939885139465332], + [-0.33103418350219727, -0.16702833771705627], + [-2.4091389179229736, -2.326857566833496], + [1.6610108613967896, -2.159703254699707], + [0.014884627424180508, 0.3887578248977661], + [0.029668325558304787, 1.8786455392837524], + [1.180362582206726, 2.699317216873169], + [1.821286678314209, -0.5960053205490112], + [-0.44835323095321655, 3.327436685562134], + [-0.3714401423931122, -2.1466753482818604], + [-1.1103475093841553, -2.4536871910095215], + [-0.39110705256462097, 0.6670510172843933], + [0.474752813577652, -1.1959707736968994], + [-0.013110585510730743, -2.52519154548645], + [-2.0836575031280518, -1.703289270401001], + [-1.1077687740325928, -0.1252644956111908], + [-0.4138077199459076, 1.1837692260742188], + [-1.977599024772644, 1.688241720199585], + [-1.659559965133667, -2.1387736797332764], + [0.03242531046271324, 0.6526556015014648], + [0.9127950072288513, 0.6099498867988586], + [-0.38478314876556396, 0.433487206697464], + [0.27454206347465515, -0.27719801664352417], + [0.10388526320457458, 2.2812814712524414], + [-0.014394169673323631, -3.177137613296509], + [-1.2871228456497192, -0.8961855173110962], + [0.5720916986465454, -0.921597957611084], + [1.1159656047821045, -0.7609877586364746], + [2.4383342266082764, -2.2983546257019043], + [-0.294057160615921, -0.9770799875259399], + [-0.9342701435089111, 1.107579231262207], + [-1.549338698387146, 3.090520143508911], + [2.6076579093933105, 2.051239013671875], + [-0.9259037375450134, 1.407211184501648], + [-0.1747353971004486, 0.540488600730896], + [-0.8963701725006104, 0.8271111249923706], + [0.6480194926261902, 1.0128909349441528], + [0.980783998966217, -0.06156221032142639], + [-0.16883476078510284, 1.0601658821105957], + [0.5839992761611938, 0.004697148688137531], + [-0.34228450059890747, -1.2423977851867676], + [2.500824451446533, 0.3665279746055603], + [-0.17641609907150269, 1.3529551029205322], + [0.05378641560673714, 2.817232847213745], + [-1.2391047477722168, 2.354328155517578], + [0.630434513092041, -0.668536365032196], + [1.7576488256454468, 0.6738647818565369], + [0.4435231387615204, 0.6000469326972961], + [-0.08794835954904556, -0.11511358618736267], + [1.6540337800979614, 0.33995017409324646], + [-0.04202975332736969, -0.5375117063522339], + [-0.4247745871543884, -0.7897617220878601], + [0.06695003807544708, 1.2000739574432373], + [-3.2508881092071533, 0.28734830021858215], + [-1.613816261291504, 0.4944162368774414], + [1.3598989248275757, 0.26117825508117676], + [2.308382511138916, 1.3462618589401245], + [-1.2137469053268433, -1.9254342317581177], + [-0.4889402985572815, 1.8136259317398071], + [-0.1870335340499878, -0.3480615019798279], + [1.0766386985778809, -1.0627082586288452], + [0.4651014506816864, 2.131748914718628], + [-0.1306295394897461, -0.7811847925186157], + [0.06433182954788208, -1.5397958755493164], + [-0.2894323468208313, -0.5789554715156555], + [-0.6081662178039551, 0.4845278263092041], + [2.697964668273926, -0.18515698611736298], + [0.1277363896369934, -0.7221432328224182], + [0.8700758218765259, 0.35042452812194824], + [0.22088994085788727, 0.495242178440094], + [-2.5843818187713623, -0.8000828623771667], + [0.6732649803161621, -1.4362232685089111], + [-1.5286413431167603, 1.0417330265045166], + [-1.1222513914108276, -0.6269875764846802], + [-0.9752035140991211, -0.8750635385513306], + [-2.6369473934173584, 0.6918523907661438], + [0.14478731155395508, -0.041986867785453796], + [-1.5629483461380005, 1.4369450807571411], + [0.38952457904815674, -2.16428804397583], + [-0.16885095834732056, 0.7976621985435486], + [-3.12416934967041, 1.256506085395813], + [0.6843105554580688, -0.4203019142150879], + [1.9345275163650513, 1.934950351715088], + [0.012184220366179943, -2.1080918312072754], + [-0.6350273489952087, 0.7358828186988831], + [-0.837304949760437, -0.6214472651481628], + [0.08211923390626907, -0.9472538232803345], + [2.9332995414733887, -1.4956780672073364], + [1.3806978464126587, -0.2916182279586792], + [0.06773144006729126, 0.9285762310028076], + [-1.1943119764328003, 1.5963770151138306], + [1.6395620107650757, -0.32285431027412415], + [-1.390851378440857, -0.08273141086101532], + [1.816330909729004, -1.2812227010726929], + [0.7921574711799622, -2.1135804653167725], + [0.5817914605140686, 1.2644577026367188], + [1.929347038269043, -0.2386285960674286], + [0.8877345323562622, 1.190008521080017], + [1.4732073545455933, 0.8935023546218872], + [-2.8518524169921875, -1.5478795766830444], + [0.2439267635345459, 0.7576767802238464], + [0.5246709585189819, -2.606659412384033], + [1.150876760482788, 1.4073830842971802], + [-0.2643202245235443, 2.0634236335754395], + [1.555483341217041, -0.0023102816194295883], + [2.0830578804016113, -1.7225427627563477], + [-0.5424830317497253, -1.070199728012085], + [0.9168899655342102, 0.8955540060997009], + [-0.8120972514152527, 2.696739912033081], + [-0.29908373951911926, -1.5310651063919067], + [1.2320337295532227, -1.556247353553772], + [1.8612544536590576, 0.08704725652933121], + [0.22133447229862213, -1.8091708421707153], + [-0.4403655230998993, -0.38571012020111084], + [-1.88539457321167, 1.192205786705017], + [2.239687919616699, 0.004709010478109121], + [1.139495611190796, 0.45733731985092163], + [-1.507995367050171, 0.19716016948223114], + [0.46986445784568787, 1.5422041416168213], + [-1.2573751211166382, -0.35984551906585693], + [-1.7415345907211304, -0.6020717024803162], + [1.0751984119415283, 0.19006384909152985], + [2.24186635017395, -0.46343153715133667], + [0.3610347509384155, -0.07658443599939346], + [-1.3111497163772583, 0.432013601064682], + [0.6164408326148987, 0.24538464844226837], + [-1.9266542196273804, -0.3256155550479889], + [-0.5870336890220642, -0.1879584938287735], + [-1.0476511716842651, 0.3677721917629242], + [-1.229940414428711, 1.2433830499649048], + [0.18550436198711395, 0.22753673791885376], + [-0.017921989783644676, 0.12625974416732788], + [1.1659504175186157, -0.5020995736122131], + [-0.5983408093452454, -1.40438973903656], + [0.7519024014472961, -0.16282692551612854], + [0.9920787811279297, -1.344896912574768], + [-0.8103678226470947, 0.3064485788345337], + [0.6956969499588013, 1.8208192586898804], + [-2.7830491065979004, -0.2299390584230423], + [-0.34681546688079834, 2.4890666007995605], + [-1.4452646970748901, -1.2216600179672241], + [-2.1872897148132324, 0.8926076292991638], + [1.706072211265564, -2.8440372943878174], + [1.1119003295898438, -2.4923460483551025], + [-2.582794666290283, 2.0973289012908936], + [0.04987720400094986, -0.2964983284473419], + [-2.063807487487793, -0.7847916483879089], + [-0.4068813621997833, 0.9135897755622864], + [-0.9814359545707703, -0.3874954879283905], + [-1.4227229356765747, 0.7337291240692139], + [0.3065044581890106, 1.3125417232513428], + [1.2160996198654175, -1.9643305540084839], + [-1.2163853645324707, 0.14608727395534515], + [-2.3030710220336914, -0.37558120489120483], + [0.9232977628707886, 2.1843791007995605], + [-0.1989777386188507, 1.651851773262024], + [-0.714374840259552, -0.39365994930267334], + [-0.7805715799331665, -2.099881887435913], + [0.9015759229660034, -1.7053706645965576], + [0.1033422127366066, 1.5256654024124146], + [-1.8773194551467896, 2.324174165725708], + [1.9227174520492554, 2.7441604137420654], + [-0.5994020104408264, 0.23984014987945557], + [1.3496100902557373, -0.9126054644584656], + [-0.8765304088592529, -3.1877026557922363], + [-1.2040035724639893, -1.5169521570205688], + [1.4261796474456787, 2.150200128555298], + [1.463774561882019, 1.6656692028045654], + [0.20364105701446533, -0.4988172650337219], + [0.5195154547691345, -0.24067887663841248], + [-1.1116786003112793, -1.1599653959274292], + [-0.8490808606147766, -0.1681060940027237], + [0.3189965784549713, -0.9641751646995544], + [-0.5664751529693604, -0.5951744318008423], + [-1.6347930431365967, -0.9137664437294006], + [0.44048091769218445, -0.47259435057640076], + [-2.147747039794922, 0.47442489862442017], + [1.834734320640564, 1.4462147951126099], + [1.1777573823928833, 1.0659226179122925], + [-0.9568989872932434, 0.09495053440332413], + [-1.838529348373413, 0.2950586676597595], + [-0.4800611734390259, 0.014894310384988785], + [-0.5235516428947449, -1.7687653303146362], + [2.0735011100769043, -0.8825281262397766], + [2.637502431869507, 0.8455678224563599], + [2.606602907180786, -0.7848446369171143], + [-1.1886937618255615, 0.9330510497093201], + [0.38082656264305115, 0.13328030705451965], + [0.6847941875457764, 0.7384101152420044], + [1.2638574838638306, -0.007309418171644211], + [0.18292222917079926, -1.22371244430542], + [0.8143821954727173, 1.4976691007614136], + [0.6571850776672363, 0.48368802666664124], + [-0.6991601586341858, 2.150190830230713], + [0.8101756572723389, 0.10206498205661774], + [-0.08768226951360703, -1.084917664527893], + [-0.7208092212677002, 0.03657956421375275], + [0.3211449086666107, 1.803687334060669], + [-0.7835946083068848, 1.6869111061096191], + ] + ) + if (p, n) == (2, 64): + return torch.tensor( + [ + [-2.7216711044311523, 0.14431366324424744], + [-0.766914427280426, 1.7193410396575928], + [-2.2575762271881104, 1.2476624250411987], + [1.233758807182312, -2.3560616970062256], + [0.8701965808868408, -0.2649352252483368], + [1.4506438970565796, 2.1776366233825684], + [-0.06305818259716034, 1.9049758911132812], + [2.536226511001587, 0.563927412033081], + [0.4599496126174927, -1.8745561838150024], + [-1.900517225265503, -0.30703988671302795], + [0.09386251866817474, 0.8755807280540466], + [1.946500539779663, -0.6743080615997314], + [2.1338934898376465, 1.4581491947174072], + [0.9429940581321716, -0.8038390278816223], + [2.0697755813598633, -1.614896535873413], + [0.772676408290863, 0.22017823159694672], + [1.0689979791641235, -1.525044322013855], + [0.6813604831695557, 1.1345642805099487], + [0.4706456661224365, 2.606626272201538], + [-1.294018030166626, -0.4372096061706543], + [-0.09134224057197571, 0.4610418677330017], + [-0.7907772064208984, -0.48412787914276123], + [0.060459110885858536, -0.9172890186309814], + [-0.5855047702789307, 2.56172513961792], + [0.11484206467866898, -2.659848213195801], + [-1.5893300771713257, 2.188580274581909], + [1.6750942468643188, 0.7089915871620178], + [-0.445697546005249, 0.7452405095100403], + [-1.8539940118789673, -1.8377939462661743], + [-1.5791912078857422, -1.017285943031311], + [-1.030419945716858, -1.5746369361877441], + [-1.9511750936508179, 0.43696075677871704], + [-0.3446580767631531, -1.8953213691711426], + [-1.4219647645950317, 0.7676230669021606], + [-0.9191089272499084, 0.5021472573280334], + [0.20464491844177246, 1.3684605360031128], + [0.5402919054031372, 0.6699410676956177], + [1.8903915882110596, 0.03638288006186485], + [0.4723062515258789, -0.6216739416122437], + [-0.41345009207725525, -0.22752176225185394], + [2.7119064331054688, -0.5111885070800781], + [1.065286636352539, 0.6950305700302124], + [0.40629103779792786, -0.14339995384216309], + [1.2815024852752686, 0.17108257114887238], + [0.01785222627222538, -0.43778058886528015], + [0.054590027779340744, -1.4225547313690186], + [0.3076786696910858, 0.30697619915008545], + [-0.9498570561408997, -0.9576997756958008], + [-2.4640724658966064, -0.9660449028015137], + [1.3714425563812256, -0.39760473370552063], + [-0.4857747256755829, 0.2386789172887802], + [1.2797833681106567, 1.3097363710403442], + [0.5508887767791748, -1.1777795553207397], + [-1.384316325187683, 0.1465839296579361], + [-0.46556955575942993, -1.2442727088928223], + [-0.3915477693080902, -0.7319604158401489], + [-1.4005504846572876, 1.3890998363494873], + [-0.8647305965423584, 1.0617644786834717], + [-0.8901953101158142, -0.01650036871433258], + [-0.9893633723258972, -2.4662880897521973], + [1.445534110069275, -1.049334168434143], + [-0.041650623083114624, 0.012734669260680676], + [-0.3302375078201294, 1.26217782497406], + [0.6934980154037476, 1.7714335918426514], + ] + ) + elif (p, n) == (2, 16): + return torch.tensor( + [ + [-0.8996632695198059, -1.6360418796539307], + [-0.961183488368988, 1.5999565124511719], + [-1.882026195526123, 0.678778350353241], + [0.36300793290138245, -1.9667866230010986], + [-0.6814072728157043, -0.576818585395813], + [0.7270012497901917, 0.6186859607696533], + [0.3359416127204895, 1.8371193408966064], + [1.859930396080017, 0.036668598651885986], + [0.17208248376846313, -0.9401724338531494], + [-1.7599700689315796, -0.6244229674339294], + [-0.8993809223175049, 0.32267823815345764], + [0.839488685131073, -0.3017036020755768], + [1.5314953327178955, 1.2942044734954834], + [-0.0011779458727687597, 0.00022069070837460458], + [1.4274526834487915, -1.207889199256897], + [-0.16123905777931213, 0.8787511587142944], + ] + ) + elif (p, n) == (1, 16): + return torch.tensor( + [ + [-2.7325894832611084], + [-2.069017171859741], + [-1.6180464029312134], + [-1.2562311887741089], + [-0.9423404335975647], + [-0.6567591428756714], + [-0.38804829120635986], + [-0.12839503586292267], + [0.12839503586292267], + [0.38804829120635986], + [0.6567591428756714], + [0.9423404335975647], + [1.2562311887741089], + [1.6180464029312134], + [2.069017171859741], + [2.7325894832611084], + ] + ) + elif (p, n) == (1, 8): + return torch.tensor( + [ + [-2.1519455909729004], + [-1.3439092636108398], + [-0.7560052871704102], + [-0.2450941801071167], + [0.2450941801071167], + [0.7560052871704102], + [1.3439092636108398], + [2.1519455909729004], + ] + ) + elif (p, n) == (1, 4): + return torch.tensor([[-1.5104175806045532], [-0.4527800381183624], [0.4527800381183624], [1.5104175806045532]]) + else: + raise NotImplementedError(f"Unsupported p={p}, n={n}") + + +def quantize_with_higgs(weight, bits: int = 4, p: int = 2, group_size: int = 256, hadamard_size: int = 1024): + assert len(weight.shape) == 2, "Only 2D weights are supported for now" + + grid = get_higgs_grid(p, 2 ** (p * bits)).to(weight.device) + grid_norm_2 = torch.linalg.norm(grid, axis=-1) ** 2 + + device = weight.device + dtype = weight.dtype + weight = weight.to(copy=True, dtype=torch.float32) + # Pad to Hadamard transform size + weight = pad_to_block(weight, [1], hadamard_size) + + # Scale and Hadamard transform + mult = weight.shape[1] // hadamard_size + weight = weight.reshape(-1, mult, hadamard_size) + scales = torch.linalg.norm(weight, axis=-1) + weight = hadamard_transform(weight, 1) / scales[:, :, None] + + # Pad to edenn_d and project + weight = pad_to_block(weight, [2], p).reshape(weight.shape[0], mult, -1, p) + + # Quantize + codes = torch.empty(weight.shape[:-1], device=device, dtype=torch.uint8) + for i in range(0, weight.shape[0], 16): + codes[i : i + 16] = torch.argmax(2 * weight[i : i + 16] @ grid.T - grid_norm_2, dim=-1).to(torch.uint8) + del weight + + codes = codes.reshape(codes.shape[0], -1) + scales = scales / sqrt(hadamard_size) + + weight, scales, tables, tables2, tune_metadata = prepare_data_transposed( + codes, + torch.repeat_interleave(scales.to(dtype), hadamard_size // group_size, dim=1), + grid.to(dtype), + num_bits=bits, + group_size=group_size, + vector_size=p, + dtype=dtype, + device=device, + check_correctness=False, + ) + + return { + "weight": weight, + "scales": scales, + "tables": tables, + "tables2": tables2.view(dtype=torch.float16), + "tune_metadata": tune_metadata, + } + + +class HiggsLinear(torch.nn.Module): + def __init__( + self, + in_features: int, + out_features: int, + num_bits: int, + bias=True, + dtype: torch.dtype = None, + device: torch.device = None, + group_size: int = 256, + hadamard_size: int = 1024, + ): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.num_bits = num_bits + self.group_size = group_size + self.hadamard_size = hadamard_size + + assert in_features % group_size == 0 + assert num_bits in [2, 3, 4] + + self.weight = nn.Parameter( + torch.empty((out_features * num_bits // 16, in_features), dtype=torch.int16, device=device), + requires_grad=False, + ) + self.scales = nn.Parameter( + torch.empty((out_features, in_features // group_size), dtype=dtype, device=device), requires_grad=False + ) + self.tables = nn.Parameter(torch.empty((2**num_bits,), dtype=dtype, device=device), requires_grad=False) + self.tables2 = nn.Parameter( + torch.empty((2**num_bits, 2**num_bits, 2), dtype=dtype, device=device), requires_grad=False + ) + + if bias: + self.bias = nn.Parameter(torch.empty(out_features, device=device, dtype=dtype), requires_grad=False) + else: + self.register_parameter("bias", None) + + self.workspace = None # must be set externally to be reused among layers + self.tune_metadata: TuneMetaData = None # must be set externally because architecture dependent + + def forward(self, x): + x = pad_to_block(x, [-1], self.hadamard_size) + + if self.workspace is None: + raise Exception("Workspace must be set before calling forward") + + return qgemm_v2( + x, + self.weight, + self.scales, + self.tables, + self.tables2.view(dtype=torch.float32), + self.workspace, + self.tune_metadata, + hadamard_size=self.hadamard_size, + ) + + +def replace_with_higgs_linear( + model, + quantization_config=None, + current_key_name=None, + has_been_replaced=False, +): + """ + Public method that recursively replaces the Linear layers of the given model with HIGGS quantized layers. + `accelerate` is needed to use this method. Returns the converted model and a boolean that indicates if the + conversion has been successfull or not. + + Args: + model (`torch.nn.Module`): + The model to convert, can be any `torch.nn.Module` instance. + quantization_config (`HiggsConfig`): + The quantization config object that contains the quantization parameters. + current_key_name (`list`, *optional*): + A list that contains the current key name. This is used for recursion and should not be passed by the user. + has_been_replaced (`bool`, *optional*): + A boolean that indicates if the conversion has been successful or not. This is used for recursion and + should not be passed by the user. + """ + + from accelerate import init_empty_weights + + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + + if isinstance(module, nn.Linear): + # Check if the current key is not in the `quantization_config.modules_to_not_convert` + current_key_name_str = ".".join(current_key_name) + if not any(current_key_name_str.endswith(key) for key in quantization_config.modules_to_not_convert): + with init_empty_weights(): + in_features = module.in_features + out_features = module.out_features + + model._modules[name] = HiggsLinear( + in_features, + out_features, + bias=module.bias is not None, + num_bits=quantization_config.bits, + hadamard_size=quantization_config.hadamard_size, + group_size=quantization_config.group_size, + ) + has_been_replaced = True + + # Store the module class in case we need to transpose the weight later + model._modules[name].source_cls = type(module) + # Force requires grad to False to avoid unexpected errors + model._modules[name].requires_grad_(False) + if len(list(module.children())) > 0: + _, has_been_replaced = replace_with_higgs_linear( + module, + quantization_config=quantization_config, + current_key_name=current_key_name, + has_been_replaced=has_been_replaced, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced + + +def dequantize_higgs(model, current_key_name=None): + """ + Dequantizes the HiggsLinear layers in the given model by replacing them with standard torch.nn.Linear layers. + Args: + model (torch.nn.Module): The model containing HiggsLinear layers to be dequantized. + current_key_name (list, optional): A list to keep track of the current module names during recursion. Defaults to None. + Returns: + torch.nn.Module: The model with HiggsLinear layers replaced by torch.nn.Linear layers. + """ + + with torch.no_grad(): + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + + if isinstance(module, HiggsLinear): + in_features = module.in_features + out_features = module.out_features + + model._modules[name] = torch.nn.Linear( + in_features, + out_features, + bias=module.bias is not None, + device=module.scales.device, + dtype=module.scales.dtype, + ) + + model._modules[name].weight.data = module( + torch.eye(in_features, device=module.scales.device, dtype=module.scales.dtype) + ).T.contiguous() + + if len(list(module.children())) > 0: + _ = dequantize_higgs( + module, + current_key_name=current_key_name, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model diff --git a/lib/python3.12/site-packages/transformers/integrations/hqq.py b/lib/python3.12/site-packages/transformers/integrations/hqq.py new file mode 100644 index 0000000000000000000000000000000000000000..4ff154ee2056c26d5e9005c7fa90a970931d41fe --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/hqq.py @@ -0,0 +1,129 @@ +# 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. +"HQQ (Half-Quadratic Quantization) integration file" + +from ..utils import is_hqq_available, is_torch_available, logging + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +# Name all modules inside the model +def autoname_modules(model): + for name, module in model.named_modules(): + module.name = name + + +# Get the linear_tag from a modul name. For example: model.layers.31.self_attn.k_proj -> self_attn.k_proj +def name_to_linear_tag(name): + return ".".join([n for n in name.split(".") if ((n not in ["model", "layers"]) and (not n.isnumeric()))]) + + +# Get all linear tags available +def get_linear_tags(model): + if is_hqq_available(): + from hqq.core.quantize import HQQLinear + + linear_tags = set() + for name, module in model.named_modules(): + if isinstance(module, (torch.nn.Linear, HQQLinear)): + linear_tags.add(name_to_linear_tag(name)) + return list(linear_tags) + + +def _prepare_for_hqq_linear(model, patch_params, has_been_replaced, current_key_name=None): + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + + if isinstance(module, torch.nn.Linear): + # Get linear tag + linear_tag = name_to_linear_tag(module.name) + + # We put the module quant_config into the nn.Linear layer so we can access it later in quantizer_hqq.create_quantized_param() + if linear_tag in patch_params: + if patch_params[linear_tag] is not None: + model._modules[name].quant_config = patch_params[linear_tag] + # Store the module class in case we need to transpose the weight later + model._modules[name].source_cls = type(module) + # Force requires grad to False to avoid unexpected errors + model._modules[name].requires_grad_(False) + + has_been_replaced = True + + # Add these fake parameters to avoid loading fail + for att in ["W_q", "meta"]: + setattr(module, att, None) + + if len(list(module.children())) > 0: + _, has_been_replaced = _prepare_for_hqq_linear( + module, + patch_params=patch_params, + has_been_replaced=has_been_replaced, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + + return model, has_been_replaced + + +def prepare_for_hqq_linear(model, quantization_config=None, modules_to_not_convert=None, has_been_replaced=False): + """ + Prepares nn.Linear layers for HQQ quantization. + Since each layer type can have separate quantization parameters, we need to do the following: + 1- tag each module with its neme via autoname_modules() + 2- Extract linear_tags (e.g. ['self_attn.q_proj', ...]) + 3- Map quantization parameters as a dictionary linear_tag -> quant_params as HQQLinear exepects it, this is referred to as patch_params + """ + + modules_to_not_convert = [] if modules_to_not_convert is None else modules_to_not_convert + + # Add name to module + autoname_modules(model) + + # Get linear tags. This allows us to use different quant params to different layer types + linear_tags = get_linear_tags(model) + + # Convert quantization_config to layer-wise config + skip_modules = quantization_config.skip_modules + quant_config = quantization_config.quant_config + linear_tags = list(set(linear_tags) - set(skip_modules) - set(modules_to_not_convert)) + + if any(key in linear_tags for key in quant_config.keys()): + # If the user doesn't specify a key from get_linear_tags, the layer is not quantized via (key, None) + patch_params = dict.fromkeys(linear_tags) + patch_params.update(quant_config) + else: + # Same quant_config for all layers + patch_params = dict.fromkeys(linear_tags, quant_config) + + model, has_been_replaced = _prepare_for_hqq_linear( + model, patch_params=patch_params, has_been_replaced=has_been_replaced + ) + + # We store quantization config as linear_tag -> hqq quant config + model.config.quantization_config = { + "quant_config": quant_config, + "quant_method": quantization_config.quant_method, + "skip_modules": skip_modules, + } + + if not has_been_replaced: + logger.warning("No linear modules were found in your model for quantization.") + + return model diff --git a/lib/python3.12/site-packages/transformers/integrations/hub_kernels.py b/lib/python3.12/site-packages/transformers/integrations/hub_kernels.py new file mode 100644 index 0000000000000000000000000000000000000000..b2ec6b53715aca2068db6e18ac4753d2720b9b09 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/hub_kernels.py @@ -0,0 +1,73 @@ +# Copyright 2025 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. +from typing import Dict, Union + + +try: + from kernels import ( + Device, + LayerRepository, + register_kernel_mapping, + replace_kernel_forward_from_hub, + use_kernel_forward_from_hub, + ) + + _hub_kernels_available = True + + _KERNEL_MAPPING: Dict[str, Dict[Union[Device, str], LayerRepository]] = { + "MultiScaleDeformableAttention": { + "cuda": LayerRepository( + repo_id="kernels-community/deformable-detr", + layer_name="MultiScaleDeformableAttention", + ) + } + } + + register_kernel_mapping(_KERNEL_MAPPING) + +except ImportError: + # Stub to make decorators int transformers work when `kernels` + # is not installed. + def use_kernel_forward_from_hub(*args, **kwargs): + def decorator(cls): + return cls + + return decorator + + class LayerRepository: + def __init__(self, *args, **kwargs): + raise RuntimeError("LayerRepository requires `kernels` to be installed. Run `pip install kernels`.") + + def replace_kernel_forward_from_hub(*args, **kwargs): + raise RuntimeError( + "replace_kernel_forward_from_hub requires `kernels` to be installed. Run `pip install kernels`." + ) + + def register_kernel_mapping(*args, **kwargs): + raise RuntimeError("register_kernel_mapping requires `kernels` to be installed. Run `pip install kernels`.") + + _hub_kernels_available = False + + +def is_hub_kernels_available(): + return _hub_kernels_available + + +__all__ = [ + "LayerRepository", + "is_hub_kernels_available", + "use_kernel_forward_from_hub", + "register_kernel_mapping", + "replace_kernel_forward_from_hub", +] diff --git a/lib/python3.12/site-packages/transformers/integrations/integration_utils.py b/lib/python3.12/site-packages/transformers/integrations/integration_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a696612c3bfd784107421f28eb227e3763ed514b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/integration_utils.py @@ -0,0 +1,2348 @@ +# Copyright 2020 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. +""" +Integrations with other Python libraries. +""" + +import functools +import importlib.metadata +import importlib.util +import json +import numbers +import os +import pickle +import shutil +import sys +import tempfile +from dataclasses import asdict, fields +from enum import Enum +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Union + +import numpy as np +import packaging.version + +from .. import PreTrainedModel, TFPreTrainedModel +from .. import __version__ as version +from ..utils import ( + PushToHubMixin, + flatten_dict, + is_datasets_available, + is_pandas_available, + is_tf_available, + is_torch_available, + logging, +) + + +logger = logging.get_logger(__name__) + +if is_torch_available(): + import torch + +# comet_ml requires to be imported before any ML frameworks +_MIN_COMET_VERSION = "3.43.2" +try: + _comet_version = importlib.metadata.version("comet_ml") + _is_comet_installed = True + + _is_comet_recent_enough = packaging.version.parse(_comet_version) >= packaging.version.parse(_MIN_COMET_VERSION) + + # Check if the Comet API Key is set + import comet_ml + + if comet_ml.config.get_config("comet.api_key") is not None: + _is_comet_configured = True + else: + _is_comet_configured = False +except (importlib.metadata.PackageNotFoundError, ImportError, ValueError, TypeError, AttributeError, KeyError): + _comet_version = None + _is_comet_installed = False + _is_comet_recent_enough = False + _is_comet_configured = False + +_has_neptune = ( + importlib.util.find_spec("neptune") is not None or importlib.util.find_spec("neptune-client") is not None +) +if TYPE_CHECKING and _has_neptune: + try: + _neptune_version = importlib.metadata.version("neptune") + logger.info(f"Neptune version {_neptune_version} available.") + except importlib.metadata.PackageNotFoundError: + try: + _neptune_version = importlib.metadata.version("neptune-client") + logger.info(f"Neptune-client version {_neptune_version} available.") + except importlib.metadata.PackageNotFoundError: + _has_neptune = False + +from .. import modelcard # noqa: E402 +from ..trainer_callback import ProgressCallback, TrainerCallback # noqa: E402 +from ..trainer_utils import PREFIX_CHECKPOINT_DIR, BestRun, IntervalStrategy # noqa: E402 +from ..training_args import ParallelMode # noqa: E402 +from ..utils import ENV_VARS_TRUE_VALUES, is_torch_xla_available # noqa: E402 + + +# Integration functions: +def is_wandb_available(): + # any value of WANDB_DISABLED disables wandb + if os.getenv("WANDB_DISABLED", "").upper() in ENV_VARS_TRUE_VALUES: + logger.warning( + "Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the " + "--report_to flag to control the integrations used for logging result (for instance --report_to none)." + ) + return False + return importlib.util.find_spec("wandb") is not None + + +def is_clearml_available(): + return importlib.util.find_spec("clearml") is not None + + +def is_comet_available(): + if os.getenv("COMET_MODE", "").upper() == "DISABLED": + logger.warning( + "Using the `COMET_MODE=DISABLED` environment variable is deprecated and will be removed in v5. Use the " + "--report_to flag to control the integrations used for logging result (for instance --report_to none)." + ) + return False + + if _is_comet_installed is False: + return False + + if _is_comet_recent_enough is False: + logger.warning( + "comet_ml version %s is installed, but version %s or higher is required. " + "Please update comet_ml to the latest version to enable Comet logging with pip install 'comet-ml>=%s'.", + _comet_version, + _MIN_COMET_VERSION, + _MIN_COMET_VERSION, + ) + return False + + if _is_comet_configured is False: + logger.warning( + "comet_ml is installed but the Comet API Key is not configured. " + "Please set the `COMET_API_KEY` environment variable to enable Comet logging. " + "Check out the documentation for other ways of configuring it: " + "https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#set-the-api-key" + ) + return False + + return True + + +def is_tensorboard_available(): + return importlib.util.find_spec("tensorboard") is not None or importlib.util.find_spec("tensorboardX") is not None + + +def is_optuna_available(): + return importlib.util.find_spec("optuna") is not None + + +def is_ray_available(): + return importlib.util.find_spec("ray") is not None + + +def is_ray_tune_available(): + if not is_ray_available(): + return False + return importlib.util.find_spec("ray.tune") is not None + + +def is_sigopt_available(): + return importlib.util.find_spec("sigopt") is not None + + +def is_azureml_available(): + if importlib.util.find_spec("azureml") is None: + return False + if importlib.util.find_spec("azureml.core") is None: + return False + return importlib.util.find_spec("azureml.core.run") is not None + + +def is_mlflow_available(): + if os.getenv("DISABLE_MLFLOW_INTEGRATION", "FALSE").upper() == "TRUE": + return False + return importlib.util.find_spec("mlflow") is not None + + +def is_dagshub_available(): + return None not in [importlib.util.find_spec("dagshub"), importlib.util.find_spec("mlflow")] + + +def is_neptune_available(): + return _has_neptune + + +def is_codecarbon_available(): + return importlib.util.find_spec("codecarbon") is not None + + +def is_flytekit_available(): + return importlib.util.find_spec("flytekit") is not None + + +def is_flyte_deck_standard_available(): + if not is_flytekit_available(): + return False + return importlib.util.find_spec("flytekitplugins.deck") is not None + + +def is_dvclive_available(): + return importlib.util.find_spec("dvclive") is not None + + +def is_swanlab_available(): + return importlib.util.find_spec("swanlab") is not None + + +def hp_params(trial): + if is_optuna_available(): + import optuna + + if isinstance(trial, optuna.trial.BaseTrial): + return trial.params + if is_ray_tune_available(): + if isinstance(trial, dict): + return trial + + if is_sigopt_available(): + if isinstance(trial, dict): + return trial + + if is_wandb_available(): + if isinstance(trial, dict): + return trial + + raise RuntimeError(f"Unknown type for trial {trial.__class__}") + + +def run_hp_search_optuna(trainer, n_trials: int, direction: str, **kwargs) -> BestRun: + import optuna + from accelerate.utils.memory import release_memory + + if trainer.args.process_index == 0: + + def _objective(trial: optuna.Trial, checkpoint_dir=None): + checkpoint = None + if checkpoint_dir: + for subdir in os.listdir(checkpoint_dir): + if subdir.startswith(PREFIX_CHECKPOINT_DIR): + checkpoint = os.path.join(checkpoint_dir, subdir) + trainer.objective = None + if trainer.args.world_size > 1: + if trainer.args.parallel_mode != ParallelMode.DISTRIBUTED: + raise RuntimeError("only support DDP optuna HPO for ParallelMode.DISTRIBUTED currently.") + trainer.hp_space(trial) + fixed_trial = optuna.trial.FixedTrial(trial.params, trial.number) + trial_main_rank_list = [fixed_trial] + torch.distributed.broadcast_object_list(trial_main_rank_list, src=0) + trainer.train(resume_from_checkpoint=checkpoint, trial=trial) + else: + trainer.train(resume_from_checkpoint=checkpoint, trial=trial) + # If there hasn't been any evaluation during the training loop. + if getattr(trainer, "objective", None) is None: + metrics = trainer.evaluate() + trainer.objective = trainer.compute_objective(metrics) + + # Free GPU memory + trainer.model_wrapped, trainer.model = release_memory(trainer.model_wrapped, trainer.model) + trainer.accelerator.clear() + + return trainer.objective + + timeout = kwargs.pop("timeout", None) + n_jobs = kwargs.pop("n_jobs", 1) + gc_after_trial = kwargs.pop("gc_after_trial", False) + directions = direction if isinstance(direction, list) else None + direction = None if directions is not None else direction + study = optuna.create_study(direction=direction, directions=directions, **kwargs) + study.optimize(_objective, n_trials=n_trials, timeout=timeout, n_jobs=n_jobs, gc_after_trial=gc_after_trial) + if not study._is_multi_objective(): + best_trial = study.best_trial + return BestRun(str(best_trial.number), best_trial.value, best_trial.params) + else: + best_trials = study.best_trials + return [BestRun(str(best.number), best.values, best.params) for best in best_trials] + else: + for i in range(n_trials): + trainer.objective = None + trial_main_rank_list = [None] + if trainer.args.parallel_mode != ParallelMode.DISTRIBUTED: + raise RuntimeError("only support DDP optuna HPO for ParallelMode.DISTRIBUTED currently.") + torch.distributed.broadcast_object_list(trial_main_rank_list, src=0) + trainer.train(resume_from_checkpoint=None, trial=trial_main_rank_list[0]) + # If there hasn't been any evaluation during the training loop. + if getattr(trainer, "objective", None) is None: + metrics = trainer.evaluate() + trainer.objective = trainer.compute_objective(metrics) + return None + + +def run_hp_search_ray(trainer, n_trials: int, direction: str, **kwargs) -> BestRun: + import ray + import ray.train + + def _objective(trial: dict, local_trainer): + try: + from transformers.utils.notebook import NotebookProgressCallback + + if local_trainer.pop_callback(NotebookProgressCallback): + local_trainer.add_callback(ProgressCallback) + except ModuleNotFoundError: + pass + + local_trainer.objective = None + + checkpoint = ray.train.get_checkpoint() + if checkpoint: + # Upon trial resume, the local_trainer's objective gets reset to None. + # If `local_trainer.train` is a noop (training has already reached + # the target number of epochs/steps), then this would + # trigger an unnecessary extra checkpoint at the end of training. + # -> Set the objective to a dummy value upon resume as a workaround. + local_trainer.objective = "objective" + + with checkpoint.as_directory() as checkpoint_dir: + checkpoint_path = next(Path(checkpoint_dir).glob(f"{PREFIX_CHECKPOINT_DIR}*")).as_posix() + local_trainer.train(resume_from_checkpoint=checkpoint_path, trial=trial) + else: + local_trainer.train(trial=trial) + + # If there hasn't been any evaluation during the training loop. + if getattr(local_trainer, "objective", None) is None: + metrics = local_trainer.evaluate() + local_trainer.objective = local_trainer.compute_objective(metrics) + + metrics.update({"objective": local_trainer.objective, "done": True}) + + with tempfile.TemporaryDirectory() as temp_checkpoint_dir: + local_trainer._tune_save_checkpoint(checkpoint_dir=temp_checkpoint_dir) + checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir) + ray.train.report(metrics, checkpoint=checkpoint) + + if not trainer._memory_tracker.skip_memory_metrics: + from ..trainer_utils import TrainerMemoryTracker + + logger.warning( + "Memory tracking for your Trainer is currently " + "enabled. Automatically disabling the memory tracker " + "since the memory tracker is not serializable." + ) + trainer._memory_tracker = TrainerMemoryTracker(skip_memory_metrics=True) + + # The model and TensorBoard writer do not pickle so we have to remove them (if they exists) + # while doing the ray hp search. + _tb_writer = trainer.pop_callback(TensorBoardCallback) + trainer.model = None + + # Setup default `resources_per_trial`. + if "resources_per_trial" not in kwargs: + # Default to 1 CPU and 1 GPU (if applicable) per trial. + kwargs["resources_per_trial"] = {"cpu": 1} + if trainer.args.n_gpu > 0: + kwargs["resources_per_trial"]["gpu"] = 1 + resource_msg = "1 CPU" + (" and 1 GPU" if trainer.args.n_gpu > 0 else "") + logger.info( + "No `resources_per_trial` arg was passed into " + "`hyperparameter_search`. Setting it to a default value " + f"of {resource_msg} for each trial." + ) + # Make sure each trainer only uses GPUs that were allocated per trial. + gpus_per_trial = kwargs["resources_per_trial"].get("gpu", 0) + trainer.args._n_gpu = gpus_per_trial + + # Setup default `progress_reporter`. + if "progress_reporter" not in kwargs: + from ray.tune import CLIReporter + + kwargs["progress_reporter"] = CLIReporter(metric_columns=["objective"]) + + if "scheduler" in kwargs: + from ray.tune.schedulers import ASHAScheduler, HyperBandForBOHB, MedianStoppingRule, PopulationBasedTraining + + # Check for `do_eval` and `eval_during_training` for schedulers that require intermediate reporting. + if isinstance( + kwargs["scheduler"], (ASHAScheduler, MedianStoppingRule, HyperBandForBOHB, PopulationBasedTraining) + ) and (not trainer.args.do_eval or trainer.args.eval_strategy == IntervalStrategy.NO): + raise RuntimeError( + "You are using {cls} as a scheduler but you haven't enabled evaluation during training. " + "This means your trials will not report intermediate results to Ray Tune, and " + "can thus not be stopped early or used to exploit other trials parameters. " + "If this is what you want, do not use {cls}. If you would like to use {cls}, " + "make sure you pass `do_eval=True` and `eval_strategy='steps'` in the " + "Trainer `args`.".format(cls=type(kwargs["scheduler"]).__name__) + ) + + trainable = ray.tune.with_parameters(_objective, local_trainer=trainer) + + @functools.wraps(trainable) + def dynamic_modules_import_trainable(*args, **kwargs): + """ + Wrapper around `tune.with_parameters` to ensure datasets_modules are loaded on each Actor. + + Without this, an ImportError will be thrown. See https://github.com/huggingface/transformers/issues/11565. + + Assumes that `_objective`, defined above, is a function. + """ + if is_datasets_available(): + import datasets.load + + dynamic_modules_path = os.path.join(datasets.load.init_dynamic_modules(), "__init__.py") + # load dynamic_modules from path + spec = importlib.util.spec_from_file_location("datasets_modules", dynamic_modules_path) + datasets_modules = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = datasets_modules + spec.loader.exec_module(datasets_modules) + return trainable(*args, **kwargs) + + # special attr set by tune.with_parameters + if hasattr(trainable, "__mixins__"): + dynamic_modules_import_trainable.__mixins__ = trainable.__mixins__ + + analysis = ray.tune.run( + dynamic_modules_import_trainable, + config=trainer.hp_space(None), + num_samples=n_trials, + **kwargs, + ) + best_trial = analysis.get_best_trial(metric="objective", mode=direction[:3], scope=trainer.args.ray_scope) + best_run = BestRun(best_trial.trial_id, best_trial.last_result["objective"], best_trial.config, analysis) + if _tb_writer is not None: + trainer.add_callback(_tb_writer) + return best_run + + +def run_hp_search_sigopt(trainer, n_trials: int, direction: str, **kwargs) -> BestRun: + import sigopt + + if trainer.args.process_index == 0: + if importlib.metadata.version("sigopt") >= "8.0.0": + sigopt.set_project("huggingface") + + experiment = sigopt.create_experiment( + name="huggingface-tune", + type="offline", + parameters=trainer.hp_space(None), + metrics=[{"name": "objective", "objective": direction, "strategy": "optimize"}], + parallel_bandwidth=1, + budget=n_trials, + ) + + logger.info(f"created experiment: https://app.sigopt.com/experiment/{experiment.id}") + + for run in experiment.loop(): + with run: + trainer.objective = None + if trainer.args.world_size > 1: + if trainer.args.parallel_mode != ParallelMode.DISTRIBUTED: + raise RuntimeError("only support DDP Sigopt HPO for ParallelMode.DISTRIBUTED currently.") + trainer._hp_search_setup(run.run) + torch.distributed.broadcast_object_list(pickle.dumps(trainer.args), src=0) + trainer.train(resume_from_checkpoint=None) + else: + trainer.train(resume_from_checkpoint=None, trial=run.run) + # If there hasn't been any evaluation during the training loop. + if getattr(trainer, "objective", None) is None: + metrics = trainer.evaluate() + trainer.objective = trainer.compute_objective(metrics) + run.log_metric("objective", trainer.objective) + + best = list(experiment.get_best_runs())[0] + best_run = BestRun(best.id, best.values["objective"].value, best.assignments) + else: + from sigopt import Connection + + conn = Connection() + proxies = kwargs.pop("proxies", None) + if proxies is not None: + conn.set_proxies(proxies) + + experiment = conn.experiments().create( + name="huggingface-tune", + parameters=trainer.hp_space(None), + metrics=[{"name": "objective", "objective": direction, "strategy": "optimize"}], + parallel_bandwidth=1, + observation_budget=n_trials, + project="huggingface", + ) + logger.info(f"created experiment: https://app.sigopt.com/experiment/{experiment.id}") + + while experiment.progress.observation_count < experiment.observation_budget: + suggestion = conn.experiments(experiment.id).suggestions().create() + trainer.objective = None + if trainer.args.world_size > 1: + if trainer.args.parallel_mode != ParallelMode.DISTRIBUTED: + raise RuntimeError("only support DDP Sigopt HPO for ParallelMode.DISTRIBUTED currently.") + trainer._hp_search_setup(suggestion) + torch.distributed.broadcast_object_list(pickle.dumps(trainer.args), src=0) + trainer.train(resume_from_checkpoint=None) + else: + trainer.train(resume_from_checkpoint=None, trial=suggestion) + # If there hasn't been any evaluation during the training loop. + if getattr(trainer, "objective", None) is None: + metrics = trainer.evaluate() + trainer.objective = trainer.compute_objective(metrics) + + values = [{"name": "objective", "value": trainer.objective}] + obs = conn.experiments(experiment.id).observations().create(suggestion=suggestion.id, values=values) + logger.info(f"[suggestion_id, observation_id]: [{suggestion.id}, {obs.id}]") + experiment = conn.experiments(experiment.id).fetch() + + best = list(conn.experiments(experiment.id).best_assignments().fetch().iterate_pages())[0] + best_run = BestRun(best.id, best.value, best.assignments) + return best_run + else: + for i in range(n_trials): + trainer.objective = None + args_main_rank = list(pickle.dumps(trainer.args)) + if trainer.args.parallel_mode != ParallelMode.DISTRIBUTED: + raise RuntimeError("only support DDP Sigopt HPO for ParallelMode.DISTRIBUTED currently.") + torch.distributed.broadcast_object_list(args_main_rank, src=0) + args = pickle.loads(bytes(args_main_rank)) + for key, value in asdict(args).items(): + if key != "local_rank": + setattr(trainer.args, key, value) + trainer.train(resume_from_checkpoint=None) + # If there hasn't been any evaluation during the training loop. + if getattr(trainer, "objective", None) is None: + metrics = trainer.evaluate() + trainer.objective = trainer.compute_objective(metrics) + return None + + +def run_hp_search_wandb(trainer, n_trials: int, direction: str, **kwargs) -> BestRun: + from ..integrations import is_wandb_available + + if not is_wandb_available(): + raise ImportError("This function needs wandb installed: `pip install wandb`") + import wandb + + # add WandbCallback if not already added in trainer callbacks + reporting_to_wandb = False + for callback in trainer.callback_handler.callbacks: + if isinstance(callback, WandbCallback): + reporting_to_wandb = True + break + if not reporting_to_wandb: + trainer.add_callback(WandbCallback()) + trainer.args.report_to = ["wandb"] + best_trial = {"run_id": None, "objective": None, "hyperparameters": None} + sweep_id = kwargs.pop("sweep_id", None) + project = kwargs.pop("project", None) + name = kwargs.pop("name", None) + entity = kwargs.pop("entity", None) + metric = kwargs.pop("metric", "eval/loss") + + sweep_config = trainer.hp_space(None) + sweep_config["metric"]["goal"] = direction + sweep_config["metric"]["name"] = metric + if name: + sweep_config["name"] = name + + def _objective(): + run = wandb.run if wandb.run else wandb.init() + trainer.state.trial_name = run.name + run.config.update({"assignments": {}, "metric": metric}) + config = wandb.config + + trainer.objective = None + + trainer.train(resume_from_checkpoint=None, trial=vars(config)["_items"]) + # If there hasn't been any evaluation during the training loop. + if getattr(trainer, "objective", None) is None: + metrics = trainer.evaluate() + trainer.objective = trainer.compute_objective(metrics) + format_metrics = rewrite_logs(metrics) + if metric not in format_metrics: + logger.warning( + f"Provided metric {metric} not found. This might result in unexpected sweeps charts. The available" + f" metrics are {format_metrics.keys()}" + ) + best_score = False + if best_trial["run_id"] is not None: + if direction == "minimize": + best_score = trainer.objective < best_trial["objective"] + elif direction == "maximize": + best_score = trainer.objective > best_trial["objective"] + + if best_score or best_trial["run_id"] is None: + best_trial["run_id"] = run.id + best_trial["objective"] = trainer.objective + best_trial["hyperparameters"] = dict(config) + + return trainer.objective + + if not sweep_id: + sweep_id = wandb.sweep(sweep_config, project=project, entity=entity) + else: + import wandb.env + + if entity: + wandb.env.set_entity(entity) + wandb.env.set_project(project) + + logger.info(f"wandb sweep id - {sweep_id}") + wandb.agent(sweep_id, function=_objective, count=n_trials) + + return BestRun(best_trial["run_id"], best_trial["objective"], best_trial["hyperparameters"], sweep_id) + + +def get_available_reporting_integrations(): + integrations = [] + if is_azureml_available() and not is_mlflow_available(): + integrations.append("azure_ml") + if is_comet_available(): + integrations.append("comet_ml") + if is_dagshub_available(): + integrations.append("dagshub") + if is_dvclive_available(): + integrations.append("dvclive") + if is_mlflow_available(): + integrations.append("mlflow") + if is_neptune_available(): + integrations.append("neptune") + if is_tensorboard_available(): + integrations.append("tensorboard") + if is_wandb_available(): + integrations.append("wandb") + if is_codecarbon_available(): + integrations.append("codecarbon") + if is_clearml_available(): + integrations.append("clearml") + if is_swanlab_available(): + integrations.append("swanlab") + return integrations + + +def rewrite_logs(d): + new_d = {} + eval_prefix = "eval_" + eval_prefix_len = len(eval_prefix) + test_prefix = "test_" + test_prefix_len = len(test_prefix) + for k, v in d.items(): + if k.startswith(eval_prefix): + new_d["eval/" + k[eval_prefix_len:]] = v + elif k.startswith(test_prefix): + new_d["test/" + k[test_prefix_len:]] = v + else: + new_d["train/" + k] = v + return new_d + + +class TensorBoardCallback(TrainerCallback): + """ + A [`TrainerCallback`] that sends the logs to [TensorBoard](https://www.tensorflow.org/tensorboard). + + Args: + tb_writer (`SummaryWriter`, *optional*): + The writer to use. Will instantiate one if not set. + """ + + def __init__(self, tb_writer=None): + has_tensorboard = is_tensorboard_available() + if not has_tensorboard: + raise RuntimeError( + "TensorBoardCallback requires tensorboard to be installed. Either update your PyTorch version or" + " install tensorboardX." + ) + if has_tensorboard: + try: + from torch.utils.tensorboard import SummaryWriter # noqa: F401 + + self._SummaryWriter = SummaryWriter + except ImportError: + try: + from tensorboardX import SummaryWriter + + self._SummaryWriter = SummaryWriter + except ImportError: + self._SummaryWriter = None + else: + self._SummaryWriter = None + self.tb_writer = tb_writer + + def _init_summary_writer(self, args, log_dir=None): + log_dir = log_dir or args.logging_dir + if self._SummaryWriter is not None: + self.tb_writer = self._SummaryWriter(log_dir=log_dir) + + def on_train_begin(self, args, state, control, **kwargs): + if not state.is_world_process_zero: + return + + log_dir = None + + if state.is_hyper_param_search: + trial_name = state.trial_name + if trial_name is not None: + log_dir = os.path.join(args.logging_dir, trial_name) + + if self.tb_writer is None: + self._init_summary_writer(args, log_dir) + + if self.tb_writer is not None: + self.tb_writer.add_text("args", args.to_json_string()) + if "model" in kwargs: + model = kwargs["model"] + if hasattr(model, "config") and model.config is not None: + model_config_json = model.config.to_json_string() + self.tb_writer.add_text("model_config", model_config_json) + + def on_log(self, args, state, control, logs=None, **kwargs): + if not state.is_world_process_zero: + return + + if self.tb_writer is None: + self._init_summary_writer(args) + + if self.tb_writer is not None: + logs = rewrite_logs(logs) + for k, v in logs.items(): + if isinstance(v, (int, float)): + self.tb_writer.add_scalar(k, v, state.global_step) + elif isinstance(v, str): + self.tb_writer.add_text(k, v, state.global_step) + else: + logger.warning( + "Trainer is attempting to log a value of " + f'"{v}" of type {type(v)} for key "{k}" as a scalar. ' + "This invocation of Tensorboard's writer.add_scalar() " + "is incorrect so we dropped this attribute." + ) + self.tb_writer.flush() + + def on_train_end(self, args, state, control, **kwargs): + if self.tb_writer: + self.tb_writer.close() + self.tb_writer = None + + +def save_model_architecture_to_file(model: Any, output_dir: str): + with open(f"{output_dir}/model_architecture.txt", "w+") as f: + if isinstance(model, PreTrainedModel): + print(model, file=f) + elif is_tf_available() and isinstance(model, TFPreTrainedModel): + + def print_to_file(s): + print(s, file=f) + + model.summary(print_fn=print_to_file) + elif is_torch_available() and ( + isinstance(model, (torch.nn.Module, PushToHubMixin)) and hasattr(model, "base_model") + ): + print(model, file=f) + + +class WandbLogModel(str, Enum): + """Enum of possible log model values in W&B.""" + + CHECKPOINT = "checkpoint" + END = "end" + FALSE = "false" + + @property + def is_enabled(self) -> bool: + """Check if the value corresponds to a state where the `WANDB_LOG_MODEL` setting is enabled.""" + return self in (WandbLogModel.CHECKPOINT, WandbLogModel.END) + + @classmethod + def _missing_(cls, value: Any) -> "WandbLogModel": + if not isinstance(value, str): + raise ValueError(f"Expecting to have a string `WANDB_LOG_MODEL` setting, but got {type(value)}") + if value.upper() in ENV_VARS_TRUE_VALUES: + raise DeprecationWarning( + f"Setting `WANDB_LOG_MODEL` as {os.getenv('WANDB_LOG_MODEL')} is deprecated and will be removed in " + "version 5 of transformers. Use one of `'end'` or `'checkpoint'` instead." + ) + logger.info(f"Setting `WANDB_LOG_MODEL` from {os.getenv('WANDB_LOG_MODEL')} to `end` instead") + return WandbLogModel.END + logger.warning( + f"Received unrecognized `WANDB_LOG_MODEL` setting value={value}; so disabling `WANDB_LOG_MODEL`" + ) + return WandbLogModel.FALSE + + +class WandbCallback(TrainerCallback): + """ + A [`TrainerCallback`] that logs metrics, media, model checkpoints to [Weight and Biases](https://www.wandb.com/). + """ + + def __init__(self): + has_wandb = is_wandb_available() + if not has_wandb: + raise RuntimeError("WandbCallback requires wandb to be installed. Run `pip install wandb`.") + if has_wandb: + import wandb + + self._wandb = wandb + self._initialized = False + self._log_model = WandbLogModel(os.getenv("WANDB_LOG_MODEL", "false")) + + def setup(self, args, state, model, **kwargs): + """ + Setup the optional Weights & Biases (*wandb*) integration. + + One can subclass and override this method to customize the setup if needed. Find more information + [here](https://docs.wandb.ai/guides/integrations/huggingface). You can also override the following environment + variables: + + Environment: + - **WANDB_LOG_MODEL** (`str`, *optional*, defaults to `"false"`): + Whether to log model and checkpoints during training. Can be `"end"`, `"checkpoint"` or `"false"`. If set + to `"end"`, the model will be uploaded at the end of training. If set to `"checkpoint"`, the checkpoint + will be uploaded every `args.save_steps` . If set to `"false"`, the model will not be uploaded. Use along + with [`~transformers.TrainingArguments.load_best_model_at_end`] to upload best model. + + + + Setting `WANDB_LOG_MODEL` as `bool` will be deprecated in version 5 of 🤗 Transformers. + + + - **WANDB_WATCH** (`str`, *optional* defaults to `"false"`): + Can be `"gradients"`, `"all"`, `"parameters"`, or `"false"`. Set to `"all"` to log gradients and + parameters. + - **WANDB_PROJECT** (`str`, *optional*, defaults to `"huggingface"`): + Set this to a custom string to store results in a different project. + - **WANDB_DISABLED** (`bool`, *optional*, defaults to `False`): + Whether to disable wandb entirely. Set `WANDB_DISABLED=true` to disable. + """ + if self._wandb is None: + return + self._initialized = True + + # prepare to handle potential configuration issues during setup + from wandb.sdk.lib.config_util import ConfigError as WandbConfigError + + if state.is_world_process_zero: + logger.info( + 'Automatic Weights & Biases logging enabled, to disable set os.environ["WANDB_DISABLED"] = "true"' + ) + combined_dict = {**args.to_dict()} + + if hasattr(model, "config") and model.config is not None: + model_config = model.config if isinstance(model.config, dict) else model.config.to_dict() + combined_dict = {**model_config, **combined_dict} + if hasattr(model, "peft_config") and model.peft_config is not None: + peft_config = model.peft_config + combined_dict = {**{"peft_config": peft_config}, **combined_dict} + trial_name = state.trial_name + init_args = {} + if trial_name is not None: + init_args["name"] = trial_name + init_args["group"] = args.run_name + elif args.run_name is not None: + init_args["name"] = args.run_name + if args.run_name == args.output_dir: + self._wandb.termwarn( + "The `run_name` is currently set to the same value as `TrainingArguments.output_dir`. If this was " + "not intended, please specify a different run name by setting the `TrainingArguments.run_name` parameter.", + repeat=False, + ) + + if self._wandb.run is None: + self._wandb.init( + project=os.getenv("WANDB_PROJECT", "huggingface"), + **init_args, + ) + # add config parameters (run may have been created manually) + self._wandb.config.update(combined_dict, allow_val_change=True) + + # define default x-axis (for latest wandb versions) + if getattr(self._wandb, "define_metric", None): + self._wandb.define_metric("train/global_step") + self._wandb.define_metric("*", step_metric="train/global_step", step_sync=True) + + # keep track of model topology and gradients, unsupported on TPU + _watch_model = os.getenv("WANDB_WATCH", "false") + if not is_torch_xla_available() and _watch_model in ("all", "parameters", "gradients"): + self._wandb.watch(model, log=_watch_model, log_freq=max(100, state.logging_steps)) + self._wandb.run._label(code="transformers_trainer") + + # add number of model parameters to wandb config + try: + self._wandb.config["model/num_parameters"] = model.num_parameters() + except AttributeError: + logger.info( + "Could not log the number of model parameters in Weights & Biases due to an AttributeError." + ) + except WandbConfigError: + logger.warning( + "A ConfigError was raised whilst setting the number of model parameters in Weights & Biases config." + ) + + # log the initial model architecture to an artifact + if self._log_model.is_enabled: + with tempfile.TemporaryDirectory() as temp_dir: + model_name = ( + f"model-{self._wandb.run.id}" + if (args.run_name is None or args.run_name == args.output_dir) + else f"model-{self._wandb.run.name}" + ) + model_artifact = self._wandb.Artifact( + name=model_name, + type="model", + metadata={ + "model_config": model.config.to_dict() if hasattr(model, "config") else None, + "num_parameters": self._wandb.config.get("model/num_parameters"), + "initial_model": True, + }, + ) + # add the architecture to a separate text file + save_model_architecture_to_file(model, temp_dir) + + for f in Path(temp_dir).glob("*"): + if f.is_file(): + with model_artifact.new_file(f.name, mode="wb") as fa: + fa.write(f.read_bytes()) + self._wandb.run.log_artifact(model_artifact, aliases=["base_model"]) + + badge_markdown = ( + f'[Visualize in Weights & Biases]({self._wandb.run.get_url()})' + ) + + modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n{badge_markdown}" + + def on_train_begin(self, args, state, control, model=None, **kwargs): + if self._wandb is None: + return + hp_search = state.is_hyper_param_search + if hp_search: + self._wandb.finish() + self._initialized = False + args.run_name = None + if not self._initialized: + self.setup(args, state, model, **kwargs) + + def on_train_end(self, args, state, control, model=None, processing_class=None, **kwargs): + if self._wandb is None: + return + if self._log_model.is_enabled and self._initialized and state.is_world_process_zero: + from ..trainer import Trainer + + fake_trainer = Trainer(args=args, model=model, processing_class=processing_class, eval_dataset=["fake"]) + with tempfile.TemporaryDirectory() as temp_dir: + fake_trainer.save_model(temp_dir) + metadata = ( + { + k: v + for k, v in dict(self._wandb.summary).items() + if isinstance(v, numbers.Number) and not k.startswith("_") + } + if not args.load_best_model_at_end + else { + f"eval/{args.metric_for_best_model}": state.best_metric, + "train/total_floss": state.total_flos, + "model/num_parameters": self._wandb.config.get("model/num_parameters"), + } + ) + metadata["final_model"] = True + logger.info("Logging model artifacts. ...") + model_name = ( + f"model-{self._wandb.run.id}" + if (args.run_name is None or args.run_name == args.output_dir) + else f"model-{self._wandb.run.name}" + ) + # add the model architecture to a separate text file + save_model_architecture_to_file(model, temp_dir) + + artifact = self._wandb.Artifact(name=model_name, type="model", metadata=metadata) + for f in Path(temp_dir).glob("*"): + if f.is_file(): + with artifact.new_file(f.name, mode="wb") as fa: + fa.write(f.read_bytes()) + self._wandb.run.log_artifact(artifact, aliases=["final_model"]) + + def on_log(self, args, state, control, model=None, logs=None, **kwargs): + single_value_scalars = [ + "train_runtime", + "train_samples_per_second", + "train_steps_per_second", + "train_loss", + "total_flos", + ] + + if self._wandb is None: + return + if not self._initialized: + self.setup(args, state, model) + if state.is_world_process_zero: + for k, v in logs.items(): + if k in single_value_scalars: + self._wandb.run.summary[k] = v + non_scalar_logs = {k: v for k, v in logs.items() if k not in single_value_scalars} + non_scalar_logs = rewrite_logs(non_scalar_logs) + self._wandb.log({**non_scalar_logs, "train/global_step": state.global_step}) + + def on_save(self, args, state, control, **kwargs): + if self._log_model == WandbLogModel.CHECKPOINT and self._initialized and state.is_world_process_zero: + checkpoint_metadata = { + k: v + for k, v in dict(self._wandb.summary).items() + if isinstance(v, numbers.Number) and not k.startswith("_") + } + checkpoint_metadata["model/num_parameters"] = self._wandb.config.get("model/num_parameters") + + ckpt_dir = f"checkpoint-{state.global_step}" + artifact_path = os.path.join(args.output_dir, ckpt_dir) + logger.info(f"Logging checkpoint artifacts in {ckpt_dir}. ...") + checkpoint_name = ( + f"model-{self._wandb.run.id}" + if (args.run_name is None or args.run_name == args.output_dir) + else f"model-{self._wandb.run.name}" + ) + artifact = self._wandb.Artifact(name=checkpoint_name, type="model", metadata=checkpoint_metadata) + artifact.add_dir(artifact_path) + self._wandb.log_artifact( + artifact, aliases=[f"epoch_{round(state.epoch, 2)}", f"checkpoint_global_step_{state.global_step}"] + ) + + def on_predict(self, args, state, control, metrics, **kwargs): + if self._wandb is None: + return + if not self._initialized: + self.setup(args, state, **kwargs) + if state.is_world_process_zero: + metrics = rewrite_logs(metrics) + self._wandb.log(metrics) + + +class CometCallback(TrainerCallback): + """ + A [`TrainerCallback`] that sends the logs to [Comet ML](https://www.comet.com/site/). + """ + + def __init__(self): + if _is_comet_installed is False or _is_comet_recent_enough is False: + raise RuntimeError( + f"CometCallback requires comet-ml>={_MIN_COMET_VERSION} to be installed. Run `pip install comet-ml>={_MIN_COMET_VERSION}`." + ) + self._initialized = False + self._log_assets = False + self._experiment = None + + def setup(self, args, state, model): + """ + Setup the optional Comet integration. + + Environment: + - **COMET_MODE** (`str`, *optional*, default to `get_or_create`): + Control whether to create and log to a new Comet experiment or append to an existing experiment. + It accepts the following values: + * `get_or_create`: Decides automatically depending if + `COMET_EXPERIMENT_KEY` is set and whether an Experiment + with that key already exists or not. + * `create`: Always create a new Comet Experiment. + * `get`: Always try to append to an Existing Comet Experiment. + Requires `COMET_EXPERIMENT_KEY` to be set. + * `ONLINE`: **deprecated**, used to create an online + Experiment. Use `COMET_START_ONLINE=1` instead. + * `OFFLINE`: **deprecated**, used to created an offline + Experiment. Use `COMET_START_ONLINE=0` instead. + * `DISABLED`: **deprecated**, used to disable Comet logging. + Use the `--report_to` flag to control the integrations used + for logging result instead. + - **COMET_PROJECT_NAME** (`str`, *optional*): + Comet project name for experiments. + - **COMET_LOG_ASSETS** (`str`, *optional*, defaults to `TRUE`): + Whether or not to log training assets (tf event logs, checkpoints, etc), to Comet. Can be `TRUE`, or + `FALSE`. + + For a number of configurable items in the environment, see + [here](https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options). + """ + self._initialized = True + log_assets = os.getenv("COMET_LOG_ASSETS", "FALSE").upper() + if log_assets in {"TRUE", "1"}: + self._log_assets = True + if state.is_world_process_zero: + comet_old_mode = os.getenv("COMET_MODE") + + mode = None + online = None + + if comet_old_mode is not None: + comet_old_mode = comet_old_mode.lower() + + if comet_old_mode == "online": + online = True + elif comet_old_mode == "offline": + online = False + elif comet_old_mode in ("get", "get_or_create", "create"): + mode = comet_old_mode + elif comet_old_mode: + logger.warning("Invalid COMET_MODE env value %r, Comet logging is disabled", comet_old_mode) + return + + # For HPO, we always create a new experiment for each trial + if state.is_hyper_param_search: + if mode is not None: + logger.warning( + "Hyperparameter Search is enabled, forcing the creation of new experimetns, COMET_MODE value %r is ignored", + comet_old_mode, + ) + mode = "create" + + import comet_ml + + # Do not use the default run_name as the experiment name + if args.run_name is not None and args.run_name != args.output_dir: + experiment_config = comet_ml.ExperimentConfig(name=args.run_name) + else: + experiment_config = comet_ml.ExperimentConfig() + + self._experiment = comet_ml.start(online=online, mode=mode, experiment_config=experiment_config) + self._experiment.__internal_api__set_model_graph__(model, framework="transformers") + + params = {"args": args.to_dict()} + + if hasattr(model, "config") and model.config is not None: + model_config = model.config.to_dict() + params["config"] = model_config + if hasattr(model, "peft_config") and model.peft_config is not None: + peft_config = model.peft_config + params["peft_config"] = peft_config + + self._experiment.__internal_api__log_parameters__( + params, framework="transformers", source="manual", flatten_nested=True + ) + + if state.is_hyper_param_search: + optimization_id = getattr(state, "trial_name", None) + optimization_params = getattr(state, "trial_params", None) + + self._experiment.log_optimization(optimization_id=optimization_id, parameters=optimization_params) + + def on_train_begin(self, args, state, control, model=None, **kwargs): + if not self._initialized: + self.setup(args, state, model) + + def on_log(self, args, state, control, model=None, logs=None, **kwargs): + if not self._initialized: + self.setup(args, state, model) + if state.is_world_process_zero: + if self._experiment is not None: + rewritten_logs = rewrite_logs(logs) + self._experiment.__internal_api__log_metrics__( + rewritten_logs, step=state.global_step, epoch=state.epoch, framework="transformers" + ) + + def on_train_end(self, args, state, control, **kwargs): + if self._initialized and state.is_world_process_zero: + if self._experiment is not None: + if self._log_assets is True: + logger.info("Logging checkpoints. This may take time.") + self._experiment.log_asset_folder( + args.output_dir, recursive=True, log_file_name=True, step=state.global_step + ) + + # We create one experiment per trial in HPO mode + if state.is_hyper_param_search: + self._experiment.clean() + self._initialized = False + + def on_predict(self, args, state, control, metrics, **kwargs): + if not self._initialized: + self.setup(args, state, model=None) + if state.is_world_process_zero and self._experiment is not None: + rewritten_metrics = rewrite_logs(metrics) + self._experiment.__internal_api__log_metrics__( + rewritten_metrics, step=state.global_step, epoch=state.epoch, framework="transformers" + ) + + +class AzureMLCallback(TrainerCallback): + """ + A [`TrainerCallback`] that sends the logs to [AzureML](https://pypi.org/project/azureml-sdk/). + """ + + def __init__(self, azureml_run=None): + if not is_azureml_available(): + raise RuntimeError("AzureMLCallback requires azureml to be installed. Run `pip install azureml-sdk`.") + self.azureml_run = azureml_run + + def on_init_end(self, args, state, control, **kwargs): + from azureml.core.run import Run + + if self.azureml_run is None and state.is_world_process_zero: + self.azureml_run = Run.get_context() + + def on_log(self, args, state, control, logs=None, **kwargs): + if self.azureml_run and state.is_world_process_zero: + for k, v in logs.items(): + if isinstance(v, (int, float)): + self.azureml_run.log(k, v, description=k) + + +class MLflowCallback(TrainerCallback): + """ + A [`TrainerCallback`] that sends the logs to [MLflow](https://www.mlflow.org/). Can be disabled by setting + environment variable `DISABLE_MLFLOW_INTEGRATION = TRUE`. + """ + + def __init__(self): + if not is_mlflow_available(): + raise RuntimeError("MLflowCallback requires mlflow to be installed. Run `pip install mlflow`.") + import mlflow + + self._MAX_PARAM_VAL_LENGTH = mlflow.utils.validation.MAX_PARAM_VAL_LENGTH + self._MAX_PARAMS_TAGS_PER_BATCH = mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH + + self._initialized = False + self._auto_end_run = False + self._log_artifacts = False + self._ml_flow = mlflow + + def setup(self, args, state, model): + """ + Setup the optional MLflow integration. + + Environment: + - **HF_MLFLOW_LOG_ARTIFACTS** (`str`, *optional*): + Whether to use MLflow `.log_artifact()` facility to log artifacts. This only makes sense if logging to a + remote server, e.g. s3 or GCS. If set to `True` or *1*, will copy each saved checkpoint on each save in + [`TrainingArguments`]'s `output_dir` to the local or remote artifact storage. Using it without a remote + storage will just copy the files to your artifact location. + - **MLFLOW_TRACKING_URI** (`str`, *optional*): + Whether to store runs at a specific path or remote server. Unset by default, which skips setting the + tracking URI entirely. + - **MLFLOW_EXPERIMENT_NAME** (`str`, *optional*, defaults to `None`): + Whether to use an MLflow experiment_name under which to launch the run. Default to `None` which will point + to the `Default` experiment in MLflow. Otherwise, it is a case sensitive name of the experiment to be + activated. If an experiment with this name does not exist, a new experiment with this name is created. + - **MLFLOW_TAGS** (`str`, *optional*): + A string dump of a dictionary of key/value pair to be added to the MLflow run as tags. Example: + `os.environ['MLFLOW_TAGS']='{"release.candidate": "RC1", "release.version": "2.2.0"}'`. + - **MLFLOW_NESTED_RUN** (`str`, *optional*): + Whether to use MLflow nested runs. If set to `True` or *1*, will create a nested run inside the current + run. + - **MLFLOW_RUN_ID** (`str`, *optional*): + Allow to reattach to an existing run which can be useful when resuming training from a checkpoint. When + `MLFLOW_RUN_ID` environment variable is set, `start_run` attempts to resume a run with the specified run ID + and other parameters are ignored. + - **MLFLOW_FLATTEN_PARAMS** (`str`, *optional*, defaults to `False`): + Whether to flatten the parameters dictionary before logging. + - **MLFLOW_MAX_LOG_PARAMS** (`int`, *optional*): + Set the maximum number of parameters to log in the run. + """ + self._log_artifacts = os.getenv("HF_MLFLOW_LOG_ARTIFACTS", "FALSE").upper() in ENV_VARS_TRUE_VALUES + self._nested_run = os.getenv("MLFLOW_NESTED_RUN", "FALSE").upper() in ENV_VARS_TRUE_VALUES + self._tracking_uri = os.getenv("MLFLOW_TRACKING_URI", None) + self._experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME", None) + self._flatten_params = os.getenv("MLFLOW_FLATTEN_PARAMS", "FALSE").upper() in ENV_VARS_TRUE_VALUES + self._run_id = os.getenv("MLFLOW_RUN_ID", None) + self._max_log_params = os.getenv("MLFLOW_MAX_LOG_PARAMS", None) + + # "synchronous" flag is only available with mlflow version >= 2.8.0 + # https://github.com/mlflow/mlflow/pull/9705 + # https://github.com/mlflow/mlflow/releases/tag/v2.8.0 + self._async_log = packaging.version.parse(self._ml_flow.__version__) >= packaging.version.parse("2.8.0") + + logger.debug( + f"MLflow experiment_name={self._experiment_name}, run_name={args.run_name}, nested={self._nested_run}," + f" tracking_uri={self._tracking_uri}" + ) + if state.is_world_process_zero: + if not self._ml_flow.is_tracking_uri_set(): + if self._tracking_uri: + self._ml_flow.set_tracking_uri(self._tracking_uri) + logger.debug(f"MLflow tracking URI is set to {self._tracking_uri}") + else: + logger.debug( + "Environment variable `MLFLOW_TRACKING_URI` is not provided and therefore will not be" + " explicitly set." + ) + else: + logger.debug(f"MLflow tracking URI is set to {self._ml_flow.get_tracking_uri()}") + + if self._ml_flow.active_run() is None or self._nested_run or self._run_id: + if self._experiment_name: + # Use of set_experiment() ensure that Experiment is created if not exists + self._ml_flow.set_experiment(self._experiment_name) + self._ml_flow.start_run(run_name=args.run_name, nested=self._nested_run) + logger.debug(f"MLflow run started with run_id={self._ml_flow.active_run().info.run_id}") + self._auto_end_run = True + combined_dict = args.to_dict() + if hasattr(model, "config") and model.config is not None: + model_config = model.config.to_dict() + combined_dict = {**model_config, **combined_dict} + combined_dict = flatten_dict(combined_dict) if self._flatten_params else combined_dict + # remove params that are too long for MLflow + for name, value in list(combined_dict.items()): + # internally, all values are converted to str in MLflow + if len(str(value)) > self._MAX_PARAM_VAL_LENGTH: + logger.warning( + f'Trainer is attempting to log a value of "{value}" for key "{name}" as a parameter. MLflow\'s' + " log_param() only accepts values no longer than 250 characters so we dropped this attribute." + " You can use `MLFLOW_FLATTEN_PARAMS` environment variable to flatten the parameters and" + " avoid this message." + ) + del combined_dict[name] + # MLflow cannot log more than 100 values in one go, so we have to split it + combined_dict_items = list(combined_dict.items()) + if self._max_log_params and self._max_log_params.isdigit(): + max_log_params = int(self._max_log_params) + if max_log_params < len(combined_dict_items): + logger.debug( + f"Reducing the number of parameters to log from {len(combined_dict_items)} to {max_log_params}." + ) + combined_dict_items = combined_dict_items[:max_log_params] + for i in range(0, len(combined_dict_items), self._MAX_PARAMS_TAGS_PER_BATCH): + if self._async_log: + self._ml_flow.log_params( + dict(combined_dict_items[i : i + self._MAX_PARAMS_TAGS_PER_BATCH]), synchronous=False + ) + else: + self._ml_flow.log_params(dict(combined_dict_items[i : i + self._MAX_PARAMS_TAGS_PER_BATCH])) + mlflow_tags = os.getenv("MLFLOW_TAGS", None) + if mlflow_tags: + mlflow_tags = json.loads(mlflow_tags) + self._ml_flow.set_tags(mlflow_tags) + self._initialized = True + + def on_train_begin(self, args, state, control, model=None, **kwargs): + if not self._initialized: + self.setup(args, state, model) + + def on_log(self, args, state, control, logs, model=None, **kwargs): + if not self._initialized: + self.setup(args, state, model) + if state.is_world_process_zero: + metrics = {} + for k, v in logs.items(): + if isinstance(v, (int, float)): + metrics[k] = v + elif isinstance(v, torch.Tensor) and v.numel() == 1: + metrics[k] = v.item() + else: + logger.warning( + f'Trainer is attempting to log a value of "{v}" of type {type(v)} for key "{k}" as a metric. ' + "MLflow's log_metric() only accepts float and int types so we dropped this attribute." + ) + + if self._async_log: + self._ml_flow.log_metrics(metrics=metrics, step=state.global_step, synchronous=False) + else: + self._ml_flow.log_metrics(metrics=metrics, step=state.global_step) + + def on_train_end(self, args, state, control, **kwargs): + if self._initialized and state.is_world_process_zero: + if self._auto_end_run and self._ml_flow.active_run(): + self._ml_flow.end_run() + + def on_save(self, args, state, control, **kwargs): + if self._initialized and state.is_world_process_zero and self._log_artifacts: + ckpt_dir = f"checkpoint-{state.global_step}" + artifact_path = os.path.join(args.output_dir, ckpt_dir) + logger.info(f"Logging checkpoint artifacts in {ckpt_dir}. This may take time.") + self._ml_flow.pyfunc.log_model( + ckpt_dir, + artifacts={"model_path": artifact_path}, + python_model=self._ml_flow.pyfunc.PythonModel(), + ) + + def __del__(self): + # if the previous run is not terminated correctly, the fluent API will + # not let you start a new run before the previous one is killed + if ( + self._auto_end_run + and callable(getattr(self._ml_flow, "active_run", None)) + and self._ml_flow.active_run() is not None + ): + self._ml_flow.end_run() + + +class DagsHubCallback(MLflowCallback): + """ + A [`TrainerCallback`] that logs to [DagsHub](https://dagshub.com/). Extends [`MLflowCallback`] + """ + + def __init__(self): + super().__init__() + if not is_dagshub_available(): + raise ImportError("DagsHubCallback requires dagshub to be installed. Run `pip install dagshub`.") + + from dagshub.upload import Repo + + self.Repo = Repo + + def setup(self, *args, **kwargs): + """ + Setup the DagsHub's Logging integration. + + Environment: + - **HF_DAGSHUB_LOG_ARTIFACTS** (`str`, *optional*): + Whether to save the data and model artifacts for the experiment. Default to `False`. + """ + + self.log_artifacts = os.getenv("HF_DAGSHUB_LOG_ARTIFACTS", "FALSE").upper() in ENV_VARS_TRUE_VALUES + self.name = os.getenv("HF_DAGSHUB_MODEL_NAME") or "main" + self.remote = os.getenv("MLFLOW_TRACKING_URI") + self.repo = self.Repo( + owner=self.remote.split(os.sep)[-2], + name=self.remote.split(os.sep)[-1].split(".")[0], + branch=os.getenv("BRANCH") or "main", + ) + self.path = Path("artifacts") + + if self.remote is None: + raise RuntimeError( + "DagsHubCallback requires the `MLFLOW_TRACKING_URI` environment variable to be set. Did you run" + " `dagshub.init()`?" + ) + + super().setup(*args, **kwargs) + + def on_train_end(self, args, state, control, **kwargs): + if self.log_artifacts: + if getattr(self, "train_dataloader", None): + torch.save(self.train_dataloader.dataset, os.path.join(args.output_dir, "dataset.pt")) + + self.repo.directory(str(self.path)).add_dir(args.output_dir) + + +class NeptuneMissingConfiguration(Exception): + def __init__(self): + super().__init__( + """ + ------ Unsupported ---- We were not able to create new runs. You provided a custom Neptune run to + `NeptuneCallback` with the `run` argument. For the integration to work fully, provide your `api_token` and + `project` by saving them as environment variables or passing them to the callback. + """ + ) + + +class NeptuneCallback(TrainerCallback): + """TrainerCallback that sends the logs to [Neptune](https://app.neptune.ai). + + Args: + api_token (`str`, *optional*): Neptune API token obtained upon registration. + You can leave this argument out if you have saved your token to the `NEPTUNE_API_TOKEN` environment + variable (strongly recommended). See full setup instructions in the + [docs](https://docs.neptune.ai/setup/installation). + project (`str`, *optional*): Name of an existing Neptune project, in the form "workspace-name/project-name". + You can find and copy the name in Neptune from the project settings -> Properties. If None (default), the + value of the `NEPTUNE_PROJECT` environment variable is used. + name (`str`, *optional*): Custom name for the run. + base_namespace (`str`, *optional*, defaults to "finetuning"): In the Neptune run, the root namespace + that will contain all of the metadata logged by the callback. + log_parameters (`bool`, *optional*, defaults to `True`): + If True, logs all Trainer arguments and model parameters provided by the Trainer. + log_checkpoints (`str`, *optional*): If "same", uploads checkpoints whenever they are saved by the Trainer. + If "last", uploads only the most recently saved checkpoint. If "best", uploads the best checkpoint (among + the ones saved by the Trainer). If `None`, does not upload checkpoints. + run (`Run`, *optional*): Pass a Neptune run object if you want to continue logging to an existing run. + Read more about resuming runs in the [docs](https://docs.neptune.ai/logging/to_existing_object). + **neptune_run_kwargs (*optional*): + Additional keyword arguments to be passed directly to the + [`neptune.init_run()`](https://docs.neptune.ai/api/neptune#init_run) function when a new run is created. + + For instructions and examples, see the [Transformers integration + guide](https://docs.neptune.ai/integrations/transformers) in the Neptune documentation. + """ + + integration_version_key = "source_code/integrations/transformers" + model_parameters_key = "model_parameters" + trial_name_key = "trial" + trial_params_key = "trial_params" + trainer_parameters_key = "trainer_parameters" + flat_metrics = {"train/epoch"} + + def __init__( + self, + *, + api_token: Optional[str] = None, + project: Optional[str] = None, + name: Optional[str] = None, + base_namespace: str = "finetuning", + run=None, + log_parameters: bool = True, + log_checkpoints: Optional[str] = None, + **neptune_run_kwargs, + ): + if not is_neptune_available(): + raise ValueError( + "NeptuneCallback requires the Neptune client library to be installed. " + "To install the library, run `pip install neptune`." + ) + + try: + from neptune import Run + from neptune.internal.utils import verify_type + except ImportError: + from neptune.new.internal.utils import verify_type + from neptune.new.metadata_containers.run import Run + + verify_type("api_token", api_token, (str, type(None))) + verify_type("project", project, (str, type(None))) + verify_type("name", name, (str, type(None))) + verify_type("base_namespace", base_namespace, str) + verify_type("run", run, (Run, type(None))) + verify_type("log_parameters", log_parameters, bool) + verify_type("log_checkpoints", log_checkpoints, (str, type(None))) + + self._base_namespace_path = base_namespace + self._log_parameters = log_parameters + self._log_checkpoints = log_checkpoints + self._initial_run: Optional[Run] = run + + self._run = None + self._is_monitoring_run = False + self._run_id = None + self._force_reset_monitoring_run = False + self._init_run_kwargs = {"api_token": api_token, "project": project, "name": name, **neptune_run_kwargs} + + self._volatile_checkpoints_dir = None + self._should_upload_checkpoint = self._log_checkpoints is not None + self._recent_checkpoint_path = None + + if self._log_checkpoints in {"last", "best"}: + self._target_checkpoints_namespace = f"checkpoints/{self._log_checkpoints}" + self._should_clean_recently_uploaded_checkpoint = True + else: + self._target_checkpoints_namespace = "checkpoints" + self._should_clean_recently_uploaded_checkpoint = False + + def _stop_run_if_exists(self): + if self._run: + self._run.stop() + del self._run + self._run = None + + def _initialize_run(self, **additional_neptune_kwargs): + try: + from neptune import init_run + from neptune.exceptions import NeptuneMissingApiTokenException, NeptuneMissingProjectNameException + except ImportError: + from neptune.new import init_run + from neptune.new.exceptions import NeptuneMissingApiTokenException, NeptuneMissingProjectNameException + + self._stop_run_if_exists() + + try: + run_params = additional_neptune_kwargs.copy() + run_params.update(self._init_run_kwargs) + self._run = init_run(**run_params) + self._run_id = self._run["sys/id"].fetch() + except (NeptuneMissingProjectNameException, NeptuneMissingApiTokenException) as e: + raise NeptuneMissingConfiguration() from e + + def _use_initial_run(self): + self._run = self._initial_run + self._is_monitoring_run = True + self._run_id = self._run["sys/id"].fetch() + self._initial_run = None + + def _ensure_run_with_monitoring(self): + if self._initial_run is not None: + self._use_initial_run() + else: + if not self._force_reset_monitoring_run and self._is_monitoring_run: + return + + if self._run and not self._is_monitoring_run and not self._force_reset_monitoring_run: + self._initialize_run(with_id=self._run_id) + self._is_monitoring_run = True + else: + self._initialize_run() + self._force_reset_monitoring_run = False + + def _ensure_at_least_run_without_monitoring(self): + if self._initial_run is not None: + self._use_initial_run() + else: + if not self._run: + self._initialize_run( + with_id=self._run_id, + capture_stdout=False, + capture_stderr=False, + capture_hardware_metrics=False, + capture_traceback=False, + ) + self._is_monitoring_run = False + + @property + def run(self): + if self._run is None: + self._ensure_at_least_run_without_monitoring() + return self._run + + @property + def _metadata_namespace(self): + return self.run[self._base_namespace_path] + + def _log_integration_version(self): + self.run[NeptuneCallback.integration_version_key] = version + + def _log_trainer_parameters(self, args): + self._metadata_namespace[NeptuneCallback.trainer_parameters_key] = args.to_sanitized_dict() + + def _log_model_parameters(self, model): + from neptune.utils import stringify_unsupported + + if model and hasattr(model, "config") and model.config is not None: + self._metadata_namespace[NeptuneCallback.model_parameters_key] = stringify_unsupported( + model.config.to_dict() + ) + + def _log_hyper_param_search_parameters(self, state): + if state and hasattr(state, "trial_name"): + self._metadata_namespace[NeptuneCallback.trial_name_key] = state.trial_name + + if state and hasattr(state, "trial_params") and state.trial_params is not None: + self._metadata_namespace[NeptuneCallback.trial_params_key] = state.trial_params + + def _log_model_checkpoint(self, source_directory: str, checkpoint: str): + target_path = relative_path = os.path.join(source_directory, checkpoint) + + if self._volatile_checkpoints_dir is not None: + consistent_checkpoint_path = os.path.join(self._volatile_checkpoints_dir, checkpoint) + try: + # Remove leading ../ from a relative path. + cpkt_path = relative_path.replace("..", "").lstrip(os.path.sep) + copy_path = os.path.join(consistent_checkpoint_path, cpkt_path) + shutil.copytree(relative_path, copy_path) + target_path = consistent_checkpoint_path + except IOError as e: + logger.warning( + "NeptuneCallback was unable to made a copy of checkpoint due to I/O exception: '{}'. " + "Could fail trying to upload.".format(e) + ) + + self._metadata_namespace[self._target_checkpoints_namespace].upload_files(target_path) + + if self._should_clean_recently_uploaded_checkpoint and self._recent_checkpoint_path is not None: + self._metadata_namespace[self._target_checkpoints_namespace].delete_files(self._recent_checkpoint_path) + + self._recent_checkpoint_path = relative_path + + def on_init_end(self, args, state, control, **kwargs): + self._volatile_checkpoints_dir = None + if self._log_checkpoints and (args.overwrite_output_dir or args.save_total_limit is not None): + self._volatile_checkpoints_dir = tempfile.TemporaryDirectory().name + + if self._log_checkpoints == "best" and not args.load_best_model_at_end: + raise ValueError("To save the best model checkpoint, the load_best_model_at_end argument must be enabled.") + + def on_train_begin(self, args, state, control, model=None, **kwargs): + if not state.is_world_process_zero: + return + + self._ensure_run_with_monitoring() + self._force_reset_monitoring_run = True + + self._log_integration_version() + if self._log_parameters: + self._log_trainer_parameters(args) + self._log_model_parameters(model) + + if state.is_hyper_param_search: + self._log_hyper_param_search_parameters(state) + + def on_train_end(self, args, state, control, **kwargs): + self._stop_run_if_exists() + + def __del__(self): + if self._volatile_checkpoints_dir is not None: + shutil.rmtree(self._volatile_checkpoints_dir, ignore_errors=True) + + self._stop_run_if_exists() + + def on_save(self, args, state, control, **kwargs): + if self._should_upload_checkpoint: + self._log_model_checkpoint(args.output_dir, f"checkpoint-{state.global_step}") + + def on_evaluate(self, args, state, control, metrics=None, **kwargs): + if self._log_checkpoints == "best": + best_metric_name = args.metric_for_best_model + if not best_metric_name.startswith("eval_"): + best_metric_name = f"eval_{best_metric_name}" + + metric_value = metrics.get(best_metric_name) + + operator = np.greater if args.greater_is_better else np.less + + self._should_upload_checkpoint = state.best_metric is None or operator(metric_value, state.best_metric) + + @classmethod + def get_run(cls, trainer): + for callback in trainer.callback_handler.callbacks: + if isinstance(callback, cls): + return callback.run + + raise Exception("The trainer doesn't have a NeptuneCallback configured.") + + def on_log(self, args, state, control, logs: Optional[Dict[str, float]] = None, **kwargs): + if not state.is_world_process_zero: + return + + if logs is not None: + for name, value in rewrite_logs(logs).items(): + if isinstance(value, (int, float)): + if name in NeptuneCallback.flat_metrics: + self._metadata_namespace[name] = value + else: + self._metadata_namespace[name].log(value, step=state.global_step) + + +class CodeCarbonCallback(TrainerCallback): + """ + A [`TrainerCallback`] that tracks the CO2 emission of training. + """ + + def __init__(self): + if not is_codecarbon_available(): + raise RuntimeError( + "CodeCarbonCallback requires `codecarbon` to be installed. Run `pip install codecarbon`." + ) + elif torch.version.hip: + raise RuntimeError( + "CodeCarbonCallback requires `codecarbon` package, which is not compatible with AMD ROCm (https://github.com/mlco2/codecarbon/pull/490). When using the Trainer, please specify the `report_to` argument (https://huggingface.co/docs/transformers/v4.39.3/en/main_classes/trainer#transformers.TrainingArguments.report_to) to disable CodeCarbonCallback." + ) + + import codecarbon + + self._codecarbon = codecarbon + self.tracker = None + + def on_init_end(self, args, state, control, **kwargs): + if self.tracker is None and state.is_local_process_zero: + # CodeCarbon will automatically handle environment variables for configuration + self.tracker = self._codecarbon.EmissionsTracker(output_dir=args.output_dir) + + def on_train_begin(self, args, state, control, model=None, **kwargs): + if self.tracker and state.is_local_process_zero: + self.tracker.start() + + def on_train_end(self, args, state, control, **kwargs): + if self.tracker and state.is_local_process_zero: + self.tracker.stop() + + +class ClearMLCallback(TrainerCallback): + """ + A [`TrainerCallback`] that sends the logs to [ClearML](https://clear.ml/). + + Environment: + - **CLEARML_PROJECT** (`str`, *optional*, defaults to `HuggingFace Transformers`): + ClearML project name. + - **CLEARML_TASK** (`str`, *optional*, defaults to `Trainer`): + ClearML task name. + - **CLEARML_LOG_MODEL** (`bool`, *optional*, defaults to `False`): + Whether to log models as artifacts during training. + """ + + log_suffix = "" + + _hparams_section = "Transformers" + _model_config_section = "Model Configuration" + _ignore_hparams_overrides = "_ignore_hparams_ui_overrides_" + _ignoge_model_config_overrides = "_ignore_model_config_ui_overrides_" + _model_config_description = "The configuration of model number {}." + _model_config_description_note = ( + "Note that, when cloning this task and running it remotely," + " the configuration might be applied to another model instead of this one." + " To avoid this, initialize the task externally by calling `Task.init`" + " before the `ClearMLCallback` is instantiated." + ) + _train_run_counter = 0 + _model_connect_counter = 0 + _task_created_in_callback = False + _should_close_on_train_end = None + + def __init__(self): + if is_clearml_available(): + import clearml + + self._clearml = clearml + else: + raise RuntimeError("ClearMLCallback requires 'clearml' to be installed. Run `pip install clearml`.") + + self._initialized = False + self._clearml_task = None + + self._log_model = False + self._checkpoints_saved = [] + + def setup(self, args, state, model, processing_class, **kwargs): + if self._clearml is None: + return + if self._initialized: + return + ClearMLCallback._train_run_counter += 1 + ClearMLCallback._model_connect_counter += 1 + ClearMLCallback.log_suffix = ( + "" if ClearMLCallback._train_run_counter == 1 else "_" + str(ClearMLCallback._train_run_counter) + ) + if state.is_world_process_zero: + logger.info("Automatic ClearML logging enabled.") + if self._clearml_task is None: + if ClearMLCallback._should_close_on_train_end is None: + if not self._clearml.Task.running_locally() or self._clearml.Task.current_task(): + ClearMLCallback._should_close_on_train_end = False + else: + ClearMLCallback._should_close_on_train_end = True + + # This might happen when running inside of a pipeline, where the task is already initialized + # from outside of Hugging Face + if self._clearml.Task.running_locally() and self._clearml.Task.current_task(): + self._clearml_task = self._clearml.Task.current_task() + self._log_model = os.getenv( + "CLEARML_LOG_MODEL", + "FALSE" if not ClearMLCallback._task_created_in_callback else "TRUE", + ).upper() in ENV_VARS_TRUE_VALUES.union({"TRUE"}) + logger.info("External ClearML Task has been connected.") + else: + self._clearml_task = self._clearml.Task.init( + project_name=os.getenv("CLEARML_PROJECT", "HuggingFace Transformers"), + task_name=os.getenv("CLEARML_TASK", "Trainer"), + auto_connect_frameworks={"tensorboard": False, "pytorch": False}, + output_uri=True, + ) + self._log_model = os.getenv("CLEARML_LOG_MODEL", "TRUE").upper() in ENV_VARS_TRUE_VALUES.union( + {"TRUE"} + ) + ClearMLCallback._task_created_in_callback = True + logger.info("ClearML Task has been initialized.") + self._initialized = True + + suffixed_hparams_section = ClearMLCallback._hparams_section + ClearMLCallback.log_suffix + ignore_hparams_config_section = suffixed_hparams_section + "/" + ClearMLCallback._ignore_hparams_overrides + if self._clearml.Task.running_locally(): + self._copy_training_args_as_hparams(args, suffixed_hparams_section) + self._clearml_task.set_parameter( + name=ignore_hparams_config_section, + value=True, + value_type=bool, + description=( + "If True, ignore Transformers hyperparameters overrides done in the UI/backend " + + "when running remotely. Otherwise, the overrides will be applied when running remotely" + ), + ) + elif not self._clearml_task.get_parameter(ignore_hparams_config_section, default=True, cast=True): + self._clearml_task.connect(args, suffixed_hparams_section) + else: + self._copy_training_args_as_hparams( + args, ClearMLCallback._hparams_section + ClearMLCallback.log_suffix + ) + + if getattr(model, "config", None) is not None: + ignore_model_config_section = ( + suffixed_hparams_section + "/" + ClearMLCallback._ignoge_model_config_overrides + ) + configuration_object_description = ClearMLCallback._model_config_description.format( + ClearMLCallback._model_connect_counter + ) + if ClearMLCallback._model_connect_counter != ClearMLCallback._train_run_counter: + configuration_object_description += " " + ClearMLCallback._model_config_description_note + if self._clearml.Task.running_locally(): + self._clearml_task.set_parameter( + name=ignore_model_config_section, + value=True, + value_type=bool, + description=( + "If True, ignore Transformers model configuration overrides done in the UI/backend " + + "when running remotely. Otherwise, the overrides will be applied when running remotely" + ), + ) + self._clearml_task.set_configuration_object( + name=ClearMLCallback._model_config_section + ClearMLCallback.log_suffix, + config_dict=model.config.to_dict(), + description=configuration_object_description, + ) + elif not self._clearml_task.get_parameter(ignore_model_config_section, default=True, cast=True): + model.config = model.config.from_dict( + self._clearml_task.get_configuration_object_as_dict( + ClearMLCallback._model_config_section + ClearMLCallback.log_suffix + ) + ) + else: + self._clearml_task.set_configuration_object( + name=ClearMLCallback._model_config_section + ClearMLCallback.log_suffix, + config_dict=model.config.to_dict(), + description=configuration_object_description, + ) + + def on_train_begin(self, args, state, control, model=None, processing_class=None, **kwargs): + if self._clearml is None: + return + self._checkpoints_saved = [] + if state.is_hyper_param_search: + self._initialized = False + if not self._initialized: + self.setup(args, state, model, processing_class, **kwargs) + + def on_train_end(self, args, state, control, **kwargs): + if ClearMLCallback._should_close_on_train_end: + self._clearml_task.close() + ClearMLCallback._train_run_counter = 0 + + def on_log(self, args, state, control, model=None, processing_class=None, logs=None, **kwargs): + if self._clearml is None: + return + if not self._initialized: + self.setup(args, state, model, processing_class, **kwargs) + if state.is_world_process_zero: + eval_prefix = "eval_" + eval_prefix_len = len(eval_prefix) + test_prefix = "test_" + test_prefix_len = len(test_prefix) + single_value_scalars = [ + "train_runtime", + "train_samples_per_second", + "train_steps_per_second", + "train_loss", + "total_flos", + "epoch", + ] + for k, v in logs.items(): + if isinstance(v, (int, float)): + if k in single_value_scalars: + self._clearml_task.get_logger().report_single_value( + name=k + ClearMLCallback.log_suffix, value=v + ) + elif k.startswith(eval_prefix): + self._clearml_task.get_logger().report_scalar( + title="eval" + ClearMLCallback.log_suffix, + series=k[eval_prefix_len:], + value=v, + iteration=state.global_step, + ) + elif k.startswith(test_prefix): + self._clearml_task.get_logger().report_scalar( + title="test" + ClearMLCallback.log_suffix, + series=k[test_prefix_len:], + value=v, + iteration=state.global_step, + ) + else: + self._clearml_task.get_logger().report_scalar( + title="train" + ClearMLCallback.log_suffix, + series=k, + value=v, + iteration=state.global_step, + ) + else: + logger.warning( + "Trainer is attempting to log a value of " + f'"{v}" of type {type(v)} for key "{k}" as a scalar. ' + "This invocation of ClearML logger's report_scalar() " + "is incorrect so we dropped this attribute." + ) + + def on_save(self, args, state, control, **kwargs): + if self._log_model and self._clearml_task and state.is_world_process_zero: + ckpt_dir = f"checkpoint-{state.global_step}" + artifact_path = os.path.join(args.output_dir, ckpt_dir) + name = ckpt_dir + ClearMLCallback.log_suffix + logger.info(f"Logging checkpoint artifact `{name}`. This may take some time.") + output_model = self._clearml.OutputModel(task=self._clearml_task, name=name) + output_model.connect(task=self._clearml_task, name=name) + output_model.update_weights_package( + weights_path=artifact_path, + target_filename=ckpt_dir, + iteration=state.global_step, + auto_delete_file=False, + ) + self._checkpoints_saved.append(output_model) + while args.save_total_limit and args.save_total_limit < len(self._checkpoints_saved): + try: + self._clearml.model.Model.remove( + self._checkpoints_saved[0], + delete_weights_file=True, + force=True, + raise_on_errors=True, + ) + except Exception as e: + logger.warning( + "Could not remove checkpoint `{}` after going over the `save_total_limit`. Error is: {}".format( + self._checkpoints_saved[0].name, e + ) + ) + break + self._checkpoints_saved = self._checkpoints_saved[1:] + + def _copy_training_args_as_hparams(self, training_args, prefix): + as_dict = { + field.name: getattr(training_args, field.name) + for field in fields(training_args) + if field.init and not field.name.endswith("_token") + } + flat_dict = {str(k): v for k, v in self._clearml.utilities.proxy_object.flatten_dictionary(as_dict).items()} + self._clearml_task._arguments.copy_from_dict(flat_dict, prefix=prefix) + + +class FlyteCallback(TrainerCallback): + """A [`TrainerCallback`] that sends the logs to [Flyte](https://flyte.org/). + NOTE: This callback only works within a Flyte task. + + Args: + save_log_history (`bool`, *optional*, defaults to `True`): + When set to True, the training logs are saved as a Flyte Deck. + + sync_checkpoints (`bool`, *optional*, defaults to `True`): + When set to True, checkpoints are synced with Flyte and can be used to resume training in the case of an + interruption. + + Example: + + ```python + # Note: This example skips over some setup steps for brevity. + from flytekit import current_context, task + + + @task + def train_hf_transformer(): + cp = current_context().checkpoint + trainer = Trainer(..., callbacks=[FlyteCallback()]) + output = trainer.train(resume_from_checkpoint=cp.restore()) + ``` + """ + + def __init__(self, save_log_history: bool = True, sync_checkpoints: bool = True): + super().__init__() + if not is_flytekit_available(): + raise ImportError("FlyteCallback requires flytekit to be installed. Run `pip install flytekit`.") + + if not is_flyte_deck_standard_available() or not is_pandas_available(): + logger.warning( + "Syncing log history requires both flytekitplugins-deck-standard and pandas to be installed. " + "Run `pip install flytekitplugins-deck-standard pandas` to enable this feature." + ) + save_log_history = False + + from flytekit import current_context + + self.cp = current_context().checkpoint + self.save_log_history = save_log_history + self.sync_checkpoints = sync_checkpoints + + def on_save(self, args, state, control, **kwargs): + if self.sync_checkpoints and state.is_world_process_zero: + ckpt_dir = f"checkpoint-{state.global_step}" + artifact_path = os.path.join(args.output_dir, ckpt_dir) + + logger.info(f"Syncing checkpoint in {ckpt_dir} to Flyte. This may take time.") + self.cp.save(artifact_path) + + def on_train_end(self, args, state, control, **kwargs): + if self.save_log_history: + import pandas as pd + from flytekit import Deck + from flytekitplugins.deck.renderer import TableRenderer + + log_history_df = pd.DataFrame(state.log_history) + Deck("Log History", TableRenderer().to_html(log_history_df)) + + +class DVCLiveCallback(TrainerCallback): + """ + A [`TrainerCallback`] that sends the logs to [DVCLive](https://www.dvc.org/doc/dvclive). + + Use the environment variables below in `setup` to configure the integration. To customize this callback beyond + those environment variables, see [here](https://dvc.org/doc/dvclive/ml-frameworks/huggingface). + + Args: + live (`dvclive.Live`, *optional*, defaults to `None`): + Optional Live instance. If None, a new instance will be created using **kwargs. + log_model (Union[Literal["all"], bool], *optional*, defaults to `None`): + Whether to use `dvclive.Live.log_artifact()` to log checkpoints created by [`Trainer`]. If set to `True`, + the final checkpoint is logged at the end of training. If set to `"all"`, the entire + [`TrainingArguments`]'s `output_dir` is logged at each checkpoint. + """ + + def __init__( + self, + live: Optional[Any] = None, + log_model: Optional[Union[Literal["all"], bool]] = None, + **kwargs, + ): + if not is_dvclive_available(): + raise RuntimeError("DVCLiveCallback requires dvclive to be installed. Run `pip install dvclive`.") + from dvclive import Live + + self._initialized = False + self.live = None + if isinstance(live, Live): + self.live = live + elif live is not None: + raise RuntimeError(f"Found class {live.__class__} for live, expected dvclive.Live") + + self._log_model = log_model + if self._log_model is None: + log_model_env = os.getenv("HF_DVCLIVE_LOG_MODEL", "FALSE") + if log_model_env.upper() in ENV_VARS_TRUE_VALUES: + self._log_model = True + elif log_model_env.lower() == "all": + self._log_model = "all" + + def setup(self, args, state, model): + """ + Setup the optional DVCLive integration. To customize this callback beyond the environment variables below, see + [here](https://dvc.org/doc/dvclive/ml-frameworks/huggingface). + + Environment: + - **HF_DVCLIVE_LOG_MODEL** (`str`, *optional*): + Whether to use `dvclive.Live.log_artifact()` to log checkpoints created by [`Trainer`]. If set to `True` or + *1*, the final checkpoint is logged at the end of training. If set to `all`, the entire + [`TrainingArguments`]'s `output_dir` is logged at each checkpoint. + """ + from dvclive import Live + + self._initialized = True + if state.is_world_process_zero: + if not self.live: + self.live = Live() + self.live.log_params(args.to_dict()) + + def on_train_begin(self, args, state, control, model=None, **kwargs): + if not self._initialized: + self.setup(args, state, model) + + def on_log(self, args, state, control, model=None, logs=None, **kwargs): + if not self._initialized: + self.setup(args, state, model) + if state.is_world_process_zero: + from dvclive.plots import Metric + from dvclive.utils import standardize_metric_name + + for key, value in logs.items(): + if Metric.could_log(value): + self.live.log_metric(standardize_metric_name(key, "dvclive.huggingface"), value) + else: + logger.warning( + "Trainer is attempting to log a value of " + f'"{value}" of type {type(value)} for key "{key}" as a scalar. ' + "This invocation of DVCLive's Live.log_metric() " + "is incorrect so we dropped this attribute." + ) + self.live.next_step() + + def on_save(self, args, state, control, **kwargs): + if self._log_model == "all" and self._initialized and state.is_world_process_zero: + self.live.log_artifact(args.output_dir) + + def on_train_end(self, args, state, control, **kwargs): + if self._initialized and state.is_world_process_zero: + from transformers.trainer import Trainer + + if self._log_model is True: + fake_trainer = Trainer( + args=args, + model=kwargs.get("model"), + processing_class=kwargs.get("processing_class"), + eval_dataset=["fake"], + ) + name = "best" if args.load_best_model_at_end else "last" + output_dir = os.path.join(args.output_dir, name) + fake_trainer.save_model(output_dir) + self.live.log_artifact(output_dir, name=name, type="model", copy=True) + self.live.end() + + +class SwanLabCallback(TrainerCallback): + """ + A [`TrainerCallback`] that logs metrics, media, model checkpoints to [SwanLab](https://swanlab.cn/). + """ + + def __init__(self): + if not is_swanlab_available(): + raise RuntimeError("SwanLabCallback requires swanlab to be installed. Run `pip install swanlab`.") + import swanlab + + self._swanlab = swanlab + self._initialized = False + self._log_model = os.getenv("SWANLAB_LOG_MODEL", None) + + def setup(self, args, state, model, **kwargs): + """ + Setup the optional SwanLab (*swanlab*) integration. + + One can subclass and override this method to customize the setup if needed. Find more information + [here](https://docs.swanlab.cn/guide_cloud/integration/integration-huggingface-transformers.html). + + You can also override the following environment variables. Find more information about environment + variables [here](https://docs.swanlab.cn/en/api/environment-variable.html#environment-variables) + + Environment: + - **SWANLAB_API_KEY** (`str`, *optional*, defaults to `None`): + Cloud API Key. During login, this environment variable is checked first. If it doesn't exist, the system + checks if the user is already logged in. If not, the login process is initiated. + + - If a string is passed to the login interface, this environment variable is ignored. + - If the user is already logged in, this environment variable takes precedence over locally stored + login information. + + - **SWANLAB_PROJECT** (`str`, *optional*, defaults to `None`): + Set this to a custom string to store results in a different project. If not specified, the name of the current + running directory is used. + + - **SWANLAB_LOG_DIR** (`str`, *optional*, defaults to `swanlog`): + This environment variable specifies the storage path for log files when running in local mode. + By default, logs are saved in a folder named swanlog under the working directory. + + - **SWANLAB_MODE** (`Literal["local", "cloud", "disabled"]`, *optional*, defaults to `cloud`): + SwanLab's parsing mode, which involves callbacks registered by the operator. Currently, there are three modes: + local, cloud, and disabled. Note: Case-sensitive. Find more information + [here](https://docs.swanlab.cn/en/api/py-init.html#swanlab-init) + + - **SWANLAB_LOG_MODEL** (`str`, *optional*, defaults to `None`): + SwanLab does not currently support the save mode functionality.This feature will be available in a future + release + + - **SWANLAB_WEB_HOST** (`str`, *optional*, defaults to `None`): + Web address for the SwanLab cloud environment for private version (its free) + + - **SWANLAB_API_HOST** (`str`, *optional*, defaults to `None`): + API address for the SwanLab cloud environment for private version (its free) + + """ + self._initialized = True + + if state.is_world_process_zero: + logger.info('Automatic SwanLab logging enabled, to disable set os.environ["SWANLAB_MODE"] = "disabled"') + combined_dict = {**args.to_dict()} + + if hasattr(model, "config") and model.config is not None: + model_config = model.config if isinstance(model.config, dict) else model.config.to_dict() + combined_dict = {**model_config, **combined_dict} + if hasattr(model, "peft_config") and model.peft_config is not None: + peft_config = model.peft_config + combined_dict = {**{"peft_config": peft_config}, **combined_dict} + trial_name = state.trial_name + init_args = {} + if trial_name is not None: + init_args["experiment_name"] = f"{args.run_name}-{trial_name}" + elif args.run_name is not None: + init_args["experiment_name"] = args.run_name + init_args["project"] = os.getenv("SWANLAB_PROJECT", None) + + if self._swanlab.get_run() is None: + self._swanlab.init( + **init_args, + ) + # show transformers logo! + self._swanlab.config["FRAMEWORK"] = "🤗transformers" + # add config parameters (run may have been created manually) + self._swanlab.config.update(combined_dict) + + # add number of model parameters to swanlab config + try: + self._swanlab.config.update({"model_num_parameters": model.num_parameters()}) + # get peft model parameters + if type(model).__name__ == "PeftModel" or type(model).__name__ == "PeftMixedModel": + trainable_params, all_param = model.get_nb_trainable_parameters() + self._swanlab.config.update({"peft_model_trainable_params": trainable_params}) + self._swanlab.config.update({"peft_model_all_param": all_param}) + except AttributeError: + logger.info("Could not log the number of model parameters in SwanLab due to an AttributeError.") + + # log the initial model architecture to an artifact + if self._log_model is not None: + logger.warning( + "SwanLab does not currently support the save mode functionality. " + "This feature will be available in a future release." + ) + badge_markdown = ( + f'[Visualize in SwanLab]({self._swanlab.get_run().public.cloud.exp_url})' + ) + + modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n{badge_markdown}" + + def on_train_begin(self, args, state, control, model=None, **kwargs): + if not self._initialized: + self.setup(args, state, model, **kwargs) + + def on_train_end(self, args, state, control, model=None, processing_class=None, **kwargs): + if self._log_model is not None and self._initialized and state.is_world_process_zero: + logger.warning( + "SwanLab does not currently support the save mode functionality. " + "This feature will be available in a future release." + ) + + def on_log(self, args, state, control, model=None, logs=None, **kwargs): + single_value_scalars = [ + "train_runtime", + "train_samples_per_second", + "train_steps_per_second", + "train_loss", + "total_flos", + ] + + if not self._initialized: + self.setup(args, state, model) + if state.is_world_process_zero: + for k, v in logs.items(): + if k in single_value_scalars: + self._swanlab.log({f"single_value/{k}": v}, step=state.global_step) + non_scalar_logs = {k: v for k, v in logs.items() if k not in single_value_scalars} + non_scalar_logs = rewrite_logs(non_scalar_logs) + self._swanlab.log({**non_scalar_logs, "train/global_step": state.global_step}, step=state.global_step) + + def on_save(self, args, state, control, **kwargs): + if self._log_model is not None and self._initialized and state.is_world_process_zero: + logger.warning( + "SwanLab does not currently support the save mode functionality. " + "This feature will be available in a future release." + ) + + def on_predict(self, args, state, control, metrics, **kwargs): + if not self._initialized: + self.setup(args, state, **kwargs) + if state.is_world_process_zero: + metrics = rewrite_logs(metrics) + self._swanlab.log(metrics) + + +INTEGRATION_TO_CALLBACK = { + "azure_ml": AzureMLCallback, + "comet_ml": CometCallback, + "mlflow": MLflowCallback, + "neptune": NeptuneCallback, + "tensorboard": TensorBoardCallback, + "wandb": WandbCallback, + "codecarbon": CodeCarbonCallback, + "clearml": ClearMLCallback, + "dagshub": DagsHubCallback, + "flyte": FlyteCallback, + "dvclive": DVCLiveCallback, + "swanlab": SwanLabCallback, +} + + +def get_reporting_integration_callbacks(report_to): + if report_to is None: + return [] + + if isinstance(report_to, str): + if "none" == report_to: + return [] + elif "all" == report_to: + report_to = get_available_reporting_integrations() + else: + report_to = [report_to] + + for integration in report_to: + if integration not in INTEGRATION_TO_CALLBACK: + raise ValueError( + f"{integration} is not supported, only {', '.join(INTEGRATION_TO_CALLBACK.keys())} are supported." + ) + + return [INTEGRATION_TO_CALLBACK[integration] for integration in report_to] diff --git a/lib/python3.12/site-packages/transformers/integrations/mistral.py b/lib/python3.12/site-packages/transformers/integrations/mistral.py new file mode 100644 index 0000000000000000000000000000000000000000..78172329277e4f5ac0b01795a123bdebe7bd8f51 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/mistral.py @@ -0,0 +1,105 @@ +from tokenizers import Regex, Tokenizer, decoders, pre_tokenizers, processors +from tokenizers.models import BPE + +from transformers import LlamaTokenizerFast +from transformers.convert_slow_tokenizer import bytes_to_unicode + + +class MistralConverter: + """ + A general tiktoken converter. + """ + + def __init__( + self, + vocab=None, + pattern=r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""", + add_prefix_space=False, + additional_special_tokens=None, + *args, + **kwargs, + ): + super().__init__(*args) + self.vocab = vocab + self.pattern = pattern + self.add_prefix_space = add_prefix_space + self.additional_special_tokens = additional_special_tokens + + def extract_vocab_merges_from_model(self, vocab: str): + bpe_ranks = vocab + byte_encoder = bytes_to_unicode() + + def token_bytes_to_string(b): + return "".join([byte_encoder[ord(char)] for char in b.decode("latin-1")]) + + merges = [] + vocab = {} + for idx, (token, rank) in enumerate(bpe_ranks.items()): + if token not in self.additional_special_tokens: + vocab[token_bytes_to_string(token)] = idx + if len(token) == 1: + continue + local = [] + for index in range(1, len(token)): + piece_l, piece_r = token[:index], token[index:] + if piece_l in bpe_ranks and piece_r in bpe_ranks and (piece_l + piece_r) in bpe_ranks: + local.append((piece_l, piece_r, rank)) + local = sorted(local, key=lambda x: (bpe_ranks[x[0]], bpe_ranks[x[1]]), reverse=False) + merges.extend(local) + else: + vocab[token] = idx + merges = sorted(merges, key=lambda val: val[2], reverse=False) + merges = [(token_bytes_to_string(val[0]), token_bytes_to_string(val[1])) for val in merges] + return vocab, merges + + def tokenizer(self): + vocab_scores, merges = self.extract_vocab_merges_from_model(self.vocab) + tokenizer = Tokenizer(BPE(vocab_scores, merges, fuse_unk=False)) + if hasattr(tokenizer.model, "ignore_merges"): + tokenizer.model.ignore_merges = True + return tokenizer + + def converted(self) -> Tokenizer: + tokenizer = self.tokenizer() + tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.Split(Regex(self.pattern), behavior="isolated", invert=False), + pre_tokenizers.ByteLevel(add_prefix_space=self.add_prefix_space, use_regex=False), + ] + ) + tokenizer.decoder = decoders.ByteLevel() + tokenizer.add_special_tokens(self.additional_special_tokens) + + tokenizer.post_processor = processors.ByteLevel(trim_offsets=False) + + return tokenizer + + +def convert_tekken_tokenizer(tokenizer_file: str): + """Convert a "tekken" tokenizer to a fast Tokenizer.""" + # Tekken format -- need to use the Converter + + from mistral_common.tokens.tokenizers.mistral import MistralTokenizer + + # Load directly using their lib + mistral_tokenizer = MistralTokenizer.from_file(tokenizer_file) + + # Extract vocab and special tokens + vocab = mistral_tokenizer.instruct_tokenizer.tokenizer._tekken_token2id_nospecial + all_special = [ + token.value if hasattr(token, "value") else token + for token in mistral_tokenizer.instruct_tokenizer.tokenizer._all_special_tokens + ] + specials_tokens = {token: all_special.index(token) for token in all_special} + specials_tokens.update(vocab) + vocab = specials_tokens + + # Convert + tokenizer = LlamaTokenizerFast( + tokenizer_object=MistralConverter(vocab=vocab, additional_special_tokens=all_special).converted(), + ) + + # Post-process + tokenizer.add_special_tokens({"additional_special_tokens": all_special}) + + return tokenizer diff --git a/lib/python3.12/site-packages/transformers/integrations/npu_flash_attention.py b/lib/python3.12/site-packages/transformers/integrations/npu_flash_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..fa26379126fc09f9f56972741292c052a4895371 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/npu_flash_attention.py @@ -0,0 +1,233 @@ +# 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. + +import os + +import torch +import torch.nn.functional as F + +from ..utils.import_utils import is_torch_npu_available + + +if is_torch_npu_available(): + import torch_npu + from einops import rearrange, repeat + + +# FlashAttention2 is supported on Ascend NPU with down-right aligned causal mask by default. +# Set environment variable `NPU_FA2_SPARSE_MODE` to 2 when using top-left aligned causal mask. +TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE = 2 +DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE = 3 + +SPARSE_MODE = int(os.getenv("NPU_FA2_SPARSE_MODE", default=DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE)) +if SPARSE_MODE not in [TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE, DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE]: + raise ValueError( + "Environment variable `NPU_FA2_SPARSE_MODE` can only be set as 2 (top-left aligned causal mask) " + "or 3 (down-right aligned causal mask)." + ) + + +def is_npu_fa2_top_left_aligned_causal_mask(): + return SPARSE_MODE == TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE if is_torch_npu_available() else False + + +# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py +class IndexFirstAxis(torch.autograd.Function): + @staticmethod + def forward(ctx, input, indices): + ctx.save_for_backward(indices) + assert input.ndim >= 2 + ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] + second_dim = other_shape.numel() + # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. + # return input[indices] + return torch.gather( + rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim) + ).reshape(-1, *other_shape) + + @staticmethod + def backward(ctx, grad_output): + (indices,) = ctx.saved_tensors + assert grad_output.ndim >= 2 + other_shape = grad_output.shape[1:] + grad_output = rearrange(grad_output, "b ... -> b (...)") + grad_input = torch.zeros( + [ctx.first_axis_dim, grad_output.shape[1]], + device=grad_output.device, + dtype=grad_output.dtype, + ) + # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. + # grad_input[indices] = grad_output + grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) + return grad_input.reshape(ctx.first_axis_dim, *other_shape), None + + +index_first_axis = IndexFirstAxis.apply + + +# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py +class IndexPutFirstAxis(torch.autograd.Function): + @staticmethod + def forward(ctx, values, indices, first_axis_dim): + ctx.save_for_backward(indices) + assert indices.ndim == 1 + assert values.ndim >= 2 + output = torch.zeros(first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype) + # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. + output[indices] = values + # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values) + return output + + @staticmethod + def backward(ctx, grad_output): + (indices,) = ctx.saved_tensors + # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. + grad_values = grad_output[indices] + # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1])) + return grad_values, None, None + + +index_put_first_axis = IndexPutFirstAxis.apply + + +# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py +def pad_input(hidden_states, indices, batch, seqlen): + """ + Arguments: + hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. + indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. + batch: int, batch size for the padded sequence. + seqlen: int, maximum sequence length for the padded sequence. + Return: + hidden_states: (batch, seqlen, ...) + """ + # dim = hidden_states.shape[-1] + # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype) + # output[indices] = hidden_states + output = index_put_first_axis(hidden_states, indices, batch * seqlen) + return rearrange(output, "(b s) ... -> b s ...", b=batch) + + +# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py +def unpad_input(hidden_states, attention_mask, unused_mask=None): + """ + Arguments: + hidden_states: (batch, seqlen, ...) + attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. + unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. + Return: + hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. + indices: (total_nnz), the indices of masked tokens from the flattened input sequence. + cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. + max_seqlen_in_batch: int + seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. + """ + all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask + seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) + used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the + # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim + # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to + # index with integer indices. Moreover, torch's index is a bit slower than it needs to be, + # so we write custom forward and backward to make it a bit faster. + return ( + index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), + indices, + cu_seqlens, + max_seqlen_in_batch, + used_seqlens_in_batch, + ) + + +def npu_flash_attn_func( + q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=False, + **kwargs, +): + keep_prob = 1.0 - dropout_p + + if not causal: + head_num = q.shape[2] + output = torch_npu.npu_fusion_attention(q, k, v, head_num, "BSND", keep_prob=keep_prob, scale=softmax_scale)[0] + else: + attn_mask_npu = torch.triu(torch.ones([2048, 2048]), diagonal=1).bool().to(q.device) + head_num = q.shape[2] + output = torch_npu.npu_fusion_attention( + q, + k, + v, + head_num, + "BSND", + keep_prob=keep_prob, + scale=softmax_scale, + atten_mask=attn_mask_npu, + sparse_mode=SPARSE_MODE, + )[0] + + return output + + +def npu_flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + dropout_p=0.0, + softmax_scale=None, + causal=False, + **kwargs, +): + keep_prob = 1.0 - dropout_p + + if not causal: + head_num = q.shape[1] + output = torch_npu.npu_fusion_attention( + q, + k, + v, + head_num, + pse=None, + atten_mask=None, + scale=softmax_scale, + keep_prob=keep_prob, + input_layout="TND", + actual_seq_qlen=tuple(cu_seqlens_q[1:].cpu().numpy().tolist()), + actual_seq_kvlen=tuple(cu_seqlens_k[1:].cpu().numpy().tolist()), + )[0] + else: + attn_mask_npu = torch.triu(torch.ones([2048, 2048]), diagonal=1).bool().to(q.device) + head_num = q.shape[1] + output = torch_npu.npu_fusion_attention( + q, + k, + v, + head_num, + pse=None, + padding_mask=None, + atten_mask=attn_mask_npu, + scale=softmax_scale, + keep_prob=keep_prob, + input_layout="TND", + actual_seq_qlen=tuple(cu_seqlens_q[1:].cpu().numpy().tolist()), + actual_seq_kvlen=tuple(cu_seqlens_k[1:].cpu().numpy().tolist()), + sparse_mode=SPARSE_MODE, + )[0] + + return output diff --git a/lib/python3.12/site-packages/transformers/integrations/peft.py b/lib/python3.12/site-packages/transformers/integrations/peft.py new file mode 100644 index 0000000000000000000000000000000000000000..6aa3b137b1914254c1b9ac49d5a8b413e29acb0b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/peft.py @@ -0,0 +1,593 @@ +# Copyright 2023 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. + +import importlib +import inspect +import warnings +from typing import Any, Dict, List, Optional, Union + +from packaging import version + +from ..utils import ( + check_peft_version, + find_adapter_config_file, + is_accelerate_available, + is_peft_available, + is_torch_available, + logging, +) + + +if is_torch_available(): + import torch + +if is_accelerate_available(): + from accelerate import dispatch_model + from accelerate.utils import get_balanced_memory, infer_auto_device_map + +# Minimum PEFT version supported for the integration +MIN_PEFT_VERSION = "0.5.0" + + +logger = logging.get_logger(__name__) + + +class PeftAdapterMixin: + """ + A class containing all functions for loading and using adapters weights that are supported in PEFT library. For + more details about adapters and injecting them on a transformer-based model, check out the documentation of PEFT + library: https://huggingface.co/docs/peft/index + + Currently supported PEFT methods are all non-prefix tuning methods. Below is the list of supported PEFT methods + that anyone can load, train and run with this mixin class: + - Low Rank Adapters (LoRA): https://huggingface.co/docs/peft/conceptual_guides/lora + - IA3: https://huggingface.co/docs/peft/conceptual_guides/ia3 + - AdaLora: https://arxiv.org/abs/2303.10512 + + Other PEFT models such as prompt tuning, prompt learning are out of scope as these adapters are not "injectable" + into a torch module. For using these methods, please refer to the usage guide of PEFT library. + + With this mixin, if the correct PEFT version is installed, it is possible to: + + - Load an adapter stored on a local path or in a remote Hub repository, and inject it in the model + - Attach new adapters in the model and train them with Trainer or by your own. + - Attach multiple adapters and iteratively activate / deactivate them + - Activate / deactivate all adapters from the model. + - Get the `state_dict` of the active adapter. + """ + + _hf_peft_config_loaded = False + + def load_adapter( + self, + peft_model_id: Optional[str] = None, + adapter_name: Optional[str] = None, + revision: Optional[str] = None, + token: Optional[str] = None, + device_map: Optional[str] = "auto", + max_memory: Optional[str] = None, + offload_folder: Optional[str] = None, + offload_index: Optional[int] = None, + peft_config: Dict[str, Any] = None, + adapter_state_dict: Optional[Dict[str, "torch.Tensor"]] = None, + low_cpu_mem_usage: bool = False, + is_trainable: bool = False, + adapter_kwargs: Optional[Dict[str, Any]] = None, + ) -> None: + """ + Load adapter weights from file or remote Hub folder. If you are not familiar with adapters and PEFT methods, we + invite you to read more about them on PEFT official documentation: https://huggingface.co/docs/peft + + Requires peft as a backend to load the adapter weights. + + Args: + peft_model_id (`str`, *optional*): + The identifier of the model to look for on the Hub, or a local path to the saved adapter config file + and adapter weights. + adapter_name (`str`, *optional*): + The adapter name to use. If not set, will use the default adapter. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`. + + + + token (`str`, `optional`): + Whether to use authentication token to load the remote folder. Useful to load private repositories + that are on HuggingFace Hub. You might need to call `huggingface-cli login` and paste your tokens to + cache it. + device_map (`str` or `Dict[str, Union[int, str, torch.device]]` or `int` or `torch.device`, *optional*): + A map that specifies where each submodule should go. It doesn't need to be refined to each + parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the + same device. If we only pass the device (*e.g.*, `"cpu"`, `"cuda:1"`, `"mps"`, or a GPU ordinal rank + like `1`) on which the model will be allocated, the device map will map the entire model to this + device. Passing `device_map = 0` means put the whole model on GPU 0. + + To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For + more information about each option see [designing a device + map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). + max_memory (`Dict`, *optional*): + A dictionary device identifier to maximum memory. Will default to the maximum memory available for each + GPU and the available CPU RAM if unset. + offload_folder (`str` or `os.PathLike`, `optional`): + If the `device_map` contains any value `"disk"`, the folder where we will offload weights. + offload_index (`int`, `optional`): + `offload_index` argument to be passed to `accelerate.dispatch_model` method. + peft_config (`Dict[str, Any]`, *optional*): + The configuration of the adapter to add, supported adapters are non-prefix tuning and adaption prompts + methods. This argument is used in case users directly pass PEFT state dicts + adapter_state_dict (`Dict[str, torch.Tensor]`, *optional*): + The state dict of the adapter to load. This argument is used in case users directly pass PEFT state + dicts + low_cpu_mem_usage (`bool`, *optional*, defaults to `False`): + Reduce memory usage while loading the PEFT adapter. This should also speed up the loading process. + Requires PEFT version 0.13.0 or higher. + is_trainable (`bool`, *optional*, defaults to `False`): + Whether the adapter should be trainable or not. If `False`, the adapter will be frozen and can only be + used for inference. + adapter_kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments passed along to the `from_pretrained` method of the adapter config and + `find_adapter_config_file` method. + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + # peft only supports low_cpu_mem_usage starting from v0.13.0 + peft_load_kwargs = {} + if low_cpu_mem_usage: + min_version_lcmu = "0.13.0" + if version.parse(importlib.metadata.version("peft")) >= version.parse(min_version_lcmu): + peft_load_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage + else: + raise ValueError( + "The version of PEFT you are using does not support `low_cpu_mem_usage` yet, " + f"please install PEFT >= {min_version_lcmu}." + ) + + adapter_name = adapter_name if adapter_name is not None else "default" + if adapter_kwargs is None: + adapter_kwargs = {} + + from peft import PeftConfig, inject_adapter_in_model, load_peft_weights + from peft.utils import set_peft_model_state_dict + + if self._hf_peft_config_loaded and adapter_name in self.peft_config: + raise ValueError(f"Adapter with name {adapter_name} already exists. Please use a different name.") + + if peft_model_id is None and (adapter_state_dict is None and peft_config is None): + raise ValueError( + "You should either pass a `peft_model_id` or a `peft_config` and `adapter_state_dict` to load an adapter." + ) + + if "device" not in adapter_kwargs: + device = self.device if not hasattr(self, "hf_device_map") else list(self.hf_device_map.values())[0] + else: + device = adapter_kwargs.pop("device") + + # To avoid PEFT errors later on with safetensors. + if isinstance(device, torch.device): + device = str(device) + + # We keep `revision` in the signature for backward compatibility + if revision is not None and "revision" not in adapter_kwargs: + adapter_kwargs["revision"] = revision + elif revision is not None and "revision" in adapter_kwargs and revision != adapter_kwargs["revision"]: + logger.error( + "You passed a `revision` argument both in `adapter_kwargs` and as a standalone argument. " + "The one in `adapter_kwargs` will be used." + ) + + # Override token with adapter_kwargs' token + if "token" in adapter_kwargs: + token = adapter_kwargs.pop("token") + + if peft_config is None: + adapter_config_file = find_adapter_config_file( + peft_model_id, + token=token, + **adapter_kwargs, + ) + + if adapter_config_file is None: + raise ValueError( + f"adapter model file not found in {peft_model_id}. Make sure you are passing the correct path to the " + "adapter model." + ) + + peft_config = PeftConfig.from_pretrained( + peft_model_id, + token=token, + **adapter_kwargs, + ) + peft_config.inference_mode = not is_trainable + + # Create and add fresh new adapters into the model. + inject_adapter_in_model(peft_config, self, adapter_name, **peft_load_kwargs) + + if not self._hf_peft_config_loaded: + self._hf_peft_config_loaded = True + + if peft_model_id is not None: + adapter_state_dict = load_peft_weights(peft_model_id, token=token, device=device, **adapter_kwargs) + + # We need to pre-process the state dict to remove unneeded prefixes - for backward compatibility + processed_adapter_state_dict = {} + prefix = "base_model.model." + for key, value in adapter_state_dict.items(): + if key.startswith(prefix): + new_key = key[len(prefix) :] + else: + new_key = key + processed_adapter_state_dict[new_key] = value + + # Load state dict + incompatible_keys = set_peft_model_state_dict( + self, processed_adapter_state_dict, adapter_name, **peft_load_kwargs + ) + + if incompatible_keys is not None: + err_msg = "" + origin_name = peft_model_id if peft_model_id is not None else "state_dict" + # Check for unexpected keys. + if hasattr(incompatible_keys, "unexpected_keys") and len(incompatible_keys.unexpected_keys) > 0: + err_msg = ( + f"Loading adapter weights from {origin_name} led to unexpected keys not found in the model: " + f"{', '.join(incompatible_keys.unexpected_keys)}. " + ) + + # Check for missing keys. + missing_keys = getattr(incompatible_keys, "missing_keys", None) + if missing_keys: + # Filter missing keys specific to the current adapter, as missing base model keys are expected. + lora_missing_keys = [k for k in missing_keys if "lora_" in k and adapter_name in k] + if lora_missing_keys: + err_msg += ( + f"Loading adapter weights from {origin_name} led to missing keys in the model: " + f"{', '.join(lora_missing_keys)}" + ) + + if err_msg: + logger.warning(err_msg) + + if peft_config.inference_mode: + self.eval() + + # Re-dispatch model and hooks in case the model is offloaded to CPU / Disk. + if ( + (getattr(self, "hf_device_map", None) is not None) + and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0) + and len(self.peft_config) == 1 + ): + self._dispatch_accelerate_model( + device_map=device_map, + max_memory=max_memory, + offload_folder=offload_folder, + offload_index=offload_index, + ) + + def add_adapter(self, adapter_config, adapter_name: Optional[str] = None) -> None: + r""" + If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT + official documentation: https://huggingface.co/docs/peft + + Adds a fresh new adapter to the current model for training purpose. If no adapter name is passed, a default + name is assigned to the adapter to follow the convention of PEFT library (in PEFT we use "default" as the + default adapter name). + + Args: + adapter_config (`~peft.PeftConfig`): + The configuration of the adapter to add, supported adapters are non-prefix tuning and adaption prompts + methods + adapter_name (`str`, *optional*, defaults to `"default"`): + The name of the adapter to add. If no name is passed, a default name is assigned to the adapter. + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + from peft import PeftConfig, inject_adapter_in_model + + adapter_name = adapter_name or "default" + + if not self._hf_peft_config_loaded: + self._hf_peft_config_loaded = True + elif adapter_name in self.peft_config: + raise ValueError(f"Adapter with name {adapter_name} already exists. Please use a different name.") + + if not isinstance(adapter_config, PeftConfig): + raise TypeError(f"adapter_config should be an instance of PeftConfig. Got {type(adapter_config)} instead.") + + # Retrieve the name or path of the model, one could also use self.config._name_or_path + # but to be consistent with what we do in PEFT: https://github.com/huggingface/peft/blob/6e783780ca9df3a623992cc4d1d665001232eae0/src/peft/mapping.py#L100 + adapter_config.base_model_name_or_path = self.__dict__.get("name_or_path", None) + inject_adapter_in_model(adapter_config, self, adapter_name) + + self.set_adapter(adapter_name) + + def set_adapter(self, adapter_name: Union[List[str], str]) -> None: + """ + If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT + official documentation: https://huggingface.co/docs/peft + + Sets a specific adapter by forcing the model to use a that adapter and disable the other adapters. + + Args: + adapter_name (`Union[List[str], str]`): + The name of the adapter to set. Can be also a list of strings to set multiple adapters. + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + if not self._hf_peft_config_loaded: + raise ValueError("No adapter loaded. Please load an adapter first.") + elif isinstance(adapter_name, list): + missing = set(adapter_name) - set(self.peft_config) + if len(missing) > 0: + raise ValueError( + f"Following adapter(s) could not be found: {', '.join(missing)}. Make sure you are passing the correct adapter name(s)." + f" current loaded adapters are: {list(self.peft_config.keys())}" + ) + elif adapter_name not in self.peft_config: + raise ValueError( + f"Adapter with name {adapter_name} not found. Please pass the correct adapter name among {list(self.peft_config.keys())}" + ) + + from peft.tuners.tuners_utils import BaseTunerLayer + from peft.utils import ModulesToSaveWrapper + + _adapters_has_been_set = False + + for _, module in self.named_modules(): + if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)): + # For backward compatbility with previous PEFT versions + if hasattr(module, "set_adapter"): + module.set_adapter(adapter_name) + else: + module.active_adapter = adapter_name + _adapters_has_been_set = True + + if not _adapters_has_been_set: + raise ValueError( + "Did not succeeded in setting the adapter. Please make sure you are using a model that supports adapters." + ) + + def disable_adapters(self) -> None: + r""" + If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT + official documentation: https://huggingface.co/docs/peft + + Disable all adapters that are attached to the model. This leads to inferring with the base model only. + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + if not self._hf_peft_config_loaded: + raise ValueError("No adapter loaded. Please load an adapter first.") + + from peft.tuners.tuners_utils import BaseTunerLayer + from peft.utils import ModulesToSaveWrapper + + for _, module in self.named_modules(): + if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)): + # The recent version of PEFT need to call `enable_adapters` instead + if hasattr(module, "enable_adapters"): + module.enable_adapters(enabled=False) + else: + module.disable_adapters = True + + def enable_adapters(self) -> None: + """ + If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT + official documentation: https://huggingface.co/docs/peft + + Enable adapters that are attached to the model. + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + if not self._hf_peft_config_loaded: + raise ValueError("No adapter loaded. Please load an adapter first.") + + from peft.tuners.tuners_utils import BaseTunerLayer + + for _, module in self.named_modules(): + if isinstance(module, BaseTunerLayer): + # The recent version of PEFT need to call `enable_adapters` instead + if hasattr(module, "enable_adapters"): + module.enable_adapters(enabled=True) + else: + module.disable_adapters = False + + def active_adapters(self) -> List[str]: + """ + If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT + official documentation: https://huggingface.co/docs/peft + + Gets the current active adapters of the model. In case of multi-adapter inference (combining multiple adapters + for inference) returns the list of all active adapters so that users can deal with them accordingly. + + For previous PEFT versions (that does not support multi-adapter inference), `module.active_adapter` will return + a single string. + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + if not is_peft_available(): + raise ImportError("PEFT is not available. Please install PEFT to use this function: `pip install peft`.") + + if not self._hf_peft_config_loaded: + raise ValueError("No adapter loaded. Please load an adapter first.") + + from peft.tuners.tuners_utils import BaseTunerLayer + + for _, module in self.named_modules(): + if isinstance(module, BaseTunerLayer): + active_adapters = module.active_adapter + break + + # For previous PEFT versions + if isinstance(active_adapters, str): + active_adapters = [active_adapters] + + return active_adapters + + def active_adapter(self) -> str: + warnings.warn( + "The `active_adapter` method is deprecated and will be removed in a future version.", FutureWarning + ) + + return self.active_adapters()[0] + + def get_adapter_state_dict(self, adapter_name: Optional[str] = None, state_dict: Optional[dict] = None) -> dict: + """ + If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT + official documentation: https://huggingface.co/docs/peft + + Gets the adapter state dict that should only contain the weights tensors of the specified adapter_name adapter. + If no adapter_name is passed, the active adapter is used. + + Args: + adapter_name (`str`, *optional*): + The name of the adapter to get the state dict from. If no name is passed, the active adapter is used. + state_dict (nested dictionary of `torch.Tensor`, *optional*) + The state dictionary of the model. Will default to `self.state_dict()`, but can be used if special + precautions need to be taken when recovering the state dictionary of a model (like when using model + parallelism). + """ + check_peft_version(min_version=MIN_PEFT_VERSION) + + if not self._hf_peft_config_loaded: + raise ValueError("No adapter loaded. Please load an adapter first.") + + from peft import get_peft_model_state_dict + + if adapter_name is None: + adapter_name = self.active_adapters()[0] + + adapter_state_dict = get_peft_model_state_dict(self, state_dict=state_dict, adapter_name=adapter_name) + return adapter_state_dict + + def _dispatch_accelerate_model( + self, + device_map: str, + max_memory: Optional[int] = None, + offload_folder: Optional[str] = None, + offload_index: Optional[int] = None, + ) -> None: + """ + Optional re-dispatch the model and attach new hooks to the model in case the model has been loaded with + accelerate (i.e. with `device_map=xxx`) + + Args: + device_map (`str` or `Dict[str, Union[int, str, torch.device]]` or `int` or `torch.device`, *optional*): + A map that specifies where each submodule should go. It doesn't need to be refined to each + parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the + same device. If we only pass the device (*e.g.*, `"cpu"`, `"cuda:1"`, `"mps"`, or a GPU ordinal rank + like `1`) on which the model will be allocated, the device map will map the entire model to this + device. Passing `device_map = 0` means put the whole model on GPU 0. + + To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For + more information about each option see [designing a device + map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). + max_memory (`Dict`, *optional*): + A dictionary device identifier to maximum memory. Will default to the maximum memory available for each + GPU and the available CPU RAM if unset. + offload_folder (`str` or `os.PathLike`, *optional*): + If the `device_map` contains any value `"disk"`, the folder where we will offload weights. + offload_index (`int`, *optional*): + The offload_index argument to be passed to `accelerate.dispatch_model` method. + """ + dispatch_model_kwargs = {} + # Safety checker for previous `accelerate` versions + # `offload_index` was introduced in https://github.com/huggingface/accelerate/pull/873/ + if "offload_index" in inspect.signature(dispatch_model).parameters: + dispatch_model_kwargs["offload_index"] = offload_index + + no_split_module_classes = self._no_split_modules + + if device_map != "sequential": + max_memory = get_balanced_memory( + self, + max_memory=max_memory, + no_split_module_classes=no_split_module_classes, + low_zero=(device_map == "balanced_low_0"), + ) + if isinstance(device_map, str): + device_map = infer_auto_device_map( + self, max_memory=max_memory, no_split_module_classes=no_split_module_classes + ) + dispatch_model( + self, + device_map=device_map, + offload_dir=offload_folder, + **dispatch_model_kwargs, + ) + + def delete_adapter(self, adapter_names: Union[List[str], str]) -> None: + """ + Delete an adapter's LoRA layers from the underlying model. + + Args: + adapter_names (`Union[List[str], str]`): + The name(s) of the adapter(s) to delete. + + Example: + + ```py + from diffusers import AutoPipelineForText2Image + import torch + + pipeline = AutoPipelineForText2Image.from_pretrained( + "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16 + ).to("cuda") + pipeline.load_lora_weights( + "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_names="cinematic" + ) + pipeline.delete_adapters("cinematic") + ``` + """ + + check_peft_version(min_version=MIN_PEFT_VERSION) + + if not self._hf_peft_config_loaded: + raise ValueError("No adapter loaded. Please load an adapter first.") + + from peft.tuners.tuners_utils import BaseTunerLayer + + if isinstance(adapter_names, str): + adapter_names = [adapter_names] + + # Check that all adapter names are present in the config + missing_adapters = [name for name in adapter_names if name not in self.peft_config] + if missing_adapters: + raise ValueError( + f"The following adapter(s) are not present and cannot be deleted: {', '.join(missing_adapters)}" + ) + + for adapter_name in adapter_names: + for module in self.modules(): + if isinstance(module, BaseTunerLayer): + if hasattr(module, "delete_adapter"): + module.delete_adapter(adapter_name) + else: + raise ValueError( + "The version of PEFT you are using is not compatible, please use a version that is greater than 0.6.1" + ) + + # For transformers integration - we need to pop the adapter from the config + if getattr(self, "_hf_peft_config_loaded", False) and hasattr(self, "peft_config"): + self.peft_config.pop(adapter_name, None) + + # In case all adapters are deleted, we need to delete the config + # and make sure to set the flag to False + if len(self.peft_config) == 0: + del self.peft_config + self._hf_peft_config_loaded = False diff --git a/lib/python3.12/site-packages/transformers/integrations/quanto.py b/lib/python3.12/site-packages/transformers/integrations/quanto.py new file mode 100644 index 0000000000000000000000000000000000000000..1c5702321937da5f1238427156452cbccdbcea4a --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/quanto.py @@ -0,0 +1,98 @@ +# 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. + +from ..utils import is_optimum_quanto_available, is_torch_available, logging + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +def replace_with_quanto_layers( + model, + quantization_config=None, + modules_to_not_convert=None, + current_key_name=None, + has_been_replaced=False, +): + """ + Public method that recursively replaces the Linear layers of the given model with Quanto quantized layers. + Returns the converted model and a boolean that indicates if the conversion has been successfull or not. + + Args: + model (`torch.nn.Module`): + The model to convert, can be any `torch.nn.Module` instance. + quantization_config (`AqlmConfig`, defaults to `None`): + The quantization config object that contains the quantization parameters. + modules_to_not_convert (`list`, *optional*, defaults to `None`): + A list of modules to not convert. If a module name is in the list (e.g. `lm_head`), it will not be + converted. + current_key_name (`list`, *optional*, defaults to `None`): + A list that contains the current key name. This is used for recursion and should not be passed by the user. + has_been_replaced (`bool`, *optional*, defaults to `None`): + A boolean that indicates if the conversion has been successful or not. This is used for recursion and + should not be passed by the user. + """ + from accelerate import init_empty_weights + + if is_optimum_quanto_available(): + from optimum.quanto import QLayerNorm, QLinear, qfloat8, qint2, qint4, qint8 + + w_mapping = {"float8": qfloat8, "int8": qint8, "int4": qint4, "int2": qint2} + a_mapping = {None: None, "float8": qfloat8, "int8": qint8} + + if modules_to_not_convert is None: + modules_to_not_convert = [] + + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + + if not any(key in ".".join(current_key_name) for key in modules_to_not_convert): + with init_empty_weights(): + if isinstance(module, torch.nn.Linear): + model._modules[name] = QLinear( + in_features=module.in_features, + out_features=module.out_features, + bias=module.bias is not None, + dtype=module.weight.dtype, + weights=w_mapping[quantization_config.weights], + activations=a_mapping[quantization_config.activations], + ) + model._modules[name].requires_grad_(False) + has_been_replaced = True + elif isinstance(module, torch.nn.LayerNorm): + if quantization_config.activations is not None: + model._modules[name] = QLayerNorm( + module.normalized_shape, + module.eps, + module.elementwise_affine, + module.bias is not None, + activations=a_mapping[quantization_config.activations], + ) + has_been_replaced = True + if len(list(module.children())) > 0: + _, has_been_replaced = replace_with_quanto_layers( + module, + quantization_config=quantization_config, + modules_to_not_convert=modules_to_not_convert, + current_key_name=current_key_name, + has_been_replaced=has_been_replaced, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced diff --git a/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py b/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..9c924c048ad52929a2d0f890a22295d5a56ef505 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py @@ -0,0 +1,65 @@ +from typing import Optional, Tuple + +import torch + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def sdpa_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + dropout: float = 0.0, + scaling: Optional[float] = None, + is_causal: Optional[bool] = None, + **kwargs, +) -> Tuple[torch.Tensor, None]: + if hasattr(module, "num_key_value_groups"): + key = repeat_kv(key, module.num_key_value_groups) + value = repeat_kv(value, module.num_key_value_groups) + + causal_mask = attention_mask + if attention_mask is not None and causal_mask.ndim == 4: + causal_mask = causal_mask[:, :, :, : key.shape[-2]] + + # SDPA with memory-efficient backend is bugged with non-contiguous inputs and custom attn_mask for some torch versions + # Reference: https://github.com/pytorch/pytorch/issues/112577. + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + # Note that it is important to check first for the shape, otherwise compile will fail with `argument 'is_causal' must be bool, not SymBool` + if is_causal is None: + is_causal = query.shape[2] > 1 and causal_mask is None + + # Shapes (e.g. query.shape[2]) are tensors during jit tracing, resulting in `is_causal` being a tensor. + # We convert it to a bool for the SDPA kernel that only accepts bools. + if torch.jit.is_tracing() and isinstance(is_causal, torch.Tensor): + is_causal = is_causal.item() + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=causal_mask, + dropout_p=dropout, + scale=scaling, + is_causal=is_causal, + ) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, None diff --git a/lib/python3.12/site-packages/transformers/integrations/spqr.py b/lib/python3.12/site-packages/transformers/integrations/spqr.py new file mode 100644 index 0000000000000000000000000000000000000000..58b71740d37c77f24fe60ff65fe179b2737c4aa7 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/spqr.py @@ -0,0 +1,122 @@ +# 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. +"SpQR (Sparse-Quantized Representation) integration file" + +from ..utils import is_accelerate_available, is_spqr_available, is_torch_available + + +if is_torch_available(): + import torch.nn as nn + + +def replace_with_spqr_linear( + model, + quantization_config=None, + modules_to_not_convert=None, + current_key_name=None, + has_been_replaced=False, +): + """ + Public method that recursively replaces the Linear layers of the given model with SpQR quantized layers. + `accelerate` is needed to use this method. Returns the converted model and a boolean that indicates if the + conversion has been successful or not. + + Args: + model (`torch.nn.Module`): + The model to convert, can be any `torch.nn.Module` instance. + quantization_config (`SpQRConfig`): + The quantization config object that contains the quantization parameters. + modules_to_not_convert (`list[str]`, *optional*): + A list of nn.Linear weights to not convert. If a parameter path is in the list (e.g. `lm_head.weight`), the corresponding module will not be + converted. + current_key_name (`list`, *optional*): + A list that contains the current key name. This is used for recursion and should not be passed by the user. + has_been_replaced (`bool`, *optional*): + A boolean that indicates if the conversion has been successful or not. This is used for recursion and + should not be passed by the user. + """ + if modules_to_not_convert is None: + modules_to_not_convert = [] + + if is_accelerate_available(): + from accelerate import init_empty_weights + if is_spqr_available(): + from spqr_quant import QuantizedLinear + + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + + if isinstance(module, nn.Linear): + # Check if the current key is not in the `modules_to_not_convert` + if ".".join(current_key_name) + ".weight" not in modules_to_not_convert: + with init_empty_weights(): + tensor_name = ".".join(current_key_name) + + shapes = quantization_config.shapes + shapes_keys = shapes.keys() + + shapes_valid = ( + f"{tensor_name}.dense_weights.shape" in shapes_keys + and f"{tensor_name}.row_offsets.shape" in shapes_keys + and f"{tensor_name}.col_vals.shape" in shapes_keys + and f"{tensor_name}.in_perm.shape" in shapes_keys + ) + + if not shapes_valid: + raise ValueError( + f"The SpQR quantization config does not contain the shape " + f"configuration for {tensor_name}. This indicates that the " + f"configuration is either invalid or corrupted." + ) + + dense_weights_shape = shapes[f"{tensor_name}.dense_weights.shape"] + row_offsets_shape = shapes[f"{tensor_name}.row_offsets.shape"] + col_vals_shape = shapes[f"{tensor_name}.col_vals.shape"] + in_perm_shape = shapes[f"{tensor_name}.in_perm.shape"] + + in_features = module.in_features + out_features = module.out_features + + model._modules[name] = QuantizedLinear.create_placehodler( + rows=out_features, + cols=in_features, + bits=quantization_config.bits, + beta1=quantization_config.beta1, + beta2=quantization_config.beta2, + dense_weights_shape=dense_weights_shape, + row_offsets_shape=row_offsets_shape, + col_vals_shape=col_vals_shape, + in_perm_shape=in_perm_shape, + ) + has_been_replaced = True + + # Store the module class in case we need to transpose the weight later + model._modules[name].source_cls = type(module) + # Force requires grad to False to avoid unexpected errors + model._modules[name].requires_grad_(False) + else: + pass + if len(list(module.children())) > 0: + _, has_been_replaced = replace_with_spqr_linear( + module, + quantization_config=quantization_config, + modules_to_not_convert=modules_to_not_convert, + current_key_name=current_key_name, + has_been_replaced=has_been_replaced, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced diff --git a/lib/python3.12/site-packages/transformers/integrations/tensor_parallel.py b/lib/python3.12/site-packages/transformers/integrations/tensor_parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..34b21444fe35faed6487caf8f655c14e660487d4 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/tensor_parallel.py @@ -0,0 +1,672 @@ +# 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. +from __future__ import annotations + +import re +from functools import lru_cache, partial +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn + +from ..utils import is_torch_greater_or_equal, logging + + +ALL_LAYERNORM_LAYERS = [nn.LayerNorm] + +logger = logging.get_logger(__name__) + +# Cache this result has it's a C FFI call which can be pretty time-consuming +_torch_distributed_available = torch.distributed.is_available() + + +if is_torch_greater_or_equal("2.5") and _torch_distributed_available: + from torch.distributed.tensor import DTensor, Placement, Replicate, Shard + + +def _blocks_to_block_sizes(total_size: int, blocks: Union[int, List[int]]) -> List[int]: + """ + Convert block count or proportions to block sizes. + + This function accepts + + - The number of blocks (int), in which case the block size is + total_size//blocks; or + - A list of block sizes (List[int]). + + In the second case, if sum(blocks) < total_size, the ratios between + the block sizes will be preserved. For instance, if blocks is + [2, 1, 1] and total_size is 1024, the returned block sizes are + [512, 256, 256]. + """ + if isinstance(blocks, list): + total_blocks = sum(blocks) + assert total_size % total_blocks == 0, f"Cannot split {total_size} in proportional blocks: {blocks}" + part_size = total_size // total_blocks + return [part_size * block for block in blocks] + else: + assert total_size % blocks == 0, f"Prepacked is not divisible by {blocks}" + single_size = total_size // blocks + return [single_size] * blocks + + +str_to_torch_dtype = { + "BOOL": torch.bool, + "U8": torch.uint8, + "I8": torch.int8, + "I16": torch.int16, + "F16": torch.float16, + "BF16": torch.bfloat16, + "I32": torch.int32, + "F32": torch.float32, + "F64": torch.float64, + "I64": torch.int64, + "F8_E4M3": torch.float8_e4m3fn, +} + + +def get_packed_weights(param, empty_param, device_mesh, rank, dim): + """ + When weights are packed (gate_up_proj), we need to make sure each shard gets its correct share. + So if you have: gate_proj ( 16, 5120, 8190) + and up_proj ( 16, 5120, 8190) + packed as gate_up_proj ( 16, 5120, 2 * 8190) + And you shard along the last dimension, you need to interleave the gate and up values: + + Now, if we shard along the last dimension across TP_size (Tensor Parallelism size), we must interleave the values from gate and up projections correctly. + + Let's take TP_size = 4 for an example: + + Packed tensor `gate_up_proj` + --------------------------------------------------------------- + [ G0 G1 G2 G3 | G4 G5 G6 G7 | ... | U0 U1 U2 U3 | U4 U5 U6 U7 | ... ] + ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ + Gate Slice 0 Gate Slice 1 Up Slice 0 Up Slice 1 + + Explanation: + - The first half of the tensor (left of the center) holds the gate_proj values. + - The second half (right of the center) holds the up_proj values. + - For TP=4, we divide each half into 4 slices. In this example, we show two slices for brevity. + - Each shard receives one slice from the gate part and the corresponding slice from the up part. + + For instance: + • Shard 0 gets: [ Gate Slice 0, Up Slice 0 ] = [ G0, G1, G2, G3, U0, U1, U2, U3 ] + • Shard 1 gets: [ Gate Slice 1, Up Slice 1 ] = [ G4, G5, G6, G7, U4, U5, U6, U7 ] + • … and so on. + + This ensures that each shard receives an equal portion of both gate and up projections, maintaining consistency across tensor parallelism. + """ + slice_ = param + total_size = empty_param.shape[dim] + world_size = device_mesh.size() + block_sizes = _blocks_to_block_sizes(total_size=total_size, blocks=2) + + tensors_slices = [] + block_offset = 0 + for block_size in block_sizes: + shard_block_size = block_size // world_size + start = rank * shard_block_size + stop = (rank + 1) * shard_block_size + tensors_slices += range(block_offset + start, block_offset + stop) + block_offset += block_size + + slice_dtype = slice_.get_dtype() + # Handle F8_E4M3 dtype by converting to float16 before slicing + # Without upcasting, the slicing causes : RuntimeError: "index_cpu" not implemented for 'Float8_e4m3fn' + if slice_dtype == "F8_E4M3": + slice_ = slice_[...].to(torch.float16) + + if dim == 0: + tensor = slice_[tensors_slices, ...] + elif dim == 1 or dim == -2: + tensor = slice_[:, tensors_slices, ...] + elif dim == 2 or dim == -1: + tensor = slice_[..., tensors_slices] + else: + raise ValueError(f"Unsupported dim {dim}, only dim 0, 1 or 2 are supported") + return tensor.to(str_to_torch_dtype[slice_dtype]) + + +def get_tensor_shard(param, empty_param, device_mesh, rank, dim): + if dim == 0: + size_ = empty_param.shape[0] + param = param[rank * (size_ // device_mesh.size()) : (rank + 1) * (size_ // device_mesh.size()), ...] + elif dim == 1 or dim == -2: + size_ = empty_param.shape[-2] + param = param[..., rank * (size_ // device_mesh.size()) : (rank + 1) * (size_ // device_mesh.size()), :] + elif dim == 2 or dim == -1: + size_ = empty_param.shape[-1] + param = param[..., rank * (size_ // device_mesh.size()) : (rank + 1) * (size_ // device_mesh.size())] + else: + raise ValueError(f"Unsupported dim {dim}, only dim 0, 1 or 2 are supported") + return param + + +def distribute_module( + module: nn.Module, + device_mesh=None, + input_fn=None, + output_fn=None, +) -> nn.Module: + """ + Copy pasted from torch's function but we remove the communications (partitionning) + as well as buffer registering that is similarly not efficient. + """ + if len(module._forward_pre_hooks) == 0: + if input_fn is not None: + module.register_forward_pre_hook(lambda mod, inputs: input_fn(mod, inputs, device_mesh)) + if output_fn is not None: + module.register_forward_hook(lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh)) + return module + + +class TensorParallelLayer: + """ + General tensor parallel layer for transformers. + """ + + use_dtensor = True + + @staticmethod + def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh): ... + + @staticmethod + def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): ... + + def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh): + raise NotImplementedError + + def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module: + if self.use_dtensor: + distribute_module( + module, + device_mesh, + partial(self._prepare_input_fn, self.input_layouts, self.desired_input_layouts), + partial(self._prepare_output_fn, self.output_layouts, self.use_local_output), + ) + + +# use_dtensor needs to be set to false for nn.Parameter when you want to view, chunk, slice +# you name it. Whatever you want to do that is a bit unconventional, you need local tensors +class GatherParallel(TensorParallelLayer): + """ + Simple class used to define the hooks to add to a layer when we just want to gather the outputs + """ + + def __init__( + self, + *, + input_layouts: Optional[Placement] = None, + output_layouts: Optional[Placement] = None, + use_local_output: bool = True, + ): + super().__init__() + self.input_layouts = (input_layouts or Replicate(),) + self.output_layouts = output_layouts + self.desired_input_layouts = (Replicate(),) + self.use_local_output = use_local_output + + @staticmethod + def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh): + if isinstance(inputs[0], DTensor): + inputs = inputs[0].to_local() + return inputs + + @staticmethod + def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): + # this op cannot be asynch, otherwise it completely breaks the outputs of models + torch.distributed.all_reduce(outputs[0], op=torch.distributed.ReduceOp.SUM, async_op=False) + return outputs + + +class IsolatedParallel(TensorParallelLayer): + """ + This class is used to isolate computation in a TP layer from the rest of the world. + Parameters need to be LOCAL, so not dtensors + """ + + @staticmethod + def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh=None): + # annotate module input placements/sharding with input_layouts + input_tensor = inputs[0] + if isinstance(input_tensor, DTensor): + input_tensor = input_tensor.to_local() + return input_tensor + + @staticmethod + def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh=None): + # TODO: figure out dynamo support for instance method and switch this to instance method + return outputs + + def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module: + distribute_module( + module, + device_mesh, + partial(self._prepare_input_fn, None, None), + partial(self._prepare_output_fn, None, None), + ) + + +class ColwiseParallel(TensorParallelLayer): + """ + General tensor parallel layer for transformers. + """ + + def __init__( + self, + *, + input_layouts: Optional[Placement] = None, + output_layouts: Optional[Placement] = None, + use_local_output: bool = True, + use_dtensor=True, + ): + super().__init__() + self.input_layouts = (input_layouts or Replicate(),) + self.output_layouts = (output_layouts or Shard(-1),) + self.desired_input_layouts = (Replicate(),) + self.use_local_output = use_local_output + self.use_dtensor = use_dtensor + + @staticmethod + def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh): + # TODO: figure out dynamo support for instance method and switch this to instance method + # annotate module input placements/sharding with input_layouts + input_tensor = inputs[0] + if not isinstance(input_tensor, DTensor): + input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False) + + # transform the input layouts to the desired layouts of ColwiseParallel + if input_layouts != desired_input_layouts: + input_tensor = input_tensor.redistribute(placements=desired_input_layouts, async_op=False) + return input_tensor + + def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh): + # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only) + # means Colwise as Linear is input * weight^T + bias, where + # weight would become Shard(1) + if param_type == "bias": + parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1) + shard = [Shard(-1)] + else: + shard = [Shard(-2)] + parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -2) + + parameter = parameter.to(param_casting_dtype) + if to_contiguous: + parameter = parameter.contiguous() + if self.use_dtensor: + parameter = DTensor.from_local(parameter, device_mesh, shard, run_check=False) + return nn.Parameter(parameter) + + @staticmethod + def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): + # outputs is a shard on last dimension DTensor, i.e. Shard(-1) + if outputs.placements != output_layouts: + outputs = outputs.redistribute(placements=output_layouts, async_op=False) + # back to local tensor + return outputs.to_local() if use_local_output else outputs + + +class PackedColwiseParallel(ColwiseParallel): + def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh): + # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only) + # means Colwise as Linear is input * weight^T + bias, where + # weight would become Shard(1) + parameter = get_packed_weights(param, empty_param, device_mesh, rank, -2) + parameter = parameter.to(param_casting_dtype) + if to_contiguous: + parameter = parameter.contiguous() + if self.use_dtensor: + parameter = DTensor.from_local(parameter, device_mesh, [Shard(-2)], run_check=False) + return nn.Parameter(parameter) + + +class RowwiseParallel(TensorParallelLayer): + """ + Partition a compatible nn.Module in a row-wise fashion. Currently supports nn.Linear and nn.Embedding. + Users can compose it with ColwiseParallel to achieve the sharding of more complicated modules. + (i.e. MLP, Attention) + + Keyword Args: + input_layouts (Placement, optional): + The DTensor layout of input tensor for the nn.Module, this is used to annotate the input tensor to + become a DTensor. If not specified, we assume the input tensor to be sharded on the last dimension. + output_layouts (Placement, optional): + The DTensor layout of the output for the nn.Module, this is used to ensure the output of the nn.Module + with the user desired layout. If not specified, the output tensor is replicated. + use_local_output (bool, optional): + Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module output, default: True. + Returns: + A :class:`ParallelStyle` object that represents Rowwise sharding of the nn.Module. + """ + + def __init__( + self, + *, + input_layouts: Optional[Placement] = None, + output_layouts: Optional[Placement] = None, + use_local_output: bool = True, + use_dtensor=True, + ): + super().__init__() + self.input_layouts = (input_layouts or Shard(-1),) + self.output_layouts = (output_layouts or Replicate(),) + self.use_local_output = use_local_output + self.use_dtensor = use_dtensor + + def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh): + # Rowwise shard weight to Shard(1), bias to Replicate(), weight be Shard(1) + # means Rowwise as nn.Linear is input * weight^T + bias, where + # weight would become Shard(0) + if param_type != "bias": + parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1) + shard = [Shard(-1)] + else: + shard = [Replicate()] + parameter = param[:] + + parameter = parameter.to(param_casting_dtype) + if to_contiguous: + parameter = parameter.contiguous() + if self.use_dtensor: + parameter = DTensor.from_local(parameter, device_mesh, shard, run_check=False) + return nn.Parameter(parameter) + + @staticmethod + def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh): + if hasattr(mod, "bias") and mod.bias is not None: + mod._bias = mod.bias + mod.bias = None + + input_tensor = inputs[0] + if not isinstance(input_tensor, DTensor): + input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False) + + if input_layouts != desired_input_layouts: + input_tensor = input_tensor.redistribute(placements=desired_input_layouts, async_op=True) + return input_tensor + + @staticmethod + def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): + # Rowwise sharding produces partial output, depending on output layouts: + # 1. to replicate -> allreduce + # 2. to shard -> reduce_scatter + if outputs.placements != output_layouts: + outputs = outputs.redistribute(placements=output_layouts, async_op=True) + if hasattr(mod, "_bias"): + outputs += mod._bias + # back to local tensor if use_local_output is True + return outputs.to_local() if use_local_output else outputs + + def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module: + module._distribute_module_applied = True + if self.use_dtensor: + if isinstance(module, nn.Linear): + # rowwise linear runtime sharding requires input tensor shard on last dim + self.desired_input_layouts: Tuple[Placement, ...] = (Shard(-1),) + elif isinstance(module, nn.Embedding): + # rowwise embedding runtime sharding requires input tensor replicated + self.desired_input_layouts = (Replicate(),) + elif isinstance(module, nn.Parameter): + # rowwise embedding runtime sharding requires input tensor replicated + self.desired_input_layouts = (Shard(-1),) + else: + raise NotImplementedError("RowwiseParallel currently only support nn.Linear and nn.Embedding!") + + distribute_module( + module, + device_mesh, + partial(self._prepare_input_fn, self.input_layouts, self.desired_input_layouts), + partial(self._prepare_output_fn, self.output_layouts, self.use_local_output), + ) + + +class PackedRowwiseParallel(RowwiseParallel): + def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh): + # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only) + # means Colwise as Linear is input * weight^T + bias, where + # weight would become Shard(1) + parameter = get_packed_weights(param, empty_param, device_mesh, rank, -1) + parameter = parameter.to(param_casting_dtype) + if to_contiguous: + parameter = parameter.contiguous() + if self.use_dtensor: + parameter = DTensor.from_local(parameter, device_mesh, [Shard(-1)], run_check=False) + return nn.Parameter(parameter) + + +class SequenceParallel(TensorParallelLayer): + """ + SequenceParallel replicates a compatible ``nn.Module`` parameters and runs the sharded computation with + input sharded on the sequence dimension. This currently supports ``nn.LayerNorm``, ``nn.Dropout``, and the + `RMSNorm python implementation `__ + + This style implements the operation that is described in the paper + `Reducing Activation Recomputation in Large Transformer Models `__ + + If the input passed in to this ``nn.Module`` is a :class:`torch.Tensor`, it assumes that the input is already sharded + on the sequence dimension and converts the input to a :class:`DTensor` sharded on the sequence dimension. If the input + passed in to this ``nn.Module`` is already a :class:`DTensor` but is not sharded on the sequence dimension, it would + redistribute the input to be sharded on the sequence dimension. + + The output of the ``nn.Module`` will be sharded on the sequence dimension. + + Keyword Args: + sequence_dim (int, optional): + The sequence dimension of the input tensor for the ``nn.Module``, this is used to annotate the input tensor to + become a DTensor that is sharded on the sequence dimension, default: 1. + use_local_output (bool, optional): + Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module output, default: False. + Returns: + A :class:`ParallelStyle` object that represents Sequence Parallel of the ``nn.Module``. + + Example:: + >>> # xdoctest: +SKIP(failing) + >>> from torch.distributed.tensor.parallel import parallelize_module, SequenceParallel + >>> from torch.distributed.device_mesh import init_device_mesh + >>> ... + >>> m = Model(...) # m is a nn.Module that contains a "norm" nn.LayerNorm submodule + >>> tp_mesh = init_device_mesh("cuda", (8,)) + >>> + >>> # By default, the input of the "norm" will be converted to DTensor that shards on the sequence dim + >>> # and the output of "norm" will return a sharded on sequence dimension :class:`DTensor`. + >>> + >>> sharded_mod = parallelize_module(m, tp_mesh, {"norm": SequenceParallel()}), + >>> ... + + .. note:: SequenceParallel style assumes ones initialization if there are weights in the nn.Module (i.e. + ``nn.LayerNorm`` or ``RMSNorm``, and they by default have ones initialization). If you have custom + inits for the weights on those modules, you need to broadcast the weights before/after parallelizing + to ensure that they are replicated. + """ + + def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = False, use_dtensor=False): + super().__init__() + self.input_layouts = (Replicate(),) + self.desired_input_layouts = (Shard(1),) + self.output_layouts = (Replicate(),) + self.use_local_output = use_local_output + self.use_dtensor = True + self.sequence_sharding = (Shard(sequence_dim),) + self.use_local_output = use_local_output + + @staticmethod + def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh): + input_tensor = inputs[0] + if not isinstance(input_tensor, DTensor): + input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False) + if input_layouts != desired_input_layouts: + input_tensor = input_tensor.redistribute(placements=desired_input_layouts, async_op=True) + return input_tensor + + @staticmethod + def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): + outputs = outputs.redistribute( + placements=(Replicate(),), async_op=True + ) # maybe we have to replicate ? because next layer is not sharded + return outputs.to_local() # if use_local_output else outputs + + def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh): + # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only) + # means Colwise as Linear is input * weight^T + bias, where + # weight would become Shard(1) + parameter = param[:] + parameter = parameter.to(param_casting_dtype) + if to_contiguous: + parameter = parameter.contiguous() + if self.use_dtensor: + parameter = DTensor.from_local(parameter, device_mesh, [Replicate()], run_check=False) + return nn.Parameter(parameter) + + +SUPPORTED_TP_STYLES = { + "colwise", + "rowwise", + "colwise_rep", + "rowwise_rep", + "local_colwise", + "local_rowwise", + "local", + "gather", + "local_packed_rowwise", + "sequence_parallel", +} + + +@lru_cache +def translate_to_torch_parallel_style(style: str): + """ + In model configurations, we use a neutral type (string) to specify parallel + styles, here we translate them into torch.distributed tensor-parallel + types. + """ + if not isinstance(style, str): + raise ValueError(f"Unsupported parallel style type {type(style)}, expected str") + + if style == "colwise": + return ColwiseParallel() + elif style == "rowwise": + return RowwiseParallel() + elif style == "colwise_rep": + return ColwiseParallel(output_layouts=Replicate()) + elif style == "rowwise_rep": + return RowwiseParallel(input_layouts=Replicate()) + elif style == "local_colwise": + return ColwiseParallel(use_dtensor=False) + elif style == "local_rowwise": + return RowwiseParallel(use_dtensor=False) + elif style == "local": + return IsolatedParallel() + elif style == "gather": + return GatherParallel() + elif style == "local_packed_rowwise": + return PackedRowwiseParallel(use_dtensor=False) + elif style == "sequence_parallel": + return SequenceParallel() + else: + raise ValueError(f"Unsupported parallel style value: {style}") + + +def add_tensor_parallel_hooks_to_module(model, module, tp_plan, layer_name, current_module_plan, device_mesh): + """ + Add hooks to the module holding the layer. Meaning: + ``` + class MyModel(nn.Module): + def __init__(self): + self.layer = nn.Linear(10, 10) + ``` + has state_dict like: + ``` + { + "layer.weight": torch.Tensor, + "layer.bias": torch.Tensor + } + ``` + we add hooks to `MyModel` as well as `layer` to make sure that the tensors are correctly sharded and gathered. + """ + + # 1. We add hooks to the layer being loaded: + if current_module_plan is not None: + tp_layer = translate_to_torch_parallel_style(current_module_plan) + try: + tp_layer.prepare_module_tp(module, device_mesh) + except NotImplementedError as e: + print( + f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP plan: {e}" + ) + + # 2. We add hooks to the parrent module if needed + if "." in layer_name: + parrent_layer_name = layer_name.rsplit(".", 1)[0] + generic_name = re.sub(r"\d+", "*", parrent_layer_name) + # The module itself needs hooks + if module_plan := tp_plan.get(generic_name, False): + tp_layer = translate_to_torch_parallel_style(module_plan) + module_to_tp_ = model.get_submodule(parrent_layer_name) + tp_layer.prepare_module_tp(module_to_tp_, device_mesh) + + +def shard_and_distribute_module( + model, param, empty_param, parameter_name, param_casting_dtype, is_contiguous, rank, device_mesh +): + r""" + Main uses cases: + - column / rowise parallelism, you just shard all the weights of the layer (weight and bias) + - packed layers: you slice the weights, then shard like above + - custom operation: + - you want to add an all-gather at the end of a local layer. + - you want to have a layer that is isolated from the rest of the world (because torch.DTensor does not work well with `.view` for instance) + + """ + param_name, param_type = parameter_name.rsplit(".", 1) if "." in parameter_name else parameter_name + tp_plan = model._tp_plan + module_to_tp = model.get_submodule(param_name) + current_module_plan = None + rank = int(rank) + generic_param_name = re.sub(r"\d+", "*", parameter_name) + if generic_param_name in tp_plan: + current_module_plan = tp_plan[generic_param_name] + elif "." in generic_param_name and generic_param_name.rsplit(".", 1)[0] in tp_plan: + current_module_plan = tp_plan[generic_param_name.rsplit(".", 1)[0]] + + # Add hooks to the module if not done yet + # add_tensor_parallel_hooks_to_module(model, module_to_tp, tp_plan, param_name, current_module_plan, device_mesh) + if not getattr(module_to_tp, "_is_hooked", False): + add_tensor_parallel_hooks_to_module(model, module_to_tp, tp_plan, param_name, current_module_plan, device_mesh) + module_to_tp._is_hooked = True + + if current_module_plan is not None: + try: + tp_layer = translate_to_torch_parallel_style(current_module_plan) + param = tp_layer.partition_tensor( + param, empty_param, param_type, param_casting_dtype, is_contiguous, rank, device_mesh + ) + except NotImplementedError as e: + print( + f"Trying to prepare {parameter_name}, but it's not supported. Corresponding module: {module_to_tp} Fix it's TP plan, current layer: {tp_layer} : {e}" + ) + else: + # TODO log no plan modules in set + # print("No plan for", parameter_name,end ="\n") + param = param[...].to(param_casting_dtype) + if is_contiguous: + param = param.contiguous() + + # SUPER IMPORTANT we have to use setattr + # otherwise loading is crazy slow + if not isinstance(param, torch.nn.Parameter): + param = torch.nn.Parameter(param) + setattr(module_to_tp, param_type, param) + # module_to_tp.load_state_dict({param_type: param}, strict=False, assign=True) + return param diff --git a/lib/python3.12/site-packages/transformers/integrations/tiktoken.py b/lib/python3.12/site-packages/transformers/integrations/tiktoken.py new file mode 100644 index 0000000000000000000000000000000000000000..6f200bc355933bcb9ec119e08b560241fc1433c5 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/tiktoken.py @@ -0,0 +1,43 @@ +from pathlib import Path +from typing import Any + +from transformers.convert_slow_tokenizer import TikTokenConverter +from transformers.tokenization_utils_fast import TIKTOKEN_VOCAB_FILE, TOKENIZER_FILE + + +def convert_tiktoken_to_fast(encoding: Any, output_dir: str): + """ + Converts given `tiktoken` encoding to `PretrainedTokenizerFast` and saves the configuration of converted tokenizer + on disk. + + Args: + encoding (`str` or `tiktoken.Encoding`): + Tokenizer from `tiktoken` library. If `encoding` is `str`, the tokenizer will be loaded with + `tiktoken.get_encoding(encoding)`. + output_dir (`str`): + Save path for converted tokenizer configuration file. + """ + output_dir = Path(output_dir) + output_dir.mkdir(exist_ok=True) + + save_file = output_dir / "tiktoken" / TIKTOKEN_VOCAB_FILE + tokenizer_file = output_dir / TOKENIZER_FILE + + save_file_absolute = str(save_file.absolute()) + output_file_absolute = str(tokenizer_file.absolute()) + + try: + from tiktoken import get_encoding + from tiktoken.load import dump_tiktoken_bpe + + if isinstance(encoding, str): + encoding = get_encoding(encoding) + + dump_tiktoken_bpe(encoding._mergeable_ranks, save_file_absolute) + except ImportError: + raise ValueError("`tiktoken` is required to save a `tiktoken` file. Install it with `pip install tiktoken`.") + + tokenizer = TikTokenConverter( + vocab_file=save_file_absolute, pattern=encoding._pat_str, additional_special_tokens=encoding._special_tokens + ).converted() + tokenizer.save(output_file_absolute) diff --git a/lib/python3.12/site-packages/transformers/integrations/tpu.py b/lib/python3.12/site-packages/transformers/integrations/tpu.py new file mode 100644 index 0000000000000000000000000000000000000000..1f7b0df3cd3b905805581fe2cf9c0314c1dda72a --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/tpu.py @@ -0,0 +1,36 @@ +# 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. + +from torch.utils.data import DataLoader + +from ..utils import is_torch_xla_available + + +def tpu_spmd_dataloader(dataloader: DataLoader): + if is_torch_xla_available(): + import torch_xla.distributed.parallel_loader as pl + + assert isinstance(dataloader, pl.MpDeviceLoader), ( + "The dataloader must be a `torch_xla.distributed.parallel_loader.MpDeviceLoader`." + ) + + # This is to support PyTorch/XLA FSDP via SPMD. + # Here we shard the input data's 0th dim across the fsdp axis. + import torch_xla.distributed.spmd as xs + + sharding_spec = xs.ShardingSpec(xs.get_global_mesh(), ("fsdp", None)) + dataloader._parallel_loader_kwargs["input_sharding"] = sharding_spec + return dataloader + else: + return dataloader diff --git a/lib/python3.12/site-packages/transformers/integrations/vptq.py b/lib/python3.12/site-packages/transformers/integrations/vptq.py new file mode 100644 index 0000000000000000000000000000000000000000..aa435517e81ebe3f9e624b45d34a601e512f5b1b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/integrations/vptq.py @@ -0,0 +1,101 @@ +# 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. +"VPTQ (Vector Post-Training Quantization) integration file" + +import torch.nn as nn +from accelerate import init_empty_weights +from vptq import VQuantLinear + + +def replace_with_vptq_linear( + model, + quantization_config=None, + modules_to_not_convert=None, + current_key_name=None, + has_been_replaced=False, +): + """ + Public method that recursively replaces the Linear layers of the given model with VPTQ quantized layers. + `accelerate` is needed to use this method. Returns the converted model and a boolean that indicates if the + conversion has been successfull or not. + + Args: + model (`torch.nn.Module`): + The model to convert, can be any `torch.nn.Module` instance. + quantization_config (`VptqConfig`): + The quantization config object that contains the quantization parameters. + modules_to_not_convert (`List[`str`]`, *optional*, defaults to `["lm_head"]`): + Names of the modules to not convert in `VQuantLinear`. In practice we keep the `lm_head` in full precision + for numerical stability reasons. + current_key_name (`list`, *optional*): + A list that contains the current key name. This is used for recursion and should not be passed by the user. + has_been_replaced (`bool`, *optional*): + A boolean that indicates if the conversion has been successful or not. This is used for recursion and + should not be passed by the user. + """ + + modules_to_not_convert = ["lm_head"] if not modules_to_not_convert else modules_to_not_convert + + for name, module in model.named_children(): + if current_key_name is None: + current_key_name = [] + current_key_name.append(name) + layer_name = ".".join(current_key_name) + shared_layer_config = quantization_config.shared_layer_config + config_for_layers = quantization_config.config_for_layers + + if ( + isinstance(module, nn.Linear) + and layer_name not in modules_to_not_convert + and ((layer_name in config_for_layers) or (current_key_name[-1] in shared_layer_config)) + ): + layer_params = config_for_layers.get(layer_name, None) or shared_layer_config.get( + current_key_name[-1], None + ) + + with init_empty_weights(): + in_features = module.in_features + out_features = module.out_features + + model._modules[name] = VQuantLinear( + in_features, + out_features, + vector_lens=layer_params["vector_lens"], + num_centroids=layer_params["num_centroids"], + num_res_centroids=layer_params["num_res_centroids"], + group_num=layer_params["group_num"], + group_size=layer_params["group_size"], + outlier_size=layer_params["outlier_size"], + indices_as_float=layer_params["indices_as_float"], + enable_norm=layer_params["enable_norm"], + enable_perm=layer_params["enable_perm"], + is_indice_packed=True, + enable_proxy_error=False, + bias=module.bias is not None, + ) + has_been_replaced = True + + # Force requires grad to False to avoid unexpected errors + model._modules[name].requires_grad_(False) + if len(list(module.children())) > 0: + _, has_been_replaced = replace_with_vptq_linear( + module, + quantization_config=quantization_config, + modules_to_not_convert=modules_to_not_convert, + current_key_name=current_key_name, + has_been_replaced=has_been_replaced, + ) + # Remove the last key for recursion + current_key_name.pop(-1) + return model, has_been_replaced diff --git a/lib/python3.12/site-packages/transformers/keras_callbacks.py b/lib/python3.12/site-packages/transformers/keras_callbacks.py new file mode 100644 index 0000000000000000000000000000000000000000..57e72aea7e6cf22a0ff420267ccca1703a21cdf2 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/keras_callbacks.py @@ -0,0 +1,413 @@ +import logging +import os +from pathlib import Path +from time import sleep +from typing import Callable, Optional, Union + +import numpy as np +import tensorflow as tf +from huggingface_hub import Repository, create_repo +from packaging.version import parse + +from . import IntervalStrategy, PreTrainedTokenizerBase +from .modelcard import TrainingSummary +from .modeling_tf_utils import keras + + +logger = logging.getLogger(__name__) + + +class KerasMetricCallback(keras.callbacks.Callback): + """ + Callback to compute metrics at the end of every epoch. Unlike normal Keras metrics, these do not need to be + compilable by TF. It is particularly useful for common NLP metrics like BLEU and ROUGE that require string + operations or generation loops that cannot be compiled. Predictions (or generations) will be computed on the + `eval_dataset` before being passed to the `metric_fn` in `np.ndarray` format. The `metric_fn` should compute + metrics and return a dict mapping metric names to metric values. + + We provide an example of a suitable metric_fn that computes ROUGE scores for a summarization model below. Note that + this example skips some post-processing for readability and simplicity, and should probably not be used as-is! + + ```py + from datasets import load_metric + + rouge_metric = load_metric("rouge") + + + def rouge_fn(predictions, labels): + decoded_predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True) + decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True) + result = rouge_metric.compute(predictions=decoded_predictions, references=decoded_labels) + return {key: value.mid.fmeasure * 100 for key, value in result.items()} + ``` + + The above function will return a dict containing values which will be logged like any other Keras metric: + + ``` + {'rouge1': 37.4199, 'rouge2': 13.9768, 'rougeL': 34.361, 'rougeLsum': 35.0781 + ``` + + Args: + metric_fn (`Callable`): + Metric function provided by the user. It will be called with two arguments - `predictions` and `labels`. + These contain the model's outputs and matching labels from the dataset. It should return a dict mapping + metric names to numerical values. + eval_dataset (`tf.data.Dataset` or `dict` or `tuple` or `np.ndarray` or `tf.Tensor`): + Validation data to be used to generate predictions for the `metric_fn`. + output_cols (`List[str], *optional*): + A list of columns to be retained from the model output as the predictions. Defaults to all. + label_cols ('`List[str]`, *optional*'): + A list of columns to be retained from the input dataset as the labels. Will be autodetected if this is not + supplied. + batch_size (`int`, *optional*): + Batch size. Only used when the data is not a pre-batched `tf.data.Dataset`. + predict_with_generate (`bool`, *optional*, defaults to `False`): + Whether we should use `model.generate()` to get outputs for the model. + use_xla_generation (`bool`, *optional*, defaults to `False`): + If we're generating, whether to compile model generation with XLA. This can massively increase the speed of + generation (up to 100X speedup) but will require a new XLA compilation for each input shape. When using XLA + generation, it's a good idea to pad your inputs to the same size, or to use the `pad_to_multiple_of` + argument in your `tokenizer` or `DataCollator`, which will reduce the number of unique input shapes and + save a lot of compilation time. This option has no effect is `predict_with_generate` is `False`. + generate_kwargs (`dict`, *optional*): + Keyword arguments to pass to `model.generate()` when generating. Has no effect if `predict_with_generate` + is `False`. + + """ + + def __init__( + self, + metric_fn: Callable, + eval_dataset: Union[tf.data.Dataset, np.ndarray, tf.Tensor, tuple, dict], + output_cols: Optional[list[str]] = None, + label_cols: Optional[list[str]] = None, + batch_size: Optional[int] = None, + predict_with_generate: bool = False, + use_xla_generation: bool = False, + generate_kwargs: Optional[dict] = None, + ): + super().__init__() + self.metric_fn = metric_fn + self.batch_size = batch_size + if not isinstance(eval_dataset, tf.data.Dataset): + if batch_size is None: + raise ValueError( + "When passing data to KerasMetricCallback that is not a pre-batched tf.data.Dataset " + "the batch_size argument must be set." + ) + # Wrap a tf.data.Dataset around it + eval_dataset = tf.data.Dataset.from_tensor_slices(eval_dataset).batch(batch_size, drop_remainder=False) + self.eval_dataset = eval_dataset + self.predict_with_generate = predict_with_generate + self.output_cols = output_cols + + # This next block attempts to parse out which elements of the dataset should be appended to the labels list + # that is passed to the metric_fn + if isinstance(eval_dataset.element_spec, tuple) and len(eval_dataset.element_spec) == 2: + input_spec, label_spec = eval_dataset.element_spec + else: + input_spec = eval_dataset.element_spec + label_spec = None + if label_cols is not None: + for label in label_cols: + if label not in input_spec: + raise ValueError(f"Label {label} is in label_cols but could not be found in the dataset inputs!") + self.label_cols = label_cols + self.use_keras_label = False + elif label_spec is not None: + # If the dataset inputs are split into a 2-tuple of inputs and labels, + # assume the second element is the labels + self.label_cols = None + self.use_keras_label = True + elif "labels" in input_spec: + self.label_cols = ["labels"] + self.use_keras_label = False + logging.warning("No label_cols specified for KerasMetricCallback, assuming you want the 'labels' key.") + elif "start_positions" in input_spec and "end_positions" in input_spec: + self.label_cols = ["start_positions", "end_positions"] + self.use_keras_label = False + logging.warning( + "No label_cols specified for KerasMetricCallback, assuming you want the " + "start_positions and end_positions keys." + ) + else: + raise ValueError("Could not autodetect label_cols for KerasMetricCallback, please specify them!") + if parse(tf.__version__) < parse("2.7"): + logging.warning("TF versions less than 2.7 may encounter issues with KerasMetricCallback!") + + self.use_xla_generation = use_xla_generation + self.generate_kwargs = {} if generate_kwargs is None else generate_kwargs + + self.generation_function = None + + @staticmethod + def _concatenate_batches(batches, padding_index=-100): + # If all batches are unidimensional or same length, do a simple concatenation + if batches[0].ndim == 1 or all(batch.shape[1] == batches[0].shape[1] for batch in batches): + return np.concatenate(batches, axis=0) + + # Welp, they're not the same length. Let's do some padding + max_len = max([batch.shape[1] for batch in batches]) + num_samples = sum([batch.shape[0] for batch in batches]) + output = np.full_like( + batches[0], fill_value=padding_index, shape=[num_samples, max_len] + list(batches[0].shape[2:]) + ) + # i keeps track of which part of the concatenated array we're writing the next batch to + i = 0 + for batch in batches: + output[i : i + len(batch), : batch.shape[1]] = batch + i += len(batch) + return output + + def _postprocess_predictions_or_labels(self, inputs): + if isinstance(inputs[0], dict): + outputs = {} + for key in inputs[0].keys(): + outputs[key] = self._concatenate_batches([batch[key] for batch in inputs]) + # If it's a dict with only one key, just return the array + if len(outputs) == 1: + outputs = list(outputs.values())[0] + elif isinstance(inputs[0], list) or isinstance(inputs[0], tuple): + outputs = [] + for input_list in zip(*inputs): + outputs.append(self._concatenate_batches(input_list)) + if len(outputs) == 1: + outputs = outputs[0] # If it's a list with only one element, just return the array + elif isinstance(inputs[0], np.ndarray): + outputs = self._concatenate_batches(inputs) + elif isinstance(inputs[0], tf.Tensor): + outputs = self._concatenate_batches([tensor.numpy() for tensor in inputs]) + else: + raise TypeError(f"Couldn't handle batch of type {type(inputs[0])}!") + return outputs + + def on_epoch_end(self, epoch, logs=None): + if hasattr(self.model, "config"): + ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + main_input_name = None + if self.predict_with_generate: + # This dense conditional recognizes the case where we have an encoder-decoder model, but + # avoids getting tangled up when we just have a model with a layer called 'encoder' + if hasattr(self.model, "encoder") and hasattr(self.model.encoder, "main_input_name"): + main_input_name = self.model.encoder.main_input_name + else: + main_input_name = getattr(self.model, "main_input_name", "input_ids") + + if self.use_xla_generation and self.generation_function is None: + + def generation_function(inputs, attention_mask): + return self.model.generate(inputs, attention_mask=attention_mask, **self.generate_kwargs) + + self.generation_function = tf.function(generation_function, jit_compile=True) + + prediction_list = [] + label_list = [] + + # The whole predict/generate loop is handled inside this method + for batch in self.eval_dataset: + if isinstance(batch, tuple): + batch, labels = batch + else: + labels = None + if self.predict_with_generate: + if isinstance(batch, dict): + generation_inputs = batch[main_input_name] + attention_mask = batch.get("attention_mask", None) + else: + generation_inputs = batch + attention_mask = None + if self.use_xla_generation: + predictions = self.generation_function(generation_inputs, attention_mask=attention_mask) + else: + predictions = self.model.generate( + generation_inputs, attention_mask=attention_mask, **self.generate_kwargs + ) + else: + predictions = self.model.predict_on_batch(batch) + if isinstance(predictions, dict): + # This converts any dict-subclass to a regular dict + # Keras REALLY doesn't like it when we pass around a BatchEncoding or other derived class + predictions = dict(predictions) + if self.output_cols is not None: + predictions = {key: predictions[key] for key in self.output_cols} + else: + predictions = { + key: val for key, val in predictions.items() if key not in ignore_keys + ["loss"] + } + prediction_list.append(predictions) + if not self.use_keras_label: + labels = {key: batch[key].numpy() for key in self.label_cols} + elif isinstance(labels, dict): + labels = {key: array.numpy() for key, array in labels.items()} + elif isinstance(labels, list) or isinstance(labels, tuple): + labels = [array.numpy() for array in labels] + elif isinstance(labels, tf.Tensor): + labels = labels.numpy() + else: + raise TypeError(f"Confused by labels of type {type(labels)}") + label_list.append(labels) + + all_preds = self._postprocess_predictions_or_labels(prediction_list) + all_labels = self._postprocess_predictions_or_labels(label_list) + + metric_output = self.metric_fn((all_preds, all_labels)) + if not isinstance(metric_output, dict): + raise TypeError( + f"metric_fn should return a dict mapping metric names to values but instead returned {metric_output}" + ) + # This is the critical bit - Keras passes a dict containing the loss and standard metric values for this epoch + # in the logs argument. Ordinarily, this is so the callback can read them, but in this case we write a bunch of + # new keys in there, which will then get read by the History callback and treated like any other metric value. + # I promise that I have it in writing from Chollet that this is okay. + logs.update(metric_output) + + +class PushToHubCallback(keras.callbacks.Callback): + """ + Callback that will save and push the model to the Hub regularly. By default, it pushes once per epoch, but this can + be changed with the `save_strategy` argument. Pushed models can be accessed like any other model on the hub, such + as with the `from_pretrained` method. + + ```py + from transformers.keras_callbacks import PushToHubCallback + + push_to_hub_callback = PushToHubCallback( + output_dir="./model_save", + tokenizer=tokenizer, + hub_model_id="gpt5-7xlarge", + ) + + model.fit(train_dataset, callbacks=[push_to_hub_callback]) + ``` + + Args: + output_dir (`str`): + The output directory where the model predictions and checkpoints will be written and synced with the + repository on the Hub. + save_strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"epoch"`): + The checkpoint save strategy to adopt during training. Possible values are: + + - `"no"`: Save is done at the end of training. + - `"epoch"`: Save is done at the end of each epoch. + - `"steps"`: Save is done every `save_steps` + save_steps (`int`, *optional*): + The number of steps between saves when using the "steps" `save_strategy`. + tokenizer (`PreTrainedTokenizerBase`, *optional*): + The tokenizer used by the model. If supplied, will be uploaded to the repo alongside the weights. + hub_model_id (`str`, *optional*): + The name of the repository to keep in sync with the local `output_dir`. It can be a simple model ID in + which case the model will be pushed in your namespace. Otherwise it should be the whole repository name, + for instance `"user_name/model"`, which allows you to push to an organization you are a member of with + `"organization_name/model"`. + + Will default to the name of `output_dir`. + hub_token (`str`, *optional*): + The token to use to push the model to the Hub. Will default to the token in the cache folder obtained with + `huggingface-cli login`. + checkpoint (`bool`, *optional*, defaults to `False`): + Whether to save full training checkpoints (including epoch and optimizer state) to allow training to be + resumed. Only usable when `save_strategy` is `"epoch"`. + """ + + def __init__( + self, + output_dir: Union[str, Path], + save_strategy: Union[str, IntervalStrategy] = "epoch", + save_steps: Optional[int] = None, + tokenizer: Optional[PreTrainedTokenizerBase] = None, + hub_model_id: Optional[str] = None, + hub_token: Optional[str] = None, + checkpoint: bool = False, + **model_card_args, + ): + super().__init__() + if checkpoint and save_strategy != "epoch": + raise ValueError("Cannot save checkpoints when save_strategy is not 'epoch'!") + if isinstance(save_strategy, str): + save_strategy = IntervalStrategy(save_strategy.lower()) + self.save_strategy = save_strategy + if self.save_strategy == IntervalStrategy.STEPS and (not isinstance(save_steps, int) or save_steps <= 0): + raise ValueError("Please supply a positive integer argument for save_steps when save_strategy == 'steps'!") + self.save_steps = save_steps + output_dir = Path(output_dir) + + # Create repo and retrieve repo_id + if hub_model_id is None: + hub_model_id = output_dir.absolute().name + self.hub_model_id = create_repo(repo_id=hub_model_id, exist_ok=True, token=hub_token).repo_id + + self.output_dir = output_dir + self.repo = Repository(str(self.output_dir), clone_from=self.hub_model_id, token=hub_token) + + self.tokenizer = tokenizer + self.last_job = None + self.checkpoint = checkpoint + self.training_history = None + self.model_card_args = model_card_args + + def on_train_begin(self, logs=None): + # Although we can access model.history, we have no guarantees that the History callback will fire before this + # one, so we keep track of it here too + self.training_history = [] + + def on_train_batch_end(self, batch, logs=None): + if self.save_strategy == IntervalStrategy.STEPS and (batch + 1) % self.save_steps == 0: + if self.last_job is not None and not self.last_job.is_done: + return # The last upload is still running, don't start another + self.model.save_pretrained(self.output_dir) + if self.tokenizer is not None: + self.tokenizer.save_pretrained(self.output_dir) + _, self.last_job = self.repo.push_to_hub( + commit_message=f"Training in progress steps {batch}", blocking=False + ) + + def on_epoch_end(self, epoch, logs=None): + logs = logs.copy() # Don't accidentally write things that Keras will read later + if "epoch" not in logs: + logs["epoch"] = epoch + self.training_history.append(logs) + if self.save_strategy == IntervalStrategy.EPOCH: + if self.last_job is not None and not self.last_job.is_done: + return # The last upload is still running, don't start another + self.model.save_pretrained(self.output_dir) + if self.tokenizer is not None: + self.tokenizer.save_pretrained(self.output_dir) + if self.checkpoint: + checkpoint_dir = os.path.join(self.output_dir, "checkpoint") + self.model._save_checkpoint(checkpoint_dir, epoch) + train_summary = TrainingSummary.from_keras( + model=self.model, + model_name=self.hub_model_id, + keras_history=self.training_history, + **self.model_card_args, + ) + model_card = train_summary.to_model_card() + with (self.output_dir / "README.md").open("w") as f: + f.write(model_card) + _, self.last_job = self.repo.push_to_hub( + commit_message=f"Training in progress epoch {epoch}", blocking=False + ) + + def on_train_end(self, logs=None): + # Makes sure the latest version of the model is uploaded + if self.last_job is not None and not self.last_job.is_done: + logging.info("Pushing the last epoch to the Hub, this may take a while...") + while not self.last_job.is_done: + sleep(1) + else: + self.model.save_pretrained(self.output_dir) + if self.tokenizer is not None: + self.tokenizer.save_pretrained(self.output_dir) + train_summary = TrainingSummary.from_keras( + model=self.model, + model_name=self.hub_model_id, + keras_history=self.training_history, + **self.model_card_args, + ) + model_card = train_summary.to_model_card() + with (self.output_dir / "README.md").open("w") as f: + f.write(model_card) + self.repo.push_to_hub(commit_message="End of training", blocking=True) diff --git a/lib/python3.12/site-packages/transformers/model_debugging_utils.py b/lib/python3.12/site-packages/transformers/model_debugging_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0992636a141ccea2a2ad5cad26eb4fd4ecfb1f0b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/model_debugging_utils.py @@ -0,0 +1,329 @@ +# Copyright 2025 The HuggingFace Inc. 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. + +import functools +import json +import os +import re +from contextlib import contextmanager +from typing import Optional + +from transformers.utils.import_utils import export + +from .utils import is_torch_available + + +if is_torch_available(): + import torch + import torch.distributed.tensor + from torch import nn + + from .modeling_utils import PreTrainedModel + +from .utils import logging + + +logger = logging.get_logger(__name__) + +# Note to code inspectors: this toolbox is intended for people who add models to `transformers`. +_torch_distributed_available = torch.distributed.is_available() + + +def _is_rank_zero(): + """Return True if rank=0 or we aren't running distributed.""" + if not (_torch_distributed_available and torch.distributed.is_initialized()): + return True + return torch.distributed.get_rank() == 0 + + +MEMORY_ADDRESS_REGEX = re.compile(r"object at 0x[0-9A-Fa-f]+") + + +def _sanitize_repr_for_diff(x_str: str) -> str: + """ + Replace memory addresses in an object's repr with a stable placeholder + so that beautiful JSON diffs won't be ruined by ephemeral addresses. + """ + return MEMORY_ADDRESS_REGEX.sub("object at 0xXXXXXXXX", x_str) + + +def _dtensor_repr(x): + """Return a stable string representation for a DTensor-like object.""" + if _is_rank_zero(): + return f"DTensor (rank0) -> {repr(x._local_tensor)}" + return "DTensor(non-rank0)" + + +def _serialize_io(value): + """ + Recursively build a JSON-serializable Python structure from `value`. + Tensors and DTensors become sanitized repr strings. + Lists/tuples/dicts are recursed into. + All memory addresses are replaced with a stable placeholder. + + Args: + value: Any Python object, often including torch Tensors, lists, dicts, etc. + + Returns: + A nested Python structure (list, dict, or sanitized string) that is safe to json.dump. + """ + if isinstance(value, (list, tuple)): + return [_serialize_io(v) for v in value] + + if isinstance(value, dict): + return {k: _serialize_io(v) for k, v in value.items()} + + if hasattr(value, "_local_tensor"): + # DTensor-like handling, just use local tensor attribute + return { + "shape": repr(value._local_tensor.shape), + "dtype": repr(value._local_tensor.dtype), + "value": _sanitize_repr_for_diff(repr(value)), + } + + if isinstance(value, torch.Tensor): + # standard PyTorch Tensor + # return also the shape of such + return {"shape": repr(value.shape), "dtype": repr(value.dtype), "value": _sanitize_repr_for_diff(repr(value))} + + # fallback for everything else (bool, int, float, None, or custom class) + return _sanitize_repr_for_diff(repr(value)) + + +def prune_outputs_if_children(node): + # if there are children, remove this node's "outputs" + # so we only see outputs at the leaf level + if node.get("children"): + node.pop("outputs", None) + for child in node["children"]: + prune_outputs_if_children(child) + + +def log_model_debug_trace(debug_path, model): + if debug_path: + try: + os.makedirs(debug_path, exist_ok=False) + output_path = os.path.join(debug_path, model._debugger_module_dump_name + "_debug_tree.json") + except Exception as e: + raise ValueError(f"Unexpected or existing debug_path={debug_path}. {e}") + else: + output_path = model._debugger_module_dump_name + "_debug_tree.json" + logger.info(f"Writing model trace at {output_path}") + with open(output_path, "w") as outfile: + prune_outputs_if_children(model._call_tree) + json.dump(model._call_tree, outfile, indent=2) + + +def _attach_debugger_logic(model, class_name, debug_path: str): + # Prepare data structures on the model object + model._call_tree = {"module_path": class_name, "inputs": None, "outputs": None, "children": []} + model._debugger_model_call_stack = [] + model._debugger_module_dump_name = class_name # used for final JSON filename + + def wrap_forward(module, full_path): + orig_forward = module.forward + + @functools.wraps(orig_forward) + def wrapped_forward(*inps, **kws): + if _is_rank_zero(): + dict_inputs = {"args": inps, "kwargs": kws} + dict_inputs = {k: dict_inputs[k] for k in dict_inputs if len(dict_inputs[k]) > 0} + node = { + "module_path": full_path, + "inputs": _serialize_io(dict_inputs), + "outputs": None, + "children": [], + } + model._debugger_model_call_stack.append(node) + with torch.inference_mode(): + out = orig_forward(*inps, **kws) + + if _is_rank_zero(): + if sum(1 for _ in module.named_children()) > 0: + node["outputs"] = None + else: + node["outputs"] = _serialize_io(out) + + finished = model._debugger_model_call_stack.pop() + # prune empty vertices here as well (mostly empty children nodes) + if not finished["children"]: + finished.pop("children") + + if model._debugger_model_call_stack: + model._debugger_model_call_stack[-1]["children"].append(finished) + return out + + module.forward = wrapped_forward + + # wrap all submodules + for name, submodule in model.named_modules(): + if name == "": + continue + wrap_forward(submodule, f"{class_name}.{name}") + + # wrap top-level forward + real_top_forward = model.forward + + @functools.wraps(real_top_forward) + def top_wrapped_forward(*inps, **kws): + if _is_rank_zero(): + top_node = { + "module_path": f"{class_name} (top-level)", + "inputs": _serialize_io({"args": inps, "kwargs": kws}), + "outputs": None, + "children": [], + } + model._debugger_model_call_stack.append(top_node) + + out = real_top_forward(*inps, **kws) + + if _is_rank_zero() and model._debugger_model_call_stack: + top_node["outputs"] = _serialize_io(out) + finished = model._debugger_model_call_stack.pop() + model._call_tree["inputs"] = finished["inputs"] + model._call_tree["outputs"] = finished["outputs"] + model._call_tree["children"] = finished["children"] + # prune empty stuff for visibility + [model._call_tree.pop(k, None) for k in list(model._call_tree.keys()) if not model._call_tree[k]] + + return out + + model.forward = top_wrapped_forward + + # Final hook for writing JSON on forward-end + def final_hook(_, inputs, outputs): + if _is_rank_zero() and model._debugger_model_call_stack: + finished = model._debugger_model_call_stack.pop() + model._call_tree["inputs"] = finished["inputs"] + model._call_tree["outputs"] = finished["outputs"] + model._call_tree["children"] = finished["children"] + + if _is_rank_zero(): + log_model_debug_trace(debug_path=debug_path, model=model) + + model.register_forward_hook(final_hook) + # Optionally also for a couple possible hooks that have specific names. It should be just one. + # This means modules that are not typically called "forward" within the model. But we should not need to recurse + # through them. + possible_model_calls = ["language_model", "model"] + for model_call in possible_model_calls: + this_model_call = getattr(model, model_call, None) + if this_model_call and isinstance(this_model_call, (nn.Module, PreTrainedModel)): + this_model_call.register_forward_hook(final_hook) + break # exit the loop after finding one (unsure, but should be just one call.) + + +@export(backends=("torch",)) +def model_addition_debugger(cls): + """ + # Model addition debugger - a model adder tracer + This decorator is a power user tool intended for model adders. + It tracks all forward calls within a model forward and logs a slice of each input and output on a nested Json. + To note, this decorator enforces `torch.inference_mode()`. + ## Usage + + add decorator to your model class + ```python + from ...modeling_utils import model_addition_debugger + + @model_addition_debugger + class MyModel(nn.Module) # Can inherit from PreTrainedModel too + # ... nothing else changes + ``` + Then, in a separate script (example is for Llava) + + ```python + import torch + from PIL import Image + import requests + from transformers import LlavaProcessor, LlavaForConditionalGeneration + torch.random.manual_seed(673) + + # load pretrained model and processor + model_id = "llava-hf/llava-1.5-7b-hf" + processor = LlavaProcessor.from_pretrained(model_id) + model = LlavaForConditionalGeneration.from_pretrained(model_id, low_cpu_mem_usage=True) + + # create random image input + random_image = Image.fromarray(torch.randint(0, 256, (224, 224, 3), dtype=torch.uint8).numpy()) + + # prompt + prompt = "Describe this image." + + # process inputs + inputs = processor(text=prompt, images=random_image, return_tensors="pt") + + # call forward method (not .generate!) + with torch.no_grad(): + output = model.forward(**inputs) + ``` + + """ + orig_init = cls.__init__ + + @functools.wraps(cls.__init__) + def wrapped_init(self, *args, **kwargs): + orig_init(self, *args, **kwargs) + _attach_debugger_logic(self, cls.__name__) + + cls.__init__ = wrapped_init + return cls + + +@export(backends=("torch",)) +@contextmanager +def model_addition_debugger_context(model, debug_path: Optional[str] = None): + """ + # Model addition debugger - context manager for model adders + This context manager is a power user tool intended for model adders. + It tracks all forward calls within a model forward and logs a slice of each input and output on a nested Json. + To note, this context manager enforces `torch.inference_mode()`. + + ## Usage + + add the context manager to a model to debug + + ```python + import torch + from PIL import Image + import requests + from transformers import LlavaProcessor, LlavaForConditionalGeneration + torch.random.manual_seed(673) + + # load pretrained model and processor + model_id = "llava-hf/llava-1.5-7b-hf" + processor = LlavaProcessor.from_pretrained(model_id) + model = LlavaForConditionalGeneration.from_pretrained(model_id, low_cpu_mem_usage=True) + + # create random image input + random_image = Image.fromarray(torch.randint(0, 256, (224, 224, 3), dtype=torch.uint8).numpy()) + + # prompt + prompt = "Describe this image." + + # process inputs + inputs = processor(text=prompt, images=random_image, return_tensors="pt") + + # call forward method (not .generate!) + with model_addition_debugger_context(model): + output = model.forward(**inputs) + ``` + + """ + _attach_debugger_logic(model, model.__class__.__name__, debug_path) + try: + yield model + finally: + pass diff --git a/lib/python3.12/site-packages/transformers/modelcard.py b/lib/python3.12/site-packages/transformers/modelcard.py new file mode 100644 index 0000000000000000000000000000000000000000..da5c4077785a5d681e3d658ecd25f72c5e097c91 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modelcard.py @@ -0,0 +1,909 @@ +# Copyright 2018 The HuggingFace Inc. team. +# +# 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. +"""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, Optional, Union + +import requests +import yaml +from huggingface_hub import model_info +from huggingface_hub.utils import HFValidationError + +from . import __version__ +from .models.auto.modeling_auto import ( + MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES, + MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, + MODEL_FOR_CTC_MAPPING_NAMES, + MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES, + MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, + MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES, + MODEL_FOR_MASKED_LM_MAPPING_NAMES, + MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES, + MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES, + MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES, + MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES, + MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES, + MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES, + MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES, + MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES, +) +from .training_args import ParallelMode +from .utils import ( + MODEL_CARD_NAME, + cached_file, + is_datasets_available, + is_offline_mode, + is_tf_available, + is_tokenizers_available, + is_torch_available, + logging, +) + + +TASK_MAPPING = { + "text-generation": MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, + "image-classification": MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES, + "image-segmentation": MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES, + "fill-mask": MODEL_FOR_MASKED_LM_MAPPING_NAMES, + "object-detection": MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES, + "question-answering": MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES, + "text2text-generation": MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES, + "text-classification": MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES, + "table-question-answering": MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES, + "token-classification": MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES, + "audio-classification": MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES, + "automatic-speech-recognition": {**MODEL_FOR_CTC_MAPPING_NAMES, **MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES}, + "zero-shot-image-classification": MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES, + "image-text-to-text": MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES, +} + +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://arxiv.org/abs/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://arxiv.org/abs/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 = """ + +""" + +AUTOGENERATED_KERAS_COMMENT = """ + +""" + + +TASK_TAG_TO_NAME_MAPPING = { + "fill-mask": "Masked Language Modeling", + "image-classification": "Image Classification", + "image-segmentation": "Image Segmentation", + "multiple-choice": "Multiple Choice", + "object-detection": "Object Detection", + "question-answering": "Question Answering", + "summarization": "Summarization", + "table-question-answering": "Table Question Answering", + "text-classification": "Text Classification", + "text-generation": "Causal Language Modeling", + "text2text-generation": "Sequence-to-sequence Language Modeling", + "token-classification": "Token Classification", + "translation": "Translation", + "zero-shot-classification": "Zero Shot Classification", + "automatic-speech-recognition": "Automatic Speech Recognition", + "audio-classification": "Audio Classification", +} + + +METRIC_TAGS = [ + "accuracy", + "bleu", + "f1", + "matthews_correlation", + "pearsonr", + "precision", + "recall", + "rouge", + "sacrebleu", + "spearmanr", + "wer", +] + + +def _listify(obj): + if obj is None: + return [] + elif isinstance(obj, str): + return [obj] + else: + return obj + + +def _insert_values_as_list(metadata, name, values): + if values is None: + return metadata + if isinstance(values, str): + values = [values] + values = [v for v in values if v is not None] + if len(values) == 0: + return metadata + metadata[name] = values + return metadata + + +def infer_metric_tags_from_eval_results(eval_results): + if eval_results is None: + return {} + result = {} + for key in eval_results.keys(): + if key.lower().replace(" ", "_") in METRIC_TAGS: + result[key.lower().replace(" ", "_")] = key + elif key.lower() == "rouge1": + result["rouge"] = key + return result + + +def _insert_value(metadata, name, value): + if value is None: + return metadata + metadata[name] = value + return metadata + + +def is_hf_dataset(dataset): + if not is_datasets_available(): + return False + + from datasets import Dataset, IterableDataset + + return isinstance(dataset, (Dataset, IterableDataset)) + + +def _get_mapping_values(mapping): + result = [] + for v in mapping.values(): + if isinstance(v, (tuple, list)): + result += list(v) + else: + result.append(v) + return result + + +@dataclass +class TrainingSummary: + model_name: str + 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. + if ( + self.license is None + and not is_offline_mode() + and self.finetuned_from is not None + and len(self.finetuned_from) > 0 + ): + try: + info = model_info(self.finetuned_from) + for tag in info.tags: + if tag.startswith("license:"): + self.license = tag[8:] + except (requests.exceptions.HTTPError, requests.exceptions.ConnectionError, HFValidationError): + pass + + def create_model_index(self, metric_mapping): + model_index = {"name": self.model_name} + + # Dataset mapping tag -> name + dataset_names = _listify(self.dataset) + dataset_tags = _listify(self.dataset_tags) + dataset_args = _listify(self.dataset_args) + dataset_metadata = _listify(self.dataset_metadata) + if len(dataset_args) < len(dataset_tags): + dataset_args = dataset_args + [None] * (len(dataset_tags) - len(dataset_args)) + dataset_mapping = dict(zip(dataset_tags, dataset_names)) + dataset_arg_mapping = dict(zip(dataset_tags, dataset_args)) + dataset_metadata_mapping = dict(zip(dataset_tags, dataset_metadata)) + + task_mapping = { + task: TASK_TAG_TO_NAME_MAPPING[task] for task in _listify(self.tasks) if task in TASK_TAG_TO_NAME_MAPPING + } + + model_index["results"] = [] + + if len(task_mapping) == 0 and len(dataset_mapping) == 0: + return [model_index] + if len(task_mapping) == 0: + task_mapping = {None: None} + if len(dataset_mapping) == 0: + dataset_mapping = {None: None} + + # One entry per dataset and per task + all_possibilities = [(task_tag, ds_tag) for task_tag in task_mapping for ds_tag in dataset_mapping] + for task_tag, ds_tag in all_possibilities: + result = {} + if task_tag is not None: + result["task"] = {"name": task_mapping[task_tag], "type": task_tag} + + if ds_tag is not None: + metadata = dataset_metadata_mapping.get(ds_tag, {}) + result["dataset"] = { + "name": dataset_mapping[ds_tag], + "type": ds_tag, + **metadata, + } + if dataset_arg_mapping[ds_tag] is not None: + result["dataset"]["args"] = dataset_arg_mapping[ds_tag] + + if len(metric_mapping) > 0: + result["metrics"] = [] + for metric_tag, metric_name in metric_mapping.items(): + result["metrics"].append( + { + "name": metric_name, + "type": metric_tag, + "value": self.eval_results[metric_name], + } + ) + + # Remove partial results to avoid the model card being rejected. + if "task" in result and "dataset" in result and "metrics" in result: + model_index["results"].append(result) + else: + logger.info(f"Dropping the following result as it does not have all the necessary fields:\n{result}") + + return [model_index] + + def create_metadata(self): + metric_mapping = infer_metric_tags_from_eval_results(self.eval_results) + + metadata = {} + metadata = _insert_value(metadata, "library_name", "transformers") + metadata = _insert_values_as_list(metadata, "language", self.language) + metadata = _insert_value(metadata, "license", self.license) + if self.finetuned_from is not None and isinstance(self.finetuned_from, str) and len(self.finetuned_from) > 0: + metadata = _insert_value(metadata, "base_model", self.finetuned_from) + metadata = _insert_values_as_list(metadata, "tags", self.tags) + metadata = _insert_values_as_list(metadata, "datasets", self.dataset_tags) + metadata = _insert_values_as_list(metadata, "metrics", list(metric_mapping.keys())) + metadata["model-index"] = self.create_model_index(metric_mapping) + + return metadata + + def to_model_card(self): + model_card = "" + + metadata = yaml.dump(self.create_metadata(), sort_keys=False) + if len(metadata) > 0: + model_card = f"---\n{metadata}---\n" + + # Now the model card for realsies. + if self.source == "trainer": + model_card += AUTOGENERATED_TRAINER_COMMENT + else: + model_card += AUTOGENERATED_KERAS_COMMENT + + model_card += f"\n# {self.model_name}\n\n" + + if self.finetuned_from is None: + model_card += "This model was trained from scratch on " + else: + model_card += ( + "This model is a fine-tuned version of" + f" [{self.finetuned_from}](https://huggingface.co/{self.finetuned_from}) on " + ) + + if self.dataset is None or (isinstance(self.dataset, list) and len(self.dataset) == 0): + model_card += "an unknown dataset." + else: + if isinstance(self.dataset, str): + model_card += f"the {self.dataset} dataset." + elif isinstance(self.dataset, (tuple, list)) and len(self.dataset) == 1: + model_card += f"the {self.dataset[0]} dataset." + else: + model_card += ( + ", ".join([f"the {ds}" for ds in self.dataset[:-1]]) + f" and the {self.dataset[-1]} datasets." + ) + + if self.eval_results is not None: + model_card += "\nIt achieves the following results on the evaluation set:\n" + model_card += "\n".join([f"- {name}: {_maybe_round(value)}" for name, value in self.eval_results.items()]) + model_card += "\n" + + model_card += "\n## Model description\n\nMore information needed\n" + model_card += "\n## Intended uses & limitations\n\nMore information needed\n" + model_card += "\n## Training and evaluation data\n\nMore information needed\n" + + model_card += "\n## Training procedure\n" + model_card += "\n### Training hyperparameters\n" + if self.hyperparameters is not None: + model_card += "\nThe following hyperparameters were used during training:\n" + model_card += "\n".join([f"- {name}: {value}" for name, value in self.hyperparameters.items()]) + model_card += "\n" + else: + model_card += "\nMore information needed\n" + + if self.eval_lines is not None: + model_card += "\n### Training results\n\n" + model_card += make_markdown_table(self.eval_lines) + model_card += "\n" + + model_card += "\n### Framework versions\n\n" + model_card += f"- Transformers {__version__}\n" + + if self.source == "trainer" and is_torch_available(): + import torch + + model_card += f"- Pytorch {torch.__version__}\n" + elif self.source == "keras" and is_tf_available(): + import tensorflow as tf + + model_card += f"- TensorFlow {tf.__version__}\n" + if is_datasets_available(): + import datasets + + model_card += f"- Datasets {datasets.__version__}\n" + if is_tokenizers_available(): + import tokenizers + + model_card += f"- Tokenizers {tokenizers.__version__}\n" + + return model_card + + @classmethod + def from_trainer( + cls, + trainer, + language=None, + license=None, + tags=None, + model_name=None, + finetuned_from=None, + tasks=None, + dataset_tags=None, + dataset_metadata=None, + dataset=None, + dataset_args=None, + ): + # Infer default from dataset + one_dataset = trainer.eval_dataset if trainer.eval_dataset is not None else trainer.train_dataset + if is_hf_dataset(one_dataset) and (dataset_tags is None or dataset_args is None or dataset_metadata is None): + default_tag = one_dataset.builder_name + # Those are not real datasets from the Hub so we exclude them. + if default_tag not in ["csv", "json", "pandas", "parquet", "text"]: + if dataset_metadata is None: + dataset_metadata = [{"config": one_dataset.config_name, "split": str(one_dataset.split)}] + if dataset_tags is None: + dataset_tags = [default_tag] + if dataset_args is None: + dataset_args = [one_dataset.config_name] + + if dataset is None and dataset_tags is not None: + dataset = dataset_tags + + # Infer default finetuned_from + if ( + finetuned_from is None + and hasattr(trainer.model.config, "_name_or_path") + and not os.path.isdir(trainer.model.config._name_or_path) + ): + finetuned_from = trainer.model.config._name_or_path + + # Infer default task tag: + if tasks is None: + model_class_name = trainer.model.__class__.__name__ + for task, mapping in TASK_MAPPING.items(): + if model_class_name in _get_mapping_values(mapping): + tasks = task + + if model_name is None: + model_name = Path(trainer.args.output_dir).name + if len(model_name) == 0: + model_name = finetuned_from + + # Add `generated_from_trainer` to the tags + if tags is None: + tags = ["generated_from_trainer"] + elif isinstance(tags, str) and tags != "generated_from_trainer": + tags = [tags, "generated_from_trainer"] + elif "generated_from_trainer" not in tags: + tags.append("generated_from_trainer") + + _, eval_lines, eval_results = parse_log_history(trainer.state.log_history) + hyperparameters = extract_hyperparameters_from_trainer(trainer) + + return cls( + language=language, + license=license, + tags=tags, + model_name=model_name, + finetuned_from=finetuned_from, + tasks=tasks, + dataset=dataset, + dataset_tags=dataset_tags, + dataset_args=dataset_args, + dataset_metadata=dataset_metadata, + eval_results=eval_results, + eval_lines=eval_lines, + hyperparameters=hyperparameters, + ) + + @classmethod + def from_keras( + cls, + model, + model_name, + keras_history=None, + language=None, + license=None, + tags=None, + finetuned_from=None, + tasks=None, + dataset_tags=None, + dataset=None, + dataset_args=None, + ): + # Infer default from dataset + if dataset is not None: + if is_hf_dataset(dataset) and (dataset_tags is None or dataset_args is None): + default_tag = dataset.builder_name + # Those are not real datasets from the Hub so we exclude them. + if default_tag not in ["csv", "json", "pandas", "parquet", "text"]: + if dataset_tags is None: + dataset_tags = [default_tag] + if dataset_args is None: + dataset_args = [dataset.config_name] + + if dataset is None and dataset_tags is not None: + dataset = dataset_tags + + # Infer default finetuned_from + if ( + finetuned_from is None + and hasattr(model.config, "_name_or_path") + and not os.path.isdir(model.config._name_or_path) + ): + finetuned_from = model.config._name_or_path + + # Infer default task tag: + if tasks is None: + model_class_name = model.__class__.__name__ + for task, mapping in TASK_MAPPING.items(): + if model_class_name in _get_mapping_values(mapping): + tasks = task + + # Add `generated_from_keras_callback` to the tags + if tags is None: + tags = ["generated_from_keras_callback"] + elif isinstance(tags, str) and tags != "generated_from_keras_callback": + tags = [tags, "generated_from_keras_callback"] + elif "generated_from_keras_callback" not in tags: + tags.append("generated_from_keras_callback") + + if keras_history is not None: + _, eval_lines, eval_results = parse_keras_history(keras_history) + else: + eval_lines = [] + eval_results = {} + hyperparameters = extract_hyperparameters_from_keras(model) + + return cls( + language=language, + license=license, + tags=tags, + model_name=model_name, + finetuned_from=finetuned_from, + tasks=tasks, + dataset_tags=dataset_tags, + dataset=dataset, + dataset_args=dataset_args, + eval_results=eval_results, + eval_lines=eval_lines, + hyperparameters=hyperparameters, + source="keras", + ) + + +def parse_keras_history(logs): + """ + Parse the `logs` of either a `keras.History` object returned by `model.fit()` or an accumulated logs `dict` + passed to the `PushToHubCallback`. Returns lines and logs compatible with those returned by `parse_log_history`. + """ + if hasattr(logs, "history"): + # This looks like a `History` object + if not hasattr(logs, "epoch"): + # This history looks empty, return empty results + return None, [], {} + logs.history["epoch"] = logs.epoch + logs = logs.history + else: + # Training logs is a list of dicts, let's invert it to a dict of lists to match a History object + logs = {log_key: [single_dict[log_key] for single_dict in logs] for log_key in logs[0]} + + lines = [] + for i in range(len(logs["epoch"])): + epoch_dict = {log_key: log_value_list[i] for log_key, log_value_list in logs.items()} + values = {} + for k, v in epoch_dict.items(): + if k.startswith("val_"): + k = "validation_" + k[4:] + elif k != "epoch": + k = "train_" + k + splits = k.split("_") + name = " ".join([part.capitalize() for part in splits]) + values[name] = v + lines.append(values) + + eval_results = lines[-1] + + return logs, lines, eval_results + + +def parse_log_history(log_history): + """ + Parse the `log_history` of a Trainer to get the intermediate and final evaluation results. + """ + idx = 0 + while idx < len(log_history) and "train_runtime" not in log_history[idx]: + idx += 1 + + # If there are no training logs + if idx == len(log_history): + idx -= 1 + while idx >= 0 and "eval_loss" not in log_history[idx]: + idx -= 1 + + if idx >= 0: + return None, None, log_history[idx] + else: + return None, None, None + + # From now one we can assume we have training logs: + train_log = log_history[idx] + lines = [] + training_loss = "No log" + for i in range(idx): + if "loss" in log_history[i]: + training_loss = log_history[i]["loss"] + if "eval_loss" in log_history[i]: + metrics = log_history[i].copy() + _ = metrics.pop("total_flos", None) + epoch = metrics.pop("epoch", None) + step = metrics.pop("step", None) + _ = metrics.pop("eval_runtime", None) + _ = metrics.pop("eval_samples_per_second", None) + _ = metrics.pop("eval_steps_per_second", None) + _ = metrics.pop("eval_jit_compilation_time", None) + values = {"Training Loss": training_loss, "Epoch": epoch, "Step": step} + for k, v in metrics.items(): + if k == "eval_loss": + values["Validation Loss"] = v + else: + splits = k.split("_") + name = " ".join([part.capitalize() for part in splits[1:]]) + values[name] = v + lines.append(values) + + idx = len(log_history) - 1 + while idx >= 0 and "eval_loss" not in log_history[idx]: + idx -= 1 + + if idx > 0: + eval_results = {} + for key, value in log_history[idx].items(): + if key.startswith("eval_"): + key = key[5:] + if key not in ["runtime", "samples_per_second", "steps_per_second", "epoch", "step"]: + camel_cased_key = " ".join([part.capitalize() for part in key.split("_")]) + eval_results[camel_cased_key] = value + return train_log, lines, eval_results + else: + return train_log, lines, None + + +def extract_hyperparameters_from_keras(model): + from .modeling_tf_utils import keras + + hyperparameters = {} + if hasattr(model, "optimizer") and model.optimizer is not None: + hyperparameters["optimizer"] = model.optimizer.get_config() + else: + hyperparameters["optimizer"] = None + hyperparameters["training_precision"] = keras.mixed_precision.global_policy().name + + return hyperparameters + + +def _maybe_round(v, decimals=4): + if isinstance(v, float) and len(str(v).split(".")) > 1 and len(str(v).split(".")[1]) > decimals: + return f"{v:.{decimals}f}" + return str(v) + + +def _regular_table_line(values, col_widths): + values_with_space = [f"| {v}" + " " * (w - len(v) + 1) for v, w in zip(values, col_widths)] + return "".join(values_with_space) + "|\n" + + +def _second_table_line(col_widths): + values = ["|:" + "-" * w + ":" for w in col_widths] + return "".join(values) + "|\n" + + +def make_markdown_table(lines): + """ + Create a nice Markdown table from the results in `lines`. + """ + if lines is None or len(lines) == 0: + return "" + col_widths = {key: len(str(key)) for key in lines[0].keys()} + for line in lines: + for key, value in line.items(): + if col_widths[key] < len(_maybe_round(value)): + col_widths[key] = len(_maybe_round(value)) + + table = _regular_table_line(list(lines[0].keys()), list(col_widths.values())) + table += _second_table_line(list(col_widths.values())) + for line in lines: + table += _regular_table_line([_maybe_round(v) for v in line.values()], list(col_widths.values())) + return table + + +_TRAINING_ARGS_KEYS = [ + "learning_rate", + "train_batch_size", + "eval_batch_size", + "seed", +] + + +def extract_hyperparameters_from_trainer(trainer): + hyperparameters = {k: getattr(trainer.args, k) for k in _TRAINING_ARGS_KEYS} + + if trainer.args.parallel_mode not in [ParallelMode.NOT_PARALLEL, ParallelMode.NOT_DISTRIBUTED]: + hyperparameters["distributed_type"] = ( + "multi-GPU" if trainer.args.parallel_mode == ParallelMode.DISTRIBUTED else trainer.args.parallel_mode.value + ) + if trainer.args.world_size > 1: + hyperparameters["num_devices"] = trainer.args.world_size + if trainer.args.gradient_accumulation_steps > 1: + hyperparameters["gradient_accumulation_steps"] = trainer.args.gradient_accumulation_steps + + total_train_batch_size = ( + trainer.args.train_batch_size * trainer.args.world_size * trainer.args.gradient_accumulation_steps + ) + if total_train_batch_size != hyperparameters["train_batch_size"]: + hyperparameters["total_train_batch_size"] = total_train_batch_size + total_eval_batch_size = trainer.args.eval_batch_size * trainer.args.world_size + if total_eval_batch_size != hyperparameters["eval_batch_size"]: + hyperparameters["total_eval_batch_size"] = total_eval_batch_size + + if trainer.args.optim: + optimizer_name = trainer.args.optim + optimizer_args = trainer.args.optim_args if trainer.args.optim_args else "No additional optimizer arguments" + + if "adam" in optimizer_name.lower(): + hyperparameters["optimizer"] = ( + f"Use {optimizer_name} with betas=({trainer.args.adam_beta1},{trainer.args.adam_beta2}) and" + f" epsilon={trainer.args.adam_epsilon} and optimizer_args={optimizer_args}" + ) + else: + hyperparameters["optimizer"] = f"Use {optimizer_name} and the args are:\n{optimizer_args}" + + hyperparameters["lr_scheduler_type"] = trainer.args.lr_scheduler_type.value + if trainer.args.warmup_ratio != 0.0: + hyperparameters["lr_scheduler_warmup_ratio"] = trainer.args.warmup_ratio + if trainer.args.warmup_steps != 0.0: + hyperparameters["lr_scheduler_warmup_steps"] = trainer.args.warmup_steps + if trainer.args.max_steps != -1: + hyperparameters["training_steps"] = trainer.args.max_steps + else: + hyperparameters["num_epochs"] = trainer.args.num_train_epochs + + if trainer.args.fp16: + if trainer.use_apex: + hyperparameters["mixed_precision_training"] = f"Apex, opt level {trainer.args.fp16_opt_level}" + else: + hyperparameters["mixed_precision_training"] = "Native AMP" + + if trainer.args.label_smoothing_factor != 0.0: + hyperparameters["label_smoothing_factor"] = trainer.args.label_smoothing_factor + + return hyperparameters diff --git a/lib/python3.12/site-packages/transformers/modeling_attn_mask_utils.py b/lib/python3.12/site-packages/transformers/modeling_attn_mask_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..dfdd976f0156139024c6f7788870a71e4a41754d --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_attn_mask_utils.py @@ -0,0 +1,481 @@ +# Copyright 2023 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. +from dataclasses import dataclass +from typing import Optional, Union + +import torch + +from .utils.import_utils import is_torchdynamo_compiling + + +@dataclass +class AttentionMaskConverter: + """ + A utility attention mask class that allows one to: + - Create a causal 4d mask + - Create a causal 4d mask with slided window + - Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length, + key_value_length) that can be multiplied with attention scores + + Examples: + + ```python + >>> import torch + >>> from transformers.modeling_attn_mask_utils import AttentionMaskConverter + + >>> converter = AttentionMaskConverter(True) + >>> converter.to_4d(torch.tensor([[0, 0, 0, 1, 1]]), 5, key_value_length=5, dtype=torch.float32) + tensor([[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, -3.4028e+38], + [-3.4028e+38, -3.4028e+38, -3.4028e+38, 0.0000e+00, 0.0000e+00]]]]) + ``` + + Parameters: + is_causal (`bool`): + Whether the attention mask should be a uni-directional (causal) or bi-directional mask. + + sliding_window (`int`, *optional*): + Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer. + """ + + is_causal: bool + sliding_window: int + + def __init__(self, is_causal: bool, sliding_window: Optional[int] = None): + self.is_causal = is_causal + self.sliding_window = sliding_window + + if self.sliding_window is not None and self.sliding_window <= 0: + raise ValueError( + f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`" + ) + + def to_causal_4d( + self, + batch_size: int, + query_length: int, + key_value_length: int, + dtype: torch.dtype, + device: Union[torch.device, "str"] = "cpu", + ) -> 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). + """ + if not self.is_causal: + raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.") + + # If shape is not cached, create a new causal mask and cache it + input_shape = (batch_size, query_length) + past_key_values_length = key_value_length - query_length + + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + causal_4d_mask = None + if input_shape[-1] > 1 or self.sliding_window is not None: + causal_4d_mask = self._make_causal_mask( + input_shape, + dtype, + device=device, + past_key_values_length=past_key_values_length, + sliding_window=self.sliding_window, + ) + + return causal_4d_mask + + def to_4d( + self, + attention_mask_2d: torch.Tensor, + query_length: int, + dtype: torch.dtype, + 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, + key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is + causal, a causal mask will be added. + """ + input_shape = (attention_mask_2d.shape[0], query_length) + + # create causal mask + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + causal_4d_mask = None + if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal: + if key_value_length is None: + raise ValueError( + "This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask." + ) + + past_key_values_length = key_value_length - query_length + causal_4d_mask = self._make_causal_mask( + input_shape, + dtype, + device=attention_mask_2d.device, + past_key_values_length=past_key_values_length, + sliding_window=self.sliding_window, + ) + elif self.sliding_window is not None: + raise NotImplementedError("Sliding window is currently only implemented for causal masking") + + # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len] + expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to( + attention_mask_2d.device + ) + + if causal_4d_mask is not None: + expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min) + + # expanded_attn_mask + causal_4d_mask can cause some overflow + expanded_4d_mask = expanded_attn_mask + + return expanded_4d_mask + + @staticmethod + def _make_causal_mask( + input_ids_shape: torch.Size, + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, + sliding_window: Optional[int] = None, + ): + """ + Make causal mask used for bi-directional self-attention. + """ + bsz, tgt_len = input_ids_shape + mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device) + mask_cond = torch.arange(mask.size(-1), device=device) + mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0) + + mask = mask.to(dtype) + + if past_key_values_length > 0: + mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1) + + # add lower triangular sliding window mask if necessary + if sliding_window is not None: + diagonal = past_key_values_length - sliding_window - 1 + + context_mask = torch.tril(torch.ones_like(mask, dtype=torch.bool), diagonal=diagonal) + # Recent changes in PyTorch prevent mutations on tensors converted with aten::_to_copy + # See https://github.com/pytorch/pytorch/issues/127571 + if is_torchdynamo_compiling(): + mask = mask.clone() + mask.masked_fill_(context_mask, torch.finfo(dtype).min) + + 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: Optional[int] = None): + """ + Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. + """ + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) + + inverted_mask = 1.0 - expanded_mask + + return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min) + + @staticmethod + def _unmask_unattended( + expanded_mask: torch.FloatTensor, + min_dtype: float, + ): + # fmt: off + """ + Attend to all tokens in masked rows from the expanded attention mask, for example the relevant first rows when + using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + Details: https://github.com/pytorch/pytorch/issues/110213 + + `expanded_mask` is [bsz, num_masks, tgt_seq_len, src_seq_len] or [bsz, tgt_seq_len, src_seq_len]. + `attention_mask` is [bsz, src_seq_len]. + + The dimension num_masks of `expanded_mask` is most often 1, but it can also be the number of heads in the case of alibi attention bias. + + For example, if `expanded_mask` is (e.g. here left-padding case) + ``` + [[[[0, 0, 0], + [0, 0, 0], + [0, 0, 1]]], + [[[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]], + [[[0, 0, 0], + [0, 1, 0], + [0, 1, 1]]]] + ``` + then the modified `expanded_mask` will be + ``` + [[[[1, 1, 1], <-- modified + [1, 1, 1], <-- modified + [0, 0, 1]]], + [[[1, 0, 0], + [1, 1, 0], + [1, 1, 1]]], + [[[1, 1, 1], <-- modified + [0, 1, 0], + [0, 1, 1]]]] + ``` + """ + # fmt: on + if expanded_mask.dtype == torch.bool: + raise ValueError( + "AttentionMaskConverter._unmask_unattended expects a float `expanded_mask`, got a BoolTensor." + ) + + return expanded_mask.mul(~torch.all(expanded_mask == min_dtype, dim=-1, keepdim=True)) + + @staticmethod + def _ignore_causal_mask_sdpa( + attention_mask: Optional[torch.Tensor], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, + is_training: bool = False, + ) -> bool: + """ + Detects whether the optional user-specified attention_mask & the automatically created causal mask can be + ignored in case PyTorch's SDPA is used, rather relying on SDPA's `is_causal` argument. + + In case no token is masked in the `attention_mask` argument, if `query_length == 1` or + `key_value_length == query_length`, we rather rely on SDPA `is_causal` argument to use causal/non-causal masks, + allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is + passed). + """ + + _, query_length = inputs_embeds.shape[0], inputs_embeds.shape[1] + key_value_length = query_length + past_key_values_length + + is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy) or is_torchdynamo_compiling() + + ignore_causal_mask = False + + if attention_mask is None: + # TODO: When tracing with TorchDynamo with fullgraph=True, the model is recompiled depending on the input + # shape, thus SDPA's `is_causal` argument is rightfully updated + # (see https://gist.github.com/fxmarty/1313f39037fc1c112508989628c57363). However, when using + # `torch.export` or `torch.onnx.dynamo_export`, we must pass an example input, and `is_causal` behavior is + # hard-coded. If a user exports a model with q_len > 1, the exported model will hard-code `is_causal=True` + # which is in general wrong (see https://github.com/pytorch/pytorch/issues/108108). + # Thus, we only set `ignore_causal_mask = True` if the model is set to training. + # + # Besides, jit.trace can not handle the `q_len > 1` condition for `is_causal` + # ("TypeError: scaled_dot_product_attention(): argument 'is_causal' must be bool, not Tensor"). + if ( + (is_training or not is_tracing) + and (query_length == 1 or key_value_length == query_length) + and (sliding_window is None or key_value_length < sliding_window) + ): + ignore_causal_mask = True + elif sliding_window is None or key_value_length < sliding_window: + if len(attention_mask.shape) == 4: + return False + elif not is_tracing and torch.all(attention_mask == 1): + if query_length == 1 or key_value_length == query_length: + # For query_length == 1, causal attention and bi-directional attention are the same. + ignore_causal_mask = True + + # Unfortunately, for query_length > 1 and key_value_length != query_length, we cannot generally ignore + # the attention mask, as SDPA causal mask generation may be wrong. We will set `is_causal=False` in + # SDPA and rely on Transformers attention_mask instead, hence not setting it to None here. + # Reference: https://github.com/pytorch/pytorch/issues/108108 + # TODO: maybe revisit this with https://github.com/pytorch/pytorch/pull/114823 in PyTorch 2.3. + + return ignore_causal_mask + + +def _prepare_4d_causal_attention_mask( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, tuple, list], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + 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 + `(batch_size, key_value_length)` + + Args: + attention_mask (`torch.Tensor` or `None`): + A 2D attention mask of shape `(batch_size, key_value_length)` + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + inputs_embeds (`torch.Tensor`): + The embedded inputs as a torch Tensor. + past_key_values_length (`int`): + The length of the key value cache. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = input_shape[-1] + past_key_values_length + + # 4d mask is passed through the layers + if attention_mask is not None and len(attention_mask.shape) == 2: + attention_mask = attn_mask_converter.to_4d( + attention_mask, input_shape[-1], key_value_length=key_value_length, dtype=inputs_embeds.dtype + ) + elif attention_mask is not None and len(attention_mask.shape) == 4: + expected_shape = (input_shape[0], 1, input_shape[1], key_value_length) + if tuple(attention_mask.shape) != expected_shape: + raise ValueError( + f"Incorrect 4D attention_mask shape: {tuple(attention_mask.shape)}; expected: {expected_shape}." + ) + else: + # if the 4D mask has correct shape - invert it and fill with negative infinity + inverted_mask = 1.0 - attention_mask + attention_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(inputs_embeds.dtype).min + ) + else: + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + + return attention_mask + + +# Adapted from _prepare_4d_causal_attention_mask +def _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask: Optional[torch.Tensor], + input_shape: Union[torch.Size, tuple, list], + inputs_embeds: torch.Tensor, + past_key_values_length: int, + sliding_window: Optional[int] = None, +): + """ + Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`. + + In case no token is masked in the `attention_mask` argument, we simply set it to `None` for the cases `query_length == 1` and + `key_value_length == query_length`, and rely instead on SDPA `is_causal` argument to use causal/non-causal masks, + allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is passed). + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = input_shape[-1] + past_key_values_length + + # torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1` + # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing. + # TODO: For dynamo, rather use a check on fullgraph=True once this is possible (https://github.com/pytorch/pytorch/pull/120400). + is_tracing = torch.jit.is_tracing() or isinstance(inputs_embeds, torch.fx.Proxy) or is_torchdynamo_compiling() + + ignore_causal_mask = AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + sliding_window=sliding_window, + ) + + if ignore_causal_mask: + expanded_4d_mask = None + elif attention_mask is None: + expanded_4d_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device + ) + else: + if attention_mask.dim() == 4: + expanded_4d_mask = attention_mask + else: + expanded_4d_mask = attn_mask_converter.to_4d( + attention_mask, + input_shape[-1], + dtype=inputs_embeds.dtype, + key_value_length=key_value_length, + ) + + # Attend to all tokens in masked rows from the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + if not is_tracing and expanded_4d_mask.device.type == "cuda": + expanded_4d_mask = AttentionMaskConverter._unmask_unattended( + expanded_4d_mask, min_dtype=torch.finfo(inputs_embeds.dtype).min + ) + + return expanded_4d_mask + + +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)` + + Args: + mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + tgt_len (`int`): + The target length or query length the created mask shall have. + """ + 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: 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)` + + Args: + mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + tgt_len (`int`): + The target length or query length the created mask shall have. + """ + _, key_value_length = mask.shape + tgt_len = tgt_len if tgt_len is not None else key_value_length + + is_tracing = torch.jit.is_tracing() or isinstance(mask, torch.fx.Proxy) or is_torchdynamo_compiling() + + # torch.jit.trace, symbolic_trace and torchdynamo with fullgraph=True are unable to capture data-dependent controlflows. + if not is_tracing and torch.all(mask == 1): + return None + else: + return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len) + + +def _create_4d_causal_attention_mask( + input_shape: Union[torch.Size, tuple, list], + dtype: torch.dtype, + device: torch.device, + past_key_values_length: int = 0, + sliding_window: Optional[int] = None, +) -> Optional[torch.Tensor]: + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` + + Args: + input_shape (`tuple(int)` or `list(int)` or `torch.Size`): + The input shape should be a tuple that defines `(batch_size, query_length)`. + dtype (`torch.dtype`): + The torch dtype the created mask shall have. + device (`int`): + The torch device the created mask shall have. + sliding_window (`int`, *optional*): + If the model uses windowed attention, a sliding window should be passed. + """ + attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window) + + key_value_length = past_key_values_length + input_shape[-1] + attention_mask = attn_mask_converter.to_causal_4d( + input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device + ) + + return attention_mask diff --git a/lib/python3.12/site-packages/transformers/modeling_flash_attention_utils.py b/lib/python3.12/site-packages/transformers/modeling_flash_attention_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c7d54dc415146e0f289750c86caffa5021399a92 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_flash_attention_utils.py @@ -0,0 +1,437 @@ +# Copyright 2024 The Fairseq Authors and the HuggingFace Inc. 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. + +import inspect +import os +from typing import Optional, TypedDict + +import torch +import torch.nn.functional as F + +from .utils import ( + is_flash_attn_2_available, + is_flash_attn_greater_or_equal, + is_flash_attn_greater_or_equal_2_10, + is_torch_npu_available, + logging, +) + + +logger = logging.get_logger(__name__) +flash_attn_func = None + + +if is_flash_attn_2_available(): + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + from flash_attn import flash_attn_func, flash_attn_varlen_func + from flash_attn.layers.rotary import apply_rotary_emb # noqa + + +# patch functions in package `flash-attn` when using flash-attention on Ascend NPU. +if is_torch_npu_available(): + from torch_npu import npu_rotary_mul as apply_rotary_emb # noqa + + from .integrations.npu_flash_attention import index_first_axis, pad_input, unpad_input + from .integrations.npu_flash_attention import npu_flash_attn_func as flash_attn_func + from .integrations.npu_flash_attention import npu_flash_attn_varlen_func as flash_attn_varlen_func + + +if flash_attn_func: + _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters) + + +def is_flash_attn_available(): + """Determine whether flash-attention can be used or not.""" + + # if package `flash-attn` is available, flash-attention can be used natively. + if is_flash_attn_2_available(): + return True + + # flash-attention can be used on Ascend NPU without package `flash-attn` + if is_torch_npu_available(): + return True + + return False + + +def flash_attn_supports_top_left_mask(): + """Determine whether flash-attention uses top-left or down-right mask""" + + if is_flash_attn_2_available(): + # top-left mask is used in package `flash-attn` with version lower than 2.1.0 + return not is_flash_attn_greater_or_equal_2_10() + + if is_torch_npu_available(): + # down-right mask is used on Ascend NPU by default, set env `NPU_FA2_SPARSE_MODE=2` to activate top-left mask. + from .integrations.npu_flash_attention import is_npu_fa2_top_left_aligned_causal_mask + + return is_npu_fa2_top_left_aligned_causal_mask() + + return False + + +def _get_unpad_data(attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]: + """ + Retrieves indexing data required to repad unpadded (ragged) tensors. + + Arguments: + attention_mask (`torch.Tensor`): + Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. + + Return: + indices (`torch.Tensor`): + The indices of non-masked tokens from the flattened input sequence. + cu_seqlens (`torch.Tensor`): + The cumulative sequence lengths, used to index into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,). + max_seqlen_in_batch (`int`): + Maximum sequence length in batch. + """ + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return ( + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def _upad_input( + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + attention_mask: torch.Tensor, + query_length: int, +): + """ + Unpads query, key, and values tensors, using a single dimension for all tokens even though they belong to different batches. + + This function is used instead of `flash_attn.bert_padding.unpad_input` in order to avoid the recomputation of the same intermediary + tensors for query, key, value tensors. + + Arguments: + query_layer (`torch.Tensor`): + Query state with padding. Shape: (batch_size, query_length, num_heads, head_dim). + key_layer (`torch.Tensor`): + Key state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim). + value_layer (`torch.Tensor`): + Value state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim). + attention_mask (`torch.Tensor`): + Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. + query_length (`int`): + Target length. + + Return: + query_layer (`torch.Tensor`): + Query state without padding. Shape: (total_target_length, num_heads, head_dim). + key_layer (`torch.Tensor`): + Key state with padding. Shape: (total_source_length, num_key_value_heads, head_dim). + value_layer (`torch.Tensor`): + Value state with padding. Shape: (total_source_length, num_key_value_heads, head_dim). + indices_q (`torch.Tensor`): + The indices of non-masked tokens from the flattened input target sequence. + (cu_seqlens_q, cu_seqlens_k) (`Tuple[int]`): + The cumulative sequence lengths for the target (query) and source (key, value), used to index into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,). + (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`Tuple[int]`): + Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value). + """ + indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) + batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape + + key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k) + value_layer = index_first_axis( + value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k + ) + if query_length == kv_seq_len: + query_layer = index_first_axis(query_layer.reshape(batch_size * kv_seq_len, -1, head_dim), indices_k) + cu_seqlens_q = cu_seqlens_k + max_seqlen_in_batch_q = max_seqlen_in_batch_k + indices_q = indices_k + elif query_length == 1: + max_seqlen_in_batch_q = 1 + cu_seqlens_q = torch.arange( + batch_size + 1, dtype=torch.int32, device=query_layer.device + ) # There is a memcpy here, that is very bad. + indices_q = cu_seqlens_q[:-1] + query_layer = query_layer.squeeze(1) + else: + # The -q_len: slice assumes left padding. + attention_mask = attention_mask[:, -query_length:] + query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q, *_ = unpad_input(query_layer, attention_mask) + + return ( + query_layer, + key_layer, + value_layer, + indices_q, + (cu_seqlens_q, cu_seqlens_k), + (max_seqlen_in_batch_q, max_seqlen_in_batch_k), + ) + + +def prepare_fa2_from_position_ids(query, key, value, position_ids): + """ + This function returns necessary arguments to call `flash_attn_varlen_func`. + All three query, key, value states will be flattened. + Cumulative lengths of each examples in the batch will be extracted from position_ids. + + NOTE: ideally cumulative lengths should be prepared at the data collator stage + + Arguments: + query (`torch.Tensor`): + Query state with padding. Shape: (batch_size, query_length, num_heads, head_dim). + key (`torch.Tensor`): + Key state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim). + value (`torch.Tensor`): + Value state with padding. Shape: (batch_size, kv_seq_len, num_key_value_heads, head_dim). + position_ids (`torch.Tensor`): + Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. + + Return: + query (`torch.Tensor`): + Query state without padding. Shape: (total_target_length, num_heads, head_dim). + key (`torch.Tensor`): + Key state with padding. Shape: (total_source_length, num_key_value_heads, head_dim). + value (`torch.Tensor`): + Value state with padding. Shape: (total_source_length, num_key_value_heads, head_dim). + indices_q (`torch.Tensor`): + The indices of non-masked tokens from the flattened input target sequence. + (cu_seqlens_q, cu_seqlens_k) (`Tuple[int]`): + The cumulative sequence lengths for the target (query) and source (key, value), used to index into ragged (unpadded) tensors. `cu_seqlens` shape is (batch_size + 1,). + (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`Tuple[int]`): + Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value). + """ + query = query.view(-1, query.size(-2), query.size(-1)) + key = key.contiguous().view(-1, key.size(-2), key.size(-1)) + value = value.contiguous().view(-1, value.size(-2), value.size(-1)) + position_ids = position_ids.flatten() + indices_q = torch.arange(position_ids.size(0), device=position_ids.device, dtype=torch.int32) + + cu_seq_lens = torch.cat( + ( + indices_q[position_ids == 0], + torch.tensor(position_ids.size(), device=position_ids.device, dtype=torch.int32), + ) + ) + + max_length = position_ids.max() + 1 + + return (query, key, value, indices_q, (cu_seq_lens, cu_seq_lens), (max_length, max_length)) + + +def fa_peft_integration_check( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + target_dtype: Optional[torch.dtype] = None, +): + """ + PEFT usually casts the layer norms in float32 for training stability reasons + therefore the input hidden states gets silently casted in float32. Hence, we need + cast them back in float16 / bfloat16 just to be sure everything works as expected. + This might slowdown training & inference so it is recommended to not cast the LayerNorms! + + Args: + query (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value (`torch.Tensor`): + Input value states to be passed to Flash Attention API + target_dtype (`torch.dtype`, *optional*): + The dtype to convert the attention tensors to. Conversion can be ignored by + not providing the target dtype. + """ + if target_dtype is None: + return query, key, value + + input_dtype = query.dtype + if input_dtype == torch.float32: + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query = query.to(target_dtype) + key = key.to(target_dtype) + value = value.to(target_dtype) + + return query, key, value + + +flash_241 = is_flash_attn_greater_or_equal("2.4.1") +deterministic_g = os.environ.get("FLASH_ATTENTION_DETERMINISTIC", "0") == "1" + + +def _flash_attention_forward( + query_states: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + attention_mask: Optional[torch.Tensor], + query_length: int, + is_causal: bool, + dropout: float = 0.0, + position_ids: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + sliding_window: Optional[int] = None, + use_top_left_mask: bool = False, + softcap: Optional[float] = None, + deterministic: Optional[bool] = None, + cu_seq_lens_q: Optional[torch.LongTensor] = None, + cu_seq_lens_k: Optional[torch.LongTensor] = None, + max_length_q: Optional[int] = None, + max_length_k: Optional[int] = None, + target_dtype: Optional[torch.dtype] = None, + **kwargs, +): + """ + Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token + first unpad the input, then computes the attention scores and pad the final attention scores. + + Args: + query_states (`torch.Tensor`): + Input query states to be passed to Flash Attention API + key_states (`torch.Tensor`): + Input key states to be passed to Flash Attention API + value_states (`torch.Tensor`): + Input value states to be passed to Flash Attention API + attention_mask (`torch.Tensor`, *optional*): + The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the + position of padding tokens and 1 for the position of non-padding tokens. + dropout (`float`): + Attention dropout + softmax_scale (`float`, *optional*): + The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) + use_top_left_mask (`bool`, defaults to `False`): + flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. + softcap (`float`, *optional*): + Softcap for the attention logits, used e.g. in gemma2. + deterministic (`bool`, *optional*): + Determines if the deterministic option introduced in flash_attn>=2.4.1 is enabled. + """ + if not use_top_left_mask: + causal = is_causal + else: + # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. + causal = is_causal and query_length != 1 + + # Assuming 4D tensors, key_states.shape[1] is the key/value sequence length (source length). + use_sliding_windows = ( + _flash_supports_window_size and sliding_window is not None and key_states.shape[1] > sliding_window + ) + flash_kwargs = {"window_size": (sliding_window, sliding_window)} if use_sliding_windows else {} + + if flash_241: + if deterministic is None: + deterministic = deterministic_g + flash_kwargs["deterministic"] = deterministic + + if softcap is not None: + flash_kwargs["softcap"] = softcap + + # PEFT possibly silently casts tensors to fp32, this potentially reconverts to correct dtype or is a no op + query_states, key_states, value_states = fa_peft_integration_check( + query_states, key_states, value_states, target_dtype + ) + + # Contains at least one padding token in the sequence + if attention_mask is not None: + batch_size = query_states.shape[0] + query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = _upad_input( + query_states, key_states, value_states, attention_mask, query_length + ) + cu_seqlens_q, cu_seqlens_k = cu_seq_lens + max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens + + attn_output_unpad = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_in_batch_q, + max_seqlen_k=max_seqlen_in_batch_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + **flash_kwargs, + ) + attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) + + # If position_ids is provided and check all examples do not contain only 1 sequence, If tensor in increasing + # then we probably have one sequence, otherwise it is packed. Additionally check we are in pre-fill/training stage. + # Use `flash_attn_varlen_func` to prevent cross-example attention and also allow padding free approach + elif position_ids is not None and ( + max_length_q is not None or (query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all()) + ): + batch_size = query_states.size(0) + + if cu_seq_lens_q is None or cu_seq_lens_k is None: + query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = ( + prepare_fa2_from_position_ids(query_states, key_states, value_states, position_ids) + ) + + cu_seq_lens_q, cu_seq_lens_k = cu_seq_lens + max_length_q, max_length_k = max_seq_lens + + else: + query_states = query_states.reshape(-1, query_states.size(-2), query_states.size(-1)) + key_states = key_states.reshape(-1, key_states.size(-2), key_states.size(-1)) + value_states = value_states.reshape(-1, value_states.size(-2), value_states.size(-1)) + + attn_output = flash_attn_varlen_func( + query_states, + key_states, + value_states, + cu_seqlens_q=cu_seq_lens_q, + cu_seqlens_k=cu_seq_lens_k, + max_seqlen_q=max_length_q, + max_seqlen_k=max_length_k, + dropout_p=dropout, + softmax_scale=softmax_scale, + causal=causal, + **flash_kwargs, + ) + + attn_output = attn_output.view(batch_size, -1, attn_output.size(-2), attn_output.size(-1)) + + else: + attn_output = flash_attn_func( + query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal, **flash_kwargs + ) + + return attn_output + + +class FlashAttentionKwargs(TypedDict, total=False): + """ + Keyword arguments for Flash Attention with Compile. + + Attributes: + cu_seq_lens_q (`torch.LongTensor`, *optional*) + Gets cumulative sequence length for query state. + cu_seq_lens_k (`torch.LongTensor`, *optional*) + Gets cumulative sequence length for key state. + max_length_q (`int`, *optional*): + Maximum sequence length for query state. + max_length_k (`int`, *optional*): + Maximum sequence length for key state. + """ + + cu_seq_lens_q: Optional[torch.LongTensor] + cu_seq_lens_k: Optional[torch.LongTensor] + max_length_q: Optional[int] + max_length_k: Optional[int] diff --git a/lib/python3.12/site-packages/transformers/modeling_flax_outputs.py b/lib/python3.12/site-packages/transformers/modeling_flax_outputs.py new file mode 100644 index 0000000000000000000000000000000000000000..325f968571ed400903b6c3bb15a4faf10f7e56c9 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_flax_outputs.py @@ -0,0 +1,700 @@ +# Copyright 2021 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. +from typing import Dict, Optional, Tuple + +import flax +import jax.numpy as jnp + +from .utils import ModelOutput + + +@flax.struct.dataclass +class FlaxBaseModelOutput(ModelOutput): + """ + Base class for model's outputs, with potential hidden states and attentions. + + Args: + last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + last_hidden_state: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxBaseModelOutputWithNoAttention(ModelOutput): + """ + Base class for model's outputs, with potential hidden states. + + Args: + last_hidden_state (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings, if the model has an embedding layer, + one + for the output of each layer) of shape `(batch_size, num_channels, height, width)`. Hidden-states of the + model at the output of each layer plus the optional initial embedding outputs. + """ + + last_hidden_state: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxBaseModelOutputWithPoolingAndNoAttention(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + last_hidden_state (`jnp.ndarray` of shape `(batch_size, num_channels, height, width)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`jnp.ndarray` of shape `(batch_size, hidden_size)`): + Last layer hidden-state after a pooling operation on the spatial dimensions. + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings, if the model has an embedding layer, + one + for the output of each layer) of shape `(batch_size, num_channels, height, width)`. Hidden-states of the + model at the output of each layer plus the optional initial embedding outputs. + """ + + last_hidden_state: Optional[jnp.ndarray] = None + pooler_output: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxImageClassifierOutputWithNoAttention(ModelOutput): + """ + Base class for outputs of image classification models. + + Args: + logits (`jnp.ndarray` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when + `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings, if the model has an embedding layer, + one + for the output of each stage) of shape `(batch_size, num_channels, height, width)`. Hidden-states (also + called feature maps) of the model at the output of each stage. + """ + + logits: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxBaseModelOutputWithPast(ModelOutput): + """ + Base class for model's outputs, with potential hidden states and attentions. + + Args: + last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + past_key_values (`Dict[str, jnp.ndarray]`): + Dictionary of pre-computed hidden-states (key and values in the attention blocks) that can be used for fast + auto-regressive decoding. Pre-computed key and value hidden-states are of shape *[batch_size, max_length]*. + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + last_hidden_state: Optional[jnp.ndarray] = None + past_key_values: Optional[Dict[str, jnp.ndarray]] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxBaseModelOutputWithPooling(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`jnp.ndarray` of shape `(batch_size, hidden_size)`): + Last layer hidden-state of the first token of the sequence (classification token) further processed by a + Linear layer and a Tanh activation function. The Linear layer weights are trained from the next sentence + prediction (classification) objective during pretraining. + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + last_hidden_state: Optional[jnp.ndarray] = None + pooler_output: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxBaseModelOutputWithPoolingAndCrossAttentions(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`jnp.ndarray` of shape `(batch_size, hidden_size)`): + Last layer hidden-state of the first token of the sequence (classification token) after further processing + through the layers used for the auxiliary pretraining task. E.g. for BERT-family of models, this returns + the classification token after processing through a linear layer and a tanh activation function. The linear + layer weights are trained from the next sentence prediction (classification) objective during pretraining. + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings, if the model has an embedding layer, + one + for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + past_key_values (`tuple(tuple(jnp.ndarray))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(jnp.ndarray)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if + `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads, + encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if + `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` + input) to speed up sequential decoding. + """ + + last_hidden_state: Optional[jnp.ndarray] = None + pooler_output: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + past_key_values: Optional[Tuple[Tuple[jnp.ndarray]]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + cross_attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxBaseModelOutputWithPastAndCrossAttentions(ModelOutput): + """ + Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding). + + Args: + last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`tuple(tuple(jnp.ndarray))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(jnp.ndarray)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if + `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads, + encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if + `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` + input) to speed up sequential decoding. + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + """ + + last_hidden_state: Optional[jnp.ndarray] = None + past_key_values: Optional[Tuple[Tuple[jnp.ndarray]]] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + cross_attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxSeq2SeqModelOutput(ModelOutput): + """ + Base class for model encoder's outputs that also contains : pre-computed hidden states that can speed up sequential + decoding. + + Args: + last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the decoder of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`tuple(tuple(jnp.ndarray))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(jnp.ndarray)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + last_hidden_state: Optional[jnp.ndarray] = None + past_key_values: Optional[Tuple[Tuple[jnp.ndarray]]] = None + decoder_hidden_states: Optional[Tuple[jnp.ndarray]] = None + decoder_attentions: Optional[Tuple[jnp.ndarray]] = None + cross_attentions: Optional[Tuple[jnp.ndarray]] = None + encoder_last_hidden_state: Optional[jnp.ndarray] = None + encoder_hidden_states: Optional[Tuple[jnp.ndarray]] = None + encoder_attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxCausalLMOutputWithCrossAttentions(ModelOutput): + """ + Base class for causal language model (or autoregressive) outputs. + + Args: + logits (`jnp.ndarray` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Cross attentions weights after the attention softmax, used to compute the weighted average in the + cross-attention heads. + past_key_values (`tuple(tuple(jnp.ndarray))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `jnp.ndarray` tuples of length `config.n_layers`, with each tuple containing the cached key, value + states of the self-attention and the cross-attention layers if model is used in encoder-decoder setting. + Only relevant if `config.is_decoder = True`. + + Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + """ + + logits: Optional[jnp.ndarray] = None + past_key_values: Optional[Tuple[Tuple[jnp.ndarray]]] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + cross_attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxMaskedLMOutput(ModelOutput): + """ + Base class for masked language models outputs. + + Args: + logits (`jnp.ndarray` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + logits: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + + +FlaxCausalLMOutput = FlaxMaskedLMOutput + + +@flax.struct.dataclass +class FlaxSeq2SeqLMOutput(ModelOutput): + """ + Base class for sequence-to-sequence language models outputs. + + Args: + logits (`jnp.ndarray` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(jnp.ndarray))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(jnp.ndarray)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + logits: Optional[jnp.ndarray] = None + past_key_values: Optional[Tuple[Tuple[jnp.ndarray]]] = None + decoder_hidden_states: Optional[Tuple[jnp.ndarray]] = None + decoder_attentions: Optional[Tuple[jnp.ndarray]] = None + cross_attentions: Optional[Tuple[jnp.ndarray]] = None + encoder_last_hidden_state: Optional[jnp.ndarray] = None + encoder_hidden_states: Optional[Tuple[jnp.ndarray]] = None + encoder_attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxNextSentencePredictorOutput(ModelOutput): + """ + Base class for outputs of models predicting if two sentences are consecutive or not. + + Args: + logits (`jnp.ndarray` of shape `(batch_size, 2)`): + Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation + before SoftMax). + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + logits: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxSequenceClassifierOutput(ModelOutput): + """ + Base class for outputs of sentence classification models. + + Args: + logits (`jnp.ndarray` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + logits: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxSeq2SeqSequenceClassifierOutput(ModelOutput): + """ + Base class for outputs of sequence-to-sequence sentence classification models. + + Args: + logits (`jnp.ndarray` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + past_key_values (`tuple(tuple(jnp.ndarray))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(jnp.ndarray)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + logits: Optional[jnp.ndarray] = None + past_key_values: Optional[Tuple[Tuple[jnp.ndarray]]] = None + decoder_hidden_states: Optional[Tuple[jnp.ndarray]] = None + decoder_attentions: Optional[Tuple[jnp.ndarray]] = None + cross_attentions: Optional[Tuple[jnp.ndarray]] = None + encoder_last_hidden_state: Optional[jnp.ndarray] = None + encoder_hidden_states: Optional[Tuple[jnp.ndarray]] = None + encoder_attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxMultipleChoiceModelOutput(ModelOutput): + """ + Base class for outputs of multiple choice models. + + Args: + logits (`jnp.ndarray` of shape `(batch_size, num_choices)`): + *num_choices* is the second dimension of the input tensors. (see *input_ids* above). + + Classification scores (before SoftMax). + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + logits: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxTokenClassifierOutput(ModelOutput): + """ + Base class for outputs of token classification models. + + Args: + logits (`jnp.ndarray` of shape `(batch_size, sequence_length, config.num_labels)`): + Classification scores (before SoftMax). + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + logits: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxQuestionAnsweringModelOutput(ModelOutput): + """ + Base class for outputs of question answering models. + + Args: + start_logits (`jnp.ndarray` of shape `(batch_size, sequence_length)`): + Span-start scores (before SoftMax). + end_logits (`jnp.ndarray` of shape `(batch_size, sequence_length)`): + Span-end scores (before SoftMax). + hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + start_logits: Optional[jnp.ndarray] = None + end_logits: Optional[jnp.ndarray] = None + hidden_states: Optional[Tuple[jnp.ndarray]] = None + attentions: Optional[Tuple[jnp.ndarray]] = None + + +@flax.struct.dataclass +class FlaxSeq2SeqQuestionAnsweringModelOutput(ModelOutput): + """ + Base class for outputs of sequence-to-sequence question answering models. + + Args: + start_logits (`jnp.ndarray` of shape `(batch_size, sequence_length)`): + Span-start scores (before SoftMax). + end_logits (`jnp.ndarray` of shape `(batch_size, sequence_length)`): + Span-end scores (before SoftMax). + past_key_values (`tuple(tuple(jnp.ndarray))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(jnp.ndarray)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`jnp.ndarray` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(jnp.ndarray)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `jnp.ndarray` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(jnp.ndarray)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `jnp.ndarray` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + start_logits: Optional[jnp.ndarray] = None + end_logits: Optional[jnp.ndarray] = None + past_key_values: Optional[Tuple[Tuple[jnp.ndarray]]] = None + decoder_hidden_states: Optional[Tuple[jnp.ndarray]] = None + decoder_attentions: Optional[Tuple[jnp.ndarray]] = None + cross_attentions: Optional[Tuple[jnp.ndarray]] = None + encoder_last_hidden_state: Optional[jnp.ndarray] = None + encoder_hidden_states: Optional[Tuple[jnp.ndarray]] = None + encoder_attentions: Optional[Tuple[jnp.ndarray]] = None diff --git a/lib/python3.12/site-packages/transformers/modeling_flax_pytorch_utils.py b/lib/python3.12/site-packages/transformers/modeling_flax_pytorch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..072850657725fb41e54eb33d0a7ad9d489968b00 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_flax_pytorch_utils.py @@ -0,0 +1,490 @@ +# coding=utf-8 +# Copyright 2021 The HuggingFace Inc. team. +# +# 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. +"""PyTorch - Flax general utilities.""" + +import os +from pickle import UnpicklingError +from typing import Dict, Tuple + +import jax +import jax.numpy as jnp +import numpy as np +from flax.serialization import from_bytes +from flax.traverse_util import flatten_dict, unflatten_dict + +import transformers + +from . import is_safetensors_available, is_torch_available +from .utils import logging + + +if is_torch_available(): + import torch + +if is_safetensors_available(): + from safetensors import safe_open + from safetensors.flax import load_file as safe_load_file + + +logger = logging.get_logger(__name__) + + +##################### +# PyTorch => Flax # +##################### + + +def load_pytorch_checkpoint_in_flax_state_dict( + flax_model, pytorch_checkpoint_path, is_sharded, allow_missing_keys=False +): + """Load pytorch checkpoints in a flax model""" + + if not is_sharded: + pt_path = os.path.abspath(pytorch_checkpoint_path) + logger.info(f"Loading PyTorch weights from {pt_path}") + + if pt_path.endswith(".safetensors"): + pt_state_dict = {} + with safe_open(pt_path, framework="flax") as f: + for k in f.keys(): + pt_state_dict[k] = f.get_tensor(k) + else: + try: + import torch # noqa: F401 + except (ImportError, ModuleNotFoundError): + logger.error( + "Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see" + " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation" + " instructions." + ) + raise + + pt_state_dict = torch.load(pt_path, map_location="cpu", weights_only=True) + logger.info(f"PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()):,} parameters.") + + flax_state_dict = convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model) + else: + # model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files + flax_state_dict = convert_pytorch_sharded_state_dict_to_flax(pytorch_checkpoint_path, flax_model) + return flax_state_dict + + +def rename_key_and_reshape_tensor( + pt_tuple_key: Tuple[str], + pt_tensor: np.ndarray, + random_flax_state_dict: Dict[str, jnp.ndarray], + model_prefix: str, +) -> (Tuple[str], np.ndarray): + """Rename PT weight names to corresponding Flax weight names and reshape tensor if necessary""" + + def is_key_or_prefix_key_in_dict(key: Tuple[str]) -> bool: + """Checks if `key` of `(prefix,) + key` is in random_flax_state_dict""" + return len(set(random_flax_state_dict) & {key, (model_prefix,) + key}) > 0 + + # layer norm + renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",) + if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(renamed_pt_tuple_key): + return renamed_pt_tuple_key, pt_tensor + + # batch norm layer mean + renamed_pt_tuple_key = pt_tuple_key[:-1] + ("mean",) + if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(pt_tuple_key): + return renamed_pt_tuple_key, pt_tensor + + # batch norm layer var + renamed_pt_tuple_key = pt_tuple_key[:-1] + ("var",) + if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(pt_tuple_key): + return renamed_pt_tuple_key, pt_tensor + + # embedding + renamed_pt_tuple_key = pt_tuple_key[:-1] + ("embedding",) + if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(renamed_pt_tuple_key): + return renamed_pt_tuple_key, pt_tensor + + # conv layer + renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",) + if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(pt_tuple_key): + pt_tensor = pt_tensor.transpose(2, 3, 1, 0) + return renamed_pt_tuple_key, pt_tensor + + # linear layer + renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",) + if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(pt_tuple_key): + pt_tensor = pt_tensor.T + return renamed_pt_tuple_key, pt_tensor + + # old PyTorch layer norm weight + renamed_pt_tuple_key = pt_tuple_key[:-1] + ("weight",) + if pt_tuple_key[-1] == "gamma": + return renamed_pt_tuple_key, pt_tensor + + # old PyTorch layer norm bias + renamed_pt_tuple_key = pt_tuple_key[:-1] + ("bias",) + if pt_tuple_key[-1] == "beta": + return renamed_pt_tuple_key, pt_tensor + + # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030 + name = None + if pt_tuple_key[-3::2] == ("parametrizations", "original0"): + name = pt_tuple_key[-2] + "_g" + elif pt_tuple_key[-3::2] == ("parametrizations", "original1"): + name = pt_tuple_key[-2] + "_v" + if name is not None: + renamed_pt_tuple_key = pt_tuple_key[:-3] + (name,) + return renamed_pt_tuple_key, pt_tensor + + return pt_tuple_key, pt_tensor + + +def convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model): + # convert pytorch tensor to numpy + from_bin = is_torch_available() and isinstance(next(iter(pt_state_dict.values())), torch.Tensor) + bfloat16 = torch.bfloat16 if from_bin else "bfloat16" + + weight_dtypes = {k: v.dtype for k, v in pt_state_dict.items()} + + if from_bin: + for k, v in pt_state_dict.items(): + # numpy currently does not support bfloat16, need to go over float32 in this case to not lose precision + if v.dtype == bfloat16: + v = v.float() + pt_state_dict[k] = v.cpu().numpy() + + model_prefix = flax_model.base_model_prefix + + # use params dict if the model contains batch norm layers + if "params" in flax_model.params: + flax_model_params = flax_model.params["params"] + else: + flax_model_params = flax_model.params + random_flax_state_dict = flatten_dict(flax_model_params) + + # add batch_stats keys,values to dict + if "batch_stats" in flax_model.params: + flax_batch_stats = flatten_dict(flax_model.params["batch_stats"]) + random_flax_state_dict.update(flax_batch_stats) + + flax_state_dict = {} + + load_model_with_head_into_base_model = (model_prefix not in flax_model_params) and ( + model_prefix in {k.split(".")[0] for k in pt_state_dict.keys()} + ) + load_base_model_into_model_with_head = (model_prefix in flax_model_params) and ( + model_prefix not in {k.split(".")[0] for k in pt_state_dict.keys()} + ) + + # Need to change some parameters name to match Flax names + for pt_key, pt_tensor in pt_state_dict.items(): + pt_tuple_key = tuple(pt_key.split(".")) + is_bfloat_16 = weight_dtypes[pt_key] == bfloat16 + + # remove base model prefix if necessary + has_base_model_prefix = pt_tuple_key[0] == model_prefix + if load_model_with_head_into_base_model and has_base_model_prefix: + pt_tuple_key = pt_tuple_key[1:] + + # Correctly rename weight parameters + flax_key, flax_tensor = rename_key_and_reshape_tensor( + pt_tuple_key, pt_tensor, random_flax_state_dict, model_prefix + ) + + # add model prefix if necessary + require_base_model_prefix = (model_prefix,) + flax_key in random_flax_state_dict + if load_base_model_into_model_with_head and require_base_model_prefix: + flax_key = (model_prefix,) + flax_key + + if flax_key in random_flax_state_dict: + if flax_tensor.shape != random_flax_state_dict[flax_key].shape: + raise ValueError( + f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape " + f"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}." + ) + + # add batch stats if the model contains batchnorm layers + if "batch_stats" in flax_model.params: + if "mean" in flax_key[-1] or "var" in flax_key[-1]: + flax_state_dict[("batch_stats",) + flax_key] = jnp.asarray(flax_tensor) + continue + # remove num_batches_tracked key + if "num_batches_tracked" in flax_key[-1]: + flax_state_dict.pop(flax_key, None) + continue + + # also add unexpected weight so that warning is thrown + flax_state_dict[("params",) + flax_key] = ( + jnp.asarray(flax_tensor) if not is_bfloat_16 else jnp.asarray(flax_tensor, dtype=jnp.bfloat16) + ) + else: + # also add unexpected weight so that warning is thrown + flax_state_dict[flax_key] = ( + jnp.asarray(flax_tensor) if not is_bfloat_16 else jnp.asarray(flax_tensor, dtype=jnp.bfloat16) + ) + + return unflatten_dict(flax_state_dict) + + +############################ +# Sharded Pytorch => Flax # +############################ + + +def convert_pytorch_sharded_state_dict_to_flax(shard_filenames, flax_model): + import torch + + # Load the index + flax_state_dict = {} + for shard_file in shard_filenames: + # load using msgpack utils + pt_state_dict = torch.load(shard_file, weights_only=True) + weight_dtypes = {k: v.dtype for k, v in pt_state_dict.items()} + pt_state_dict = { + k: v.numpy() if v.dtype != torch.bfloat16 else v.float().numpy() for k, v in pt_state_dict.items() + } + + model_prefix = flax_model.base_model_prefix + + # use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict + if "batch_stats" in flax_model.params: + flax_model_params = flax_model.params["params"] + + random_flax_state_dict = flatten_dict(flax_model_params) + random_flax_state_dict.update(flatten_dict(flax_model.params["batch_stats"])) + else: + flax_model_params = flax_model.params + random_flax_state_dict = flatten_dict(flax_model_params) + + load_model_with_head_into_base_model = (model_prefix not in flax_model_params) and ( + model_prefix in {k.split(".")[0] for k in pt_state_dict.keys()} + ) + load_base_model_into_model_with_head = (model_prefix in flax_model_params) and ( + model_prefix not in {k.split(".")[0] for k in pt_state_dict.keys()} + ) + # Need to change some parameters name to match Flax names + for pt_key, pt_tensor in pt_state_dict.items(): + pt_tuple_key = tuple(pt_key.split(".")) + is_bfloat_16 = weight_dtypes[pt_key] == torch.bfloat16 + + # remove base model prefix if necessary + has_base_model_prefix = pt_tuple_key[0] == model_prefix + if load_model_with_head_into_base_model and has_base_model_prefix: + pt_tuple_key = pt_tuple_key[1:] + + # Correctly rename weight parameters + flax_key, flax_tensor = rename_key_and_reshape_tensor( + pt_tuple_key, pt_tensor, random_flax_state_dict, model_prefix + ) + # add model prefix if necessary + require_base_model_prefix = (model_prefix,) + flax_key in random_flax_state_dict + if load_base_model_into_model_with_head and require_base_model_prefix: + flax_key = (model_prefix,) + flax_key + + if flax_key in random_flax_state_dict: + if flax_tensor.shape != random_flax_state_dict[flax_key].shape: + raise ValueError( + f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape " + f"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}." + ) + + # add batch stats if the model contains batchnorm layers + if "batch_stats" in flax_model.params: + if "mean" in flax_key[-1]: + flax_state_dict[("batch_stats",) + flax_key] = jnp.asarray(flax_tensor) + continue + if "var" in flax_key[-1]: + flax_state_dict[("batch_stats",) + flax_key] = jnp.asarray(flax_tensor) + continue + # remove num_batches_tracked key + if "num_batches_tracked" in flax_key[-1]: + flax_state_dict.pop(flax_key, None) + continue + + # also add unexpected weight so that warning is thrown + flax_state_dict[("params",) + flax_key] = ( + jnp.asarray(flax_tensor) if not is_bfloat_16 else jnp.asarray(flax_tensor, dtype=jnp.bfloat16) + ) + + else: + # also add unexpected weight so that warning is thrown + flax_state_dict[flax_key] = ( + jnp.asarray(flax_tensor) if not is_bfloat_16 else jnp.asarray(flax_tensor, dtype=jnp.bfloat16) + ) + return unflatten_dict(flax_state_dict) + + +##################### +# Flax => PyTorch # +##################### + + +def load_flax_checkpoint_in_pytorch_model(model, flax_checkpoint_path): + """Load flax checkpoints in a PyTorch model""" + flax_checkpoint_path = os.path.abspath(flax_checkpoint_path) + logger.info(f"Loading Flax weights from {flax_checkpoint_path}") + + # import correct flax class + flax_cls = getattr(transformers, "Flax" + model.__class__.__name__) + + # load flax weight dict + if flax_checkpoint_path.endswith(".safetensors"): + flax_state_dict = safe_load_file(flax_checkpoint_path) + flax_state_dict = unflatten_dict(flax_state_dict, sep=".") + else: + with open(flax_checkpoint_path, "rb") as state_f: + try: + flax_state_dict = from_bytes(flax_cls, state_f.read()) + except UnpicklingError: + raise EnvironmentError(f"Unable to convert {flax_checkpoint_path} to Flax deserializable object. ") + + return load_flax_weights_in_pytorch_model(model, flax_state_dict) + + +def load_flax_weights_in_pytorch_model(pt_model, flax_state): + """Load flax checkpoints in a PyTorch model""" + + try: + import torch # noqa: F401 + except (ImportError, ModuleNotFoundError): + logger.error( + "Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see" + " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation" + " instructions." + ) + raise + + # check if we have bf16 weights + is_type_bf16 = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype == jnp.bfloat16, flax_state)).values() + if any(is_type_bf16): + # convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16 + # and bf16 is not fully supported in PT yet. + logger.warning( + "Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` " + "before loading those in PyTorch model." + ) + flax_state = jax.tree_util.tree_map( + lambda params: params.astype(np.float32) if params.dtype == jnp.bfloat16 else params, flax_state + ) + + flax_state_dict = flatten_dict(flax_state) + pt_model_dict = pt_model.state_dict() + + load_model_with_head_into_base_model = (pt_model.base_model_prefix in flax_state) and ( + pt_model.base_model_prefix not in {k.split(".")[0] for k in pt_model_dict.keys()} + ) + load_base_model_into_model_with_head = (pt_model.base_model_prefix not in flax_state) and ( + pt_model.base_model_prefix in {k.split(".")[0] for k in pt_model_dict.keys()} + ) + + # keep track of unexpected & missing keys + unexpected_keys = [] + missing_keys = set(pt_model_dict.keys()) + + for flax_key_tuple, flax_tensor in flax_state_dict.items(): + has_base_model_prefix = flax_key_tuple[0] == pt_model.base_model_prefix + require_base_model_prefix = ".".join((pt_model.base_model_prefix,) + flax_key_tuple) in pt_model_dict + + # adapt flax_key to prepare for loading from/to base model only + if load_model_with_head_into_base_model and has_base_model_prefix: + flax_key_tuple = flax_key_tuple[1:] + elif load_base_model_into_model_with_head and require_base_model_prefix: + flax_key_tuple = (pt_model.base_model_prefix,) + flax_key_tuple + + # rename flax weights to PyTorch format + if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(flax_key_tuple) not in pt_model_dict: + # conv layer + flax_key_tuple = flax_key_tuple[:-1] + ("weight",) + flax_tensor = jnp.transpose(flax_tensor, (3, 2, 0, 1)) + elif flax_key_tuple[-1] == "kernel" and ".".join(flax_key_tuple) not in pt_model_dict: + # linear layer + flax_key_tuple = flax_key_tuple[:-1] + ("weight",) + flax_tensor = flax_tensor.T + elif flax_key_tuple[-1] in ["scale", "embedding"]: + flax_key_tuple = flax_key_tuple[:-1] + ("weight",) + + # adding batch stats from flax batch norm to pt + elif "mean" in flax_key_tuple[-1]: + flax_key_tuple = flax_key_tuple[:-1] + ("running_mean",) + elif "var" in flax_key_tuple[-1]: + flax_key_tuple = flax_key_tuple[:-1] + ("running_var",) + + if "batch_stats" in flax_state: + flax_key = ".".join(flax_key_tuple[1:]) # Remove the params/batch_stats header + else: + flax_key = ".".join(flax_key_tuple) + + # We also need to look at `pt_model_dict` and see if there are keys requiring further transformation. + special_pt_names = {} + # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030 + for key in pt_model_dict: + key_components = key.split(".") + name = None + if key_components[-3::2] == ["parametrizations", "original0"]: + name = key_components[-2] + "_g" + elif key_components[-3::2] == ["parametrizations", "original1"]: + name = key_components[-2] + "_v" + if name is not None: + key_components = key_components[:-3] + [name] + key_to_check = ".".join(key_components) + special_pt_names[key_to_check] = key + + if flax_key in special_pt_names: + flax_key = special_pt_names[flax_key] + + if flax_key in pt_model_dict: + if flax_tensor.shape != pt_model_dict[flax_key].shape: + raise ValueError( + f"Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected " + f"to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}." + ) + else: + # add weight to pytorch dict + flax_tensor = np.asarray(flax_tensor) if not isinstance(flax_tensor, np.ndarray) else flax_tensor + pt_model_dict[flax_key] = torch.from_numpy(flax_tensor) + # remove from missing keys + missing_keys.remove(flax_key) + else: + # weight is not expected by PyTorch model + unexpected_keys.append(flax_key) + + pt_model.load_state_dict(pt_model_dict) + + # re-transform missing_keys to list + missing_keys = list(missing_keys) + + if len(unexpected_keys) > 0: + logger.warning( + "Some weights of the Flax model were not used when initializing the PyTorch model" + f" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing" + f" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture" + " (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This" + f" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect" + " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a" + " FlaxBertForSequenceClassification model)." + ) + else: + logger.warning(f"All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n") + if len(missing_keys) > 0: + logger.warning( + f"Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly" + f" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to" + " use it for predictions and inference." + ) + else: + logger.warning( + f"All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n" + "If your task is similar to the task the model of the checkpoint was trained on, " + f"you can already use {pt_model.__class__.__name__} for predictions without further training." + ) + + return pt_model diff --git a/lib/python3.12/site-packages/transformers/modeling_flax_utils.py b/lib/python3.12/site-packages/transformers/modeling_flax_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c775ee85bbefb6a70481863bb8c546ef3efa43d3 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_flax_utils.py @@ -0,0 +1,1273 @@ +# coding=utf-8 +# Copyright 2021 The Google Flax Team Authors and The HuggingFace Inc. team. +# +# 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. + + +import gc +import json +import os +import warnings +from functools import partial +from pickle import UnpicklingError +from typing import Any, Dict, Optional, Set, Tuple, Union + +import flax.linen as nn +import jax +import jax.numpy as jnp +import msgpack.exceptions +from flax.core.frozen_dict import FrozenDict, unfreeze +from flax.serialization import from_bytes, to_bytes +from flax.traverse_util import flatten_dict, unflatten_dict +from jax.random import PRNGKey + +from .configuration_utils import PretrainedConfig +from .dynamic_module_utils import custom_object_save +from .generation import FlaxGenerationMixin, GenerationConfig +from .modeling_flax_pytorch_utils import load_pytorch_checkpoint_in_flax_state_dict +from .utils import ( + FLAX_WEIGHTS_INDEX_NAME, + FLAX_WEIGHTS_NAME, + SAFE_WEIGHTS_INDEX_NAME, + SAFE_WEIGHTS_NAME, + WEIGHTS_INDEX_NAME, + WEIGHTS_NAME, + PushToHubMixin, + add_code_sample_docstrings, + add_start_docstrings_to_model_forward, + cached_file, + copy_func, + download_url, + has_file, + is_offline_mode, + is_remote_url, + logging, + replace_return_docstrings, +) +from .utils.hub import convert_file_size_to_int, get_checkpoint_shard_files +from .utils.import_utils import is_safetensors_available + + +if is_safetensors_available(): + from safetensors import safe_open + from safetensors.flax import load_file as safe_load_file + from safetensors.flax import save_file as safe_save_file + +logger = logging.get_logger(__name__) + + +def quick_gelu(x): + return x * jax.nn.sigmoid(1.702 * x) + + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.swish, + "swish": nn.swish, + "gelu_new": partial(nn.gelu, approximate=True), + "quick_gelu": quick_gelu, + "gelu_pytorch_tanh": partial(nn.gelu, approximate=True), + "tanh": nn.tanh, +} + + +def flax_shard_checkpoint(params, max_shard_size="10GB"): + """ + Splits a model state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a + given size. The sub-checkpoints are determined by iterating through the `state_dict` in the order of its keys, so + there is no optimization made to make each sub-checkpoint as close as possible to the maximum size passed. For + example, if the limit is 10GB and we have weights 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]. + + + + If one of the model's weight is bigger that `max_shard_size`, it will end up in its own sub-checkpoint which will + have a size greater than `max_shard_size`. + + + + Args: + params (`Union[Dict, FrozenDict]`): A `PyTree` of model parameters. + max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`): + The maximum size of each sub-checkpoint. If expressed as a string, needs to be digits followed by a unit + (like `"5MB"`). + """ + max_shard_size = convert_file_size_to_int(max_shard_size) + + sharded_state_dicts = [] + current_block = {} + current_block_size = 0 + total_size = 0 + + # flatten the weights to chunk + weights = flatten_dict(params, sep="/") + for item in weights: + weight_size = weights[item].size * weights[item].dtype.itemsize + + # If this weight is going to tip up over the maximal size, we split. + if current_block_size + weight_size > max_shard_size: + sharded_state_dicts.append(current_block) + current_block = {} + current_block_size = 0 + + current_block[item] = weights[item] + current_block_size += weight_size + total_size += weight_size + + # Add the last block + sharded_state_dicts.append(current_block) + + # If we only have one shard, we return it + if len(sharded_state_dicts) == 1: + return {FLAX_WEIGHTS_NAME: sharded_state_dicts[0]}, None + + # Otherwise, let's build the index + weight_map = {} + shards = {} + for idx, shard in enumerate(sharded_state_dicts): + shard_file = FLAX_WEIGHTS_NAME.replace(".msgpack", f"-{idx + 1:05d}-of-{len(sharded_state_dicts):05d}.msgpack") + shards[shard_file] = shard + for weight_name in shard.keys(): + weight_map[weight_name] = shard_file + + # Add the metadata + metadata = {"total_size": total_size} + index = {"metadata": metadata, "weight_map": weight_map} + return shards, index + + +class FlaxPreTrainedModel(PushToHubMixin, FlaxGenerationMixin): + r""" + Base class for all models. + + [`FlaxPreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading, + downloading and saving models. + + Class attributes (overridden by derived classes): + + - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class + for this model architecture. + - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived + classes of the same architecture adding modules on top of the base model. + - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP + models, `pixel_values` for vision models and `input_values` for speech models). + """ + + config_class = None + base_model_prefix = "" + main_input_name = "input_ids" + _auto_class = None + _missing_keys = set() + + def __init__( + self, + config: PretrainedConfig, + module: nn.Module, + input_shape: Tuple = (1, 1), + seed: int = 0, + dtype: jnp.dtype = jnp.float32, + _do_init: bool = True, + ): + if config is None: + raise ValueError("config cannot be None") + + if module is None: + raise ValueError("module cannot be None") + + # Those are private to be exposed as typed property on derived classes. + self._config = config + self._module = module + + # Those are public as their type is generic to every derived classes. + self.key = PRNGKey(seed) + self.dtype = dtype + self.input_shape = input_shape + self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None + + # To check if the model was initialized automatically. + self._is_initialized = _do_init + + if _do_init: + # randomly initialized parameters + random_params = self.init_weights(self.key, input_shape) + params_shape_tree = jax.eval_shape(lambda params: params, random_params) + else: + init_fn = partial(self.init_weights, input_shape=input_shape) + params_shape_tree = jax.eval_shape(init_fn, self.key) + + logger.info( + "Model weights are not initialized as `_do_init` is set to `False`. " + f"Make sure to call `{self.__class__.__name__}.init_weights` manually to initialize the weights." + ) + + # get the shape of the parameters + self._params_shape_tree = params_shape_tree + + # save required_params as set + self._required_params = set(flatten_dict(unfreeze(params_shape_tree)).keys()) + + # initialize the parameters + if _do_init: + self.params = random_params + + def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> Dict: + raise NotImplementedError(f"init method has to be implemented for {self}") + + def enable_gradient_checkpointing(self): + raise NotImplementedError(f"gradient checkpointing method has to be implemented for {self}") + + @classmethod + def _from_config(cls, config, **kwargs): + """ + All context managers that the model should be initialized under go here. + """ + return cls(config, **kwargs) + + @property + def framework(self) -> str: + """ + :str: Identifies that this is a Flax model. + """ + return "flax" + + @property + def config(self) -> PretrainedConfig: + return self._config + + @property + def module(self) -> nn.Module: + return self._module + + @property + def params(self) -> Union[Dict, FrozenDict]: + if not self._is_initialized: + raise ValueError( + "`params` cannot be accessed from model when the model is created with `_do_init=False`. " + "You must call `init_weights` manually and store the params outside of the model and " + "pass it explicitly where needed." + ) + return self._params + + @property + def required_params(self) -> Set: + return self._required_params + + @property + def params_shape_tree(self) -> Dict: + return self._params_shape_tree + + @params.setter + def params(self, params: Union[Dict, FrozenDict]): + # don't set params if the model is not initialized + if not self._is_initialized: + raise ValueError( + "`params` cannot be set from model when the model is created with `_do_init=False`. " + "You store the params outside of the model." + ) + + if isinstance(params, FrozenDict): + params = unfreeze(params) + param_keys = set(flatten_dict(params).keys()) + if len(self.required_params - param_keys) > 0: + raise ValueError( + "Some parameters are missing. Make sure that `params` include the following " + f"parameters {self.required_params - param_keys}" + ) + self._params = params + + def _cast_floating_to(self, params: Union[Dict, FrozenDict], dtype: jnp.dtype, mask: Any = None) -> Any: + """ + Helper method to cast floating-point values of given parameter `PyTree` to given `dtype`. + """ + + # taken from https://github.com/deepmind/jmp/blob/3a8318abc3292be38582794dbf7b094e6583b192/jmp/_src/policy.py#L27 + def conditional_cast(param): + if isinstance(param, jnp.ndarray) and jnp.issubdtype(param.dtype, jnp.floating): + param = param.astype(dtype) + return param + + if mask is None: + return jax.tree_util.tree_map(conditional_cast, params) + + flat_params = flatten_dict(params) + flat_mask, _ = jax.tree_util.tree_flatten(mask) + + for masked, key in zip(flat_mask, sorted(flat_params.keys())): + if masked: + flat_params[key] = conditional_cast(flat_params[key]) + + return unflatten_dict(flat_params) + + def to_bf16(self, params: Union[Dict, FrozenDict], mask: Any = None): + r""" + Cast the floating-point `params` to `jax.numpy.bfloat16`. This returns a new `params` tree and does not cast + the `params` in place. + + This method can be used on TPU to explicitly convert the model parameters to bfloat16 precision to do full + half-precision training or to save weights in bfloat16 for inference in order to save memory and improve speed. + + Arguments: + params (`Union[Dict, FrozenDict]`): + A `PyTree` of model parameters. + mask (`Union[Dict, FrozenDict]`): + A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params + you want to cast, and should be `False` for those you want to skip. + + Examples: + + ```python + >>> from transformers import FlaxBertModel + + >>> # load model + >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased") + >>> # By default, the model parameters will be in fp32 precision, to cast these to bfloat16 precision + >>> model.params = model.to_bf16(model.params) + >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale) + >>> # then pass the mask as follows + >>> from flax import traverse_util + + >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased") + >>> flat_params = traverse_util.flatten_dict(model.params) + >>> mask = { + ... path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale")) + ... for path in flat_params + ... } + >>> mask = traverse_util.unflatten_dict(mask) + >>> model.params = model.to_bf16(model.params, mask) + ```""" + return self._cast_floating_to(params, jnp.bfloat16, mask) + + def to_fp32(self, params: Union[Dict, FrozenDict], mask: Any = None): + r""" + Cast the floating-point `params` to `jax.numpy.float32`. This method can be used to explicitly convert the + model parameters to fp32 precision. This returns a new `params` tree and does not cast the `params` in place. + + Arguments: + params (`Union[Dict, FrozenDict]`): + A `PyTree` of model parameters. + mask (`Union[Dict, FrozenDict]`): + A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params + you want to cast, and should be `False` for those you want to skip + + Examples: + + ```python + >>> from transformers import FlaxBertModel + + >>> # Download model and configuration from huggingface.co + >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased") + >>> # By default, the model params will be in fp32, to illustrate the use of this method, + >>> # we'll first cast to fp16 and back to fp32 + >>> model.params = model.to_f16(model.params) + >>> # now cast back to fp32 + >>> model.params = model.to_fp32(model.params) + ```""" + return self._cast_floating_to(params, jnp.float32, mask) + + def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None): + r""" + Cast the floating-point `params` to `jax.numpy.float16`. This returns a new `params` tree and does not cast the + `params` in place. + + This method can be used on GPU to explicitly convert the model parameters to float16 precision to do full + half-precision training or to save weights in float16 for inference in order to save memory and improve speed. + + Arguments: + params (`Union[Dict, FrozenDict]`): + A `PyTree` of model parameters. + mask (`Union[Dict, FrozenDict]`): + A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params + you want to cast, and should be `False` for those you want to skip + + Examples: + + ```python + >>> from transformers import FlaxBertModel + + >>> # load model + >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased") + >>> # By default, the model params will be in fp32, to cast these to float16 + >>> model.params = model.to_fp16(model.params) + >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale) + >>> # then pass the mask as follows + >>> from flax import traverse_util + + >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased") + >>> flat_params = traverse_util.flatten_dict(model.params) + >>> mask = { + ... path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale")) + ... for path in flat_params + ... } + >>> mask = traverse_util.unflatten_dict(mask) + >>> model.params = model.to_fp16(model.params, mask) + ```""" + return self._cast_floating_to(params, jnp.float16, mask) + + @classmethod + def load_flax_weights(cls, resolved_archive_file): + try: + if resolved_archive_file.endswith(".safetensors"): + state = safe_load_file(resolved_archive_file) + state = unflatten_dict(state, sep=".") + else: + with open(resolved_archive_file, "rb") as state_f: + state = from_bytes(cls, state_f.read()) + except (UnpicklingError, msgpack.exceptions.ExtraData) as e: + try: + with open(resolved_archive_file) as f: + if f.read().startswith("version"): + raise OSError( + "You seem to have cloned a repository without having git-lfs installed. Please" + " install git-lfs and run `git lfs install` followed by `git lfs pull` in the" + " folder you cloned." + ) + else: + raise ValueError from e + except (UnicodeDecodeError, ValueError): + raise EnvironmentError(f"Unable to convert {resolved_archive_file} to Flax deserializable object. ") + + return state + + @classmethod + def load_flax_sharded_weights(cls, shard_files): + """ + This is the same as [`flax.serialization.from_bytes`] + (https:lax.readthedocs.io/en/latest/_modules/flax/serialization.html#from_bytes) but for a sharded checkpoint. + + This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being + loaded in the model. + + Args: + shard_files (`List[str]`: + The list of shard files to load. + + Returns: + `Dict`: A nested dictionary of the model parameters, in the expected format for flax models : `{'model': + {'params': {'...'}}}`. + """ + + # Load the index + state_sharded_dict = {} + + for shard_file in shard_files: + # load using msgpack utils + try: + with open(shard_file, "rb") as state_f: + state = from_bytes(cls, state_f.read()) + except (UnpicklingError, msgpack.exceptions.ExtraData) as e: + with open(shard_file) as f: + if f.read().startswith("version"): + raise OSError( + "You seem to have cloned a repository without having git-lfs installed. Please" + " install git-lfs and run `git lfs install` followed by `git lfs pull` in the" + " folder you cloned." + ) + else: + raise ValueError from e + except (UnicodeDecodeError, ValueError): + raise EnvironmentError(f"Unable to convert {shard_file} to Flax deserializable object. ") + + state = flatten_dict(state, sep="/") + state_sharded_dict.update(state) + del state + gc.collect() + + # the state dict is unflattened to the match the format of model.params + return unflatten_dict(state_sharded_dict, sep="/") + + @classmethod + def can_generate(cls) -> bool: + """ + Returns whether this model can generate sequences with `.generate()`. Returns: + `bool`: Whether this model can generate sequences with `.generate()`. + """ + # Detects whether `prepare_inputs_for_generation` has been overwritten, which is a requirement for generation. + # Alternatively, the model can also have a custom `generate` function. + if "GenerationMixin" in str(cls.prepare_inputs_for_generation) and "GenerationMixin" in str(cls.generate): + return False + return True + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Union[str, os.PathLike], + dtype: jnp.dtype = jnp.float32, + *model_args, + config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None, + cache_dir: Optional[Union[str, os.PathLike]] = None, + ignore_mismatched_sizes: bool = False, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + **kwargs, + ): + r""" + Instantiate a pretrained flax model from a pre-trained model configuration. + + The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come + pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning + task. + + The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those + weights are discarded. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`): + Can be either: + + - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co. + - A path to a *directory* containing model weights saved using + [`~FlaxPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + - A path or url to a *pt index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In this case, + `from_pt` should be set to `True`. + dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`): + The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and + `jax.numpy.bfloat16` (on TPUs). + + This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If + specified all the computation will be performed with the given `dtype`. + + **Note that this only specifies the dtype of the computation and does not influence the dtype of model + parameters.** + + If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and + [`~FlaxPreTrainedModel.to_bf16`]. + model_args (sequence of positional arguments, *optional*): + All remaining positional arguments will be passed to the underlying model's `__init__` method. + config (`Union[PretrainedConfig, str, os.PathLike]`, *optional*): + Can be either: + + - an instance of a class derived from [`PretrainedConfig`], + - a string or path valid as input to [`~PretrainedConfig.from_pretrained`]. + + Configuration for the model to use instead of an automatically loaded configuration. Configuration can + be automatically loaded when: + + - The model is a model provided by the library (loaded with the *model id* string of a pretrained + model). + - The model was saved using [`~PreTrainedModel.save_pretrained`] and is reloaded by supplying the + save directory. + - The model is loaded by supplying a local directory as `pretrained_model_name_or_path` and a + configuration JSON file named *config.json* is found in the directory. + cache_dir (`Union[str, 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. + from_pt (`bool`, *optional*, defaults to `False`): + Load the model weights from a PyTorch checkpoint save file (see docstring of + `pretrained_model_name_or_path` argument). + ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`): + Whether or not to raise an error if some of the weights from the checkpoint do not have the same size + as the weights of the model (if for instance, you are instantiating a model with 10 labels from a + checkpoint with 3 labels). + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + local_files_only(`bool`, *optional*, defaults to `False`): + Whether or not to only look at local files (i.e., do not try to download the model). + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`. + + + + 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. + kwargs (remaining dictionary of keyword arguments, *optional*): + Can be used to update the configuration object (after it being loaded) and initiate the model (e.g., + `output_attentions=True`). Behaves differently depending on whether a `config` is provided or + automatically loaded: + + - If a configuration is provided with `config`, `**kwargs` will be directly passed to the + underlying model's `__init__` method (we assume all relevant updates to the configuration have + already been done) + - If a configuration is not provided, `kwargs` will be first passed to the configuration class + initialization function ([`~PretrainedConfig.from_pretrained`]). Each key of `kwargs` that + corresponds to a configuration attribute will be used to override said attribute with the + supplied `kwargs` value. Remaining keys that do not correspond to any configuration attribute + will be passed to the underlying model's `__init__` function. + + Examples: + + ```python + >>> from transformers import BertConfig, FlaxBertModel + + >>> # Download model and configuration from huggingface.co and cache. + >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased") + >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable). + >>> model = FlaxBertModel.from_pretrained("./test/saved_model/") + >>> # Loading from a PyTorch checkpoint file instead of a PyTorch model (slower, for example purposes, not runnable). + >>> config = BertConfig.from_json_file("./pt_model/config.json") + >>> model = FlaxBertModel.from_pretrained("./pt_model/pytorch_model.bin", from_pt=True, config=config) + ```""" + from_pt = kwargs.pop("from_pt", False) + resume_download = kwargs.pop("resume_download", None) + proxies = kwargs.pop("proxies", None) + use_auth_token = kwargs.pop("use_auth_token", None) + trust_remote_code = kwargs.pop("trust_remote_code", None) + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + _do_init = kwargs.pop("_do_init", True) + subfolder = kwargs.pop("subfolder", "") + commit_hash = kwargs.pop("_commit_hash", None) + + # Not relevant for Flax Models + _ = kwargs.pop("adapter_kwargs", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if trust_remote_code is True: + logger.warning( + "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is" + " ignored." + ) + + user_agent = {"file_type": "model", "framework": "flax", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + # Load config if we don't provide a configuration + if not isinstance(config, PretrainedConfig): + config_path = config if config is not None else pretrained_model_name_or_path + config, model_kwargs = cls.config_class.from_pretrained( + config_path, + cache_dir=cache_dir, + return_unused_kwargs=True, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder=subfolder, + _from_auto=from_auto_class, + _from_pipeline=from_pipeline, + _commit_hash=commit_hash, + **kwargs, + ) + else: + model_kwargs = kwargs.copy() + + if commit_hash is None: + commit_hash = getattr(config, "_commit_hash", None) + + # Add the dtype to model_kwargs + model_kwargs["dtype"] = dtype + + # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the + # index of the files. + is_sharded = False + + # Load model + if pretrained_model_name_or_path is not None: + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + is_local = os.path.isdir(pretrained_model_name_or_path) + if os.path.isdir(pretrained_model_name_or_path): + if os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME)): + # Load from a Flax checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME) + elif os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_INDEX_NAME)): + # Load from a sharded Flax checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_INDEX_NAME) + is_sharded = True + elif is_safetensors_available() and os.path.isfile( + os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME) + ): + # Load from a safetensors checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME) + elif from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)): + # Load from a PyTorch checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME) + elif from_pt and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_INDEX_NAME) + ): + # Load from a sharded pytorch checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_INDEX_NAME) + is_sharded = True + # At this stage we don't have a weight file so we will raise an error. + elif is_safetensors_available() and os.path.isfile( + os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME) + ): + # Load from a sharded safetensors checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME) + is_sharded = True + raise NotImplementedError("Support for sharded checkpoints using safetensors is coming soon!") + elif os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)): + raise EnvironmentError( + f"Error no file named {FLAX_WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} " + "but there is a file for PyTorch weights. Use `from_pt=True` to load this model from those " + "weights." + ) + else: + raise EnvironmentError( + f"Error no file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME} found in directory " + f"{pretrained_model_name_or_path}." + ) + elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)): + archive_file = pretrained_model_name_or_path + is_local = True + elif is_remote_url(pretrained_model_name_or_path): + filename = pretrained_model_name_or_path + resolved_archive_file = download_url(pretrained_model_name_or_path) + else: + if from_pt: + filename = WEIGHTS_NAME + else: + filename = FLAX_WEIGHTS_NAME + + try: + # Load from URL or cache if already cached + cached_file_kwargs = { + "cache_dir": cache_dir, + "force_download": force_download, + "proxies": proxies, + "resume_download": resume_download, + "local_files_only": local_files_only, + "token": token, + "user_agent": user_agent, + "revision": revision, + "subfolder": subfolder, + "_raise_exceptions_for_gated_repo": False, + "_raise_exceptions_for_missing_entries": False, + "_commit_hash": commit_hash, + } + resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs) + + # Maybe the checkpoint is sharded, we try to grab the index name in this case. + if resolved_archive_file is None and filename == FLAX_WEIGHTS_NAME: + resolved_archive_file = cached_file( + pretrained_model_name_or_path, FLAX_WEIGHTS_INDEX_NAME, **cached_file_kwargs + ) + if resolved_archive_file is not None: + is_sharded = True + + # Maybe the checkpoint is pytorch sharded, we try to grab the pytorch index name in this case. + if resolved_archive_file is None and from_pt: + resolved_archive_file = cached_file( + pretrained_model_name_or_path, WEIGHTS_INDEX_NAME, **cached_file_kwargs + ) + if resolved_archive_file is not None: + is_sharded = True + + # If we still haven't found anything, look for `safetensors`. + if resolved_archive_file is None: + # No support for sharded safetensors yet, so we'll raise an error if that's all we find. + filename = SAFE_WEIGHTS_NAME + resolved_archive_file = cached_file( + pretrained_model_name_or_path, SAFE_WEIGHTS_NAME, **cached_file_kwargs + ) + + # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None + # result when internet is up, the repo and revision exist, but the file does not. + if resolved_archive_file is None: + # Otherwise, maybe there is a TF or Torch model file. We try those to give a helpful error + # message. + has_file_kwargs = { + "revision": revision, + "proxies": proxies, + "token": token, + "cache_dir": cache_dir, + "local_files_only": local_files_only, + } + if has_file(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME, **has_file_kwargs): + is_sharded = True + raise NotImplementedError( + "Support for sharded checkpoints using safetensors is coming soon!" + ) + elif has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs): + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {FLAX_WEIGHTS_NAME} but there is a file for PyTorch weights. Use `from_pt=True` to" + " load this model from those weights." + ) + elif has_file(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME, **has_file_kwargs): + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {FLAX_WEIGHTS_INDEX_NAME} but there is a sharded file for PyTorch weights. Use" + " `from_pt=True` to load this model from those weights." + ) + else: + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}." + ) + except EnvironmentError: + # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted + # to the original exception. + raise + except Exception: + # For any other exception, we throw a generic error. + raise EnvironmentError( + f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it" + " from 'https://huggingface.co/models', make sure you don't have a local directory with the" + f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a" + f" directory containing a file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}." + ) + + if is_local: + logger.info(f"loading weights file {archive_file}") + resolved_archive_file = archive_file + filename = resolved_archive_file.split(os.path.sep)[-1] + else: + logger.info(f"loading weights file {filename} from cache at {resolved_archive_file}") + else: + resolved_archive_file = None + + # We'll need to download and cache each checkpoint shard if the checkpoint is sharded. + if is_sharded: + # resolved_archive_file becomes a list of files that point to the different checkpoint shards in this case. + resolved_archive_file, _ = get_checkpoint_shard_files( + pretrained_model_name_or_path, + resolved_archive_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _commit_hash=commit_hash, + ) + + safetensors_from_pt = False + if filename == SAFE_WEIGHTS_NAME: + with safe_open(resolved_archive_file, framework="flax") as f: + safetensors_metadata = f.metadata() + if safetensors_metadata is None or safetensors_metadata.get("format") not in ["pt", "tf", "flax"]: + raise OSError( + f"The safetensors archive passed at {resolved_archive_file} does not contain the valid metadata." + " Make sure you save your model with the `save_pretrained` method." + ) + safetensors_from_pt = safetensors_metadata.get("format") == "pt" + + # init random models + model = cls(config, *model_args, _do_init=_do_init, **model_kwargs) + + if from_pt or safetensors_from_pt: + state = load_pytorch_checkpoint_in_flax_state_dict(model, resolved_archive_file, is_sharded) + else: + if is_sharded: + state = cls.load_flax_sharded_weights(resolved_archive_file) + else: + state = cls.load_flax_weights(resolved_archive_file) + # make sure all arrays are stored as jnp.arrays + # NOTE: This is to prevent a bug this will be fixed in Flax >= v0.3.4: + # https://github.com/google/flax/issues/1261 + if _do_init: + state = jax.tree_util.tree_map(jnp.array, state) + else: + # keep the params on CPU if we don't want to initialize + state = jax.tree_util.tree_map(lambda x: jax.device_put(x, jax.local_devices(backend="cpu")[0]), state) + + if "batch_stats" in state: # if flax model contains batch norm layers + # if model is base model only use model_prefix key + if ( + cls.base_model_prefix not in dict(model.params_shape_tree["params"]) + and cls.base_model_prefix in state["params"] + ): + state["params"] = state["params"][cls.base_model_prefix] + state["batch_stats"] = state["batch_stats"][cls.base_model_prefix] + + # if model is head model and we are loading weights from base model + # we initialize new params dict with base_model_prefix + if ( + cls.base_model_prefix in dict(model.params_shape_tree["params"]) + and cls.base_model_prefix not in state["params"] + ): + state = { + "params": {cls.base_model_prefix: state["params"]}, + "batch_stats": {cls.base_model_prefix: state["batch_stats"]}, + } + + else: + # if model is base model only use model_prefix key + if cls.base_model_prefix not in dict(model.params_shape_tree) and cls.base_model_prefix in state: + state = state[cls.base_model_prefix] + + # if model is head model and we are loading weights from base model + # we initialize new params dict with base_model_prefix + if cls.base_model_prefix in dict(model.params_shape_tree) and cls.base_model_prefix not in state: + state = {cls.base_model_prefix: state} + + # flatten dicts + state = flatten_dict(state) + + random_state = flatten_dict(unfreeze(model.params if _do_init else model.params_shape_tree)) + + missing_keys = model.required_params - set(state.keys()) + unexpected_keys = set(state.keys()) - model.required_params + + # Disabling warning when porting pytorch weights to flax, flax does not uses num_batches_tracked + for unexpected_key in unexpected_keys.copy(): + if "num_batches_tracked" in unexpected_key[-1]: + unexpected_keys.remove(unexpected_key) + + if missing_keys and not _do_init: + logger.warning( + f"The checkpoint {pretrained_model_name_or_path} is missing required keys: {missing_keys}. " + "Make sure to call model.init_weights to initialize the missing weights." + ) + cls._missing_keys = missing_keys + + # Mismatched keys contains tuples key/shape1/shape2 of weights in the checkpoint that have a shape not + # matching the weights in the model. + mismatched_keys = [] + for key in state.keys(): + if key in random_state and state[key].shape != random_state[key].shape: + if ignore_mismatched_sizes: + mismatched_keys.append((key, state[key].shape, random_state[key].shape)) + state[key] = random_state[key] + else: + raise ValueError( + f"Trying to load the pretrained weight for {key} failed: checkpoint has shape " + f"{state[key].shape} which is incompatible with the model shape {random_state[key].shape}. " + "Using `ignore_mismatched_sizes=True` if you really want to load this checkpoint inside this " + "model." + ) + + # add missing keys as random parameters if we are initializing + if missing_keys and _do_init: + for missing_key in missing_keys: + state[missing_key] = random_state[missing_key] + + # remove unexpected keys to not be saved again + for unexpected_key in unexpected_keys: + del state[unexpected_key] + + if len(unexpected_keys) > 0: + logger.warning( + f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" + f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are" + f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" + " with another architecture (e.g. initializing a BertForSequenceClassification model from a" + " BertForPreTraining model).\n- This IS NOT expected if you are initializing" + f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical" + " (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)." + ) + else: + logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n") + + if len(missing_keys) > 0: + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" + " TRAIN this model on a down-stream task to be able to use it for predictions and inference." + ) + elif len(mismatched_keys) == 0: + logger.info( + f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at" + f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint" + f" was trained on, you can already use {model.__class__.__name__} for predictions without further" + " training." + ) + if len(mismatched_keys) > 0: + mismatched_warning = "\n".join( + [ + f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated" + for key, shape1, shape2 in mismatched_keys + ] + ) + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not" + f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able" + " to use it for predictions and inference." + ) + + # dictionary of key: dtypes for the model params + param_dtypes = jax.tree_util.tree_map(lambda x: x.dtype, state) + # extract keys of parameters not in jnp.float32 + fp16_params = [k for k in param_dtypes if param_dtypes[k] == jnp.float16] + bf16_params = [k for k in param_dtypes if param_dtypes[k] == jnp.bfloat16] + + # raise a warning if any of the parameters are not in jnp.float32 + if len(fp16_params) > 0: + logger.warning( + f"Some of the weights of {model.__class__.__name__} were initialized in float16 precision from " + f"the model checkpoint at {pretrained_model_name_or_path}:\n{fp16_params}\n" + "You should probably UPCAST the model weights to float32 if this was not intended. " + "See [`~FlaxPreTrainedModel.to_fp32`] for further information on how to do this." + ) + + if len(bf16_params) > 0: + logger.warning( + f"Some of the weights of {model.__class__.__name__} were initialized in bfloat16 precision from " + f"the model checkpoint at {pretrained_model_name_or_path}:\n{bf16_params}\n" + "You should probably UPCAST the model weights to float32 if this was not intended. " + "See [`~FlaxPreTrainedModel.to_fp32`] for further information on how to do this." + ) + + # If it is a model with generation capabilities, attempt to load the generation config + if model.can_generate(): + try: + model.generation_config = GenerationConfig.from_pretrained( + pretrained_model_name_or_path, + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder=subfolder, + _from_auto=from_auto_class, + _from_pipeline=from_pipeline, + **kwargs, + ) + except OSError: + logger.info( + "Generation config file not found, using a generation config created from the model config." + ) + pass + + if _do_init: + # set correct parameters + model.params = unflatten_dict(state) + return model + else: + return model, unflatten_dict(state) + + def save_pretrained( + self, + save_directory: Union[str, os.PathLike], + params=None, + push_to_hub=False, + max_shard_size="10GB", + token: Optional[Union[str, bool]] = None, + safe_serialization: bool = False, + **kwargs, + ): + """ + Save a model and its configuration file to a directory, so that it can be re-loaded using the + `[`~FlaxPreTrainedModel.from_pretrained`]` class method + + Arguments: + save_directory (`str` or `os.PathLike`): + Directory to which to save. Will be created if it doesn't exist. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`): + The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size + lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`). + + + + If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard + which will be bigger than `max_shard_size`. + + + + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + kwargs (`Dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + safe_serialization (`bool`, *optional*, defaults to `False`): + Whether to save the model using `safetensors` or through msgpack. + """ + use_auth_token = kwargs.pop("use_auth_token", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if token is not None: + kwargs["token"] = token + + if os.path.isfile(save_directory): + logger.error(f"Provided path ({save_directory}) should be a directory, not a file") + return + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = self._create_repo(repo_id, **kwargs) + files_timestamps = self._get_files_timestamps(save_directory) + + # get abs dir + save_directory = os.path.abspath(save_directory) + # save config as well + self.config.architectures = [self.__class__.__name__[4:]] + + # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + custom_object_save(self, save_directory, config=self.config) + + self.config.save_pretrained(save_directory) + if self.can_generate(): + self.generation_config.save_pretrained(save_directory) + + # save model + weights_name = SAFE_WEIGHTS_NAME if safe_serialization else FLAX_WEIGHTS_NAME + output_model_file = os.path.join(save_directory, weights_name) + + shards, index = flax_shard_checkpoint(params if params is not None else self.params, max_shard_size) + # Clean the folder from a previous save + for filename in os.listdir(save_directory): + full_filename = os.path.join(save_directory, filename) + weights_no_suffix = weights_name.replace(".bin", "").replace(".safetensors", "") + if ( + filename.startswith(weights_no_suffix) + and os.path.isfile(full_filename) + and filename not in shards.keys() + ): + os.remove(full_filename) + + if index is None: + if safe_serialization: + params = params if params is not None else self.params + flat_dict = flatten_dict(params, sep=".") + safe_save_file(flat_dict, output_model_file, metadata={"format": "flax"}) + else: + with open(output_model_file, "wb") as f: + params = params if params is not None else self.params + model_bytes = to_bytes(params) + f.write(model_bytes) + + else: + save_index_file = os.path.join(save_directory, FLAX_WEIGHTS_INDEX_NAME) + # Save the index as well + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + logger.info( + f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " + f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the " + f"index located at {save_index_file}." + ) + for shard_file, shard in shards.items(): + # the shard item are unflattened, to save them we need to flatten them again + with open(os.path.join(save_directory, shard_file), mode="wb") as f: + params = unflatten_dict(shard, sep="/") + shard_bytes = to_bytes(params) + f.write(shard_bytes) + + logger.info(f"Model weights saved in {output_model_file}") + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=token, + ) + + @classmethod + def register_for_auto_class(cls, auto_class="FlaxAutoModel"): + """ + Register this class with a given auto class. This should only be used for custom models as the ones in the + library are already mapped with an auto class. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"FlaxAutoModel"`): + The auto class to register this new model with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + +# To update the docstring, we need to copy the method, otherwise we change the original docstring. +FlaxPreTrainedModel.push_to_hub = copy_func(FlaxPreTrainedModel.push_to_hub) +if FlaxPreTrainedModel.push_to_hub.__doc__ is not None: + FlaxPreTrainedModel.push_to_hub.__doc__ = FlaxPreTrainedModel.push_to_hub.__doc__.format( + object="model", object_class="FlaxAutoModel", object_files="model checkpoint" + ) + + +def overwrite_call_docstring(model_class, docstring): + # copy __call__ function to be sure docstring is changed only for this function + model_class.__call__ = copy_func(model_class.__call__) + # delete existing docstring + model_class.__call__.__doc__ = None + # set correct docstring + model_class.__call__ = add_start_docstrings_to_model_forward(docstring)(model_class.__call__) + + +def append_call_sample_docstring( + model_class, checkpoint, output_type, config_class, mask=None, revision=None, real_checkpoint=None +): + model_class.__call__ = copy_func(model_class.__call__) + model_class.__call__ = add_code_sample_docstrings( + checkpoint=checkpoint, + output_type=output_type, + config_class=config_class, + model_cls=model_class.__name__, + revision=revision, + real_checkpoint=real_checkpoint, + )(model_class.__call__) + + +def append_replace_return_docstrings(model_class, output_type, config_class): + model_class.__call__ = copy_func(model_class.__call__) + model_class.__call__ = replace_return_docstrings( + output_type=output_type, + config_class=config_class, + )(model_class.__call__) diff --git a/lib/python3.12/site-packages/transformers/modeling_gguf_pytorch_utils.py b/lib/python3.12/site-packages/transformers/modeling_gguf_pytorch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b8e6de57f5c958c978d0b4a919dd9702342fc409 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_gguf_pytorch_utils.py @@ -0,0 +1,486 @@ +# Copyright 2024 The ggml.ai team and The HuggingFace Inc. team. and pygguf author (github.com/99991) +# https://github.com/99991/pygguf +# +# 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. + +import re +from typing import NamedTuple, Optional + +import numpy as np +from tqdm.auto import tqdm + +from .integrations import ( + GGUF_CONFIG_MAPPING, + GGUF_TOKENIZER_MAPPING, + _gguf_parse_value, +) +from .utils import is_torch_available +from .utils.import_utils import is_gguf_available +from .utils.logging import get_logger + + +if is_torch_available(): + import torch + +logger = get_logger(__name__) + + +GGUF_TO_TRANSFORMERS_MAPPING = { + "ignore": { + "GGUF": { + "version": "version", + "tensor_count": "tensor_count", + "kv_count": "kv_count", + }, + "general": {"file_type": "file_type", "quantization_version": "quantization_version"}, + }, + "config": GGUF_CONFIG_MAPPING, + "tokenizer": {"tokenizer": GGUF_TOKENIZER_MAPPING["tokenizer"]}, + "tokenizer_config": {"tokenizer": GGUF_TOKENIZER_MAPPING["tokenizer_config"]}, +} + +GGUF_SUPPORTED_ARCHITECTURES = list(GGUF_TO_TRANSFORMERS_MAPPING["config"].keys()) + + +class GGUFTensor(NamedTuple): + weights: np.ndarray + name: str + metadata: dict + + +class TensorProcessor: + def __init__(self, config=None): + self.config = config or {} + + def process(self, weights, name, **kwargs): + return GGUFTensor(weights, name, {}) + + +class LlamaTensorProcessor(TensorProcessor): + def __init__(self, config=None): + super().__init__(config=config) + + def process(self, weights, name, **kwargs): + if ".attn_k." in name or ".attn_q." in name: + num_heads = self.config.get("num_attention_heads") + num_kv_heads = self.config.get("num_key_value_heads") + + if None in (num_heads, num_kv_heads): + return GGUFTensor(weights, name, {}) + if ".attn_q." in name: + weights = self._reverse_permute_weights(weights, num_heads, num_heads) + elif ".attn_k." in name: + weights = self._reverse_permute_weights(weights, num_heads, num_kv_heads) + return GGUFTensor(weights, name, {}) + + def _reverse_permute_weights( + self, weights: np.ndarray, n_head: int, num_kv_heads: Optional[int] = None + ) -> np.ndarray: + # Original permutation implementation + # https://github.com/ggerganov/llama.cpp/blob/a38b884c6c4b0c256583acfaaabdf556c62fabea/convert_hf_to_gguf.py#L1402-L1408 + if num_kv_heads is not None and n_head != num_kv_heads: + n_head = num_kv_heads + + dim = weights.shape[0] // n_head // 2 + w = weights.reshape(n_head, dim, 2, *weights.shape[1:]) + return w.swapaxes(2, 1).reshape(weights.shape) + + +class Qwen2MoeTensorProcessor(TensorProcessor): + def __init__(self, config=None): + super().__init__(config=config) + + def process(self, weights, name, **kwargs): + if "_exp" in name: + tensor_key_mapping = kwargs.get("tensor_key_mapping") + parsed_parameters = kwargs.get("parsed_parameters") + if tensor_key_mapping: + self._split_moe_expert_tensor(weights, parsed_parameters, name, tensor_key_mapping) + return GGUFTensor(weights, None, {}) + if "ffn_gate_inp_shexp" in name: + # for compatibility tensor shared_expert_gate must be (1, 2048) dim, + # quantized one is (2048) + weights = np.expand_dims(weights, axis=0) + return GGUFTensor(weights, name, {}) + + def _split_moe_expert_tensor( + self, weights: np.ndarray, parsed_parameters: dict[str, dict], name: str, tensor_key_mapping: dict + ): + # Original merge implementation + # https://github.com/ggerganov/llama.cpp/blob/master/convert_hf_to_gguf.py#L1994-L2022 + name = tensor_key_mapping[name] + w_counter = self.config.get("num_experts", 60) + for i in range(0, w_counter): + temp_name = name.replace("mlp.experts.", f"mlp.experts.{i}.") + exp_weight = weights[i] + parsed_parameters["tensors"][temp_name] = torch.from_numpy(np.copy(exp_weight)) + + +class BloomTensorProcessor(TensorProcessor): + def __init__(self, config=None): + super().__init__(config=config) + + def process(self, weights, name, **kwargs): + if "attn_qkv" in name: + num_heads = self.config["n_head"] + n_embed = self.config["hidden_size"] + if "weight" in name: + weights = self._reverse_reshape_weights(weights, num_heads, n_embed) + else: + weights = self._reverse_reshape_bias(weights, num_heads, n_embed) + return GGUFTensor(weights, name, {}) + + def _reverse_reshape_weights(self, weights: np.ndarray, n_head: int, n_embed: int): + # Original reshape implementation + # https://github.com/ggerganov/llama.cpp/blob/master/convert_hf_to_gguf.py#L972-L985 + q, k, v = np.array_split(weights, 3, axis=0) + + q = q.reshape(n_head, n_embed // n_head, n_embed) + k = k.reshape(n_head, n_embed // n_head, n_embed) + v = v.reshape(n_head, n_embed // n_head, n_embed) + qkv_weights = np.stack([q, k, v], axis=1) + + return qkv_weights.reshape(n_head * 3 * (n_embed // n_head), n_embed) + + def _reverse_reshape_bias(self, weights: np.ndarray, n_head: int, n_embed: int): + # Original reshape implementation + # https://github.com/ggerganov/llama.cpp/blob/master/convert_hf_to_gguf.py#L986-L998 + q_bias, k_bias, v_bias = np.array_split(weights, 3) + + q_bias = q_bias.reshape(n_head, n_embed // n_head) + k_bias = k_bias.reshape(n_head, n_embed // n_head) + v_bias = v_bias.reshape(n_head, n_embed // n_head) + + qkv_bias = np.stack([q_bias, k_bias, v_bias], axis=1).flatten() + return qkv_bias + + +class T5TensorProcessor(TensorProcessor): + def __init__(self, config=None): + super().__init__(config=config) + + def process(self, weights, name, **kwargs): + bid = None + for chunk in name.split("."): + if chunk.isdigit(): + bid = int(chunk) + break + return GGUFTensor(weights, name, {"bid": bid}) + + +class GPT2TensorProcessor(TensorProcessor): + def __init__(self, config=None): + super().__init__(config=config) + + def process(self, weights, name, **kwargs): + # Original transpose implementation + # https://github.com/ggerganov/llama.cpp/blob/a38b884c6c4b0c256583acfaaabdf556c62fabea/convert_hf_to_gguf.py#L2060-L2061 + if ( + "attn_qkv.weight" in name + or "ffn_down.weight" in name + or "ffn_up.weight" in name + or "attn_output.weight" in name + ): + weights = weights.T + + # Handle special case for output.weight + if name == "output.weight": + # output.weight has conflicts with attn_output.weight in name checking + # Store the tensor directly and signal to skip further processing + name = "lm_head.weight" + parsed_parameters = kwargs.get("parsed_parameters", {}) + parsed_parameters["tensors"][name] = torch.from_numpy(np.copy(weights)) + name = None # Signal to skip further processing + return GGUFTensor(weights, name, {}) + + +class MambaTensorProcessor(TensorProcessor): + def __init__(self, config=None): + super().__init__(config=config) + + def process(self, weights, name, **kwargs): + if "ssm_conv1d.weight" in name: + # for compatibility tensor ssm_conv1d must be (5120, 1, 4]) dim, + # quantized one is (5120, 4) + weights = np.expand_dims(weights, axis=1) + if "ssm_a" in name: + # Original exponential implementation + # https://github.com/ggerganov/llama.cpp/blob/master/convert_hf_to_gguf.py#L2975-L2977 + weights = np.log(-weights) + return GGUFTensor(weights, name, {}) + + +class NemotronTensorProcessor(TensorProcessor): + def __init__(self, config=None): + super().__init__(config=config) + + # ref : https://github.com/ggerganov/llama.cpp/blob/master/convert_hf_to_gguf.py#L4666 + def process(self, weights, name, **kwargs): + if "norm.weight" in name: + weights = weights - 1 + return GGUFTensor(weights, name, {}) + + +class Gemma2TensorProcessor(TensorProcessor): + def __init__(self, config=None): + super().__init__(config=config) + + # ref: https://github.com/ggerganov/llama.cpp/blob/d79d8f39b4da6deca4aea8bf130c6034c482b320/convert_hf_to_gguf.py#L3191 + # ref: https://github.com/huggingface/transformers/blob/fc37f38915372c15992b540dfcbbe00a916d4fc6/src/transformers/models/gemma/modeling_gemma.py#L89 + def process(self, weights, name, **kwargs): + if "norm.weight" in name: + weights = weights - 1 + return GGUFTensor(weights, name, {}) + + +TENSOR_PROCESSORS = { + "llama": LlamaTensorProcessor, + "qwen2moe": Qwen2MoeTensorProcessor, + "bloom": BloomTensorProcessor, + "t5": T5TensorProcessor, + "t5encoder": T5TensorProcessor, + "gpt2": GPT2TensorProcessor, + "mamba": MambaTensorProcessor, + "nemotron": NemotronTensorProcessor, + "gemma2": Gemma2TensorProcessor, +} + + +def read_field(reader, field): + value = reader.fields[field] + return [_gguf_parse_value(value.parts[_data_index], value.types) for _data_index in value.data] + + +# modified from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/model_loader/loader.py#L1115-L1147 +def get_gguf_hf_weights_map( + hf_model, + model_type: Optional[str] = None, + num_layers: Optional[int] = None, + qual_name: str = "", +): + """ + GGUF uses this naming convention for their tensors from HF checkpoint: + `blk.N.BB.weight` and `blk.N.BB.bias` + where N signifies the block number of a layer, and BB signifies the + attention/mlp layer components. + See "Standardized tensor names" in + https://github.com/ggerganov/ggml/blob/master/docs/gguf.md for details. + """ + if is_gguf_available() and is_torch_available(): + from gguf import MODEL_ARCH_NAMES, get_tensor_name_map + else: + logger.error( + "Loading a GGUF checkpoint in PyTorch, requires both PyTorch and GGUF>=0.10.0 to be installed. Please see " + "https://pytorch.org/ and https://github.com/ggerganov/llama.cpp/tree/master/gguf-py for installation instructions." + ) + raise ImportError("Please install torch and gguf>=0.10.0 to load a GGUF checkpoint in PyTorch.") + + model_type = hf_model.config.model_type if model_type is None else model_type + num_layers = hf_model.config.num_hidden_layers if num_layers is None else num_layers + # hack: ggufs have a different name for cohere + if model_type == "cohere": + model_type = "command-r" + elif model_type == "qwen2_moe": + model_type = "qwen2moe" + arch = None + for key, value in MODEL_ARCH_NAMES.items(): + if value == model_type: + arch = key + break + if arch is None: + raise NotImplementedError( + f"Unknown gguf model_type: {model_type} in gguf-py. " + "This might because you're using an outdated version of gguf-py package, " + "you can install `gguf` package from source refer to " + "https://github.com/ggerganov/llama.cpp/tree/master/gguf-py#development" + ) + name_map = get_tensor_name_map(arch, num_layers) + + # Use a dummy conversion to get the mapping, because + # hf => gguf and gguf => hf mappings are reversed + gguf_to_hf_name_map = {} + state_dict = hf_model.state_dict() + for hf_name in state_dict.keys(): + # An exception for qwen2moe model, where the expert layers are packed + if model_type == "qwen2moe" and "mlp.experts." in hf_name: + hf_name = re.sub(r"mlp.experts.\d+.", "mlp.experts.", hf_name) + + name, suffix = hf_name, "" + if hf_name.endswith(".weight") or hf_name.endswith(".bias"): + name, suffix = hf_name.rsplit(".", 1) + suffix = "." + suffix + + gguf_name = name_map.get_name(name) + if gguf_name is None: + continue + + gguf_to_hf_name_map[gguf_name + suffix] = qual_name + hf_name + + # Some model like Bloom converted from BloomModel instead of BloomForCausalLM + # Therefore, we need to check submodule as well to get a correct mapping + if named_children := hf_model.named_children(): + for name, child in named_children: + sub_map = get_gguf_hf_weights_map(child, model_type, num_layers, qual_name=f"{qual_name}{name}.") + # Ignore the keys that are already in the main map to avoid overwriting + sub_map = {k: v for k, v in sub_map.items() if k not in gguf_to_hf_name_map} + gguf_to_hf_name_map.update(sub_map) + + return gguf_to_hf_name_map + + +def load_gguf_checkpoint(gguf_checkpoint_path, return_tensors=False, model_to_load=None): + """ + Load a GGUF file and return a dictionary of parsed parameters containing tensors, the parsed + tokenizer and config attributes. + + Args: + gguf_checkpoint_path (`str`): + The path the to GGUF file to load + return_tensors (`bool`, defaults to `False`): + Whether to read the tensors from the file and return them. Not doing so is faster + and only loads the metadata in memory. + """ + if is_gguf_available() and is_torch_available(): + from gguf import GGUFReader, dequantize + else: + logger.error( + "Loading a GGUF checkpoint in PyTorch, requires both PyTorch and GGUF>=0.10.0 to be installed. Please see " + "https://pytorch.org/ and https://github.com/ggerganov/llama.cpp/tree/master/gguf-py for installation instructions." + ) + raise ImportError("Please install torch and gguf>=0.10.0 to load a GGUF checkpoint in PyTorch.") + + reader = GGUFReader(gguf_checkpoint_path) + fields = reader.fields + reader_keys = list(fields.keys()) + + parsed_parameters = {k: {} for k in GGUF_TO_TRANSFORMERS_MAPPING} + + architecture = read_field(reader, "general.architecture")[0] + model_name = read_field(reader, "general.name") + + updated_architecture = None + # in llama.cpp mistral models use the same architecture as llama. We need + # to add this patch to ensure things work correctly on our side. + if "llama" in architecture and "mistral" in model_name: + updated_architecture = "mistral" + # FIXME: Currently this implementation is only for flan-t5 architecture. + # It needs to be developed for supporting legacy t5. + elif "t5" in architecture or "t5encoder" in architecture: + parsed_parameters["config"]["is_gated_act"] = True + if "t5encoder" in architecture: + parsed_parameters["config"]["architectures"] = ["T5EncoderModel"] + updated_architecture = "t5" + else: + updated_architecture = architecture + + if "qwen2moe" in architecture: + updated_architecture = "qwen2_moe" + + # For stablelm architecture, we need to set qkv_bias and use_parallel_residual from tensors + # If `qkv_bias=True`, qkv_proj with bias will be present in the tensors + # If `use_parallel_residual=False`, ffn_norm will be present in the tensors + if "stablelm" in architecture: + attn_bias_name = {"attn_q.bias", "attn_k.bias", "attn_v.bias"} + ffn_norm_name = "ffn_norm" + qkv_bias = any(bias_name in tensor.name for tensor in reader.tensors for bias_name in attn_bias_name) + use_parallel_residual = any(ffn_norm_name in tensor.name for tensor in reader.tensors) + parsed_parameters["config"]["use_qkv_bias"] = qkv_bias + parsed_parameters["config"]["use_parallel_residual"] = not use_parallel_residual + + if architecture not in GGUF_SUPPORTED_ARCHITECTURES and updated_architecture not in GGUF_SUPPORTED_ARCHITECTURES: + raise ValueError(f"GGUF model with architecture {architecture} is not supported yet.") + + # Handle tie_word_embeddings, if lm_head.weight is not present in tensors, + # tie_word_embeddings is true otherwise false + exceptions = ["falcon", "bloom"] + parsed_parameters["config"]["tie_word_embeddings"] = ( + all("output.weight" != tensor.name for tensor in reader.tensors) or architecture in exceptions + ) + + # List all key-value pairs in a columnized format + for gguf_key, field in reader.fields.items(): + gguf_key = gguf_key.replace(architecture, updated_architecture) + split = gguf_key.split(".") + prefix = split[0] + config_key = ".".join(split[1:]) + + value = [_gguf_parse_value(field.parts[_data_index], field.types) for _data_index in field.data] + + if len(value) == 1: + value = value[0] + + if isinstance(value, str) and architecture in value: + value = value.replace(architecture, updated_architecture) + + for parameter in GGUF_TO_TRANSFORMERS_MAPPING: + parameter_renames = GGUF_TO_TRANSFORMERS_MAPPING[parameter] + if prefix in parameter_renames and config_key in parameter_renames[prefix]: + renamed_config_key = parameter_renames[prefix][config_key] + if renamed_config_key == -1: + continue + + if renamed_config_key is not None: + parsed_parameters[parameter][renamed_config_key] = value + + if gguf_key in reader_keys: + reader_keys.remove(gguf_key) + + if gguf_key in reader_keys: + logger.info(f"Some keys were not parsed and added into account {gguf_key} | {value}") + + # retrieve config vocab_size from tokenizer + # Please refer to https://github.com/huggingface/transformers/issues/32526 for more details + if "vocab_size" not in parsed_parameters["config"]: + tokenizer_parameters = parsed_parameters["tokenizer"] + if "tokens" in tokenizer_parameters: + parsed_parameters["config"]["vocab_size"] = len(tokenizer_parameters["tokens"]) + else: + logger.warning( + "Can't find a way to retrieve missing config vocab_size from tokenizer parameters. " + "This will use default value from model config class and cause unexpected behavior." + ) + + if return_tensors: + parsed_parameters["tensors"] = {} + + tensor_key_mapping = get_gguf_hf_weights_map(model_to_load) + config = parsed_parameters.get("config", {}) + + ProcessorClass = TENSOR_PROCESSORS.get(architecture, TensorProcessor) + processor = ProcessorClass(config=config) + + for tensor in tqdm(reader.tensors, desc="Converting and de-quantizing GGUF tensors..."): + name = tensor.name + weights = dequantize(tensor.data, tensor.tensor_type) + + result = processor.process( + weights=weights, + name=name, + tensor_key_mapping=tensor_key_mapping, + parsed_parameters=parsed_parameters, + ) + + weights = result.weights + name = result.name + + if name not in tensor_key_mapping: + continue + + name = tensor_key_mapping[name] + + parsed_parameters["tensors"][name] = torch.from_numpy(np.copy(weights)) + + if len(reader_keys) > 0: + logger.info(f"Some keys of the GGUF file were not considered: {reader_keys}") + + return parsed_parameters diff --git a/lib/python3.12/site-packages/transformers/modeling_outputs.py b/lib/python3.12/site-packages/transformers/modeling_outputs.py new file mode 100644 index 0000000000000000000000000000000000000000..60a3642f87c6b64771584b361ff6aef3ac06cb87 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_outputs.py @@ -0,0 +1,1753 @@ +# Copyright 2020 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. + +import warnings +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch + +from .utils import ModelOutput + + +@dataclass +class BaseModelOutput(ModelOutput): + """ + Base class for model's outputs, with potential hidden states and attentions. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class BaseModelOutputWithNoAttention(ModelOutput): + """ + Base class for model's outputs, with potential hidden states. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, num_channels, height, width)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class BaseModelOutputWithPooling(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`): + Last layer hidden-state of the first token of the sequence (classification token) after further processing + through the layers used for the auxiliary pretraining task. E.g. for BERT-family of models, this returns + the classification token after processing through a linear layer and a tanh activation function. The linear + layer weights are trained from the next sentence prediction (classification) objective during pretraining. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + pooler_output: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class BaseModelOutputWithPoolingAndNoAttention(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`): + Last layer hidden-state after a pooling operation on the spatial dimensions. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, num_channels, height, width)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + pooler_output: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class BaseModelOutputWithPast(ModelOutput): + """ + Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding). + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if + `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads, + encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if + `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` + input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class BaseModelOutputWithCrossAttentions(ModelOutput): + """ + Base class for model's outputs, with potential hidden states and attentions. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class BaseModelOutputWithPoolingAndCrossAttentions(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`): + Last layer hidden-state of the first token of the sequence (classification token) after further processing + through the layers used for the auxiliary pretraining task. E.g. for BERT-family of models, this returns + the classification token after processing through a linear layer and a tanh activation function. The linear + layer weights are trained from the next sentence prediction (classification) objective during pretraining. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if + `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads, + encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if + `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` + input) to speed up sequential decoding. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + pooler_output: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class BaseModelOutputWithPastAndCrossAttentions(ModelOutput): + """ + Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding). + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if + `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads, + encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if + `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` + input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class MoECausalLMOutputWithPast(ModelOutput): + """ + Base class for causal language model (or autoregressive) outputs as well as Mixture of Expert's router hidden + states terms, to train a MoE model. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + z_loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided): + z_loss for the sparse modules. + aux_loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided): + aux_loss for the sparse modules. + router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`. + + Router logits of the encoder model, useful to compute the auxiliary loss and the z_loss for the sparse + modules. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + z_loss: Optional[torch.FloatTensor] = None + aux_loss: Optional[torch.FloatTensor] = None + router_logits: Optional[Tuple[torch.FloatTensor]] = None + + +@dataclass +class MoEModelOutput(ModelOutput): + """ + Base class for model's outputs, with potential hidden states and attentions. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + router_probs (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_probs=True` and `config.add_router_probs=True` is passed or when `config.output_router_probs=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`. + + Raw router probabilities that are computed by MoE routers, these terms are used to compute the auxiliary + loss and the z_loss for Mixture of Experts models. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + router_probs: Optional[Tuple[torch.FloatTensor]] = None + + +@dataclass +class MoeModelOutputWithPast(ModelOutput): + """ + Base class for model's outputs, with potential hidden states and attentions. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if + `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads, + encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if + `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` + input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_probs=True` and `config.add_router_probs=True` is passed or when `config.output_router_probs=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`. + + Raw router logtis (post-softmax) that are computed by MoE routers, these terms are used to compute the auxiliary + loss for Mixture of Experts models. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + router_logits: Optional[Tuple[torch.FloatTensor]] = None + + +@dataclass +class MoeCausalLMOutputWithPast(ModelOutput): + """ + Base class for causal language model (or autoregressive) with mixture of experts outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + + aux_loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided): + aux_loss for the sparse modules. + + router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_probs=True` and `config.add_router_probs=True` is passed or when `config.output_router_probs=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`. + + Raw router logtis (post-softmax) that are computed by MoE routers, these terms are used to compute the auxiliary + loss for Mixture of Experts models. + + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + aux_loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + router_logits: Optional[Tuple[torch.FloatTensor]] = None + + +@dataclass +class MoEModelOutputWithPastAndCrossAttentions(ModelOutput): + """ + Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding) as well as + Mixture of Expert's router hidden states terms, to train a MoE model. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if + `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads, + encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if + `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` + input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + router_probs (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_probs=True` and `config.add_router_probs=True` is passed or when `config.output_router_probs=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`. + + Raw router probabilities that are computed by MoE routers, these terms are used to compute the auxiliary + loss and the z_loss for Mixture of Experts models. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + router_probs: Optional[Tuple[torch.FloatTensor]] = None + + +@dataclass +class Seq2SeqModelOutput(ModelOutput): + """ + Base class for model encoder's outputs that also contains : pre-computed hidden states that can speed up sequential + decoding. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the decoder of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the optional initial embedding outputs. + decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the optional initial embedding outputs. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_last_hidden_state: Optional[torch.FloatTensor] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class Seq2SeqMoEModelOutput(ModelOutput): + """ + Base class for model encoder's outputs that also contains : pre-computed hidden states that can speed up sequential + decoding. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the decoder of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the optional initial embedding outputs. + decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + decoder_router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`. + + Router logits of the decoder model, useful to compute the auxiliary loss for Mixture of Experts models. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the optional initial embedding outputs. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + encoder_router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`. + + Router logits of the encoder model, useful to compute the auxiliary loss and the z_loss for the sparse + modules. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_router_logits: Optional[Tuple[torch.FloatTensor]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_last_hidden_state: Optional[torch.FloatTensor] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_router_logits: Optional[Tuple[torch.FloatTensor]] = None + + +@dataclass +class CausalLMOutput(ModelOutput): + """ + Base class for causal language model (or autoregressive) outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class CausalLMOutputWithPast(ModelOutput): + """ + Base class for causal language model (or autoregressive) outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class CausalLMOutputWithCrossAttentions(ModelOutput): + """ + Base class for causal language model (or autoregressive) outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Cross attentions weights after the attention softmax, used to compute the weighted average in the + cross-attention heads. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `torch.FloatTensor` tuples of length `config.n_layers`, with each tuple containing the cached key, + value states of the self-attention and the cross-attention layers if model is used in encoder-decoder + setting. Only relevant if `config.is_decoder = True`. + + Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class SequenceClassifierOutputWithPast(ModelOutput): + """ + Base class for outputs of sentence classification models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) + + Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class MaskedLMOutput(ModelOutput): + """ + Base class for masked language models outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Masked language modeling (MLM) loss. + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class Seq2SeqLMOutput(ModelOutput): + """ + Base class for sequence-to-sequence language models outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss. + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_last_hidden_state: Optional[torch.FloatTensor] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class Seq2SeqMoEOutput(ModelOutput): + """ + Base class for sequence-to-sequence language models outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling loss. + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + decoder_router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`. + + Router logits of the decoder model, useful to compute the auxiliary loss for Mixture of Experts models. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + encoder_router_logits (`tuple(torch.FloatTensor)`, *optional*, returned when `output_router_logits=True` is passed or when `config.add_router_probs=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, num_experts)`. + + Router logits of the encoder model, useful to compute the auxiliary loss and z_loss for Mixture of Experts + models. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + encoder_z_loss: Optional[torch.FloatTensor] = None + decoder_z_loss: Optional[torch.FloatTensor] = None + encoder_aux_loss: Optional[torch.FloatTensor] = None + decoder_aux_loss: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_router_logits: Optional[Tuple[torch.FloatTensor]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_last_hidden_state: Optional[torch.FloatTensor] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_router_logits: Optional[Tuple[torch.FloatTensor]] = None + + +@dataclass +class NextSentencePredictorOutput(ModelOutput): + """ + Base class for outputs of models predicting if two sentences are consecutive or not. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `next_sentence_label` is provided): + Next sequence prediction (classification) loss. + logits (`torch.FloatTensor` of shape `(batch_size, 2)`): + Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation + before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class SequenceClassifierOutput(ModelOutput): + """ + Base class for outputs of sentence classification models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class Seq2SeqSequenceClassifierOutput(ModelOutput): + """ + Base class for outputs of sequence-to-sequence sentence classification models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `label` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_last_hidden_state: Optional[torch.FloatTensor] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class MultipleChoiceModelOutput(ModelOutput): + """ + Base class for outputs of multiple choice models. + + Args: + loss (`torch.FloatTensor` of shape *(1,)*, *optional*, returned when `labels` is provided): + Classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, num_choices)`): + *num_choices* is the second dimension of the input tensors. (see *input_ids* above). + + Classification scores (before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class TokenClassifierOutput(ModelOutput): + """ + Base class for outputs of token classification models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided) : + Classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`): + Classification scores (before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class QuestionAnsweringModelOutput(ModelOutput): + """ + Base class for outputs of question answering models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Total span extraction loss is the sum of a Cross-Entropy for the start and end positions. + start_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Span-start scores (before SoftMax). + end_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Span-end scores (before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + start_logits: Optional[torch.FloatTensor] = None + end_logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class Seq2SeqQuestionAnsweringModelOutput(ModelOutput): + """ + Base class for outputs of sequence-to-sequence question answering models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Total span extraction loss is the sum of a Cross-Entropy for the start and end positions. + start_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Span-start scores (before SoftMax). + end_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): + Span-end scores (before SoftMax). + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + loss: Optional[torch.FloatTensor] = None + start_logits: Optional[torch.FloatTensor] = None + end_logits: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_last_hidden_state: Optional[torch.FloatTensor] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class SemanticSegmenterOutput(ModelOutput): + """ + Base class for outputs of semantic segmentation models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels, logits_height, logits_width)`): + Classification scores for each pixel. + + + + The logits returned do not necessarily have the same size as the `pixel_values` passed as inputs. This is + to avoid doing two interpolations and lose some quality when a user needs to resize the logits to the + original image size as post-processing. You should always check your logits shape and resize as needed. + + + + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, patch_size, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, patch_size, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class ImageClassifierOutput(ModelOutput): + """ + Base class for outputs of image classification models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states + (also called feature maps) of the model at the output of each stage. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, patch_size, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class ImageClassifierOutputWithNoAttention(ModelOutput): + """ + Base class for outputs of image classification models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each stage) of shape `(batch_size, num_channels, height, width)`. Hidden-states (also + called feature maps) of the model at the output of each stage. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class DepthEstimatorOutput(ModelOutput): + """ + Base class for outputs of depth estimation models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + predicted_depth (`torch.FloatTensor` of shape `(batch_size, height, width)`): + Predicted depth for each pixel. + + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, num_channels, height, width)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, patch_size, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + predicted_depth: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class ImageSuperResolutionOutput(ModelOutput): + """ + Base class for outputs of image super resolution models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Reconstruction loss. + reconstruction (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Reconstructed images, possibly upscaled. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states + (also called feature maps) of the model at the output of each stage. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, patch_size, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + reconstruction: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class Wav2Vec2BaseModelOutput(ModelOutput): + """ + Base class for models that have been trained with the Wav2Vec2 loss objective. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + extract_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, conv_dim[-1])`): + Sequence of extracted feature vectors of the last convolutional layer of the model. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of + shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + extract_features: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class XVectorOutput(ModelOutput): + """ + Output type of [`Wav2Vec2ForXVector`]. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification loss. + logits (`torch.FloatTensor` of shape `(batch_size, config.xvector_output_dim)`): + Classification hidden states before AMSoftmax. + embeddings (`torch.FloatTensor` of shape `(batch_size, config.xvector_output_dim)`): + Utterance embeddings used for vector similarity-based retrieval. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of + shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: Optional[torch.FloatTensor] = None + logits: Optional[torch.FloatTensor] = None + embeddings: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class BackboneOutput(ModelOutput): + """ + Base class for outputs of backbones. + + Args: + feature_maps (`tuple(torch.FloatTensor)` of shape `(batch_size, num_channels, height, width)`): + Feature maps of the stages. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of + shape `(batch_size, sequence_length, hidden_size)` or `(batch_size, num_channels, height, width)`, + depending on the backbone. + + Hidden-states of the model at the output of each stage plus the initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. Only applicable if the backbone uses attention. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + feature_maps: Optional[Tuple[torch.FloatTensor]] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class BaseModelOutputWithPoolingAndProjection(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`): + Last layer hidden-state of the first token of the sequence (classification token) after further processing + through the layers used for the auxiliary pretraining task. E.g. for BERT-family of models, this returns + the classification token after processing through a linear layer and a tanh activation function. The linear + layer weights are trained from the next sentence prediction (classification) objective during pretraining. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + projection_state (`tuple(torch.FloatTensor)`, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` of shape `(batch_size,config.project_dim)`. + + Text embeddings before the projection layer, used to mimic the last hidden state of the teacher encoder. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + pooler_output: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + projection_state: Optional[Tuple[torch.FloatTensor]] = None + + +@dataclass +class Seq2SeqSpectrogramOutput(ModelOutput): + """ + Base class for sequence-to-sequence spectrogram outputs. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Spectrogram generation loss. + spectrogram (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_bins)`): + The predicted spectrogram. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + loss: Optional[torch.FloatTensor] = None + spectrogram: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_last_hidden_state: Optional[torch.FloatTensor] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + +@dataclass +class Seq2SeqTSModelOutput(ModelOutput): + """ + Base class for time series model's encoder outputs that also contains pre-computed hidden states that can speed up + sequential decoding. + + Args: + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the decoder of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the optional initial embedding outputs. + decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the optional initial embedding outputs. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + loc (`torch.FloatTensor` of shape `(batch_size,)` or `(batch_size, input_size)`, *optional*): + Shift values of each time series' context window which is used to give the model inputs of the same + magnitude and then used to shift back to the original magnitude. + scale (`torch.FloatTensor` of shape `(batch_size,)` or `(batch_size, input_size)`, *optional*): + Scaling values of each time series' context window which is used to give the model inputs of the same + magnitude and then used to rescale back to the original magnitude. + static_features (`torch.FloatTensor` of shape `(batch_size, feature size)`, *optional*): + Static features of each time series' in a batch which are copied to the covariates at inference time. + """ + + last_hidden_state: Optional[torch.FloatTensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_last_hidden_state: Optional[torch.FloatTensor] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + loc: Optional[torch.FloatTensor] = None + scale: Optional[torch.FloatTensor] = None + static_features: Optional[torch.FloatTensor] = None + + +@dataclass +class Seq2SeqTSPredictionOutput(ModelOutput): + """ + Base class for time series model's decoder outputs that also contain the loss as well as the parameters of the + chosen distribution. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when a `future_values` is provided): + Distributional loss. + params (`torch.FloatTensor` of shape `(batch_size, num_samples, num_params)`): + Parameters of the chosen distribution. + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape + `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape + `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + loc (`torch.FloatTensor` of shape `(batch_size,)` or `(batch_size, input_size)`, *optional*): + Shift values of each time series' context window which is used to give the model inputs of the same + magnitude and then used to shift back to the original magnitude. + scale (`torch.FloatTensor` of shape `(batch_size,)` or `(batch_size, input_size)`, *optional*): + Scaling values of each time series' context window which is used to give the model inputs of the same + magnitude and then used to rescale back to the original magnitude. + static_features (`torch.FloatTensor` of shape `(batch_size, feature size)`, *optional*): + Static features of each time series' in a batch which are copied to the covariates at inference time. + """ + + loss: Optional[torch.FloatTensor] = None + params: Optional[Tuple[torch.FloatTensor]] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + cross_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_last_hidden_state: Optional[torch.FloatTensor] = None + encoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + encoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + loc: Optional[torch.FloatTensor] = None + scale: Optional[torch.FloatTensor] = None + static_features: Optional[torch.FloatTensor] = None + + +@dataclass +class SampleTSPredictionOutput(ModelOutput): + """ + Base class for time series model's predictions outputs that contains the sampled values from the chosen + distribution. + + Args: + sequences (`torch.FloatTensor` of shape `(batch_size, num_samples, prediction_length)` or `(batch_size, num_samples, prediction_length, input_size)`): + Sampled values from the chosen distribution. + """ + + sequences: Optional[torch.FloatTensor] = None + + +@dataclass +class MaskedImageModelingOutput(ModelOutput): + """ + Base class for outputs of masked image completion / in-painting models. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `bool_masked_pos` is provided): + Reconstruction loss. + reconstruction (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Reconstructed / completed images. + hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or + when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + + one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states + (also called feature maps) of the model at the output of each stage. + attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when + `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, patch_size, + sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in + the self-attention heads. + """ + + loss: Optional[torch.FloatTensor] = None + reconstruction: Optional[torch.FloatTensor] = None + hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None + attentions: Optional[Tuple[torch.FloatTensor, ...]] = None + + @property + def logits(self): + warnings.warn( + "logits attribute is deprecated and will be removed in version 5 of Transformers." + " Please use the reconstruction attribute to retrieve the final output instead.", + FutureWarning, + ) + return self.reconstruction diff --git a/lib/python3.12/site-packages/transformers/modeling_rope_utils.py b/lib/python3.12/site-packages/transformers/modeling_rope_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9a9cfd072b2788ebde08f26e3337ef6eb933d0c2 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_rope_utils.py @@ -0,0 +1,655 @@ +# 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. + +import math +from functools import wraps +from typing import Optional + +from .configuration_utils import PretrainedConfig +from .utils import is_torch_available, logging + + +logger = logging.get_logger(__name__) + + +if is_torch_available(): + import torch + + +def dynamic_rope_update(rope_forward): + """ + Decorator function to update the RoPE parameters in the forward pass, if the model is using a dynamic RoPE + (i.e. a RoPE implementation that may recompute its frequencies in the forward pass). + + Args: + rope_forward (Callable): + The forward pass of the RoPE implementation. + + Returns: + The decorated forward pass. + """ + + def longrope_frequency_update(self, position_ids, device): + """Longrope uses long factor if sequence is larger than original pretraining length, short otherwise.""" + seq_len = torch.max(position_ids) + 1 + if hasattr(self.config, "original_max_position_embeddings"): + original_max_position_embeddings = self.config.original_max_position_embeddings + else: + original_max_position_embeddings = self.config.max_position_embeddings + if seq_len > original_max_position_embeddings: + if not hasattr(self, "long_inv_freq"): + self.long_inv_freq, _ = self.rope_init_fn( + self.config, device, seq_len=original_max_position_embeddings + 1 + ) + self.register_buffer("inv_freq", self.long_inv_freq, persistent=False) + else: + # This .to() is needed if the model has been moved to a device after being initialized (because + # the buffer is automatically moved, but not the original copy) + self.original_inv_freq = self.original_inv_freq.to(device) + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + + def dynamic_frequency_update(self, position_ids, device): + """ + dynamic RoPE layers should recompute `inv_freq` in the following situations: + 1 - growing beyond the cached sequence length (allow scaling) + 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) + """ + seq_len = torch.max(position_ids) + 1 + if seq_len > self.max_seq_len_cached: # growth + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len) + self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation + self.max_seq_len_cached = seq_len + + if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + # This .to() is needed if the model has been moved to a device after being initialized (because + # the buffer is automatically moved, but not the original copy) + self.original_inv_freq = self.original_inv_freq.to(device) + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + self.max_seq_len_cached = self.original_max_seq_len + + @wraps(rope_forward) + def wrapper(self, x, position_ids): + if "dynamic" in self.rope_type: + dynamic_frequency_update(self, position_ids, device=x.device) + elif self.rope_type == "longrope": + longrope_frequency_update(self, position_ids, device=x.device) + return rope_forward(self, x, position_ids) + + return wrapper + + +def _compute_default_rope_parameters( + config: Optional[PretrainedConfig] = None, + device: Optional["torch.device"] = None, + seq_len: Optional[int] = None, + **rope_kwargs, +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + if config is not None and len(rope_kwargs) > 0: + raise ValueError( + "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in " + f"`_compute_default_rope_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}" + ) + if len(rope_kwargs) > 0: + base = rope_kwargs["base"] + dim = rope_kwargs["dim"] + elif config is not None: + base = config.rope_theta + partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + dim = int(head_dim * partial_rotary_factor) + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)) + return inv_freq, attention_factor + + +def _compute_linear_scaling_rope_parameters( + config: Optional[PretrainedConfig] = None, + device: Optional["torch.device"] = None, + seq_len: Optional[int] = None, + **rope_kwargs, +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies with linear scaling. Credits to the Reddit user /u/kaiokendev + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + if config is not None and len(rope_kwargs) > 0: + raise ValueError( + "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in " + f"`_compute_linear_scaling_rope_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}" + ) + if len(rope_kwargs) > 0: + factor = rope_kwargs["factor"] + elif config is not None: + factor = config.rope_scaling["factor"] + + # Gets the default RoPE parameters + inv_freq, attention_factor = _compute_default_rope_parameters(config, device, seq_len, **rope_kwargs) + + # Then applies linear scaling to the frequencies. + # NOTE: originally, scaling was applied to the position_ids. However, we get `embs = inv_freq @ position_ids`, so + # applying scaling to the inverse frequencies is equivalent. + inv_freq /= factor + return inv_freq, attention_factor + + +def _compute_dynamic_ntk_parameters( + config: Optional[PretrainedConfig] = None, + device: Optional["torch.device"] = None, + seq_len: Optional[int] = None, + **rope_kwargs, +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies with NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length, used to update the dynamic RoPE at inference time. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + # TODO (joao): use the new `original_max_position_embeddings` from rope_scaling + if config is not None and len(rope_kwargs) > 0: + raise ValueError( + "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in " + f"`_compute_dynamic_ntk_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}" + ) + if len(rope_kwargs) > 0: + base = rope_kwargs["base"] + dim = rope_kwargs["dim"] + max_position_embeddings = rope_kwargs["max_position_embeddings"] + factor = rope_kwargs["factor"] + elif config is not None: + base = config.rope_theta + partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + dim = int(head_dim * partial_rotary_factor) + max_position_embeddings = config.max_position_embeddings + factor = config.rope_scaling["factor"] + + attention_factor = 1.0 # Unused in this type of RoPE + + # seq_len: default to max_position_embeddings, e.g. at init time + seq_len = seq_len if seq_len is not None and seq_len > max_position_embeddings else max_position_embeddings + + # Compute the inverse frequencies + base = base * ((factor * seq_len / max_position_embeddings) - (factor - 1)) ** (dim / (dim - 2)) + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)) + return inv_freq, attention_factor + + +def _compute_yarn_parameters( + config: PretrainedConfig, device: "torch.device", seq_len: Optional[int] = None, **rope_kwargs +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies with NTK scaling. Please refer to the + [original paper](https://arxiv.org/abs/2309.00071) + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin. + """ + # No need to keep BC with yarn, unreleased when this new pattern was created. + if len(rope_kwargs) > 0: + raise ValueError( + f"Unexpected arguments: `**rope_kwargs` should be unset in `_compute_yarn_parameters`, got {rope_kwargs}" + ) + + base = config.rope_theta + partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + dim = int(head_dim * partial_rotary_factor) + factor = config.rope_scaling["factor"] + attention_factor = config.rope_scaling.get("attention_factor") + mscale = config.rope_scaling.get("mscale") + mscale_all_dim = config.rope_scaling.get("mscale_all_dim") + + # NOTE: DeekSeek-V3 (and potentially other models) modify `max_position_embeddings` and have a + # `original_max_position_embeddings` field containing the pretrained value. They use the ratio between these two + # values to compute the default attention scaling factor, instead of using `factor`. + if "original_max_position_embeddings" in config.rope_scaling: + original_max_position_embeddings = config.rope_scaling["original_max_position_embeddings"] + factor = config.max_position_embeddings / original_max_position_embeddings + else: + original_max_position_embeddings = config.max_position_embeddings + + def get_mscale(scale, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + # Sets the attention factor as suggested in the paper + if attention_factor is None: + if mscale and mscale_all_dim: + attention_factor = float(get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dim)) + else: + attention_factor = get_mscale(factor) + + # Optional config options + # beta_fast/beta_slow: as suggested in the paper, default to 32/1 (correspondingly) + beta_fast = config.rope_scaling.get("beta_fast") or 32 + beta_slow = config.rope_scaling.get("beta_slow") or 1 + + # Compute the inverse frequencies + def find_correction_dim(num_rotations, dim, base, max_position_embeddings): + """Inverse dimension formula to find the dimension based on the number of rotations""" + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base)) + + def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings): + """Find dimension range bounds based on rotations""" + low = math.floor(find_correction_dim(low_rot, dim, base, max_position_embeddings)) + high = math.ceil(find_correction_dim(high_rot, dim, base, max_position_embeddings)) + return max(low, 0), min(high, dim - 1) + + def linear_ramp_factor(min, max, dim): + if min == max: + max += 0.001 # Prevent singularity + + linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func + + # Note on variable naming: "interpolation" comes from the original technique, where we interpolate the position IDs + # to expand the possible context length. In other words, interpolation = apply scaling factor. + pos_freqs = base ** (torch.arange(0, dim, 2).to(device=device, dtype=torch.float) / dim) + inv_freq_extrapolation = 1.0 / pos_freqs + inv_freq_interpolation = 1.0 / (factor * pos_freqs) + + low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_max_position_embeddings) + + # Get n-dimensional rotational scaling corrected for extrapolation + inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2).to(device=device, dtype=torch.float) + inv_freq = ( + inv_freq_interpolation * (1 - inv_freq_extrapolation_factor) + + inv_freq_extrapolation * inv_freq_extrapolation_factor + ) + return inv_freq, attention_factor + + +def _compute_longrope_parameters( + config: PretrainedConfig, device: "torch.device", seq_len: Optional[int] = None, **rope_kwargs +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies with LongRoPE scaling. Please refer to the + [original implementation](https://github.com/microsoft/LongRoPE) + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin. + """ + # TODO (joao): use the new `original_max_position_embeddings` from rope_scaling + # No need to keep BC with longrope, unreleased when this new pattern was created. + if len(rope_kwargs) > 0: + raise ValueError( + "Unexpected arguments: `**rope_kwargs` should be unset in `_compute_longrope_parameters`, got " + f"{rope_kwargs}" + ) + + base = config.rope_theta + partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + dim = int(head_dim * partial_rotary_factor) + long_factor = config.rope_scaling["long_factor"] + short_factor = config.rope_scaling["short_factor"] + factor = config.rope_scaling.get("factor") + attention_factor = config.rope_scaling.get("attention_factor") + + # NOTE: Phi3 (and potentially other models) modify `max_position_embeddings` and have a + # `original_max_position_embeddings` field containing the pretrained value. They use the ratio between these two + # values to compute the default attention scaling factor, instead of using `factor`. + if hasattr(config, "original_max_position_embeddings"): + original_max_position_embeddings = config.original_max_position_embeddings + factor = config.max_position_embeddings / config.original_max_position_embeddings + else: + original_max_position_embeddings = config.max_position_embeddings + + # Sets the attention factor as suggested in the paper + if attention_factor is None: + if factor <= 1.0: + attention_factor = 1.0 + else: + attention_factor = math.sqrt(1 + math.log(factor) / math.log(original_max_position_embeddings)) + + # Compute the inverse frequencies -- scaled based on the target sequence length + if seq_len and seq_len > original_max_position_embeddings: + ext_factors = torch.tensor(long_factor, dtype=torch.float32, device=device) + else: + ext_factors = torch.tensor(short_factor, dtype=torch.float32, device=device) + inv_freq_shape = torch.arange(0, dim, 2, dtype=torch.int64, device=device).float() / dim + inv_freq = 1.0 / (ext_factors * base**inv_freq_shape) + + return inv_freq, attention_factor + + +def _compute_llama3_parameters( + config: PretrainedConfig, device: "torch.device", seq_len: Optional[int] = None, **rope_kwargs +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies for llama 3.1. + + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin. + """ + # Gets the default RoPE parameters + inv_freq, attention_factor = _compute_default_rope_parameters(config, device, seq_len, **rope_kwargs) + + factor = config.rope_scaling["factor"] # `8` in the original implementation + low_freq_factor = config.rope_scaling["low_freq_factor"] # `1` in the original implementation + high_freq_factor = config.rope_scaling["high_freq_factor"] # `4` in the original implementation + old_context_len = config.rope_scaling["original_max_position_embeddings"] # `8192` in the original implementation + + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + + wavelen = 2 * math.pi / inv_freq + # wavelen < high_freq_wavelen: do nothing + # wavelen > low_freq_wavelen: divide by factor + inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq) + # otherwise: interpolate between the two, using a smooth factor + smooth_factor = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor) + smoothed_inv_freq = (1 - smooth_factor) * inv_freq_llama / factor + smooth_factor * inv_freq_llama + is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen) + inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama) + + return inv_freq_llama, attention_factor + + +# This maps the "rope_type" string field in rope config to the corresponding function to compute the RoPE parameters +# from the model config. You can append new {'rope_type': callable} pairs to this dictionary to enable custom RoPE +# parameterizations, as long as the callable has the same signature. +ROPE_INIT_FUNCTIONS = { + "default": _compute_default_rope_parameters, + "linear": _compute_linear_scaling_rope_parameters, + "dynamic": _compute_dynamic_ntk_parameters, + "yarn": _compute_yarn_parameters, + "longrope": _compute_longrope_parameters, + "llama3": _compute_llama3_parameters, +} + + +def _check_received_keys( + rope_type: str, + received_keys: set, + required_keys: set, + optional_keys: Optional[set] = None, + ignore_keys: Optional[set] = None, +): + """Compare the received keys in `config.rope_scaling` against the expected and optional keys""" + # BC: "rope_type" was originally "type" -- let's check for "rope_type" when "type" is present + if "type" in received_keys: + received_keys -= {"type"} + required_keys.add("rope_type") + + # Some models need to store model-specific keys, and we don't want to throw warning at them + if ignore_keys is not None: + received_keys -= ignore_keys + + missing_keys = required_keys - received_keys + if missing_keys: + raise KeyError(f"Missing required keys in `rope_scaling` for 'rope_type'='{rope_type}': {missing_keys}") + + if optional_keys is not None: + unused_keys = received_keys - required_keys - optional_keys + else: + unused_keys = received_keys - required_keys + if unused_keys: + logger.warning(f"Unrecognized keys in `rope_scaling` for 'rope_type'='{rope_type}': {unused_keys}") + + +def _validate_default_rope_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type" + required_keys = {"rope_type"} + received_keys = set(rope_scaling.keys()) + _check_received_keys(rope_type, received_keys, required_keys, ignore_keys=ignore_keys) + + +def _validate_linear_scaling_rope_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type" + required_keys = {"rope_type", "factor"} + received_keys = set(rope_scaling.keys()) + _check_received_keys(rope_type, received_keys, required_keys, ignore_keys=ignore_keys) + + factor = rope_scaling["factor"] + if factor is None or not isinstance(factor, float) or factor < 1.0: + logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}") + + +def _validate_dynamic_scaling_rope_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type" + required_keys = {"rope_type", "factor"} + # TODO (joao): update logic for the inclusion of `original_max_position_embeddings` + optional_keys = {"original_max_position_embeddings"} + received_keys = set(rope_scaling.keys()) + _check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys) + + factor = rope_scaling["factor"] + if factor is None or not isinstance(factor, float) or factor < 1.0: + logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}") + + +def _validate_yarn_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type" + required_keys = {"rope_type", "factor"} + optional_keys = { + "attention_factor", + "beta_fast", + "beta_slow", + "original_max_position_embeddings", + "mscale", + "mscale_all_dim", + } + received_keys = set(rope_scaling.keys()) + _check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys) + + factor = rope_scaling["factor"] + if factor is None or not isinstance(factor, float) or factor < 1.0: + logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}") + + attention_factor = rope_scaling.get("attention_factor") + if attention_factor is not None and (not isinstance(attention_factor, float) or attention_factor < 0): + logger.warning( + f"`rope_scaling`'s attention_factor field must be a float greater than 0, got {attention_factor}" + ) + beta_fast = rope_scaling.get("beta_fast") + if beta_fast is not None and not isinstance(beta_fast, float): + logger.warning(f"`rope_scaling`'s beta_fast field must be a float, got {beta_fast}") + beta_slow = rope_scaling.get("beta_slow") + if beta_slow is not None and not isinstance(beta_slow, float): + logger.warning(f"`rope_scaling`'s beta_slow field must be a float, got {beta_slow}") + + if (beta_fast or 32) < (beta_slow or 1): + logger.warning( + f"`rope_scaling`'s beta_fast field must be greater than beta_slow, got beta_fast={beta_fast} " + f"(defaults to 32 if None) and beta_slow={beta_slow} (defaults to 1 if None)" + ) + + +def _validate_longrope_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type" + required_keys = {"rope_type", "short_factor", "long_factor"} + # TODO (joao): update logic for the inclusion of `original_max_position_embeddings` + optional_keys = {"attention_factor", "factor", "original_max_position_embeddings"} + received_keys = set(rope_scaling.keys()) + _check_received_keys(rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys) + + partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + dim = int(head_dim * partial_rotary_factor) + + short_factor = rope_scaling.get("short_factor") + if not isinstance(short_factor, list) and all(isinstance(x, (int, float)) for x in short_factor): + logger.warning(f"`rope_scaling`'s short_factor field must be a list of numbers, got {short_factor}") + if not len(short_factor) == dim // 2: + logger.warning(f"`rope_scaling`'s short_factor field must have length {dim // 2}, got {len(short_factor)}") + + long_factor = rope_scaling.get("long_factor") + if not isinstance(long_factor, list) and all(isinstance(x, (int, float)) for x in long_factor): + logger.warning(f"`rope_scaling`'s long_factor field must be a list of numbers, got {long_factor}") + if not len(long_factor) == dim // 2: + logger.warning(f"`rope_scaling`'s long_factor field must have length {dim // 2}, got {len(long_factor)}") + + # Handle Phi3 divergence: prefer the use of `attention_factor` and/or `factor` over + # `original_max_position_embeddings` to compute internal variables. The latter lives outside `rope_scaling` and is + # unique to longrope (= undesirable) + if hasattr(config, "original_max_position_embeddings"): + logger.warning_once( + "This model has set a `original_max_position_embeddings` field, to be used together with " + "`max_position_embeddings` to determine a scaling factor. Please set the `factor` field of `rope_scaling`" + "with this ratio instead -- we recommend the use of this field over `original_max_position_embeddings`, " + "as it is compatible with most model architectures." + ) + else: + factor = rope_scaling.get("factor") + if factor is None: + logger.warning("Missing required keys in `rope_scaling`: 'factor'") + elif not isinstance(factor, float) or factor < 1.0: + logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}") + + attention_factor = rope_scaling.get("attention_factor") + if attention_factor is not None: + if not isinstance(attention_factor, float) or attention_factor < 0.0: + logger.warning( + f"`rope_scaling`'s attention_factor field must be a float greater than 0, got {attention_factor}" + ) + + +def _validate_llama3_parameters(config: PretrainedConfig, ignore_keys: Optional[set] = None): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None)) # BC: "rope_type" was originally "type" + required_keys = {"rope_type", "factor", "original_max_position_embeddings", "low_freq_factor", "high_freq_factor"} + received_keys = set(rope_scaling.keys()) + _check_received_keys(rope_type, received_keys, required_keys, ignore_keys=ignore_keys) + + factor = rope_scaling["factor"] + if factor is None or not isinstance(factor, float) or factor < 1.0: + logger.warning(f"`rope_scaling`'s factor field must be a float >= 1, got {factor}") + + low_freq_factor = rope_scaling["low_freq_factor"] + high_freq_factor = rope_scaling["high_freq_factor"] + if low_freq_factor is None or not isinstance(low_freq_factor, float): + logger.warning(f"`rope_scaling`'s low_freq_factor field must be a float, got {low_freq_factor}") + if high_freq_factor is None or not isinstance(high_freq_factor, float): + logger.warning(f"`rope_scaling`'s high_freq_factor field must be a float, got {high_freq_factor}") + if high_freq_factor <= low_freq_factor: + logger.warning( + "`rope_scaling`'s high_freq_factor field must be greater than low_freq_factor, got high_freq_factor=" + f"{high_freq_factor} and low_freq_factor={low_freq_factor}" + ) + + original_max_position_embeddings = rope_scaling["original_max_position_embeddings"] + if original_max_position_embeddings is None or not isinstance(original_max_position_embeddings, int): + logger.warning( + "`rope_scaling`'s original_max_position_embeddings field must be an integer, got " + f"{original_max_position_embeddings}" + ) + if original_max_position_embeddings >= config.max_position_embeddings: + logger.warning( + "`rope_scaling`'s original_max_position_embeddings field must be less than max_position_embeddings, got " + f"{original_max_position_embeddings} and max_position_embeddings={config.max_position_embeddings}" + ) + + +# Like `ROPE_INIT_FUNCTIONS`, this validation function mapping can be dynamically updated for custom RoPE types. +ROPE_VALIDATION_FUNCTIONS = { + "default": _validate_default_rope_parameters, + "linear": _validate_linear_scaling_rope_parameters, + "dynamic": _validate_dynamic_scaling_rope_parameters, + "yarn": _validate_yarn_parameters, + "longrope": _validate_longrope_parameters, + "llama3": _validate_llama3_parameters, +} + + +def rope_config_validation(config: PretrainedConfig, ignore_keys: Optional[set] = None): + """ + Validate the RoPE config arguments, given a `PretrainedConfig` object + """ + rope_scaling = getattr(config, "rope_scaling", None) # not a default parameter in `PretrainedConfig` + if rope_scaling is None: + return + + # BC: "rope_type" was originally "type" + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", "default")) + validation_fn = ROPE_VALIDATION_FUNCTIONS.get(rope_type) + if validation_fn is not None: + validation_fn(config, ignore_keys=ignore_keys) + else: + logger.warning( + f"Missing validation function mapping in `ROPE_VALIDATION_FUNCTIONS` for 'rope_type'='{rope_type}'" + ) diff --git a/lib/python3.12/site-packages/transformers/modeling_tf_outputs.py b/lib/python3.12/site-packages/transformers/modeling_tf_outputs.py new file mode 100644 index 0000000000000000000000000000000000000000..cbc8b3682af608d0db6da02ef36764d0be10ca9c --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_tf_outputs.py @@ -0,0 +1,991 @@ +# Copyright 2020 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. + +from __future__ import annotations + +import warnings +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import tensorflow as tf + +from .utils import ModelOutput + + +@dataclass +class TFBaseModelOutput(ModelOutput): + """ + Base class for model's outputs, with potential hidden states and attentions. + + Args: + last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(tf.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + last_hidden_state: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFBaseModelOutputWithNoAttention(ModelOutput): + """ + Base class for model's outputs, with potential hidden states. + + Args: + last_hidden_state (`tf.Tensor` shape `(batch_size, num_channels, height, width)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one for + the output of each layer) of shape `(batch_size, num_channels, height, width)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + """ + + last_hidden_state: Optional[tf.Tensor] = None + hidden_states: Optional[Tuple[tf.Tensor, ...]] = None + + +@dataclass +class TFBaseModelOutputWithPooling(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`tf.Tensor` of shape `(batch_size, hidden_size)`): + Last layer hidden-state of the first token of the sequence (classification token) further processed by a + Linear layer and a Tanh activation function. The Linear layer weights are trained from the next sentence + prediction (classification) objective during pretraining. + + This output is usually *not* a good summary of the semantic content of the input, you're often better with + averaging or pooling the sequence of hidden-states for the whole input sequence. + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + last_hidden_state: Optional[tf.Tensor] = None + pooler_output: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFBaseModelOutputWithPoolingAndNoAttention(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + last_hidden_state (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`tf.Tensor` of shape `(batch_size, hidden_size)`): + Last layer hidden-state after a pooling operation on the spatial dimensions. + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one for + the output of each layer) of shape `(batch_size, num_channels, height, width)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + """ + + last_hidden_state: Optional[tf.Tensor] = None + pooler_output: Optional[tf.Tensor] = None + hidden_states: Optional[Tuple[tf.Tensor, ...]] = None + + +@dataclass +class TFBaseModelOutputWithPoolingAndCrossAttentions(ModelOutput): + """ + Base class for model's outputs that also contains a pooling of the last hidden states. + + Args: + last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + pooler_output (`tf.Tensor` of shape `(batch_size, hidden_size)`): + Last layer hidden-state of the first token of the sequence (classification token) further processed by a + Linear layer and a Tanh activation function. The Linear layer weights are trained from the next sentence + prediction (classification) objective during pretraining. + + This output is usually *not* a good summary of the semantic content of the input, you're often better with + averaging or pooling the sequence of hidden-states for the whole input sequence. + past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + """ + + last_hidden_state: Optional[tf.Tensor] = None + pooler_output: Optional[tf.Tensor] = None + past_key_values: List[tf.Tensor] | None = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + cross_attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFBaseModelOutputWithPast(ModelOutput): + """ + Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding). + + Args: + last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + last_hidden_state: Optional[tf.Tensor] = None + past_key_values: List[tf.Tensor] | None = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFBaseModelOutputWithCrossAttentions(ModelOutput): + """ + Base class for model's outputs, with potential hidden states and attentions. + + Args: + last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + hidden_states (`tuple(tf.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + """ + + last_hidden_state: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + cross_attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFBaseModelOutputWithPastAndCrossAttentions(ModelOutput): + """ + Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding). + + Args: + last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(tf.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + """ + + last_hidden_state: Optional[tf.Tensor] = None + past_key_values: List[tf.Tensor] | None = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + cross_attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFSeq2SeqModelOutput(ModelOutput): + """ + Base class for model encoder's outputs that also contains : pre-computed hidden states that can speed up sequential + decoding. + + Args: + last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the decoder of the model. + + If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, + hidden_size)` is output. + past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) of the decoder that can be + used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + last_hidden_state: Optional[tf.Tensor] = None + past_key_values: List[tf.Tensor] | None = None + decoder_hidden_states: Tuple[tf.Tensor] | None = None + decoder_attentions: Tuple[tf.Tensor] | None = None + cross_attentions: Tuple[tf.Tensor] | None = None + encoder_last_hidden_state: tf.Tensor | None = None + encoder_hidden_states: Tuple[tf.Tensor] | None = None + encoder_attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFCausalLMOutput(ModelOutput): + """ + Base class for causal language model (or autoregressive) outputs. + + Args: + loss (`tf.Tensor` of shape `(n,)`, *optional*, where n is the number of non-masked labels, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFCausalLMOutputWithPast(ModelOutput): + """ + Base class for causal language model (or autoregressive) outputs. + + Args: + loss (`tf.Tensor` of shape `(n,)`, *optional*, where n is the number of non-masked labels, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + past_key_values: List[tf.Tensor] | None = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFCausalLMOutputWithCrossAttentions(ModelOutput): + """ + Base class for causal language model (or autoregressive) outputs. + + Args: + loss (`tf.Tensor` of shape `(n,)`, *optional*, where n is the number of non-masked labels, returned when `labels` is provided): + Language modeling loss (for next-token prediction). + logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + cross_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + past_key_values: List[tf.Tensor] | None = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + cross_attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFMaskedLMOutput(ModelOutput): + """ + Base class for masked language models outputs. + + Args: + loss (`tf.Tensor` of shape `(n,)`, *optional*, where n is the number of non-masked labels, returned when `labels` is provided): + Masked language modeling (MLM) loss. + logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFSeq2SeqLMOutput(ModelOutput): + """ + Base class for sequence-to-sequence language models outputs. + + Args: + loss (`tf.Tensor` of shape `(n,)`, *optional*, where n is the number of non-masked labels, returned when `labels` is provided): + Language modeling loss. + logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.vocab_size)`): + Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). + past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) of the decoder that can be + used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder's cross-attention layer, after the attention softmax, used to compute the + weighted average in the cross-attention heads. + encoder_last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + past_key_values: List[tf.Tensor] | None = None + decoder_hidden_states: Tuple[tf.Tensor] | None = None + decoder_attentions: Tuple[tf.Tensor] | None = None + cross_attentions: Tuple[tf.Tensor] | None = None + encoder_last_hidden_state: tf.Tensor | None = None + encoder_hidden_states: Tuple[tf.Tensor] | None = None + encoder_attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFNextSentencePredictorOutput(ModelOutput): + """ + Base class for outputs of models predicting if two sentences are consecutive or not. + + Args: + loss (`tf.Tensor` of shape `(n,)`, *optional*, where n is the number of non-masked labels, returned when `next_sentence_label` is provided): + Next sentence prediction loss. + logits (`tf.Tensor` of shape `(batch_size, 2)`): + Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation + before SoftMax). + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFSequenceClassifierOutput(ModelOutput): + """ + Base class for outputs of sentence classification models. + + Args: + loss (`tf.Tensor` of shape `(batch_size, )`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFSeq2SeqSequenceClassifierOutput(ModelOutput): + """ + Base class for outputs of sequence-to-sequence sentence classification models. + + Args: + loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `label` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) of the decoder that can be + used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + cross_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)` + encoder_last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + past_key_values: List[tf.Tensor] | None = None + decoder_hidden_states: Tuple[tf.Tensor] | None = None + decoder_attentions: Tuple[tf.Tensor] | None = None + cross_attentions: Tuple[tf.Tensor] | None = None + encoder_last_hidden_state: tf.Tensor | None = None + encoder_hidden_states: Tuple[tf.Tensor] | None = None + encoder_attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFSemanticSegmenterOutput(ModelOutput): + """ + Base class for outputs of semantic segmentation models. + + Args: + loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`tf.Tensor` of shape `(batch_size, config.num_labels, logits_height, logits_width)`): + Classification scores for each pixel. + + + + The logits returned do not necessarily have the same size as the `pixel_values` passed as inputs. This is + to avoid doing two interpolations and lose some quality when a user needs to resize the logits to the + original image size as post-processing. You should always check your logits shape and resize as needed. + + + + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one for + the output of each layer) of shape `(batch_size, patch_size, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, patch_size, sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFSemanticSegmenterOutputWithNoAttention(ModelOutput): + """ + Base class for outputs of semantic segmentation models that do not output attention scores. + + Args: + loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`tf.Tensor` of shape `(batch_size, config.num_labels, logits_height, logits_width)`): + Classification scores for each pixel. + + + + The logits returned do not necessarily have the same size as the `pixel_values` passed as inputs. This is + to avoid doing two interpolations and lose some quality when a user needs to resize the logits to the + original image size as post-processing. You should always check your logits shape and resize as needed. + + + + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one for + the output of each layer) of shape `(batch_size, patch_size, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFImageClassifierOutput(ModelOutput): + """ + Base class for outputs of image classification models. + + Args: + loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one for + the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states (also called + feature maps) of the model at the output of each stage. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, patch_size, sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFMultipleChoiceModelOutput(ModelOutput): + """ + Base class for outputs of multiple choice models. + + Args: + loss (`tf.Tensor` of shape *(batch_size, )*, *optional*, returned when `labels` is provided): + Classification loss. + logits (`tf.Tensor` of shape `(batch_size, num_choices)`): + *num_choices* is the second dimension of the input tensors. (see *input_ids* above). + + Classification scores (before SoftMax). + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFTokenClassifierOutput(ModelOutput): + """ + Base class for outputs of token classification models. + + Args: + loss (`tf.Tensor` of shape `(n,)`, *optional*, where n is the number of unmasked labels, returned when `labels` is provided) : + Classification loss. + logits (`tf.Tensor` of shape `(batch_size, sequence_length, config.num_labels)`): + Classification scores (before SoftMax). + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFQuestionAnsweringModelOutput(ModelOutput): + """ + Base class for outputs of question answering models. + + Args: + loss (`tf.Tensor` of shape `(batch_size, )`, *optional*, returned when `start_positions` and `end_positions` are provided): + Total span extraction loss is the sum of a Cross-Entropy for the start and end positions. + start_logits (`tf.Tensor` of shape `(batch_size, sequence_length)`): + Span-start scores (before SoftMax). + end_logits (`tf.Tensor` of shape `(batch_size, sequence_length)`): + Span-end scores (before SoftMax). + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + start_logits: Optional[tf.Tensor] = None + end_logits: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFSeq2SeqQuestionAnsweringModelOutput(ModelOutput): + """ + Base class for outputs of sequence-to-sequence question answering models. + + Args: + loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Total span extraction loss is the sum of a Cross-Entropy for the start and end positions. + start_logits (`tf.Tensor` of shape `(batch_size, sequence_length)`): + Span-start scores (before SoftMax). + end_logits (`tf.Tensor` of shape `(batch_size, sequence_length)`): + Span-end scores (before SoftMax). + past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) of the decoder that can be + used (see `past_key_values` input) to speed up sequential decoding. + decoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the decoder at the output of each layer plus the initial embedding outputs. + decoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the decoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + encoder_last_hidden_state (`tf.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder of the model. + encoder_hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the encoder at the output of each layer plus the initial embedding outputs. + encoder_attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights of the encoder, after the attention softmax, used to compute the weighted average in the + self-attention heads. + """ + + loss: tf.Tensor | None = None + start_logits: Optional[tf.Tensor] = None + end_logits: Optional[tf.Tensor] = None + past_key_values: List[tf.Tensor] | None = None + decoder_hidden_states: Tuple[tf.Tensor] | None = None + decoder_attentions: Tuple[tf.Tensor] | None = None + encoder_last_hidden_state: tf.Tensor | None = None + encoder_hidden_states: Tuple[tf.Tensor] | None = None + encoder_attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFSequenceClassifierOutputWithPast(ModelOutput): + """ + Base class for outputs of sentence classification models. + + Args: + loss (`tf.Tensor` of shape `(batch_size, )`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + past_key_values (`List[tf.Tensor]`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + List of `tf.Tensor` of length `config.n_layers`, with each tensor of shape `(2, batch_size, num_heads, + sequence_length, embed_size_per_head)`). + + Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see + `past_key_values` input) to speed up sequential decoding. + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings + one for the output of each layer) of shape + `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + past_key_values: List[tf.Tensor] | None = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + +@dataclass +class TFImageClassifierOutputWithNoAttention(ModelOutput): + """ + Base class for outputs of image classification models. + + Args: + loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Classification (or regression if config.num_labels==1) loss. + logits (`tf.Tensor` of shape `(batch_size, config.num_labels)`): + Classification (or regression if config.num_labels==1) scores (before SoftMax). + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one for + the output of each stage) of shape `(batch_size, num_channels, height, width)`. Hidden-states (also called + feature maps) of the model at the output of each stage. + """ + + loss: tf.Tensor | None = None + logits: Optional[tf.Tensor] = None + hidden_states: Optional[Tuple[tf.Tensor, ...]] = None + + +@dataclass +class TFMaskedImageModelingOutput(ModelOutput): + """ + Base class for outputs of masked image completion / in-painting models. + + Args: + loss (`tf.Tensor` of shape `(1,)`, *optional*, returned when `bool_masked_pos` is provided): + Reconstruction loss. + reconstruction (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`): + Reconstructed / completed images. + hidden_states (`tuple(tf.Tensor)`, *optional*, returned when `output_hidden_states=True` is passed or when + `config.output_hidden_states=True`): + Tuple of `tf.Tensor` (one for the output of the embeddings, if the model has an embedding layer, + one for + the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states (also called + feature maps) of the model at the output of each stage. + attentions (`tuple(tf.Tensor)`, *optional*, returned when `output_attentions=True` is passed or when + `config.output_attentions=True`): + Tuple of `tf.Tensor` (one for each layer) of shape `(batch_size, num_heads, patch_size, sequence_length)`. + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + loss: tf.Tensor | None = None + reconstruction: Optional[tf.Tensor] = None + hidden_states: Tuple[tf.Tensor] | None = None + attentions: Tuple[tf.Tensor] | None = None + + @property + def logits(self): + warnings.warn( + "logits attribute is deprecated and will be removed in version 5 of Transformers." + " Please use the reconstruction attribute to retrieve the final output instead.", + FutureWarning, + ) + return self.reconstruction diff --git a/lib/python3.12/site-packages/transformers/modeling_tf_pytorch_utils.py b/lib/python3.12/site-packages/transformers/modeling_tf_pytorch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..84a6ddaebcc46c12a2a08e4612fde2f4f20cf648 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_tf_pytorch_utils.py @@ -0,0 +1,671 @@ +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. 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. +"""PyTorch - TF 2.0 general utilities.""" + +import os +import re + +import numpy + +from .utils import ( + ExplicitEnum, + expand_dims, + is_numpy_array, + is_safetensors_available, + is_torch_tensor, + logging, + reshape, + squeeze, + tensor_size, +) +from .utils import transpose as transpose_func + + +if is_safetensors_available(): + from safetensors import safe_open + + +logger = logging.get_logger(__name__) + + +class TransposeType(ExplicitEnum): + """ + Possible ... + """ + + NO = "no" + SIMPLE = "simple" + CONV1D = "conv1d" + CONV2D = "conv2d" + + +def convert_tf_weight_name_to_pt_weight_name( + tf_name, start_prefix_to_remove="", tf_weight_shape=None, name_scope=None +): + """ + Convert a TF 2.0 model variable name in a pytorch model weight name. + + Conventions for TF2.0 scopes -> PyTorch attribute names conversions: + + - '$1___$2' is replaced by $2 (can be used to duplicate or remove layers in TF2.0 vs PyTorch) + - '_._' is replaced by a new level separation (can be used to convert TF2.0 lists in PyTorch nn.ModulesList) + + return tuple with: + + - pytorch model weight name + - transpose: `TransposeType` member indicating whether and how TF2.0 and PyTorch weights matrices should be + transposed with regards to each other + """ + if name_scope is not None: + if not tf_name.startswith(name_scope) and "final_logits_bias" not in tf_name: + raise ValueError( + f"Weight name {tf_name} does not start with name_scope {name_scope}. This is an internal error " + "in Transformers, so (unless you were doing something really evil) please open an issue to report it!" + ) + tf_name = tf_name[len(name_scope) :] + tf_name = tf_name.lstrip("/") + tf_name = tf_name.replace(":0", "") # device ids + tf_name = re.sub( + r"/[^/]*___([^/]*)/", r"/\1/", tf_name + ) # '$1___$2' is replaced by $2 (can be used to duplicate or remove layers in TF2.0 vs PyTorch) + tf_name = tf_name.replace( + "_._", "/" + ) # '_._' is replaced by a level separation (can be used to convert TF2.0 lists in PyTorch nn.ModulesList) + tf_name = re.sub(r"//+", "/", tf_name) # Remove empty levels at the end + tf_name = tf_name.split("/") # Convert from TF2.0 '/' separators to PyTorch '.' separators + # Some weights have a single name without "/" such as final_logits_bias in BART + if len(tf_name) > 1: + tf_name = tf_name[1:] # Remove level zero + + tf_weight_shape = list(tf_weight_shape) + + # When should we transpose the weights + if tf_name[-1] == "kernel" and tf_weight_shape is not None and len(tf_weight_shape) == 4: + transpose = TransposeType.CONV2D + elif tf_name[-1] == "kernel" and tf_weight_shape is not None and len(tf_weight_shape) == 3: + transpose = TransposeType.CONV1D + elif bool( + tf_name[-1] in ["kernel", "pointwise_kernel", "depthwise_kernel"] + or "emb_projs" in tf_name + or "out_projs" in tf_name + ): + transpose = TransposeType.SIMPLE + else: + transpose = TransposeType.NO + + # Convert standard TF2.0 names in PyTorch names + if tf_name[-1] == "kernel" or tf_name[-1] == "embeddings" or tf_name[-1] == "gamma": + tf_name[-1] = "weight" + if tf_name[-1] == "beta": + tf_name[-1] = "bias" + + # The SeparableConv1D TF layer contains two weights that are translated to PyTorch Conv1D here + if tf_name[-1] == "pointwise_kernel" or tf_name[-1] == "depthwise_kernel": + tf_name[-1] = tf_name[-1].replace("_kernel", ".weight") + + # Remove prefix if needed + tf_name = ".".join(tf_name) + if start_prefix_to_remove: + tf_name = tf_name.replace(start_prefix_to_remove, "", 1) + + return tf_name, transpose + + +def apply_transpose(transpose: TransposeType, weight, match_shape=None, pt_to_tf=True): + """ + Apply a transpose to some weight then tries to reshape the weight to the same shape as a given shape, all in a + framework agnostic way. + """ + if transpose is TransposeType.CONV2D: + # Conv2D weight: + # PT: (num_out_channel, num_in_channel, kernel[0], kernel[1]) + # -> TF: (kernel[0], kernel[1], num_in_channel, num_out_channel) + axes = (2, 3, 1, 0) if pt_to_tf else (3, 2, 0, 1) + weight = transpose_func(weight, axes=axes) + elif transpose is TransposeType.CONV1D: + # Conv1D weight: + # PT: (num_out_channel, num_in_channel, kernel) + # -> TF: (kernel, num_in_channel, num_out_channel) + weight = transpose_func(weight, axes=(2, 1, 0)) + elif transpose is TransposeType.SIMPLE: + weight = transpose_func(weight) + + if match_shape is None: + return weight + + if len(match_shape) < len(weight.shape): + weight = squeeze(weight) + elif len(match_shape) > len(weight.shape): + weight = expand_dims(weight, axis=0) + + if list(match_shape) != list(weight.shape): + try: + weight = reshape(weight, match_shape) + except AssertionError as e: + e.args += (match_shape, match_shape) + raise e + + return weight + + +##################### +# PyTorch => TF 2.0 # +##################### + + +def load_pytorch_checkpoint_in_tf2_model( + tf_model, + pytorch_checkpoint_path, + tf_inputs=None, + allow_missing_keys=False, + output_loading_info=False, + _prefix=None, + tf_to_pt_weight_rename=None, +): + """Load pytorch checkpoints in a TF 2.0 model""" + try: + import tensorflow as tf # noqa: F401 + import torch # noqa: F401 + from safetensors.torch import load_file as safe_load_file # noqa: F401 + except ImportError: + logger.error( + "Loading a PyTorch model in TensorFlow, requires both PyTorch and TensorFlow to be installed. Please see " + "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions." + ) + raise + + # Treats a single file as a collection of shards with 1 shard. + if isinstance(pytorch_checkpoint_path, str): + pytorch_checkpoint_path = [pytorch_checkpoint_path] + + # Loads all shards into a single state dictionary + pt_state_dict = {} + for path in pytorch_checkpoint_path: + pt_path = os.path.abspath(path) + logger.info(f"Loading PyTorch weights from {pt_path}") + if pt_path.endswith(".safetensors"): + state_dict = safe_load_file(pt_path) + else: + state_dict = torch.load(pt_path, map_location="cpu", weights_only=True) + + pt_state_dict.update(state_dict) + + logger.info(f"PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()):,} parameters") + + return load_pytorch_weights_in_tf2_model( + tf_model, + pt_state_dict, + tf_inputs=tf_inputs, + allow_missing_keys=allow_missing_keys, + output_loading_info=output_loading_info, + _prefix=_prefix, + tf_to_pt_weight_rename=tf_to_pt_weight_rename, + ) + + +def load_pytorch_model_in_tf2_model(tf_model, pt_model, tf_inputs=None, allow_missing_keys=False): + """Load pytorch checkpoints in a TF 2.0 model""" + pt_state_dict = pt_model.state_dict() + + return load_pytorch_weights_in_tf2_model( + tf_model, pt_state_dict, tf_inputs=tf_inputs, allow_missing_keys=allow_missing_keys + ) + + +def load_pytorch_weights_in_tf2_model( + tf_model, + pt_state_dict, + tf_inputs=None, + allow_missing_keys=False, + output_loading_info=False, + _prefix=None, + tf_to_pt_weight_rename=None, +): + """Load pytorch state_dict in a TF 2.0 model.""" + try: + import tensorflow as tf # noqa: F401 + import torch # noqa: F401 + except ImportError: + logger.error( + "Loading a PyTorch model in TensorFlow, requires both PyTorch and TensorFlow to be installed. Please see " + "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions." + ) + raise + + # Numpy doesn't understand bfloat16, so upcast to a dtype that doesn't lose precision + pt_state_dict = { + k: v.numpy() if v.dtype != torch.bfloat16 else v.float().numpy() for k, v in pt_state_dict.items() + } + return load_pytorch_state_dict_in_tf2_model( + tf_model, + pt_state_dict, + tf_inputs=tf_inputs, + allow_missing_keys=allow_missing_keys, + output_loading_info=output_loading_info, + _prefix=_prefix, + tf_to_pt_weight_rename=tf_to_pt_weight_rename, + ) + + +def _log_key_warnings(missing_keys, unexpected_keys, mismatched_keys, class_name): + if len(unexpected_keys) > 0: + logger.warning( + "Some weights of the PyTorch model were not used when initializing the TF 2.0 model" + f" {class_name}: {unexpected_keys}\n- This IS expected if you are initializing" + f" {class_name} from a PyTorch model trained on another task or with another architecture" + " (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).\n- This IS" + f" NOT expected if you are initializing {class_name} from a PyTorch model that you expect" + " to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a" + " BertForSequenceClassification model)." + ) + else: + logger.warning(f"All PyTorch model weights were used when initializing {class_name}.\n") + if len(missing_keys) > 0: + logger.warning( + f"Some weights or buffers of the TF 2.0 model {class_name} were not initialized from the" + f" PyTorch model and are newly initialized: {missing_keys}\nYou should probably TRAIN this model on a" + " down-stream task to be able to use it for predictions and inference." + ) + else: + logger.warning( + f"All the weights of {class_name} were initialized from the PyTorch model.\n" + "If your task is similar to the task the model of the checkpoint was trained on, " + f"you can already use {class_name} for predictions without further training." + ) + + if len(mismatched_keys) > 0: + mismatched_warning = "\n".join( + [ + f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated" + for key, shape1, shape2 in mismatched_keys + ] + ) + logger.warning( + f"Some weights of {class_name} were not initialized from the model checkpoint" + f" are newly initialized because the shapes did not" + f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able" + " to use it for predictions and inference." + ) + + +def load_pytorch_state_dict_in_tf2_model( + tf_model, + pt_state_dict, + tf_inputs=None, + allow_missing_keys=False, + output_loading_info=False, + _prefix=None, + tf_to_pt_weight_rename=None, + ignore_mismatched_sizes=False, + skip_logger_warnings=False, +): + """Load a pytorch state_dict in a TF 2.0 model. pt_state_dict can be either an actual dict or a lazy-loading + safetensors archive created with the safe_open() function.""" + import tensorflow as tf + + if tf_inputs is None: + tf_inputs = tf_model.dummy_inputs + + if _prefix is None: + _prefix = "" + if tf_inputs: + with tf.name_scope(_prefix): + tf_model(tf_inputs, training=False) # Make sure model is built + # Convert old format to new format if needed from a PyTorch state_dict + tf_keys_to_pt_keys = {} + for key in pt_state_dict.keys(): + new_key = None + if "gamma" in key: + new_key = key.replace("gamma", "weight") + if "beta" in key: + new_key = key.replace("beta", "bias") + if "running_var" in key: + new_key = key.replace("running_var", "moving_variance") + if "running_mean" in key: + new_key = key.replace("running_mean", "moving_mean") + + # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030 + key_components = key.split(".") + name = None + if key_components[-3::2] == ["parametrizations", "original0"]: + name = key_components[-2] + "_g" + elif key_components[-3::2] == ["parametrizations", "original1"]: + name = key_components[-2] + "_v" + if name is not None: + key_components = key_components[:-3] + [name] + new_key = ".".join(key_components) + + if new_key is None: + new_key = key + tf_keys_to_pt_keys[new_key] = key + + # Matt: All TF models store the actual model stem in a MainLayer class, including the base model. + # In PT, the derived models (with heads) use the base model class as the stem instead, + # and there is no MainLayer class. This means that TF base classes have one + # extra layer in their weight names, corresponding to the MainLayer class. This code block compensates for that. + start_prefix_to_remove = "" + if not any(s.startswith(tf_model.base_model_prefix) for s in tf_keys_to_pt_keys.keys()): + start_prefix_to_remove = tf_model.base_model_prefix + "." + + symbolic_weights = tf_model.trainable_weights + tf_model.non_trainable_weights + tf_loaded_numel = 0 + all_pytorch_weights = set(tf_keys_to_pt_keys.keys()) + missing_keys = [] + mismatched_keys = [] + is_safetensor_archive = hasattr(pt_state_dict, "get_tensor") + for symbolic_weight in symbolic_weights: + sw_name = symbolic_weight.name + name, transpose = convert_tf_weight_name_to_pt_weight_name( + sw_name, + start_prefix_to_remove=start_prefix_to_remove, + tf_weight_shape=symbolic_weight.shape, + name_scope=_prefix, + ) + if tf_to_pt_weight_rename is not None: + aliases = tf_to_pt_weight_rename(name) # Is a tuple to account for possible name aliasing + for alias in aliases: # The aliases are in priority order, take the first one that matches + if alias in tf_keys_to_pt_keys: + name = alias + break + else: + # If none of the aliases match, just use the first one (it'll be reported as missing) + name = aliases[0] + + # Find associated numpy array in pytorch model state dict + if name not in tf_keys_to_pt_keys: + if allow_missing_keys: + missing_keys.append(name) + continue + elif tf_model._keys_to_ignore_on_load_missing is not None: + # authorized missing keys don't have to be loaded + if any(re.search(pat, name) is not None for pat in tf_model._keys_to_ignore_on_load_missing): + continue + raise AttributeError(f"{name} not found in PyTorch model") + state_dict_name = tf_keys_to_pt_keys[name] + if is_safetensor_archive: + array = pt_state_dict.get_tensor(state_dict_name) + else: + array = pt_state_dict[state_dict_name] + try: + array = apply_transpose(transpose, array, symbolic_weight.shape) + except tf.errors.InvalidArgumentError as e: + if not ignore_mismatched_sizes: + error_msg = str(e) + error_msg += ( + "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method." + ) + raise tf.errors.InvalidArgumentError(error_msg) + else: + mismatched_keys.append((name, array.shape, symbolic_weight.shape)) + continue + + tf_loaded_numel += tensor_size(array) + + symbolic_weight.assign(tf.cast(array, symbolic_weight.dtype)) + del array # Immediately free memory to keep peak usage as low as possible + all_pytorch_weights.discard(name) + + logger.info(f"Loaded {tf_loaded_numel:,} parameters in the TF 2.0 model.") + + unexpected_keys = list(all_pytorch_weights) + + if tf_model._keys_to_ignore_on_load_missing is not None: + for pat in tf_model._keys_to_ignore_on_load_missing: + missing_keys = [k for k in missing_keys if re.search(pat, k) is None] + if tf_model._keys_to_ignore_on_load_unexpected is not None: + for pat in tf_model._keys_to_ignore_on_load_unexpected: + unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None] + if not skip_logger_warnings: + _log_key_warnings(missing_keys, unexpected_keys, mismatched_keys, class_name=tf_model.__class__.__name__) + + if output_loading_info: + loading_info = { + "missing_keys": missing_keys, + "unexpected_keys": unexpected_keys, + "mismatched_keys": mismatched_keys, + } + return tf_model, loading_info + + return tf_model + + +def load_sharded_pytorch_safetensors_in_tf2_model( + tf_model, + safetensors_shards, + tf_inputs=None, + allow_missing_keys=False, + output_loading_info=False, + _prefix=None, + tf_to_pt_weight_rename=None, + ignore_mismatched_sizes=False, +): + all_loading_infos = [] + for shard in safetensors_shards: + with safe_open(shard, framework="tf") as safetensors_archive: + tf_model, loading_info = load_pytorch_state_dict_in_tf2_model( + tf_model, + safetensors_archive, + tf_inputs=tf_inputs, + allow_missing_keys=allow_missing_keys, + output_loading_info=True, + _prefix=_prefix, + tf_to_pt_weight_rename=tf_to_pt_weight_rename, + ignore_mismatched_sizes=ignore_mismatched_sizes, + skip_logger_warnings=True, # We will emit merged warnings at the end + ) + all_loading_infos.append(loading_info) + # Now we just need to merge the loading info + # Keys are missing only if they're missing in *every* shard + missing_keys = sorted(set.intersection(*[set(info["missing_keys"]) for info in all_loading_infos])) + # Keys are unexpected/mismatched if they're unexpected/mismatched in *any* shard + unexpected_keys = sum([info["unexpected_keys"] for info in all_loading_infos], []) + mismatched_keys = sum([info["mismatched_keys"] for info in all_loading_infos], []) + + _log_key_warnings(missing_keys, unexpected_keys, mismatched_keys, class_name=tf_model.__class__.__name__) + + if output_loading_info: + loading_info = { + "missing_keys": missing_keys, + "unexpected_keys": unexpected_keys, + "mismatched_keys": mismatched_keys, + } + return tf_model, loading_info + + return tf_model + + +##################### +# TF 2.0 => PyTorch # +##################### + + +def load_tf2_checkpoint_in_pytorch_model( + pt_model, tf_checkpoint_path, tf_inputs=None, allow_missing_keys=False, output_loading_info=False +): + """ + Load TF 2.0 HDF5 checkpoint in a PyTorch model We use HDF5 to easily do transfer learning (see + https://github.com/tensorflow/tensorflow/blob/ee16fcac960ae660e0e4496658a366e2f745e1f0/tensorflow/python/keras/engine/network.py#L1352-L1357). + """ + try: + import tensorflow as tf # noqa: F401 + import torch # noqa: F401 + except ImportError: + logger.error( + "Loading a TensorFlow model in PyTorch, requires both PyTorch and TensorFlow to be installed. Please see " + "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions." + ) + raise + + import transformers + + from .modeling_tf_utils import load_tf_weights + + logger.info(f"Loading TensorFlow weights from {tf_checkpoint_path}") + + # Instantiate and load the associated TF 2.0 model + tf_model_class_name = "TF" + pt_model.__class__.__name__ # Add "TF" at the beginning + tf_model_class = getattr(transformers, tf_model_class_name) + tf_model = tf_model_class(pt_model.config) + + if tf_inputs is None: + tf_inputs = tf_model.dummy_inputs + + if tf_inputs is not None: + tf_model(tf_inputs, training=False) # Make sure model is built + + load_tf_weights(tf_model, tf_checkpoint_path) + + return load_tf2_model_in_pytorch_model( + pt_model, tf_model, allow_missing_keys=allow_missing_keys, output_loading_info=output_loading_info + ) + + +def load_tf2_model_in_pytorch_model(pt_model, tf_model, allow_missing_keys=False, output_loading_info=False): + """Load TF 2.0 model in a pytorch model""" + weights = tf_model.weights + + return load_tf2_weights_in_pytorch_model( + pt_model, weights, allow_missing_keys=allow_missing_keys, output_loading_info=output_loading_info + ) + + +def load_tf2_weights_in_pytorch_model(pt_model, tf_weights, allow_missing_keys=False, output_loading_info=False): + """Load TF2.0 symbolic weights in a PyTorch model""" + try: + import tensorflow as tf # noqa: F401 + import torch # noqa: F401 + except ImportError: + logger.error( + "Loading a TensorFlow model in PyTorch, requires both PyTorch and TensorFlow to be installed. Please see " + "https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions." + ) + raise + + tf_state_dict = {tf_weight.name: tf_weight.numpy() for tf_weight in tf_weights} + return load_tf2_state_dict_in_pytorch_model( + pt_model, tf_state_dict, allow_missing_keys=allow_missing_keys, output_loading_info=output_loading_info + ) + + +def load_tf2_state_dict_in_pytorch_model(pt_model, tf_state_dict, allow_missing_keys=False, output_loading_info=False): + import torch + + new_pt_params_dict = {} + current_pt_params_dict = dict(pt_model.named_parameters()) + + # Make sure we are able to load PyTorch base models as well as derived models (with heads) + # TF models always have a prefix, some of PyTorch models (base ones) don't + start_prefix_to_remove = "" + if not any(s.startswith(pt_model.base_model_prefix) for s in current_pt_params_dict.keys()): + start_prefix_to_remove = pt_model.base_model_prefix + "." + + # Build a map from potential PyTorch weight names to TF 2.0 Variables + tf_weights_map = {} + for name, tf_weight in tf_state_dict.items(): + pt_name, transpose = convert_tf_weight_name_to_pt_weight_name( + name, start_prefix_to_remove=start_prefix_to_remove, tf_weight_shape=tf_weight.shape + ) + tf_weights_map[pt_name] = (tf_weight, transpose) + + all_tf_weights = set(tf_weights_map.keys()) + loaded_pt_weights_data_ptr = {} + missing_keys_pt = [] + for pt_weight_name, pt_weight in current_pt_params_dict.items(): + # Handle PyTorch shared weight ()not duplicated in TF 2.0 + if pt_weight.data_ptr() in loaded_pt_weights_data_ptr: + new_pt_params_dict[pt_weight_name] = loaded_pt_weights_data_ptr[pt_weight.data_ptr()] + continue + + pt_weight_name_to_check = pt_weight_name + # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030 + key_components = pt_weight_name.split(".") + name = None + if key_components[-3::2] == ["parametrizations", "original0"]: + name = key_components[-2] + "_g" + elif key_components[-3::2] == ["parametrizations", "original1"]: + name = key_components[-2] + "_v" + if name is not None: + key_components = key_components[:-3] + [name] + pt_weight_name_to_check = ".".join(key_components) + + # Find associated numpy array in pytorch model state dict + if pt_weight_name_to_check not in tf_weights_map: + if allow_missing_keys: + missing_keys_pt.append(pt_weight_name) + continue + + raise AttributeError(f"{pt_weight_name} not found in TF 2.0 model") + + array, transpose = tf_weights_map[pt_weight_name_to_check] + + array = apply_transpose(transpose, array, pt_weight.shape, pt_to_tf=False) + + if numpy.isscalar(array): + array = numpy.array(array) + if not is_torch_tensor(array) and not is_numpy_array(array): + array = array.numpy() + if is_numpy_array(array): + # Convert to torch tensor + array = torch.from_numpy(array) + + new_pt_params_dict[pt_weight_name] = array + loaded_pt_weights_data_ptr[pt_weight.data_ptr()] = array + all_tf_weights.discard(pt_weight_name) + + missing_keys, unexpected_keys = pt_model.load_state_dict(new_pt_params_dict, strict=False) + missing_keys += missing_keys_pt + + # Some models may have keys that are not in the state by design, removing them before needlessly warning + # the user. + if pt_model._keys_to_ignore_on_load_missing is not None: + for pat in pt_model._keys_to_ignore_on_load_missing: + missing_keys = [k for k in missing_keys if re.search(pat, k) is None] + + if pt_model._keys_to_ignore_on_load_unexpected is not None: + for pat in pt_model._keys_to_ignore_on_load_unexpected: + unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None] + + if len(unexpected_keys) > 0: + logger.warning( + "Some weights of the TF 2.0 model were not used when initializing the PyTorch model" + f" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing" + f" {pt_model.__class__.__name__} from a TF 2.0 model trained on another task or with another architecture" + " (e.g. initializing a BertForSequenceClassification model from a TFBertForPreTraining model).\n- This IS" + f" NOT expected if you are initializing {pt_model.__class__.__name__} from a TF 2.0 model that you expect" + " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a" + " TFBertForSequenceClassification model)." + ) + else: + logger.warning(f"All TF 2.0 model weights were used when initializing {pt_model.__class__.__name__}.\n") + if len(missing_keys) > 0: + logger.warning( + f"Some weights of {pt_model.__class__.__name__} were not initialized from the TF 2.0 model and are newly" + f" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to" + " use it for predictions and inference." + ) + else: + logger.warning( + f"All the weights of {pt_model.__class__.__name__} were initialized from the TF 2.0 model.\n" + "If your task is similar to the task the model of the checkpoint was trained on, " + f"you can already use {pt_model.__class__.__name__} for predictions without further training." + ) + + logger.info(f"Weights or buffers not loaded from TF 2.0 model: {all_tf_weights}") + + if output_loading_info: + loading_info = {"missing_keys": missing_keys, "unexpected_keys": unexpected_keys} + return pt_model, loading_info + + return pt_model diff --git a/lib/python3.12/site-packages/transformers/modeling_tf_utils.py b/lib/python3.12/site-packages/transformers/modeling_tf_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c0add5bb6672bf77ca0fca34a5bc194ab80f89f5 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_tf_utils.py @@ -0,0 +1,3535 @@ +# coding=utf-8 +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. 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. +"""TF general model utils.""" + +from __future__ import annotations + +import functools +import gc +import inspect +import json +import os +import pickle +import re +import warnings +from collections.abc import Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union + +import h5py +import numpy as np +import tensorflow as tf +from packaging.version import parse + +from . import DataCollatorWithPadding, DefaultDataCollator +from .activations_tf import get_tf_activation +from .configuration_utils import PretrainedConfig +from .dynamic_module_utils import custom_object_save +from .generation import GenerationConfig, TFGenerationMixin +from .tf_utils import ( + convert_batch_encoding, + expand_1d, + load_attributes_from_hdf5_group, + save_attributes_to_hdf5_group, + shape_list, +) +from .utils import ( + SAFE_WEIGHTS_INDEX_NAME, + SAFE_WEIGHTS_NAME, + TF2_WEIGHTS_INDEX_NAME, + TF2_WEIGHTS_NAME, + TF_WEIGHTS_NAME, + WEIGHTS_INDEX_NAME, + WEIGHTS_NAME, + ModelOutput, + PushToHubMixin, + cached_file, + download_url, + find_labels, + has_file, + is_offline_mode, + is_remote_url, + is_safetensors_available, + is_tf_symbolic_tensor, + logging, + requires_backends, + working_or_temp_dir, +) +from .utils.hub import convert_file_size_to_int, get_checkpoint_shard_files + + +if is_safetensors_available(): + from safetensors import safe_open + from safetensors.tensorflow import save_file as safe_save_file + +if TYPE_CHECKING: + from . import PreTrainedTokenizerBase + +logger = logging.get_logger(__name__) + +if "TF_USE_LEGACY_KERAS" not in os.environ: + os.environ["TF_USE_LEGACY_KERAS"] = "1" # Compatibility fix to make sure tf.keras stays at Keras 2 +elif os.environ["TF_USE_LEGACY_KERAS"] != "1": + logger.warning( + "Transformers is only compatible with Keras 2, but you have explicitly set `TF_USE_LEGACY_KERAS` to `0`. " + "This may result in unexpected behaviour or errors if Keras 3 objects are passed to Transformers models." + ) + +try: + import tf_keras as keras + from tf_keras import backend as K +except (ModuleNotFoundError, ImportError): + import keras + from keras import backend as K + + if parse(keras.__version__).major > 2: + raise ValueError( + "Your currently installed version of Keras is Keras 3, but this is not yet supported in " + "Transformers. Please install the backwards-compatible tf-keras package with " + "`pip install tf-keras`." + ) + + +tf_logger = tf.get_logger() + +TFModelInputType = Union[ + List[tf.Tensor], + List[np.ndarray], + Dict[str, tf.Tensor], + Dict[str, np.ndarray], + tf.Tensor, + np.ndarray, +] + + +def dummy_loss(y_true, y_pred): + if y_pred.shape.rank <= 1: + return y_pred + else: + reduction_axes = list(range(1, y_pred.shape.rank)) + return tf.reduce_mean(y_pred, axis=reduction_axes) + + +class TFModelUtilsMixin: + """ + A few utilities for `keras.Model`, to be used as a mixin. + """ + + def num_parameters(self, only_trainable: bool = False) -> int: + """ + Get the number of (optionally, trainable) parameters in the model. + + Args: + only_trainable (`bool`, *optional*, defaults to `False`): + Whether or not to return only the number of trainable parameters + + Returns: + `int`: The number of parameters. + """ + if only_trainable: + return int(sum(np.prod(w.shape.as_list()) for w in self.trainable_variables)) + else: + return self.count_params() + + +def keras_serializable(cls): + """ + Decorate a Keras Layer class to support Keras serialization. + + This is done by: + + 1. Adding a `transformers_config` dict to the Keras config dictionary in `get_config` (called by Keras at + serialization time. + 2. Wrapping `__init__` to accept that `transformers_config` dict (passed by Keras at deserialization time) and + convert it to a config object for the actual layer initializer. + 3. Registering the class as a custom object in Keras (if the Tensorflow version supports this), so that it does not + need to be supplied in `custom_objects` in the call to `keras.models.load_model`. + + Args: + cls (a `keras.layers.Layers subclass`): + Typically a `TF.MainLayer` class in this project, in general must accept a `config` argument to its + initializer. + + Returns: + The same class object, with modifications for Keras deserialization. + """ + initializer = cls.__init__ + + config_class = getattr(cls, "config_class", None) + if config_class is None: + raise AttributeError("Must set `config_class` to use @keras_serializable") + + @functools.wraps(initializer) + def wrapped_init(self, *args, **kwargs): + config = args[0] if args and isinstance(args[0], PretrainedConfig) else kwargs.pop("config", None) + + if isinstance(config, dict): + config = config_class.from_dict(config) + initializer(self, config, *args, **kwargs) + elif isinstance(config, PretrainedConfig): + if len(args) > 0: + initializer(self, *args, **kwargs) + else: + initializer(self, config, *args, **kwargs) + else: + raise ValueError("Must pass either `config` (PretrainedConfig) or `config` (dict)") + + self._config = config + self._kwargs = kwargs + + cls.__init__ = wrapped_init + + if not hasattr(cls, "get_config"): + raise TypeError("Only use @keras_serializable on keras.layers.Layer subclasses") + if hasattr(cls.get_config, "_is_default"): + + def get_config(self): + cfg = super(cls, self).get_config() + cfg["config"] = self._config.to_dict() + cfg.update(self._kwargs) + return cfg + + cls.get_config = get_config + + cls._keras_serializable = True + if hasattr(keras.utils, "register_keras_serializable"): + cls = keras.utils.register_keras_serializable()(cls) + return cls + + +class TFCausalLanguageModelingLoss: + """ + Loss function suitable for causal language modeling (CLM), that is, the task of guessing the next token. + + + + Any label of -100 will be ignored (along with the corresponding logits) in the loss computation. + + + """ + + def hf_compute_loss(self, labels, logits): + loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE) + if self.config.tf_legacy_loss: + # make sure only labels that are not equal to -100 affect the loss + active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100) + reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss) + labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss) + return loss_fn(labels, reduced_logits) + + # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway + unmasked_loss = loss_fn(tf.nn.relu(labels), logits) + # make sure only labels that are not equal to -100 affect the loss + loss_mask = tf.cast(labels != -100, dtype=unmasked_loss.dtype) + masked_loss = unmasked_loss * loss_mask + reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask) + return tf.reshape(reduced_masked_loss, (1,)) + + +class TFQuestionAnsweringLoss: + """ + Loss function suitable for question answering. + """ + + def hf_compute_loss(self, labels, logits): + loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE) + start_loss = loss_fn(labels["start_position"], logits[0]) + end_loss = loss_fn(labels["end_position"], logits[1]) + + return (start_loss + end_loss) / 2.0 + + +class TFTokenClassificationLoss: + """ + Loss function suitable for token classification. + + + + Any label of -100 will be ignored (along with the corresponding logits) in the loss computation. + + + """ + + def hf_compute_loss(self, labels, logits): + loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE) + if tf.executing_eagerly(): # Data-dependent conditionals are forbidden in XLA + if tf.math.reduce_any(labels == -1): + tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.") + + if self.config.tf_legacy_loss: + # make sure only labels that are not equal to -100 + # are taken into account as loss + if tf.math.reduce_any(labels == -1): + tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.") + active_loss = tf.reshape(labels, (-1,)) != -1 + else: + active_loss = tf.reshape(labels, (-1,)) != -100 + reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss) + labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss) + + return loss_fn(labels, reduced_logits) + + # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway + unmasked_loss = loss_fn(tf.nn.relu(labels), logits) + # make sure only labels that are not equal to -100 or -1 + # are taken into account as loss + loss_mask = tf.cast(labels >= 0, dtype=unmasked_loss.dtype) + # Avoid possible division by zero later + # Masked positions will have a loss of NaN because -100 and -1 are not valid labels + masked_loss = unmasked_loss * loss_mask + reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask) + return tf.reshape(reduced_masked_loss, (1,)) + + +class TFSequenceClassificationLoss: + """ + Loss function suitable for sequence classification. + """ + + def hf_compute_loss(self, labels, logits): + if logits.shape.rank == 1 or logits.shape[1] == 1: + loss_fn = keras.losses.MeanSquaredError(reduction=keras.losses.Reduction.NONE) + if labels.shape.rank == 1: + # MeanSquaredError returns a scalar loss if the labels are 1D, so avoid that + labels = tf.expand_dims(labels, axis=-1) + else: + loss_fn = keras.losses.SparseCategoricalCrossentropy( + from_logits=True, reduction=keras.losses.Reduction.NONE + ) + + return loss_fn(labels, logits) + + +class TFMultipleChoiceLoss: + """Loss function suitable for multiple choice tasks.""" + + def hf_compute_loss(self, labels, logits): + loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE) + return loss_fn(labels, logits) + + +class TFMaskedLanguageModelingLoss(TFCausalLanguageModelingLoss): + """ + Loss function suitable for masked language modeling (MLM), that is, the task of guessing the masked tokens. + + + + Any label of -100 will be ignored (along with the corresponding logits) in the loss computation. + + + """ + + +class TFNextSentencePredictionLoss: + """ + Loss function suitable for next sentence prediction (NSP), that is, the task of guessing the next sentence. + + + + Any label of -100 will be ignored (along with the corresponding logits) in the loss computation. + + + """ + + def hf_compute_loss(self, labels, logits): + loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE) + if self.config.tf_legacy_loss: + # make sure only labels that are not equal to -100 + # are taken into account as loss + next_sentence_active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100) + next_sentence_reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, 2)), next_sentence_active_loss) + next_sentence_label = tf.boolean_mask(tf.reshape(labels, (-1,)), next_sentence_active_loss) + + return loss_fn(next_sentence_label, next_sentence_reduced_logits) + + # make sure only labels that are not equal to -100 + # are taken into account as loss + + # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway + unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels), y_pred=logits) + ns_loss_mask = tf.cast(labels != -100, dtype=unmasked_ns_loss.dtype) + # Just zero out samples where label is -100, no reduction + masked_ns_loss = unmasked_ns_loss * ns_loss_mask + + return masked_ns_loss + + +def booleans_processing(config, **kwargs): + """ + Process the input booleans of each model. + + Args: + config ([`PretrainedConfig`]): + The config of the running model. + **kwargs: + The boolean parameters + + Returns: + A dictionary with the proper values for each boolean + """ + final_booleans = {} + + # Pure conv models (such as ConvNext) do not have `output_attentions`. If the signature has + # `output_attentions`, it will be present here in `kwargs`, even if unset (in that case, as `None`) + if "output_attentions" in kwargs: + final_booleans["output_attentions"] = ( + kwargs["output_attentions"] if kwargs["output_attentions"] is not None else config.output_attentions + ) + final_booleans["output_hidden_states"] = ( + kwargs["output_hidden_states"] if kwargs["output_hidden_states"] is not None else config.output_hidden_states + ) + final_booleans["return_dict"] = kwargs["return_dict"] if kwargs["return_dict"] is not None else config.return_dict + + if "use_cache" in kwargs: + final_booleans["use_cache"] = ( + kwargs["use_cache"] if kwargs["use_cache"] is not None else getattr(config, "use_cache", None) + ) + return final_booleans + + +def unpack_inputs(func): + """ + Decorator that processes the inputs to a Keras layer, passing them to the layer as keyword arguments. This enables + downstream use of the inputs by their variable name, even if they arrive packed as a dictionary in the first input + (common case in Keras). + + Args: + func (`callable`): + The callable function of the TensorFlow model. + + + Returns: + A callable that wraps the original `func` with the behavior described above. + """ + + original_signature = inspect.signature(func) + + @functools.wraps(func) + def run_call_with_unpacked_inputs(self, *args, **kwargs): + # isolates the actual `**kwargs` for the decorated function + kwargs_call = {key: val for key, val in kwargs.items() if key not in dict(original_signature.parameters)} + fn_args_and_kwargs = {key: val for key, val in kwargs.items() if key not in kwargs_call} + fn_args_and_kwargs.update({"kwargs_call": kwargs_call}) + + # move any arg into kwargs, if they exist + fn_args_and_kwargs.update(dict(zip(func.__code__.co_varnames[1:], args))) + + # Encoder Decoder models delegate the application of the configuration options to their inner models. + if "EncoderDecoder" in self.__class__.__name__: + config = None + else: + config = self.config + + unpacked_inputs = input_processing(func, config, **fn_args_and_kwargs) + return func(self, **unpacked_inputs) + + # Keras enforces the first layer argument to be passed, and checks it through `inspect.getfullargspec()`. This + # function does not follow wrapper chains (i.e. ignores `functools.wraps()`), meaning that without the line below + # Keras would attempt to check the first argument against the literal signature of the wrapper. + run_call_with_unpacked_inputs.__signature__ = original_signature + + return run_call_with_unpacked_inputs + + +def input_processing(func, config, **kwargs): + """ + Process the input of each TensorFlow model including the booleans. In case of a list of symbolic inputs, each input + has to be named accordingly to the parameters name, i.e. `input_ids = keras.Input(shape=(128,), dtype='int32', + name="input_ids")` otherwise the order of the tensors will not be guaranteed during the training. + + Args: + func (`callable`): + The callable function of the TensorFlow model. + config ([`PretrainedConfig`]): + The config of the running model. + **kwargs: + The inputs of the model. + + Returns: + Two lists, one for the missing layers, and another one for the unexpected layers. + """ + signature = dict(inspect.signature(func).parameters) + has_kwargs = bool(signature.pop("kwargs", None)) + signature.pop("self", None) + parameter_names = list(signature.keys()) + main_input_name = parameter_names[0] + main_input = kwargs.pop(main_input_name, None) + output = {} + allowed_types = (tf.Tensor, bool, int, ModelOutput, tuple, list, dict, np.ndarray) + + if "inputs" in kwargs["kwargs_call"]: + warnings.warn( + "The `inputs` argument is deprecated and will be removed in a future version, use `input_ids` instead.", + FutureWarning, + ) + + output["input_ids"] = kwargs["kwargs_call"].pop("inputs") + + if "decoder_cached_states" in kwargs["kwargs_call"]: + warnings.warn( + "The `decoder_cached_states` argument is deprecated and will be removed in a future version, use" + " `past_key_values` instead.", + FutureWarning, + ) + output["past_key_values"] = kwargs["kwargs_call"].pop("decoder_cached_states") + + if "past" in kwargs["kwargs_call"] and "past_key_values" in parameter_names: + warnings.warn( + "The `past` argument is deprecated and will be removed in a future version, use `past_key_values`" + " instead.", + FutureWarning, + ) + kwargs["past_key_values"] = kwargs["kwargs_call"].pop("past") + elif "past_key_values" in kwargs["kwargs_call"] and "past" in parameter_names: + kwargs["past"] = kwargs["kwargs_call"].pop("past_key_values") + + if has_kwargs: + output["kwargs"] = kwargs.pop("kwargs_call", {}) + else: + if len(kwargs["kwargs_call"]) > 0: + raise ValueError( + "The following keyword arguments are not supported by this model:" + f" {list(kwargs['kwargs_call'].keys())}." + ) + kwargs.pop("kwargs_call") + + for k, v in kwargs.items(): + if isinstance(v, allowed_types) or tf.is_tensor(v) or v is None: + output[k] = v + else: + raise ValueError(f"Data of type {type(v)} is not allowed only {allowed_types} is accepted for {k}.") + + if isinstance(main_input, (tuple, list)): + for i, input in enumerate(main_input): + # EagerTensors don't allow to use the .name property so we check for a real Tensor + if is_tf_symbolic_tensor(input): + # Tensor names have always the pattern `name:id` then we check only the + # `name` part + tensor_name = input.name.split(":")[0] + + if tensor_name in parameter_names: + output[tensor_name] = input + else: + output[parameter_names[i]] = input + elif isinstance(input, allowed_types) or input is None: + output[parameter_names[i]] = input + else: + raise ValueError( + f"Data of type {type(input)} is not allowed only {allowed_types} is accepted for" + f" {parameter_names[i]}." + ) + elif isinstance(main_input, Mapping): + if "inputs" in main_input: + warnings.warn( + "The `inputs` argument is deprecated and will be removed in a future version, use `input_ids`" + " instead.", + FutureWarning, + ) + + output["input_ids"] = main_input.pop("inputs") + + if "decoder_cached_states" in main_input: + warnings.warn( + "The `decoder_cached_states` argument is deprecated and will be removed in a future version, use" + " `past_key_values` instead.", + FutureWarning, + ) + output["past_key_values"] = main_input.pop("decoder_cached_states") + + for k, v in dict(main_input).items(): + if isinstance(v, allowed_types) or v is None: + output[k] = v + elif k not in parameter_names and "args" not in parameter_names: + logger.warning( + f"The parameter {k} does not belongs to the parameter list {parameter_names} and will be ignored." + ) + continue + else: + raise ValueError(f"Data of type {type(v)} is not allowed only {allowed_types} is accepted for {k}.") + else: + if tf.is_tensor(main_input) or main_input is None: + output[main_input_name] = main_input + else: + raise ValueError( + f"Data of type {type(main_input)} is not allowed only {allowed_types} is accepted for" + f" {main_input_name}." + ) + + # Populates any unspecified argument with their default value, according to the signature. + for name in parameter_names: + if name not in list(output.keys()) and name != "args": + output[name] = kwargs.pop(name, signature[name].default) + + # When creating a SavedModel TF calls the method with LayerCall.__call__(args, **kwargs) + # So to respect the proper output we have to add this exception + if "args" in output: + if output["args"] is not None and is_tf_symbolic_tensor(output["args"]): + tensor_name = output["args"].name.split(":")[0] + output[tensor_name] = output["args"] + else: + # `args` in this case is always the first parameter, then `input_ids` + output["input_ids"] = output["args"] + + del output["args"] + + if "kwargs" in output: + del output["kwargs"] + + cast_output = {} + for key, val in output.items(): + if isinstance(val, tf.Tensor) and val.dtype == tf.int64: + cast_output[key] = tf.cast(val, tf.int32) + elif isinstance(val, np.ndarray) and val.dtype == np.int64: + cast_output[key] = val.astype(np.int32) + else: + cast_output[key] = val + + output = cast_output + del cast_output + + if config is not None: + boolean_dict = { + k: v + for k, v in output.items() + if k in ["return_dict", "output_attentions", "output_hidden_states", "use_cache"] + } + + output.update( + booleans_processing( + config=config, + **boolean_dict, + ) + ) + + return output + + +def strip_model_name_and_prefix(name, _prefix=None): + if _prefix is not None and name.startswith(_prefix): + name = name[len(_prefix) :] + if name.startswith("/"): + name = name[1:] + if "model." not in name and len(name.split("/")) > 1: + name = "/".join(name.split("/")[1:]) + return name + + +def tf_shard_checkpoint(weights, max_shard_size="10GB", weights_name: str = TF2_WEIGHTS_NAME): + """ + Splits a model state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a + given size. + + The sub-checkpoints are determined by iterating through the `state_dict` in the order of its keys, so there is no + optimization made to make each sub-checkpoint as close as possible to the maximum size passed. For example, if the + limit is 10GB and we have weights 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]. + + + + If one of the model's weight is bigger that `max_shard_size`, it will end up in its own sub-checkpoint which will + have a size greater than `max_shard_size`. + + + + Args: + weights (`Dict[str, tf.RessourceVariable]`): The list of tf.RessourceVariable of a model to save. + max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`): + The maximum size of each sub-checkpoint. If expressed as a string, needs to be digits followed by a unit + (like `"5MB"`). + """ + max_shard_size = convert_file_size_to_int(max_shard_size) + + sharded_state_dicts = [] + current_block = [] + current_block_size = 0 + total_size = 0 + + for item in weights: + weight_size = item.numpy().size * item.dtype.size + + # If this weight is going to tip up over the maximal size, we split. + if current_block_size + weight_size > max_shard_size: + sharded_state_dicts.append(current_block) + current_block = [] + current_block_size = 0 + + current_block.append(item) + current_block_size += weight_size + total_size += weight_size + + # Add the last block + sharded_state_dicts.append(current_block) + + # If we only have one shard, we return it + if len(sharded_state_dicts) == 1: + return {weights_name: sharded_state_dicts[0]}, None + + # Otherwise, let's build the index + weight_map = {} + shards = {} + for idx, shard in enumerate(sharded_state_dicts): + shard_file = weights_name.replace(".h5", f"-{idx + 1:05d}-of-{len(sharded_state_dicts):05d}.h5") + shard_file = shard_file.replace( + ".safetensors", f"-{idx + 1:05d}-of-{len(sharded_state_dicts):05d}.safetensors" + ) + shards[shard_file] = shard + for weight in shard: + weight_name = weight.name + weight_map[weight_name] = shard_file + + # Add the metadata + metadata = {"total_size": total_size} + index = {"metadata": metadata, "weight_map": weight_map} + return shards, index + + +def load_tf_sharded_weights(model, shard_files, ignore_mismatched_sizes=False, strict=False, _prefix=None): + """ + This is the same as `load_tf_weights` but for a sharded checkpoint. Detect missing and unexpected layers and load + the TF weights from the shard file accordingly to their names and shapes. + + This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being + loaded in the model. + + Args: + model (`keras.models.Model`): The model in which to load the checkpoint. + shard_files (`str` or `os.PathLike`): A list containing the sharded checkpoint names. + ignore_mismatched_sizes`bool`, *optional`, defaults to `True`): + Whether or not to ignore the mismatch between the sizes + strict (`bool`, *optional*, defaults to `True`): + Whether to strictly enforce that the keys in the model state dict match the keys in the sharded checkpoint. + + Returns: + Three lists, one for the missing layers, another one for the unexpected layers, and a last one for the + mismatched layers. + """ + + # Load the index + unexpected_keys = set() + saved_keys = set() + mismatched_keys = set() + + # Since TF adds the name of the class to its weights, and uses the index and not the name of the layer to load + # the weight, we have to get rid of the first prefix of the name of the layer. + model_keys = set() + model_layer_map = {} + for i, k in enumerate(model.weights): + layer_name = k.name + if _prefix is not None and layer_name.startswith(_prefix): + layer_name = layer_name[len(_prefix) :] + layer_name = layer_name.lstrip("/") + if not ("model." in layer_name or len(layer_name.split("/")) == 1): + layer_name = "/".join(layer_name.split("/")[1:]) + model_keys.add(layer_name) + model_layer_map[layer_name] = i + + for shard_file in shard_files: + saved_weight_names_set, unexpected_keys_set, mismatched_keys_set = load_tf_shard( + model, + model_layer_map, + shard_file, + ignore_mismatched_sizes=ignore_mismatched_sizes, + _prefix=_prefix, + ) + saved_keys.update(saved_weight_names_set) + unexpected_keys.update(unexpected_keys_set) + mismatched_keys.update(mismatched_keys_set) + gc.collect() + + missing_keys = model_keys - saved_keys + if strict and (len(missing_keys) > 0 or len(unexpected_keys) > 0): + error_message = f"Error(s) in loading state_dict for {model.__class__.__name__}" + if len(missing_keys) > 0: + str_missing_keys = ",".join([f'"{k}"' for k in missing_keys]) + error_message += f"\nMissing key(s): {str_missing_keys}." + if len(unexpected_keys) > 0: + str_unexpected_keys = ",".join([f'"{k}"' for k in unexpected_keys]) + error_message += f"\nMissing key(s): {str_unexpected_keys}." + raise RuntimeError(error_message) + + return missing_keys, unexpected_keys, mismatched_keys + + +def load_tf_shard(model, model_layer_map, resolved_archive_file, ignore_mismatched_sizes=False, _prefix=None): + """ + Loads a shard from a sharded checkpoint file. Can be either H5 or Safetensors. + Handles missing keys and unexpected keys. + + Args: + model (`keras.models.Model`): Model in which the weights are loaded + model_layer_map (`Dict`): A dictionary mapping the layer name to the index of the layer in the model. + resolved_archive_file (`str`): Path to the checkpoint file from which the weights will be loaded + ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`): Whether to ignore the mismatched keys + + Returns: + `keras.models.Model`: Three lists, one for the layers that were found and successfully restored (from the + shard file), one for the mismatched layers, and another one for the unexpected layers. + """ + saved_weight_names_set = set() + saved_weights = {} + mismatched_keys = set() + unexpected_keys = set() + # Read the H5 file + try: + with h5py.File(resolved_archive_file, "r") as sharded_checkpoint_file: + # Retrieve the name of each layer from the H5 file + saved_h5_model_layers_name = set(load_attributes_from_hdf5_group(sharded_checkpoint_file, "layer_names")) + weight_value_tuples = [] + + # Compute missing and unexpected sub layers + # Store the weights in list of tuples that looks like [(weight_object, value_of_weight),...] + for layer_name in saved_h5_model_layers_name: + h5_layer_object = sharded_checkpoint_file[layer_name] + saved_weights[layer_name] = np.asarray(h5_layer_object) + + saved_weight_names_set.add(layer_name) + + if layer_name not in model_layer_map: + unexpected_keys.add(layer_name) + else: + symbolic_weight = model.weights[model_layer_map[layer_name]] + + saved_weight_value = saved_weights[layer_name] + # If the current weight is found + if saved_weight_value is not None: + # Check if the shape of the current weight and the one from the H5 file are different + if K.int_shape(symbolic_weight) != saved_weight_value.shape: + # If yes we reshape the weight from the H5 file accordingly to the current weight + # If the two shapes are not compatible we raise an issue + try: + array = np.reshape(saved_weight_value, K.int_shape(symbolic_weight)) + except ValueError as e: + if ignore_mismatched_sizes: + mismatched_keys.add( + (layer_name, saved_weight_value.shape, K.int_shape(symbolic_weight)) + ) + continue + else: + raise e + else: + array = saved_weight_value + + # We create the tuple that will be loaded and add it to the final list + weight_value_tuples.append((symbolic_weight, array)) + + K.batch_set_value(weight_value_tuples) + + return saved_weight_names_set, unexpected_keys, mismatched_keys + + except Exception as e: + try: + with open(resolved_archive_file) as f: + if f.read().startswith("version"): + raise OSError( + "You seem to have cloned a repository without having git-lfs installed. Please install " + "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder " + "you cloned." + ) + else: + raise ValueError( + f"Unable to locate the file {resolved_archive_file} which is necessary to load this pretrained" + " model. Make sure you have saved the model properly." + ) from e + except (UnicodeDecodeError, ValueError): + raise OSError( + f"Unable to load weights from TF checkpoint file for '{resolved_archive_file}' " + f"at '{resolved_archive_file}'. " + "If you tried to load a TF model from a sharded checkpoint, you should try converting the model " + "by loading it in pytorch and saving it locally. A convertion script should be released soon." + ) + + +def load_tf_sharded_weights_from_safetensors( + model, shard_files, ignore_mismatched_sizes=False, strict=False, _prefix=None +): + """ + This is the same as `load_tf_weights_from_safetensors` but for a sharded TF-format safetensors checkpoint. + Detect missing and unexpected layers and load the TF weights from the shard file accordingly to their names and + shapes. + + This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being + loaded in the model. + + Args: + model (`keras.models.Model`): The model in which to load the checkpoint. + shard_files (`str` or `os.PathLike`): A list containing the sharded checkpoint names. + ignore_mismatched_sizes`bool`, *optional`, defaults to `True`): + Whether or not to ignore the mismatch between the sizes + strict (`bool`, *optional*, defaults to `True`): + Whether to strictly enforce that the keys in the model state dict match the keys in the sharded checkpoint. + + Returns: + Three lists, one for the missing layers, another one for the unexpected layers, and a last one for the + mismatched layers. + """ + + # Load the index + unexpected_keys = set() + all_missing_keys = [] + mismatched_keys = set() + + for shard_file in shard_files: + missing_layers, unexpected_layers, mismatched_layers = load_tf_weights_from_safetensors( + model, + shard_file, + ignore_mismatched_sizes=ignore_mismatched_sizes, + _prefix=_prefix, + ) + all_missing_keys.append(set(missing_layers)) + unexpected_keys.update(unexpected_layers) + mismatched_keys.update(mismatched_layers) + gc.collect() + missing_keys = set.intersection(*all_missing_keys) + + if strict and (len(missing_keys) > 0 or len(unexpected_keys) > 0): + error_message = f"Error(s) in loading state_dict for {model.__class__.__name__}" + if len(missing_keys) > 0: + str_missing_keys = ",".join([f'"{k}"' for k in missing_keys]) + error_message += f"\nMissing key(s): {str_missing_keys}." + if len(unexpected_keys) > 0: + str_unexpected_keys = ",".join([f'"{k}"' for k in unexpected_keys]) + error_message += f"\nMissing key(s): {str_unexpected_keys}." + raise RuntimeError(error_message) + + return missing_keys, unexpected_keys, mismatched_keys + + +def load_tf_weights(model, resolved_archive_file, ignore_mismatched_sizes=False, _prefix=None): + """ + Detect missing and unexpected layers and load the TF weights from the shard file accordingly to their names and + shapes. + + Args: + model (`keras.models.Model`): + The model to load the weights into. + resolved_archive_file (`str`): + The location of the H5 file. + ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`): + Whether or not to ignore weights with shapes that don't match between the checkpoint of the model. + + Returns: + Three lists, one for the missing layers, another one for the unexpected layers, and a last one for the + mismatched layers. + """ + if resolved_archive_file.endswith(".safetensors"): + load_function = load_tf_weights_from_safetensors + else: + load_function = load_tf_weights_from_h5 + + return load_function( + model, resolved_archive_file, ignore_mismatched_sizes=ignore_mismatched_sizes, _prefix=_prefix + ) + + +def load_tf_weights_from_h5(model, resolved_archive_file, ignore_mismatched_sizes=False, _prefix=None): + mismatched_layers = [] + + # Read the H5 file + with h5py.File(resolved_archive_file, "r") as sharded_checkpoint_file: + # Retrieve the name of each layer from the H5 file + saved_h5_model_layers_name = set(load_attributes_from_hdf5_group(sharded_checkpoint_file, "layer_names")) + + # Find the missing layers from the high level list of layers + missing_layers = list({layer.name for layer in model.layers} - saved_h5_model_layers_name) + + # Find the unexpected layers from the high level list of layers + unexpected_layers = list(saved_h5_model_layers_name - {layer.name for layer in model.layers}) + saved_weight_names_set = set() + symbolic_weights_names = set() + weight_value_tuples = [] + + # Compute missing and unexpected sub layers + # Store the weights in list of tuples that looks like [(weight_object, value_of_weight),...] + for layer in model.layers: + # if layer_name from the H5 file belongs to the layers from the instantiated model + if layer.name in saved_h5_model_layers_name: + # Get the H5 layer object from its name + h5_layer_object = sharded_checkpoint_file[layer.name] + # Get all the weights as a list from the layer object + symbolic_weights = layer.trainable_weights + layer.non_trainable_weights + saved_weights = {} + + # Create a dict from the H5 saved model that looks like {"weight_name": weight_value} + # And a set with only the names + for weight_name in load_attributes_from_hdf5_group(h5_layer_object, "weight_names"): + # TF names always start with the model name so we ignore it + name = "/".join(weight_name.split("/")[1:]) + + if _prefix is not None: + name = _prefix + "/" + name + + saved_weights[name] = np.asarray(h5_layer_object[weight_name]) + + # Add the updated name to the final list for computing missing/unexpected values + saved_weight_names_set.add(name) + + # Loop over each weights from the instantiated model and compare with the weights from the H5 file + for symbolic_weight in symbolic_weights: + # TF names always start with the model name so we ignore it + if _prefix is not None: + delimeter = len(_prefix.split("/")) + symbolic_weight_name = "/".join( + symbolic_weight.name.split("/")[:delimeter] + + symbolic_weight.name.split("/")[delimeter + 1 :] + ) + else: + symbolic_weight_name = "/".join(symbolic_weight.name.split("/")[1:]) + + # here we check if the current weight is among the weights from the H5 file + # If yes, get the weight_value of the corresponding weight from the H5 file + # If not, make the value to None + saved_weight_value = saved_weights.get(symbolic_weight_name, None) + + # Retrocompatibility patch: some embeddings are stored with the weights name (e.g. Bart's + # `model.shared/embeddings:0` are stored as `model.shared/weights:0`) + if saved_weight_value is None and symbolic_weight_name.endswith("embeddings:0"): + symbolic_weight_name = symbolic_weight_name[:-12] + "weight:0" + saved_weight_value = saved_weights.get(symbolic_weight_name, None) + + # Add the updated name to the final list for computing missing/unexpected values + symbolic_weights_names.add(symbolic_weight_name) + + # If the current weight is found + if saved_weight_value is not None: + # Check if the shape of the current weight and the one from the H5 file are different + if K.int_shape(symbolic_weight) != saved_weight_value.shape: + # If yes we reshape the weight from the H5 file accordingly to the current weight + # If the two shapes are not compatible we raise an issue + try: + array = np.reshape(saved_weight_value, K.int_shape(symbolic_weight)) + except ValueError as e: + if ignore_mismatched_sizes: + mismatched_layers.append( + (symbolic_weight_name, saved_weight_value.shape, K.int_shape(symbolic_weight)) + ) + continue + else: + raise e + else: + array = saved_weight_value + + # We create the tuple that will be loaded and add it to the final list + weight_value_tuples.append((symbolic_weight, array)) + + # Load all the weights + K.batch_set_value(weight_value_tuples) + + # Compute the missing and unexpected layers + missing_layers.extend(list(symbolic_weights_names - saved_weight_names_set)) + unexpected_layers.extend(list(saved_weight_names_set - symbolic_weights_names)) + + return missing_layers, unexpected_layers, mismatched_layers + + +def load_tf_weights_from_safetensors(model, resolved_archive_file, ignore_mismatched_sizes=False, _prefix=None): + # Read the safetensors file + with safe_open(resolved_archive_file, framework="tf") as safetensors_archive: + mismatched_layers = [] + weight_names = [strip_model_name_and_prefix(w.name, _prefix=_prefix) for w in model.weights] + loaded_weight_names = list(safetensors_archive.keys()) + # Find the missing layers from the high level list of layers + missing_layers = list(set(weight_names) - set(loaded_weight_names)) + # Find the unexpected layers from the high level list of layers + unexpected_layers = list(set(loaded_weight_names) - set(weight_names)) + + for weight in model.weights: + weight_name = strip_model_name_and_prefix(weight.name, _prefix=_prefix) + if weight_name in loaded_weight_names: + weight_value = safetensors_archive.get_tensor(weight_name) + # Check if the shape of the current weight and the one from the H5 file are different + if K.int_shape(weight) != weight_value.shape: + # If yes we reshape the weight from the H5 file accordingly to the current weight + # If the two shapes are not compatible we raise an issue + try: + weight_value = tf.reshape(weight_value, K.int_shape(weight)) + except (ValueError, tf.errors.InvalidArgumentError) as e: + if ignore_mismatched_sizes: + mismatched_layers.append((weight_name, weight_value.shape, K.int_shape(weight))) + continue + else: + raise e + + K.set_value(weight, weight_value) # weight.assign() might break if weight is a DTensor + return missing_layers, unexpected_layers, mismatched_layers + + +def init_copy_embeddings(old_embeddings, new_num_tokens): + r""" + This function aims to reduce the embeddings in case new_num_tokens < old_num_tokens or to pad with -1 in case + new_num_tokens > old_num_tokens. A mask is also computed in order to know which weight in the embeddings should be + kept or not. Example: + + - if new_num_tokens=5 and old_num_tokens=4 and old_embeddings=[w1,w2,w3,w4] + + - mask=[True,True,True,True,False] and current_weights=[w1,w2,w3,w4,-1] + - if new_num_tokens=4 and old_num_tokens=5 and old_embeddings=[w1,w2,w3,w4,w5] + + - mask=[True,True,True,True] and current_weights=[w1,w2,w3,w4] + """ + old_num_tokens, old_embedding_dim = shape_list(old_embeddings) + size_diff = new_num_tokens - old_num_tokens + + # initialize new embeddings + # Copy token embeddings from the previous ones + if tf.math.greater(size_diff, 0): + # if the new size is greater than the old one, we extend the current embeddings with a padding until getting new size + # and we create a mask to properly identify the padded values and be replaced by the values of the newly created + # embeddings + current_weights = tf.pad( + old_embeddings.value(), tf.convert_to_tensor([[0, size_diff], [0, 0]]), constant_values=-1 + ) + num_tokens_to_copy = min(old_num_tokens, new_num_tokens) + mask = tf.fill(tf.convert_to_tensor([num_tokens_to_copy, 1]), True) + mask = tf.pad(mask, tf.convert_to_tensor([[0, size_diff], [0, 0]]), constant_values=False) + else: + # if the new size if lower than the old one, we take the current embeddings until the new size + current_weights = tf.slice( + old_embeddings.value(), + tf.convert_to_tensor([0, 0]), + tf.convert_to_tensor([new_num_tokens, old_embedding_dim]), + ) + mask = tf.fill(tf.convert_to_tensor([new_num_tokens, 1]), True) + + return mask, current_weights + + +class TFPreTrainedModel(keras.Model, TFModelUtilsMixin, TFGenerationMixin, PushToHubMixin): + r""" + Base class for all TF models. + + [`TFPreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading, + downloading and saving models as well as a few methods common to all models to: + + - resize the input embeddings, + - prune heads in the self-attention heads. + + Class attributes (overridden by derived classes): + + - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class + for this model architecture. + - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived + classes of the same architecture adding modules on top of the base model. + - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP + models, `pixel_values` for vision models and `input_values` for speech models). + """ + + config_class = None + base_model_prefix = "" + main_input_name = "input_ids" + _auto_class = None + _using_dummy_loss = None + _label_to_output_map = None + + # a list of re pattern of tensor names to ignore from the model when loading the model weights + # (and avoid unnecessary warnings). + _keys_to_ignore_on_load_missing = None + # a list of re pattern of tensor names to ignore from the weights when loading the model weights + # (and avoid unnecessary warnings). + _keys_to_ignore_on_load_unexpected = None + _requires_load_weight_prefix = False + + @property + def dummy_inputs(self) -> Dict[str, tf.Tensor]: + """ + Dummy inputs to build the network. + + Returns: + `Dict[str, tf.Tensor]`: The dummy inputs. + """ + dummies = {} + for key, spec in self.input_signature.items(): + # 2 is the most correct arbitrary size. I will not be taking questions + dummy_shape = [dim if dim is not None else 2 for dim in spec.shape] + if spec.shape[0] is None: + # But let's make the batch size 1 to save memory anyway + dummy_shape[0] = 1 + dummies[key] = tf.ones(shape=dummy_shape, dtype=spec.dtype) + if key == "token_type_ids": + # Some models have token_type_ids but with a vocab_size of 1 + dummies[key] = tf.zeros_like(dummies[key]) + if self.config.add_cross_attention and "encoder_hidden_states" in inspect.signature(self.call).parameters: + if "encoder_hidden_states" not in dummies: + if self.main_input_name == "input_ids": + dummies["encoder_hidden_states"] = tf.ones( + shape=(1, 2, self.config.hidden_size), dtype=tf.float32, name="encoder_hidden_states" + ) + else: + raise NotImplementedError( + "Model has cross-attention but we couldn't infer the shape for the encoder hidden states. Please manually override dummy_inputs!" + ) + return dummies + + def build_in_name_scope(self): + with tf.name_scope(self.name): + self.build(input_shape=None) + + @property + def framework(self) -> str: + """ + :str: Identifies that this is a TensorFlow model. + """ + return "tf" + + def build(self, input_shape=None): + pass # This is just here to make sure we don't call the superclass build() + + def __init__(self, config, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + if not isinstance(config, PretrainedConfig): + raise TypeError( + f"Parameter config in `{self.__class__.__name__}(config)` should be an instance of class " + "`PretrainedConfig`. To create a model from a pretrained model use " + f"`model = {self.__class__.__name__}.from_pretrained(PRETRAINED_MODEL_NAME)`" + ) + # Save config and origin of the pretrained weights if given in model + self.config = config + self.name_or_path = config.name_or_path + self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None + self._set_save_spec(self.input_signature) + + def get_config(self): + return self.config.to_dict() + + @functools.wraps(keras.Model.fit) + def fit(self, *args, **kwargs): + args, kwargs = convert_batch_encoding(*args, **kwargs) + return super().fit(*args, **kwargs) + + @functools.wraps(keras.Model.train_on_batch) + def train_on_batch(self, *args, **kwargs): + args, kwargs = convert_batch_encoding(*args, **kwargs) + return super().train_on_batch(*args, **kwargs) + + @functools.wraps(keras.Model.test_on_batch) + def test_on_batch(self, *args, **kwargs): + args, kwargs = convert_batch_encoding(*args, **kwargs) + return super().test_on_batch(*args, **kwargs) + + @functools.wraps(keras.Model.predict_on_batch) + def predict_on_batch(self, *args, **kwargs): + args, kwargs = convert_batch_encoding(*args, **kwargs) + return super().predict_on_batch(*args, **kwargs) + + @functools.wraps(keras.Model.predict) + def predict(self, *args, **kwargs): + args, kwargs = convert_batch_encoding(*args, **kwargs) + return super().predict(*args, **kwargs) + + @functools.wraps(keras.Model.evaluate) + def evaluate(self, *args, **kwargs): + args, kwargs = convert_batch_encoding(*args, **kwargs) + return super().evaluate(*args, **kwargs) + + @classmethod + def from_config(cls, config, **kwargs): + if isinstance(config, PretrainedConfig): + return cls._from_config(config, **kwargs) + return cls._from_config(cls.config_class.from_dict(config, **kwargs)) + + @classmethod + def _from_config(cls, config, **kwargs): + """ + All context managers that the model should be initialized under go here. + """ + return cls(config, **kwargs) + + def get_head_mask(self, head_mask: tf.Tensor | None, num_hidden_layers: int) -> tf.Tensor: + """ + Prepare the head mask if needed. + + Args: + head_mask (`tf.Tensor` with shape `[num_heads]` or `[num_hidden_layers x num_heads]`, *optional*): + The mask indicating if we should keep the heads or not (1.0 for keep, 0.0 for discard). + num_hidden_layers (`int`): + The number of hidden layers in the model. + + Returns: + `tf.Tensor` with shape `[num_hidden_layers x batch x num_heads x seq_length x seq_length]` or list with + `[None]` for each layer. + """ + if head_mask is not None: + head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers) + else: + head_mask = [None] * num_hidden_layers + + return head_mask + + def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers): + """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]""" + if head_mask.shape.rank == 1: + head_mask = head_mask[None, None, :, None, None] + head_mask = tf.repeat(head_mask, repeats=num_hidden_layers, axis=0) + elif head_mask.shape.rank == 2: + head_mask = head_mask[:, None, :, None, None] + assert head_mask.shape.rank == 5, f"head_mask.dim != 5, instead {head_mask.dim()}" + head_mask = tf.cast(head_mask, tf.float32) # switch to float if need + fp16 compatibility + return head_mask + + @tf.function + def serving(self, inputs): + """ + Args: + Method used for serving the model. Does not have a specific signature, but will be specialized as concrete + functions when saving with `save_pretrained`. + inputs (`Dict[str, tf.Tensor]`): + The input of the saved model as a dictionary of tensors. + """ + output = self.call(inputs) + + return self.serving_output(output) + + @property + def input_signature(self) -> Dict[str, tf.TensorSpec]: + """ + This property should return a dict mapping input names to tf.TensorSpec objects, representing the expected + shape and dtype for model inputs. It is used for both serving and for generating dummy inputs. + """ + model_inputs = list(inspect.signature(self.call).parameters) + sig = {} + if "input_ids" in model_inputs: + if self.__class__.__name__.endswith("ForMultipleChoice"): + text_dims = 3 + else: + text_dims = 2 + for input_name in ( + "input_ids", + "attention_mask", + "token_type_ids", + "decoder_input_ids", + "decoder_attention_mask", + ): + if input_name in model_inputs: + sig[input_name] = tf.TensorSpec([None] * text_dims, tf.int32, name=input_name) + if "pixel_values" in model_inputs: + pixel_values_shape = [None, None, None, None] + if hasattr(self.config, "vision_config"): + vision_config = self.config.vision_config + else: + vision_config = self.config + if hasattr(vision_config, "num_channels"): + pixel_values_shape[1] = vision_config.num_channels + else: + raise NotImplementedError( + "Could not infer number of channels from config, please override input_signature to specify input shapes." + ) + if hasattr(vision_config, "image_size"): + pixel_values_shape[2] = pixel_values_shape[3] = vision_config.image_size + elif hasattr(vision_config, "input_size"): + pixel_values_shape[2] = pixel_values_shape[3] = vision_config.input_size + else: + raise NotImplementedError( + "Could not infer input image shape from config, please override input_signature to specify input shapes." + ) + sig["pixel_values"] = tf.TensorSpec(pixel_values_shape, tf.float32, name="pixel_values") + if "input_features" in model_inputs: + raise NotImplementedError("Audio models need a manually defined input_signature") + return sig + + def serving_output(self, output): + """ + Prepare the output of the saved model. Can be overridden if specific serving modifications are required. + """ + if not isinstance(output, ModelOutput): + return output + for key in output: + if key.endswith("hidden_states") and not getattr(self.config, "output_hidden_states", False): + output[key] = None + elif key.endswith("attentions") and not getattr(self.config, "output_attentions", False): + output[key] = None + elif key == "past_key_values" and not getattr(self.config, "use_cache", False): + output[key] = None + elif key == "cross_attentions" and not ( + getattr(self.config, "output_attentions", False) and getattr(self.config, "add_cross_attention", False) + ): + output[key] = None + if isinstance(output[key], (tuple, list)): + try: + output[key] = tf.convert_to_tensor(output[key]) + except (ValueError, tf.errors.InvalidArgumentError): + pass # Layers may not have the same dimensions + return output + + @classmethod + def can_generate(cls) -> bool: + """ + Returns whether this model can generate sequences with `.generate()`. + + Returns: + `bool`: Whether this model can generate sequences with `.generate()`. + """ + # Detects whether `prepare_inputs_for_generation` has been overwritten, which is a requirement for generation. + # Alternatively, the model can also have a custom `generate` function. + if "GenerationMixin" in str(cls.prepare_inputs_for_generation) and "GenerationMixin" in str(cls.generate): + return False + return True + + def get_input_embeddings(self) -> keras.layers.Layer: + """ + Returns the model's input embeddings layer. + + Returns: + `tf.Variable`: The embeddings layer mapping vocabulary to hidden states. + """ + main_layer = getattr(self, self.base_model_prefix, self) + + if main_layer is not self: + return main_layer.get_input_embeddings() + else: + raise NotImplementedError + + def _save_checkpoint(self, checkpoint_dir, epoch): + if not os.path.isdir(checkpoint_dir): + os.mkdir(checkpoint_dir) + # We avoid tf.train.checkpoint or saving weights in TF format, even though that includes optimizer + # state for us, because it requires special handling for objects like custom losses, which we use + # internally and which users are likely to use too + weights_path = os.path.join(checkpoint_dir, "weights.h5") + self.save_weights(weights_path) + extra_data = {"epoch": epoch, "optimizer_state": self.optimizer.get_weights()} + extra_data_path = os.path.join(checkpoint_dir, "extra_data.pickle") + with open(extra_data_path, "wb") as f: + pickle.dump(extra_data, f) + + def prepare_tf_dataset( + self, + dataset: "datasets.Dataset", # noqa:F821 + batch_size: int = 8, + shuffle: bool = True, + tokenizer: Optional["PreTrainedTokenizerBase"] = None, + collate_fn: Optional[Callable] = None, + collate_fn_args: Optional[Dict[str, Any]] = None, + drop_remainder: Optional[bool] = None, + prefetch: bool = True, + ): + """ + Wraps a HuggingFace [`~datasets.Dataset`] as a `tf.data.Dataset` with collation and batching. This method is + designed to create a "ready-to-use" dataset that can be passed directly to Keras methods like `fit()` without + further modification. The method will drop columns from the dataset if they don't match input names for the + model. If you want to specify the column names to return rather than using the names that match this model, we + recommend using `Dataset.to_tf_dataset()` instead. + + Args: + dataset (`Any`): + A [~`datasets.Dataset`] to be wrapped as a `tf.data.Dataset`. + batch_size (`int`, *optional*, defaults to 8): + The size of batches to return. + shuffle (`bool`, defaults to `True`): + Whether to return samples from the dataset in random order. Usually `True` for training datasets and + `False` for validation/test datasets. + tokenizer ([`PreTrainedTokenizerBase`], *optional*): + A `PreTrainedTokenizer` that will be used to pad samples to create batches. Has no effect if a specific + `collate_fn` is passed instead. + collate_fn (`Callable`, *optional*): + A function that collates samples from the dataset into a single batch. Defaults to + `DefaultDataCollator` if no `tokenizer` is supplied or `DataCollatorWithPadding` if a `tokenizer` is + passed. + collate_fn_args (`Dict[str, Any]`, *optional*): + A dict of arguments to pass to the `collate_fn` alongside the list of samples. + drop_remainder (`bool`, *optional*): + Whether to drop the final batch, if the batch_size does not evenly divide the dataset length. Defaults + to the same setting as `shuffle`. + prefetch (`bool`, defaults to `True`): + Whether to add prefetching to the end of the `tf.data` pipeline. This is almost always beneficial for + performance, but can be disabled in edge cases. + + + Returns: + `Dataset`: A `tf.data.Dataset` which is ready to pass to the Keras API. + """ + requires_backends(self, ["datasets"]) + import datasets + + if collate_fn is None: + if tokenizer is None: + collate_fn = DefaultDataCollator(return_tensors="np") + else: + collate_fn = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="np") + if collate_fn_args is None: + collate_fn_args = {} + + if not isinstance(dataset, datasets.Dataset): + raise TypeError("Dataset argument should be a datasets.Dataset!") + model_inputs = list(inspect.signature(self.call).parameters) + model_labels = find_labels(self.__class__) + if "cols_to_retain" in list(inspect.signature(dataset._get_output_signature).parameters.keys()): + output_signature, _ = dataset._get_output_signature( + dataset, + batch_size=None, + collate_fn=collate_fn, + collate_fn_args=collate_fn_args, + cols_to_retain=model_inputs, + ) + else: + # TODO Matt: This is a workaround for older versions of datasets that are missing the `cols_to_retain` + # argument. We should remove this once the minimum supported version of datasets is > 2.3.2 + unwanted_columns = [ + feature + for feature in dataset.features + if feature not in model_inputs and feature not in ("label_ids", "label") + ] + dataset = dataset.remove_columns(unwanted_columns) + output_signature, _ = dataset._get_output_signature( + dataset, batch_size=None, collate_fn=collate_fn, collate_fn_args=collate_fn_args + ) + output_columns = list(output_signature.keys()) + feature_cols = [col for col in output_columns if col in model_inputs and col not in model_labels] + label_cols = [col for col in output_columns if col in model_labels] + + # Backwards compatibility for older versions of datasets. Previously, if `columns` or `label_cols` + # were a single element list, the returned element spec would be a single element. Now, passing [feature] + # will return a dict structure {"feature": feature}, and passing a single string will return a single element. + feature_cols = feature_cols[0] if len(feature_cols) == 1 else feature_cols + label_cols = label_cols[0] if len(label_cols) == 1 else label_cols + + if drop_remainder is None: + drop_remainder = shuffle + tf_dataset = dataset.to_tf_dataset( + columns=feature_cols, + label_cols=label_cols, + batch_size=batch_size, + shuffle=shuffle, + drop_remainder=drop_remainder, + collate_fn=collate_fn, + collate_fn_args=collate_fn_args, + prefetch=prefetch, + ) + return tf_dataset + + def compile( + self, + optimizer="rmsprop", + loss="auto_with_warning", + metrics=None, + loss_weights=None, + weighted_metrics=None, + run_eagerly=None, + steps_per_execution=None, + **kwargs, + ): + """ + This is a thin wrapper that sets the model's loss output head as the loss if the user does not specify a loss + function themselves. + """ + if loss in ("auto_with_warning", "passthrough"): # "passthrough" for workflow backward compatibility + logger.info( + "No loss specified in compile() - the model's internal loss computation will be used as the " + "loss. Don't panic - this is a common way to train TensorFlow models in Transformers! " + "To disable this behaviour please pass a loss argument, or explicitly pass " + "`loss=None` if you do not want your model to compute a loss. You can also specify `loss='auto'` to " + "get the internal loss without printing this info string." + ) + loss = "auto" + if loss == "auto": + loss = dummy_loss + self._using_dummy_loss = True + else: + self._using_dummy_loss = False + parent_args = list(inspect.signature(keras.Model.compile).parameters.keys()) + # This argument got renamed, we need to support both versions + if "steps_per_execution" in parent_args: + super().compile( + optimizer=optimizer, + loss=loss, + metrics=metrics, + loss_weights=loss_weights, + weighted_metrics=weighted_metrics, + run_eagerly=run_eagerly, + steps_per_execution=steps_per_execution, + **kwargs, + ) + else: + super().compile( + optimizer=optimizer, + loss=loss, + metrics=metrics, + loss_weights=loss_weights, + weighted_metrics=weighted_metrics, + run_eagerly=run_eagerly, + experimental_steps_per_execution=steps_per_execution, + **kwargs, + ) + + def compute_loss(self, *args, **kwargs): + if hasattr(keras.Model, "compute_loss"): + # This will be true in TF 2.8 or greater + return super().compute_loss(*args, **kwargs) + else: + warnings.warn( + "The old compute_loss method is deprecated as it conflicts with the Keras compute_loss " + "method added in TF 2.8. If you want the original HF compute_loss, please call " + "hf_compute_loss() instead. From TF versions >= 2.8, or Transformers versions >= 5, " + "calling compute_loss() will get the Keras method instead.", + FutureWarning, + ) + return self.hf_compute_loss(*args, **kwargs) + + def get_label_to_output_name_mapping(self): + arg_names = list(inspect.signature(self.call).parameters) + if self._label_to_output_map is not None: + return self._label_to_output_map + elif "start_positions" in arg_names: + return {"start_positions": "start_logits", "end_positions": "end_logits"} + elif "sentence_order_label" in arg_names: + return {"labels": "prediction_logits", "sentence_order_label": "sop_logits"} + elif "next_sentence_label" in arg_names: + return {"labels": "prediction_logits", "next_sentence_label": "seq_relationship_logits"} + elif "mc_labels" in arg_names: + return {"labels": "logits", "mc_labels": "mc_logits"} + else: + return {} + + def train_step(self, data): + """ + A modification of Keras's default `train_step` that correctly handles matching outputs to labels for our models + and supports directly training on the loss output head. In addition, it ensures input keys are copied to the + labels where appropriate. It will also copy label keys into the input dict when using the dummy loss, to ensure + that they are available to the model during the forward pass. + """ + + # We hardcode the most common renamings; models with weirder names can set `self._label_to_output_map` + arg_names = list(inspect.signature(self.call).parameters) + label_kwargs = find_labels(self.__class__) + label_to_output = self.get_label_to_output_name_mapping() + output_to_label = {val: key for key, val in label_to_output.items()} + if not self._using_dummy_loss and parse(tf.__version__) < parse("2.11.0"): + # Newer TF train steps leave this out + data = expand_1d(data) + x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(data) + # If the inputs are mutable dictionaries, make a shallow copy of them because we will modify + # them during input/label pre-processing. This avoids surprising the user by wrecking their data. + # In addition, modifying mutable Python inputs makes XLA compilation impossible. + if isinstance(x, dict): + x = x.copy() + if isinstance(y, dict): + y = y.copy() + + # When using a dummy loss, we ensure that separate labels are copied to the correct model arguments, + # if those keys are not already present in the input dict + if self._using_dummy_loss and y is not None: + # If y is a tensor and the model only has one label-like input, map y to that input + if len(label_kwargs) == 1 and isinstance(y, tf.Tensor): + if isinstance(x, tf.Tensor): + x = {arg_names[0]: x} + label_kwarg = next(iter(label_kwargs)) + if label_kwarg not in x: + x[label_kwarg] = y + # Otherwise, copy keys from y to x as long as they weren't already present in x + elif isinstance(y, dict): + if isinstance(x, tf.Tensor): + x = {arg_names[0]: x} + for key, val in y.items(): + if key in arg_names and key not in x: + x[key] = val + elif output_to_label.get(key, None) in arg_names and key not in x: + x[output_to_label[key]] = val + if y is None: + y = {key: val for key, val in x.items() if key in label_kwargs} + if not y and not self._using_dummy_loss: + raise ValueError("Could not find label column(s) in input dict and no separate labels were provided!") + + if isinstance(y, dict): + # Rename labels at this point to match output heads + y = {label_to_output.get(key, key): val for key, val in y.items()} + + # Run forward pass. + with tf.GradientTape() as tape: + if self._using_dummy_loss and "return_loss" in arg_names: + y_pred = self(x, training=True, return_loss=True) + else: + y_pred = self(x, training=True) + if self._using_dummy_loss: + loss = self.compiled_loss(y_pred.loss, y_pred.loss, sample_weight, regularization_losses=self.losses) + else: + loss = None + + # This next block matches outputs to label keys. Tensorflow's standard method for doing this + # can get very confused if any of the keys contain nested values (e.g. lists/tuples of Tensors) + if isinstance(y, dict) and len(y) == 1: + if list(y.keys())[0] in y_pred.keys(): + y_pred = y_pred[list(y.keys())[0]] + elif list(y_pred.keys())[0] == "loss": + y_pred = y_pred[1] + else: + y_pred = y_pred[0] + _, y = y.popitem() + elif isinstance(y, dict): + # If the labels are a dict, match keys from the output by name + y_pred = {key: val for key, val in y_pred.items() if key in y} + elif isinstance(y, tuple) or isinstance(y, list): + # If the labels are a tuple/list, match keys to the output by order, skipping the loss. + if list(y_pred.keys())[0] == "loss": + y_pred = y_pred.to_tuple()[1:] + else: + y_pred = y_pred.to_tuple() + y_pred = y_pred[: len(y)] # Remove unused fields in case those cause problems + else: + # If the labels are a single tensor, match them to the first non-loss tensor in the output + if list(y_pred.keys())[0] == "loss": + y_pred = y_pred[1] + else: + y_pred = y_pred[0] + + if loss is None: + loss = self.compiled_loss(y, y_pred, sample_weight, regularization_losses=self.losses) + + # Run backwards pass. + self.optimizer.minimize(loss, self.trainable_variables, tape=tape) + + self.compiled_metrics.update_state(y, y_pred, sample_weight) + # Collect metrics to return + return_metrics = {} + for metric in self.metrics: + result = metric.result() + if isinstance(result, dict): + return_metrics.update(result) + else: + return_metrics[metric.name] = result + return return_metrics + + def test_step(self, data): + """ + A modification of Keras's default `train_step` that correctly handles matching outputs to labels for our models + and supports directly training on the loss output head. In addition, it ensures input keys are copied to the + labels where appropriate. It will also copy label keys into the input dict when using the dummy loss, to ensure + that they are available to the model during the forward pass. + """ + # We hardcode the most common renamings; models with weirder names can set `self._label_to_output_map` + arg_names = list(inspect.signature(self.call).parameters) + label_kwargs = find_labels(self.__class__) + label_to_output = self.get_label_to_output_name_mapping() + output_to_label = {val: key for key, val in label_to_output.items()} + if not self._using_dummy_loss and parse(tf.__version__) < parse("2.11.0"): + # Newer versions leave this out + data = expand_1d(data) + x, y, sample_weight = keras.utils.unpack_x_y_sample_weight(data) + # If the inputs are mutable dictionaries, make a shallow copy of them because we will modify + # them during input/label pre-processing. This avoids surprising the user by wrecking their data. + # In addition, modifying mutable Python inputs makes XLA compilation impossible. + if isinstance(x, dict): + x = x.copy() + if isinstance(y, dict): + y = y.copy() + + # When using a dummy loss, we ensure that separate labels are copied to the correct model arguments, + # if those keys are not already present in the input dict + if self._using_dummy_loss and y is not None: + arg_names = list(inspect.signature(self.call).parameters) + # If y is a tensor and the model only has one label-like input, map y to that input + if len(label_kwargs) == 1 and isinstance(y, tf.Tensor): + if isinstance(x, tf.Tensor): + x = {arg_names[0]: x} + label_kwarg = next(iter(label_kwargs)) + if label_kwarg not in x: + x[label_kwarg] = y + # Otherwise, copy keys from y to x as long as they weren't already present in x + elif isinstance(y, dict): + if isinstance(x, tf.Tensor): + x = {arg_names[0]: x} + for key, val in y.items(): + if key in arg_names and key not in x: + x[key] = val + elif output_to_label.get(key, None) in arg_names and key not in x: + x[output_to_label[key]] = val + if y is None: + y = {key: val for key, val in x.items() if key in label_kwargs} + if not y and not self._using_dummy_loss: + raise ValueError("Could not find label column(s) in input dict and no separate labels were provided!") + + if isinstance(y, dict): + # Rename labels at this point to match output heads + y = {label_to_output.get(key, key): val for key, val in y.items()} + + # Run forward pass. + if self._using_dummy_loss and "return_loss" in arg_names: + y_pred = self(x, return_loss=True, training=False) + else: + y_pred = self(x, training=False) + if self._using_dummy_loss: + loss = self.compiled_loss(y_pred.loss, y_pred.loss, sample_weight, regularization_losses=self.losses) + else: + loss = None + + # This next block matches outputs to label keys. Tensorflow's standard method for doing this + # can get very confused if any of the keys contain nested values (e.g. lists/tuples of Tensors) + if isinstance(y, dict) and len(y) == 1: + if list(y.keys())[0] in y_pred.keys(): + y_pred = y_pred[list(y.keys())[0]] + elif list(y_pred.keys())[0] == "loss": + y_pred = y_pred[1] + else: + y_pred = y_pred[0] + _, y = y.popitem() + elif isinstance(y, dict): + # If the labels are a dict, match keys from the output by name + y_pred = {key: val for key, val in y_pred.items() if key in y} + elif isinstance(y, tuple) or isinstance(y, list): + # If the labels are a tuple/list, match keys to the output by order, skipping the loss. + if list(y_pred.keys())[0] == "loss": + y_pred = y_pred.to_tuple()[1:] + else: + y_pred = y_pred.to_tuple() + y_pred = y_pred[: len(y)] # Remove unused fields in case those cause problems + else: + # If the labels are a single tensor, match them to the first non-loss tensor in the output + if list(y_pred.keys())[0] == "loss": + y_pred = y_pred[1] + else: + y_pred = y_pred[0] + + if loss is None: + loss = self.compiled_loss(y, y_pred, sample_weight, regularization_losses=self.losses) + + self.compiled_metrics.update_state(y, y_pred, sample_weight) + # Collect metrics to return + return_metrics = {} + for metric in self.metrics: + result = metric.result() + if isinstance(result, dict): + return_metrics.update(result) + else: + return_metrics[metric.name] = result + return return_metrics + + def create_model_card( + self, + output_dir, + model_name: str, + language: Optional[str] = None, + license: Optional[str] = None, + tags: Optional[str] = None, + finetuned_from: Optional[str] = None, + tasks: Optional[str] = None, + dataset_tags: Optional[Union[str, List[str]]] = None, + dataset: Optional[Union[str, List[str]]] = None, + dataset_args: Optional[Union[str, List[str]]] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + output_dir (`str` or `os.PathLike`): + The folder in which to create the model card. + model_name (`str`, *optional*): + The name of the model. + language (`str`, *optional*): + The language of the model (if applicable) + license (`str`, *optional*): + The license of the model. Will default to the license of the pretrained model used, if the original + model given to the `Trainer` comes from a repo on the Hub. + tags (`str` or `List[str]`, *optional*): + Some tags to be included in the metadata of the model card. + finetuned_from (`str`, *optional*): + The name of the model used to fine-tune this one (if applicable). Will default to the name of the repo + of the original model given to the `Trainer` (if it comes from the Hub). + tasks (`str` or `List[str]`, *optional*): + One or several task identifiers, to be included in the metadata of the model card. + dataset_tags (`str` or `List[str]`, *optional*): + One or several dataset tags, to be included in the metadata of the model card. + dataset (`str` or `List[str]`, *optional*): + One or several dataset identifiers, to be included in the metadata of the model card. + dataset_args (`str` or `List[str]`, *optional*): + One or several dataset arguments, to be included in the metadata of the model card. + """ + # Avoids a circular import by doing this when necessary. + from .modelcard import TrainingSummary # tests_ignore + + training_summary = TrainingSummary.from_keras( + self, + keras_history=self.history, + language=language, + license=license, + tags=tags, + model_name=model_name, + finetuned_from=finetuned_from, + tasks=tasks, + dataset_tags=dataset_tags, + dataset=dataset, + dataset_args=dataset_args, + ) + model_card = training_summary.to_model_card() + with open(os.path.join(output_dir, "README.md"), "w") as f: + f.write(model_card) + + def set_input_embeddings(self, value): + """ + Set model's input embeddings + + Args: + value (`tf.Variable`): + The new weights mapping hidden states to vocabulary. + """ + main_layer = getattr(self, self.base_model_prefix) + + if main_layer is None: + raise NotImplementedError("The model does not implements the base_model_prefix attribute.") + + try: + main_layer.set_input_embeddings(value) + except AttributeError: + logger.info("Building the model") + self.build_in_name_scope() + main_layer.set_input_embeddings(value) + + def get_output_embeddings(self) -> Union[None, keras.layers.Layer]: + """ + Returns the model's output embeddings + + Returns: + `tf.Variable`: The new weights mapping vocabulary to hidden states. + """ + if self.get_lm_head() is not None: + lm_head = self.get_lm_head() + + try: + return lm_head.get_output_embeddings() + except AttributeError: + logger.info("Building the model") + self.build_in_name_scope() + + return lm_head().get_output_embeddings() + + return None # Overwrite for models with output embeddings + + def set_output_embeddings(self, value): + """ + Set model's output embeddings + + Args: + value (`tf.Variable`): + The new weights mapping hidden states to vocabulary. + """ + if self.get_lm_head() is not None: + lm_head = self.get_lm_head() + try: + lm_head.set_output_embeddings(value) + except AttributeError: + logger.info("Building the model") + self.build_in_name_scope() + lm_head.set_output_embeddings(value) + + def get_output_layer_with_bias(self) -> Union[None, keras.layers.Layer]: + """ + Get the layer that handles a bias attribute in case the model has an LM head with weights tied to the + embeddings + + Return: + `keras.layers.Layer`: The layer that handles the bias, None if not an LM model. + """ + warnings.warn( + "The method get_output_layer_with_bias is deprecated. Please use `get_lm_head` instead.", FutureWarning + ) + return self.get_lm_head() + + def get_prefix_bias_name(self) -> Union[None, str]: + """ + Get the concatenated _prefix name of the bias from the model name to the parent layer + + Return: + `str`: The _prefix name of the bias. + """ + warnings.warn("The method get_prefix_bias_name is deprecated. Please use `get_bias` instead.", FutureWarning) + return None + + def get_bias(self) -> Union[None, Dict[str, tf.Variable]]: + """ + Dict of bias attached to an LM head. The key represents the name of the bias attribute. + + Return: + `tf.Variable`: The weights representing the bias, None if not an LM model. + """ + if self.get_lm_head() is not None: + lm_head = self.get_lm_head() + try: + return lm_head.get_bias() + except AttributeError: + self.build_in_name_scope() + + return lm_head.get_bias() + return None + + def set_bias(self, value): + """ + Set all the bias in the LM head. + + Args: + value (`Dict[tf.Variable]`): + All the new bias attached to an LM head. + """ + if self.get_lm_head() is not None: + lm_head = self.get_lm_head() + try: + lm_head.set_bias(value) + except AttributeError: + self.build_in_name_scope() + lm_head.set_bias(value) + + def get_lm_head(self) -> keras.layers.Layer: + """ + The LM Head layer. This method must be overwritten by all the models that have a lm head. + + Return: + `keras.layers.Layer`: The LM head layer if the model has one, None if not. + """ + return None + + def resize_token_embeddings( + self, new_num_tokens: Optional[int] = None + ) -> Union[keras.layers.Embedding, tf.Variable]: + """ + Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`. + + Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method. + + Arguments: + new_num_tokens (`int`, *optional*): + The number of new tokens in the embedding matrix. Increasing the size will add newly initialized + vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just + returns a pointer to the input tokens without doing anything. + + Return: + `tf.Variable` or `keras.layers.Embedding`: Pointer to the input tokens of the model. + """ + # TODO (joao): flagged for replacement (by `_v2_resized_token_embeddings`) due to embeddings refactor + + # Run the new code path if the model has a keras embeddings layer + if isinstance(self.get_input_embeddings(), keras.layers.Embedding): + return self._v2_resized_token_embeddings(new_num_tokens) + + if new_num_tokens is None or new_num_tokens == self.config.vocab_size: + return self._get_word_embedding_weight(self.get_input_embeddings()) + + model_embeds = self._resize_token_embeddings(new_num_tokens) + + # Update base model and current model config + self.config.vocab_size = new_num_tokens + + return model_embeds + + def _v2_resized_token_embeddings(self, new_num_tokens: Optional[int] = None) -> keras.layers.Embedding: + """ + Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`. + + Arguments: + new_num_tokens (`int`, *optional*): + The number of new tokens in the embedding matrix. Increasing the size will add newly initialized + vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just + returns a pointer to the input tokens without doing anything. + + Return: + `keras.layers.Embedding`: Pointer to the input tokens of the model. + """ + if new_num_tokens is None or new_num_tokens == self.config.vocab_size: + return self.get_input_embeddings() + + model_embeds = self._v2_resize_token_embeddings(new_num_tokens) + + # Update base model and current model config + self.config.vocab_size = new_num_tokens + + return model_embeds + + def _get_word_embedding_weight(model, embedding_layer): + # TODO (joao): flagged for delection due to embeddings refactor + + # If the variable holds the weights themselves, return them + if isinstance(embedding_layer, tf.Tensor): + return embedding_layer + # Otherwise, try to get them from the layer's attributes + + embeds = getattr(embedding_layer, "weight", None) + if embeds is not None: + return embeds + + embeds = getattr(embedding_layer, "decoder", None) + if embeds is not None: + return embeds + + # The reason why the attributes don't exist might be + # because the model is not built, so retry getting + # the argument after building the model + model.build_in_name_scope() + + embeds = getattr(embedding_layer, "weight", None) + if embeds is not None: + return embeds + + embeds = getattr(embedding_layer, "decoder", None) + if embeds is not None: + return embeds + + return None + + def _resize_token_embeddings(self, new_num_tokens): + # TODO (joao): flagged for replacement (by `_v2_resize_token_embeddings`) due to embeddings refactor + old_embeddings = self._get_word_embedding_weight(self.get_input_embeddings()) + new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens) + + # if word embeddings are not tied, make sure that lm head bias is resized as well + if self.get_bias() is not None: + old_lm_head_bias = self.get_bias() + new_lm_head_bias = self._get_resized_lm_head_bias(old_lm_head_bias, new_num_tokens) + + self.set_bias(new_lm_head_bias) + + # if word embeddings are not tied, make sure that lm head decoder is resized as well + if self.get_output_embeddings() is not None: + old_lm_head_decoder = self._get_word_embedding_weight(self.get_output_embeddings()) + new_lm_head_decoder = self._get_resized_lm_head_decoder(old_lm_head_decoder, new_num_tokens) + + self.set_output_embeddings(new_lm_head_decoder) + + self.set_input_embeddings(new_embeddings) + + return self.get_input_embeddings() + + def _v2_resize_token_embeddings(self, new_num_tokens): + old_embeddings = self.get_input_embeddings() + new_embeddings = self._v2_get_resized_embeddings(old_embeddings, new_num_tokens) + self.set_input_embeddings(new_embeddings) + + # If word embeddings are not tied, make sure that lm head bias is resized as well + if self.get_bias() is not None: + old_lm_head_bias = self.get_bias() + new_lm_head_bias = self._v2_get_resized_lm_head_bias(old_lm_head_bias, new_num_tokens) + self.set_bias(new_lm_head_bias) + + # If word embeddings are not tied, make sure that lm head decoder is resized as well. + tied_weights = self.get_input_embeddings() == self.get_output_embeddings() + if self.get_output_embeddings() is not None and not tied_weights: + old_lm_head_decoder = self._get_word_embedding_weight(self.get_output_embeddings()) + # TODO (joao): this one probably needs a v2 version with other models + new_lm_head_decoder = self._get_resized_lm_head_decoder(old_lm_head_decoder, new_num_tokens) + self.set_output_embeddings(new_lm_head_decoder) + + return self.get_input_embeddings() + + def _get_resized_lm_head_bias(self, old_lm_head_bias, new_num_tokens): + """ + Build a resized bias from the old ones. Increasing the size will add newly initialized vectors at the end. + Reducing the size will remove vectors from the end + + Args: + old_lm_head_bias (`tf.Variable`): + Old lm head bias to be resized. + new_num_tokens (`int`, *optional*): + New number of tokens in the linear matrix. + + Increasing the size will add newly initialized vectors at the end. Reducing the size will remove + vectors from the end. If not provided or `None`, just returns None + + Return: + `tf.Variable`: Pointer to the resized bias. + """ + # TODO (joao): flagged for replacement (by `_v2_get_resized_lm_head_bias`) due to embeddings refactor + new_lm_head_bias = {} + + for attr, weight in old_lm_head_bias.items(): + first_dim, old_num_tokens = (None, shape_list(weight)[0]) if tf.rank(weight) == 1 else shape_list(weight) + size_diff = new_num_tokens - old_num_tokens + final_shape = [new_num_tokens] if first_dim is None else [first_dim, new_num_tokens] + + # initialize new bias + if tf.math.greater(size_diff, 0): + padding_shape = [[0, size_diff]] if first_dim is None else [[0, 0], [0, size_diff]] + current_bias = tf.pad(weight.value(), tf.convert_to_tensor(padding_shape), constant_values=-1) + num_tokens_to_copy = min(old_num_tokens, new_num_tokens) + mask_shape = [num_tokens_to_copy] if first_dim is None else [1, num_tokens_to_copy] + bias_mask = tf.fill(tf.convert_to_tensor(mask_shape), True) + bias_mask = tf.pad(bias_mask, tf.convert_to_tensor(padding_shape), constant_values=False) + else: + slice_from = [0] if first_dim is None else [0, 0] + current_bias = tf.slice( + weight.value(), tf.convert_to_tensor(slice_from), tf.convert_to_tensor(final_shape) + ) + bias_mask = tf.fill(tf.convert_to_tensor(final_shape), True) + + new_bias = self.add_weight( + shape=final_shape, + initializer="zeros", + trainable=True, + name=weight.name.split(":")[0], + ) + init_bias = tf.where(bias_mask, current_bias, new_bias.value()) + + new_bias.assign(init_bias) + new_lm_head_bias[attr] = new_bias + + return new_lm_head_bias + + def _v2_get_resized_lm_head_bias( + self, old_lm_head_bias: Dict[str, tf.Variable], new_num_tokens: int + ) -> Dict[str, tf.Tensor]: + """ + Build a resized bias from the old ones. Increasing the size will add newly initialized vectors at the end. + Reducing the size will remove vectors from the end + + Args: + old_lm_head_bias (`Dict[str, tf.Variable]`): + Old lm head bias to be resized. + new_num_tokens (`int`): + New number of tokens in the linear matrix. Increasing the size will add newly initialized vectors at + the end. Reducing the size will remove vectors from the end. + + Return: + `tf.Tensor`: Values for the resized bias. + """ + new_lm_head_bias = {} + + for attr, weight in old_lm_head_bias.items(): + # Determine the size difference (depending on the shape) + first_dim, old_num_tokens = (None, shape_list(weight)[0]) if tf.rank(weight) == 1 else shape_list(weight) + size_diff = new_num_tokens - old_num_tokens + + # Copy the old bias values to the new bias + if old_num_tokens > new_num_tokens: + new_bias = weight.value()[..., :new_num_tokens] + else: + padding_shape = [[0, size_diff]] if first_dim is None else [[0, 0], [0, size_diff]] + new_bias = tf.pad(weight.value(), tf.convert_to_tensor(padding_shape)) + + new_lm_head_bias[attr] = new_bias + return new_lm_head_bias + + def _get_resized_lm_head_decoder(self, old_lm_head_decoder, new_num_tokens): + """ + Build a resized decoder from the old ones. Increasing the size will add newly initialized vectors at the end. + Reducing the size will remove vectors from the end + + Args: + old_lm_head_decoder (`tf.Variable`): + Old lm head decoder to be resized. + new_num_tokens (`int`, *optional*): + New number of tokens in the linear matrix. + + Increasing the size will add newly initialized vectors at the end. Reducing the size will remove + vectors from the end. If not provided or `None`, just returns None + + Return: + `tf.Variable`: Pointer to the resized decoder or None if the output embeddings are different from the input + ones. + """ + new_lm_head_decoder = old_lm_head_decoder + is_input_output_equals = tf.reduce_any( + self._get_word_embedding_weight(self.get_input_embeddings()) == old_lm_head_decoder + ) + + if old_lm_head_decoder is not None and not is_input_output_equals: + old_embedding_dim = shape_list(old_lm_head_decoder)[1] + decoder_mask, current_decoder = init_copy_embeddings(old_lm_head_decoder, new_num_tokens) + new_lm_head_decoder = self.add_weight( + shape=(new_num_tokens, old_embedding_dim), + initializer="zeros", + trainable=True, + name=old_lm_head_decoder.name.split(":")[0], + ) + init_decoder = tf.where(decoder_mask, current_decoder, new_lm_head_decoder.value()) + + new_lm_head_decoder.assign(init_decoder) + + return new_lm_head_decoder + + def _get_resized_embeddings(self, old_embeddings, new_num_tokens=None) -> tf.Variable: + """ + Build a resized Embedding weights from a provided token Embedding weights. Increasing the size will add newly + initialized vectors at the end. Reducing the size will remove vectors from the end + + Args: + old_embeddings (`tf.Variable`): + Old embeddings to be resized. + new_num_tokens (`int`, *optional*): + New number of tokens in the embedding matrix. + + Increasing the size will add newly initialized vectors at the end. Reducing the size will remove + vectors from the end. If not provided or `None`, just returns a pointer to the input tokens + `tf.Variable` module of the model without doing anything. + + Return: + `tf.Variable`: Pointer to the resized Embedding Module or the old Embedding Module if `new_num_tokens` is + `None` + """ + # TODO (joao): flagged for replacement (by `_v2_get_resized_embeddings`) due to embeddings refactor + old_embedding_dim = shape_list(old_embeddings)[1] + init_range = getattr(self.config, "initializer_range", 0.02) + embeddings_mask, current_embeddings = init_copy_embeddings(old_embeddings, new_num_tokens) + new_embeddings = self.add_weight( + name=old_embeddings.name.split(":")[0], + shape=[new_num_tokens, old_embedding_dim], + initializer=get_initializer(init_range), + dtype=tf.float32, + ) + init_embeddings = tf.where(embeddings_mask, current_embeddings, new_embeddings.value()) + + new_embeddings.assign(init_embeddings) + + return new_embeddings + + def _v2_get_resized_embeddings( + self, old_embeddings: keras.layers.Embedding, new_num_tokens: int + ) -> keras.layers.Embedding: + """ + Build a resized Embedding layer from a provided Embedding layer. Increasing the size will add newly initialized + vectors at the end. Reducing the size will remove vectors from the end. + + Args: + old_embeddings (`keras.layers.Embedding`): + Old embeddings to be resized. + new_num_tokens (`int`, *optional*): + New number of tokens in the embedding matrix. + + Return: + `keras.layers.Embedding`: Resized Embedding layer. + """ + + # Get the initialization range for the embeddings + init_range = 0.02 # default value + potential_initialization_variable_names = [ + "initializer_range", # most common + "initializer_factor", # e.g. T5 + "init_std", # e.g BART + ] + for var_name in potential_initialization_variable_names: + if hasattr(self.config, var_name): + init_range = getattr(self.config, var_name) + + # Get a new (initialized) embeddings layer + new_embeddings = keras.layers.Embedding( + input_dim=new_num_tokens, + output_dim=old_embeddings.output_dim, + embeddings_initializer=keras.initializers.TruncatedNormal(stddev=init_range), + name=old_embeddings.embeddings.name[:-13], # exact same scoped name except "/embeddings:0" + ) + new_embeddings(tf.constant([[0]])) + + # Copy the old embeddings to the new embeddings + if old_embeddings.input_dim >= new_num_tokens: + init_embeddings = old_embeddings.embeddings[:new_num_tokens] + else: + init_embeddings = tf.concat( + [old_embeddings.embeddings, new_embeddings.embeddings[old_embeddings.input_dim :]], axis=0 + ) + new_embeddings.embeddings.assign(init_embeddings) + return new_embeddings + + def prune_heads(self, heads_to_prune): + """ + Prunes heads of the base model. + + Arguments: + heads_to_prune (`Dict[int, List[int]]`): + Dictionary with keys being selected layer indices (`int`) and associated values being the list of heads + to prune in said layer (list of `int`). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on + layer 1 and heads 2 and 3 on layer 2. + """ + raise NotImplementedError + + def save_pretrained( + self, + save_directory, + saved_model=False, + version=1, + push_to_hub=False, + signatures=None, + max_shard_size: Union[int, str] = "5GB", + create_pr: bool = False, + safe_serialization: bool = False, + token: Optional[Union[str, bool]] = None, + **kwargs, + ): + """ + Save a model and its configuration file to a directory, so that it can be re-loaded using the + [`~TFPreTrainedModel.from_pretrained`] class method. + + Arguments: + save_directory (`str`): + Directory to which to save. Will be created if it doesn't exist. + saved_model (`bool`, *optional*, defaults to `False`): + If the model has to be saved in saved model format as well or not. + version (`int`, *optional*, defaults to 1): + The version of the saved model. A saved model needs to be versioned in order to be properly loaded by + TensorFlow Serving as detailed in the official documentation + https://www.tensorflow.org/tfx/serving/serving_basic + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + signatures (`dict` or `tf.function`, *optional*): + Model's signature used for serving. This will be passed to the `signatures` argument of model.save(). + max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`): + The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size + lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`). + + + + If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard + which will be bigger than `max_shard_size`. + + + + create_pr (`bool`, *optional*, defaults to `False`): + Whether or not to create a PR with the uploaded files or directly commit. + safe_serialization (`bool`, *optional*, defaults to `False`): + Whether to save the model using `safetensors` or the traditional TensorFlow way (that uses `h5`). + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + kwargs (`Dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + use_auth_token = kwargs.pop("use_auth_token", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if token is not None: + kwargs["token"] = token + + if os.path.isfile(save_directory): + logger.error(f"Provided path ({save_directory}) should be a directory, not a file") + return + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = self._create_repo(repo_id, **kwargs) + files_timestamps = self._get_files_timestamps(save_directory) + + if saved_model: + # If `torch_dtype` is in the config with a torch dtype class as the value, we need to change it to string. + # (Although TF doesn't care about this attribute, we can't just remove it or set it to `None`.) + if getattr(self.config, "torch_dtype", None) is not None and not isinstance(self.config.torch_dtype, str): + self.config.torch_dtype = str(self.config.torch_dtype).split(".")[1] + if signatures is None: + serving_default = self.serving.get_concrete_function(self.input_signature) + if any(spec.dtype == tf.int32 for spec in self.input_signature.values()): + int64_spec = { + key: tf.TensorSpec( + shape=spec.shape, dtype=tf.int64 if spec.dtype == tf.int32 else spec.dtype, name=spec.name + ) + for key, spec in self.input_signature.items() + } + int64_serving = self.serving.get_concrete_function(int64_spec) + signatures = {"serving_default": serving_default, "int64_serving": int64_serving} + else: + signatures = serving_default + saved_model_dir = os.path.join(save_directory, "saved_model", str(version)) + self.save(saved_model_dir, include_optimizer=False, signatures=signatures) + logger.info(f"Saved model created in {saved_model_dir}") + + # Save configuration file + self.config.architectures = [self.__class__.__name__[2:]] + + # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + custom_object_save(self, save_directory, config=self.config) + + self.config.save_pretrained(save_directory) + if self.can_generate(): + self.generation_config.save_pretrained(save_directory) + + # If we save using the predefined names, we can load using `from_pretrained` + weights_name = SAFE_WEIGHTS_NAME if safe_serialization else TF2_WEIGHTS_NAME + output_model_file = os.path.join(save_directory, weights_name) + + shards, index = tf_shard_checkpoint(self.weights, max_shard_size, weights_name=weights_name) + + # Clean the folder from a previous save + for filename in os.listdir(save_directory): + full_filename = os.path.join(save_directory, filename) + # If we have a shard file that is not going to be replaced, we delete it, but only from the main process + # in distributed settings to avoid race conditions. + weights_no_suffix = weights_name.replace(".bin", "").replace(".safetensors", "") + if ( + filename.startswith(weights_no_suffix) + and os.path.isfile(full_filename) + and filename not in shards.keys() + ): + os.remove(full_filename) + + if index is None: + if safe_serialization: + state_dict = {strip_model_name_and_prefix(w.name): w.value() for w in self.weights} + safe_save_file(state_dict, output_model_file, metadata={"format": "tf"}) + else: + self.save_weights(output_model_file) + logger.info(f"Model weights saved in {output_model_file}") + else: + save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else TF2_WEIGHTS_INDEX_NAME + save_index_file = os.path.join(save_directory, save_index_file) + # Save the index as well + with open(save_index_file, "w", encoding="utf-8") as index_file: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + index_file.write(content) + logger.info( + f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " + f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the " + f"index located at {save_index_file}." + ) + for shard_file, shard in shards.items(): + if safe_serialization: + shard_state_dict = {strip_model_name_and_prefix(w.name): w.value() for w in shard} + safe_save_file( + shard_state_dict, os.path.join(save_directory, shard_file), metadata={"format": "tf"} + ) + else: + with h5py.File(os.path.join(save_directory, shard_file), mode="w") as shard_file: + layers = [] + for layer in sorted(shard, key=lambda x: x.name): + if "model." in layer.name or len(layer.name.split("/")) == 1: + layer_name = layer.name + else: + layer_name = "/".join(layer.name.split("/")[1:]) + param_dset = shard_file.create_dataset( + layer_name, layer.numpy().shape, dtype=layer.numpy().dtype + ) + param_dset[:] = layer.numpy() + layers.append(layer_name.encode("utf8")) + save_attributes_to_hdf5_group(shard_file, "layer_names", layers) + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=token, + ) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], + *model_args, + config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None, + cache_dir: Optional[Union[str, os.PathLike]] = None, + ignore_mismatched_sizes: bool = False, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + use_safetensors: Optional[bool] = None, + **kwargs, + ): + r""" + Instantiate a pretrained TF 2.0 model from a pre-trained model configuration. + + The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come + pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning + task. + + The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those + weights are discarded. + + Parameters: + pretrained_model_name_or_path (`str`, *optional*): + Can be either: + + - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co. + - A path to a *directory* containing model weights saved using + [`~TFPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + - A path or url to a *PyTorch state_dict save file* (e.g, `./pt_model/pytorch_model.bin`). In this + case, `from_pt` should be set to `True` and a configuration object should be provided as `config` + argument. This loading path is slower than converting the PyTorch model in a TensorFlow model + using the provided conversion scripts and loading the TensorFlow model afterwards. + - `None` if you are both providing the configuration and state dictionary (resp. with keyword + arguments `config` and `state_dict`). + model_args (sequence of positional arguments, *optional*): + All remaining positional arguments will be passed to the underlying model's `__init__` method. + config (`Union[PretrainedConfig, str]`, *optional*): + Can be either: + + - an instance of a class derived from [`PretrainedConfig`], + - a string valid as input to [`~PretrainedConfig.from_pretrained`]. + + Configuration for the model to use instead of an automatically loaded configuration. Configuration can + be automatically loaded when: + + - The model is a model provided by the library (loaded with the *model id* string of a pretrained + model). + - The model was saved using [`~TFPreTrainedModel.save_pretrained`] and is reloaded by supplying the + save directory. + - The model is loaded by supplying a local directory as `pretrained_model_name_or_path` and a + configuration JSON file named *config.json* is found in the directory. + from_pt (`bool`, *optional*, defaults to `False`): + Load the model weights from a PyTorch state_dict save file (see docstring of + `pretrained_model_name_or_path` argument). + ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`): + Whether or not to raise an error if some of the weights from the checkpoint do not have the same size + as the weights of the model (if for instance, you are instantiating a model with 10 labels from a + checkpoint with 3 labels). + cache_dir (`str`, *optional*): + Path to a directory in which a downloaded pretrained model configuration should be cached if the + standard cache should not be used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies: + (`Dict[str, str], `optional`): 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. + output_loading_info(`bool`, *optional*, defaults to `False`): Whether ot not to also return a + dictionary containing missing keys, unexpected keys and error messages. + local_files_only(`bool`, *optional*, defaults to `False`): + Whether or not to only look at local files (e.g., not try downloading the model). + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`. + + + + mirror (`str`, *optional*): + Mirror source to accelerate downloads in China. If you are from China and have an accessibility + problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety. + Please refer to the mirror site for more information. + 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. + tf_to_pt_weight_rename (`Callable`, *optional*): + A function that is called to transform the names of weights during the PyTorch to TensorFlow + crossloading process. This is not necessary for most models, but is useful to allow composite models to + be crossloaded correctly. + use_safetensors (`bool`, *optional*, defaults to `None`): + Whether or not to use `safetensors` checkpoints. Defaults to `None`. If not specified and `safetensors` + is not installed, it will be set to `False`. + kwargs (remaining dictionary of keyword arguments, *optional*): + Can be used to update the configuration object (after it being loaded) and initiate the model (e.g., + `output_attentions=True`). Behaves differently depending on whether a `config` is provided or + automatically loaded: + + - If a configuration is provided with `config`, `**kwargs` will be directly passed to the + underlying model's `__init__` method (we assume all relevant updates to the configuration have + already been done) + - If a configuration is not provided, `kwargs` will be first passed to the configuration class + initialization function ([`~PretrainedConfig.from_pretrained`]). Each key of `kwargs` that + corresponds to a configuration attribute will be used to override said attribute with the + supplied `kwargs` value. Remaining keys that do not correspond to any configuration attribute + will be passed to the underlying model's `__init__` function. + + Examples: + + ```python + >>> from transformers import BertConfig, TFBertModel + + >>> # Download model and configuration from huggingface.co and cache. + >>> model = TFBertModel.from_pretrained("google-bert/bert-base-uncased") + >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable). + >>> model = TFBertModel.from_pretrained("./test/saved_model/") + >>> # Update configuration during loading. + >>> model = TFBertModel.from_pretrained("google-bert/bert-base-uncased", output_attentions=True) + >>> assert model.config.output_attentions == True + >>> # Loading from a Pytorch model file instead of a TensorFlow checkpoint (slower, for example purposes, not runnable). + >>> config = BertConfig.from_json_file("./pt_model/my_pt_model_config.json") + >>> model = TFBertModel.from_pretrained("./pt_model/my_pytorch_model.bin", from_pt=True, config=config) + ```""" + from_pt = kwargs.pop("from_pt", False) + resume_download = kwargs.pop("resume_download", None) + proxies = kwargs.pop("proxies", None) + output_loading_info = kwargs.pop("output_loading_info", False) + use_auth_token = kwargs.pop("use_auth_token", None) + trust_remote_code = kwargs.pop("trust_remote_code", None) + _ = kwargs.pop("mirror", None) + load_weight_prefix = kwargs.pop("load_weight_prefix", None) + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + subfolder = kwargs.pop("subfolder", "") + commit_hash = kwargs.pop("_commit_hash", None) + tf_to_pt_weight_rename = kwargs.pop("tf_to_pt_weight_rename", None) + + # Not relevant for TF models + _ = kwargs.pop("adapter_kwargs", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if trust_remote_code is True: + logger.warning( + "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is" + " ignored." + ) + + user_agent = {"file_type": "model", "framework": "tensorflow", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + if use_safetensors is None and not is_safetensors_available(): + use_safetensors = False + + # Load config if we don't provide a configuration + if not isinstance(config, PretrainedConfig): + config_path = config if config is not None else pretrained_model_name_or_path + config, model_kwargs = cls.config_class.from_pretrained( + config_path, + cache_dir=cache_dir, + return_unused_kwargs=True, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + _from_auto=from_auto_class, + _from_pipeline=from_pipeline, + _commit_hash=commit_hash, + **kwargs, + ) + else: + model_kwargs = kwargs + + if commit_hash is None: + commit_hash = getattr(config, "_commit_hash", None) + + # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the + # index of the files. + is_sharded = False + # Load model + if pretrained_model_name_or_path is not None: + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + is_local = os.path.isdir(pretrained_model_name_or_path) + if is_local: + if from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)): + # Load from a PyTorch checkpoint in priority if from_pt + archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME) + elif from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME)): + # Load from a sharded PyTorch checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME) + is_sharded = True + elif use_safetensors is not False and os.path.isfile( + os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME) + ): + # Load from a safetensors checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME) + elif use_safetensors is not False and os.path.isfile( + os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME) + ): + # Load from a sharded safetensors checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME) + is_sharded = True + elif os.path.isfile(os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)): + # Load from a TF 2.0 checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME) + elif os.path.isfile(os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_INDEX_NAME)): + # Load from a sharded TF 2.0 checkpoint + archive_file = os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_INDEX_NAME) + is_sharded = True + + # At this stage we don't have a weight file so we will raise an error. + elif use_safetensors: + raise EnvironmentError( + f"Error no file named {SAFE_WEIGHTS_NAME} or {SAFE_WEIGHTS_INDEX_NAME} found in directory {pretrained_model_name_or_path}. " + f"Please make sure that the model has been saved with `safe_serialization=True` or do not " + f"set `use_safetensors=True`." + ) + elif os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)) or os.path.isfile( + os.path.join(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME) + ): + raise EnvironmentError( + f"Error no file named {TF2_WEIGHTS_NAME} or {SAFE_WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} " + "but there is a file for PyTorch weights. Use `from_pt=True` to load this model from those " + "weights." + ) + else: + raise EnvironmentError( + f"Error no file named {TF2_WEIGHTS_NAME}, {SAFE_WEIGHTS_NAME} or {WEIGHTS_NAME} found in directory " + f"{pretrained_model_name_or_path}." + ) + elif os.path.isfile(pretrained_model_name_or_path): + archive_file = pretrained_model_name_or_path + is_local = True + elif os.path.isfile(pretrained_model_name_or_path + ".index"): + archive_file = pretrained_model_name_or_path + ".index" + is_local = True + elif is_remote_url(pretrained_model_name_or_path): + filename = pretrained_model_name_or_path + resolved_archive_file = download_url(pretrained_model_name_or_path) + else: + # set correct filename + if from_pt: + filename = WEIGHTS_NAME + elif use_safetensors is not False: + filename = SAFE_WEIGHTS_NAME + else: + filename = TF2_WEIGHTS_NAME + + try: + # Load from URL or cache if already cached + cached_file_kwargs = { + "cache_dir": cache_dir, + "force_download": force_download, + "proxies": proxies, + "resume_download": resume_download, + "local_files_only": local_files_only, + "token": token, + "user_agent": user_agent, + "revision": revision, + "subfolder": subfolder, + "_raise_exceptions_for_gated_repo": False, + "_raise_exceptions_for_missing_entries": False, + "_commit_hash": commit_hash, + } + resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs) + + # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None + # result when internet is up, the repo and revision exist, but the file does not. + if resolved_archive_file is None and filename == SAFE_WEIGHTS_NAME: + # Did not find the safetensors file, let's fallback to TF. + # No support for sharded safetensors yet, so we'll raise an error if that's all we find. + filename = TF2_WEIGHTS_NAME + resolved_archive_file = cached_file( + pretrained_model_name_or_path, TF2_WEIGHTS_NAME, **cached_file_kwargs + ) + if resolved_archive_file is None and filename == TF2_WEIGHTS_NAME: + # Maybe the checkpoint is sharded, we try to grab the index name in this case. + resolved_archive_file = cached_file( + pretrained_model_name_or_path, TF2_WEIGHTS_INDEX_NAME, **cached_file_kwargs + ) + if resolved_archive_file is not None: + is_sharded = True + if resolved_archive_file is None and filename == WEIGHTS_NAME: + # Maybe the checkpoint is sharded, we try to grab the index name in this case. + resolved_archive_file = cached_file( + pretrained_model_name_or_path, WEIGHTS_INDEX_NAME, **cached_file_kwargs + ) + if resolved_archive_file is not None: + is_sharded = True + if resolved_archive_file is None: + # Otherwise, maybe there is a PyTorch or Flax model file. We try those to give a helpful error + # message. + has_file_kwargs = { + "revision": revision, + "proxies": proxies, + "token": token, + "cache_dir": cache_dir, + "local_files_only": local_files_only, + } + if has_file(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME, **has_file_kwargs): + is_sharded = True + elif has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs): + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {TF2_WEIGHTS_NAME} but there is a file for PyTorch weights. Use `from_pt=True` to" + " load this model from those weights." + ) + else: + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named {WEIGHTS_NAME}," + f" {TF2_WEIGHTS_NAME} or {TF_WEIGHTS_NAME}" + ) + + except EnvironmentError: + # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted + # to the original exception. + raise + except Exception: + # For any other exception, we throw a generic error. + + raise EnvironmentError( + f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it" + " from 'https://huggingface.co/models', make sure you don't have a local directory with the" + f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a" + f" directory containing a file named {WEIGHTS_NAME}, {TF2_WEIGHTS_NAME} or {TF_WEIGHTS_NAME}" + ) + if is_local: + logger.info(f"loading weights file {archive_file}") + resolved_archive_file = archive_file + filename = resolved_archive_file.split(os.path.sep)[-1] + else: + logger.info(f"loading weights file {filename} from cache at {resolved_archive_file}") + else: + resolved_archive_file = None + + # We'll need to download and cache each checkpoint shard if the checkpoint is sharded. + if is_sharded: + # resolved_archive_file becomes a list of files that point to the different checkpoint shards in this case. + resolved_archive_file, sharded_metadata = get_checkpoint_shard_files( + pretrained_model_name_or_path, + resolved_archive_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + _commit_hash=commit_hash, + ) + + safetensors_from_pt = False + if filename == SAFE_WEIGHTS_NAME: + with safe_open(resolved_archive_file, framework="tf") as f: + safetensors_metadata = f.metadata() + if safetensors_metadata is None or safetensors_metadata.get("format") not in ["pt", "tf", "flax", "mlx"]: + raise OSError( + f"The safetensors archive passed at {resolved_archive_file} does not contain the valid metadata." + " Make sure you save your model with the `save_pretrained` method." + ) + safetensors_from_pt = safetensors_metadata.get("format") == "pt" + elif filename == SAFE_WEIGHTS_INDEX_NAME: + with safe_open(resolved_archive_file[0], framework="tf") as f: + safetensors_metadata = f.metadata() + if safetensors_metadata is None or safetensors_metadata.get("format") not in ["pt", "tf", "flax", "mlx"]: + raise OSError( + f"The safetensors archive passed at {resolved_archive_file} does not contain the valid metadata." + " Make sure you save your model with the `save_pretrained` method." + ) + safetensors_from_pt = safetensors_metadata.get("format") == "pt" + + config.name_or_path = pretrained_model_name_or_path + + # composed models, *e.g.* TFRag, require special treatment when it comes to loading + # pre-trained weights. + if cls._requires_load_weight_prefix and model_kwargs.get("name") is not None: + model_kwargs["load_weight_prefix"] = load_weight_prefix + "/" + model_kwargs.get("name") + + # Instantiate model. + model = cls(config, *model_args, **model_kwargs) + + if tf_to_pt_weight_rename is None and hasattr(model, "tf_to_pt_weight_rename"): + # TODO Matt: This is a temporary workaround to allow weight renaming, but requires a method + # to be defined for each class that requires a rename. We can probably just have a class-level + # dict and a single top-level method or something and cut down a lot of boilerplate code + tf_to_pt_weight_rename = model.tf_to_pt_weight_rename + + if from_pt: + from .modeling_tf_pytorch_utils import load_pytorch_checkpoint_in_tf2_model + + # Load from a PyTorch checkpoint + return load_pytorch_checkpoint_in_tf2_model( + model, + resolved_archive_file, + allow_missing_keys=True, + output_loading_info=output_loading_info, + _prefix=load_weight_prefix, + tf_to_pt_weight_rename=tf_to_pt_weight_rename, + ) + + # we might need to extend the variable scope for composite models + if load_weight_prefix is not None: + with tf.compat.v1.variable_scope(load_weight_prefix): + model.build_in_name_scope() # build the network with dummy inputs + else: + model.build_in_name_scope() # build the network with dummy inputs + + if safetensors_from_pt and not is_sharded: + from .modeling_tf_pytorch_utils import load_pytorch_state_dict_in_tf2_model + + with safe_open(resolved_archive_file, framework="tf") as safetensors_archive: + # Load from a PyTorch safetensors checkpoint + # We load in TF format here because PT weights often need to be transposed, and this is much + # faster on GPU. Loading as numpy and transposing on CPU adds several seconds to load times. + return load_pytorch_state_dict_in_tf2_model( + model, + safetensors_archive, + tf_inputs=False, # No need to build the model again + allow_missing_keys=True, + output_loading_info=output_loading_info, + _prefix=load_weight_prefix, + ignore_mismatched_sizes=ignore_mismatched_sizes, + tf_to_pt_weight_rename=tf_to_pt_weight_rename, + ) + elif safetensors_from_pt: + from .modeling_tf_pytorch_utils import load_sharded_pytorch_safetensors_in_tf2_model + + return load_sharded_pytorch_safetensors_in_tf2_model( + model, + resolved_archive_file, + tf_inputs=False, + allow_missing_keys=True, + output_loading_info=output_loading_info, + _prefix=load_weight_prefix, + ignore_mismatched_sizes=ignore_mismatched_sizes, + tf_to_pt_weight_rename=tf_to_pt_weight_rename, + ) + + # 'by_name' allow us to do transfer learning by skipping/adding layers + # see https://github.com/tensorflow/tensorflow/blob/00fad90125b18b80fe054de1055770cfb8fe4ba3/tensorflow/python/keras/engine/network.py#L1339-L1357 + try: + if is_sharded: + for file in resolved_archive_file: + os.path.isfile(file), f"Error retrieving files {file}" + if filename == SAFE_WEIGHTS_INDEX_NAME: + missing_keys, unexpected_keys, mismatched_keys = load_tf_sharded_weights_from_safetensors( + model, + resolved_archive_file, + ignore_mismatched_sizes=ignore_mismatched_sizes, + _prefix=load_weight_prefix, + ) + else: + missing_keys, unexpected_keys, mismatched_keys = load_tf_sharded_weights( + model, + resolved_archive_file, + ignore_mismatched_sizes=ignore_mismatched_sizes, + _prefix=load_weight_prefix, + ) + else: + # Handles both H5 and safetensors + missing_keys, unexpected_keys, mismatched_keys = load_tf_weights( + model, + resolved_archive_file, + ignore_mismatched_sizes=ignore_mismatched_sizes, + _prefix=load_weight_prefix, + ) + except OSError as e: + try: + with open(resolved_archive_file) as f: + if f.read().startswith("version"): + raise OSError( + "You seem to have cloned a repository without having git-lfs installed. Please install " + "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder " + "you cloned." + ) + else: + raise ValueError from e + except (UnicodeDecodeError, ValueError): + raise OSError( + "Unable to load weights from h5 file. " + "If you tried to load a TF 2.0 model from a PyTorch checkpoint, please set from_pt=True. " + ) + + if cls._keys_to_ignore_on_load_missing is not None: + for pat in cls._keys_to_ignore_on_load_missing: + missing_keys = [k for k in missing_keys if re.search(pat, k) is None] + + if cls._keys_to_ignore_on_load_unexpected is not None: + for pat in cls._keys_to_ignore_on_load_unexpected: + unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None] + + if len(unexpected_keys) > 0: + logger.warning( + f"Some layers from the model checkpoint at {pretrained_model_name_or_path} were not used when" + f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are" + f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" + " with another architecture (e.g. initializing a BertForSequenceClassification model from a" + " BertForPreTraining model).\n- This IS NOT expected if you are initializing" + f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical" + " (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)." + ) + else: + logger.warning(f"All model checkpoint layers were used when initializing {model.__class__.__name__}.\n") + + if len(missing_keys) > 0: + logger.warning( + f"Some layers of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" + " TRAIN this model on a down-stream task to be able to use it for predictions and inference." + ) + elif len(mismatched_keys) == 0: + logger.warning( + f"All the layers of {model.__class__.__name__} were initialized from the model checkpoint at" + f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint" + f" was trained on, you can already use {model.__class__.__name__} for predictions without further" + " training." + ) + if len(mismatched_keys) > 0: + mismatched_warning = "\n".join( + [ + f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated" + for key, shape1, shape2 in mismatched_keys + ] + ) + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not" + f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able" + " to use it for predictions and inference." + ) + + # If it is a model with generation capabilities, attempt to load the generation config + if model.can_generate(): + try: + model.generation_config = GenerationConfig.from_pretrained( + pretrained_model_name_or_path, + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder=subfolder, + _from_auto=from_auto_class, + _from_pipeline=from_pipeline, + **kwargs, + ) + except OSError: + logger.info( + "Generation config file not found, using a generation config created from the model config." + ) + pass + + if output_loading_info: + loading_info = { + "missing_keys": missing_keys, + "unexpected_keys": unexpected_keys, + "mismatched_keys": mismatched_keys, + } + + return model, loading_info + + return model + + def push_to_hub( + self, + repo_id: str, + use_temp_dir: Optional[bool] = None, + commit_message: Optional[str] = None, + private: Optional[bool] = None, + max_shard_size: Optional[Union[int, str]] = "10GB", + token: Optional[Union[bool, str]] = None, + # (`use_auth_token` is deprecated: we have to keep it here as we don't have **kwargs) + use_auth_token: Optional[Union[bool, str]] = None, + create_pr: bool = False, + **base_model_card_args, + ) -> str: + """ + Upload the model files to the 🤗 Model Hub while synchronizing a local clone of the repo in `repo_path_or_name`. + + Parameters: + repo_id (`str`): + The name of the repository you want to push your model to. It should contain your organization name + when pushing to a given organization. + use_temp_dir (`bool`, *optional*): + Whether or not to use a temporary directory to store the files saved before they are pushed to the Hub. + Will default to `True` if there is no directory named like `repo_id`, `False` otherwise. + commit_message (`str`, *optional*): + Message to commit while pushing. Will default to `"Upload model"`. + private (`bool`, *optional*): + Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. + token (`bool` or `str`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). Will default to `True` if `repo_url` + is not specified. + max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`): + Only applicable for models. The maximum size for a checkpoint before being sharded. Checkpoints shard + will then be each of size lower than this size. If expressed as a string, needs to be digits followed + by a unit (like `"5MB"`). + create_pr (`bool`, *optional*, defaults to `False`): + Whether or not to create a PR with the uploaded files or directly commit. + + Examples: + + ```python + from transformers import TFAutoModel + + model = TFAutoModel.from_pretrained("google-bert/bert-base-cased") + + # Push the model to your namespace with the name "my-finetuned-bert". + model.push_to_hub("my-finetuned-bert") + + # Push the model to an organization with the name "my-finetuned-bert". + model.push_to_hub("huggingface/my-finetuned-bert") + ``` + """ + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if "repo_path_or_name" in base_model_card_args: + warnings.warn( + "The `repo_path_or_name` argument is deprecated and will be removed in v5 of Transformers. Use " + "`repo_id` instead." + ) + repo_id = base_model_card_args.pop("repo_path_or_name") + # Deprecation warning will be sent after for repo_url and organization + repo_url = base_model_card_args.pop("repo_url", None) + organization = base_model_card_args.pop("organization", None) + + if os.path.isdir(repo_id): + working_dir = repo_id + repo_id = repo_id.split(os.path.sep)[-1] + else: + working_dir = repo_id.split("/")[-1] + + repo_id = self._create_repo( + repo_id, private=private, token=token, repo_url=repo_url, organization=organization + ) + + if use_temp_dir is None: + use_temp_dir = not os.path.isdir(working_dir) + + with working_or_temp_dir(working_dir=working_dir, use_temp_dir=use_temp_dir) as work_dir: + files_timestamps = self._get_files_timestamps(work_dir) + + # Save all files. + self.save_pretrained(work_dir, max_shard_size=max_shard_size) + if hasattr(self, "history") and hasattr(self, "create_model_card"): + # This is a Keras model and we might be able to fish out its History and make a model card out of it + base_model_card_args = { + "output_dir": work_dir, + "model_name": Path(repo_id).name, + } + base_model_card_args.update(base_model_card_args) + self.create_model_card(**base_model_card_args) + + self._upload_modified_files( + work_dir, + repo_id, + files_timestamps, + commit_message=commit_message, + token=token, + create_pr=create_pr, + ) + + @classmethod + def register_for_auto_class(cls, auto_class="TFAutoModel"): + """ + Register this class with a given auto class. This should only be used for custom models as the ones in the + library are already mapped with an auto class. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"TFAutoModel"`): + The auto class to register this new model with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + +class TFConv1D(keras.layers.Layer): + """ + 1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2). + + Basically works like a linear layer but the weights are transposed. + + Args: + nf (`int`): + The number of output features. + nx (`int`): + The number of input features. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation to use to initialize the weights. + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments passed along to the `__init__` of `keras.layers.Layer`. + """ + + def __init__(self, nf, nx, initializer_range=0.02, **kwargs): + super().__init__(**kwargs) + self.nf = nf + self.nx = nx + self.initializer_range = initializer_range + + def build(self, input_shape): + if self.built: + return + self.built = True + self.weight = self.add_weight( + "weight", shape=[self.nx, self.nf], initializer=get_initializer(self.initializer_range) + ) + self.bias = self.add_weight("bias", shape=[1, self.nf], initializer=tf.zeros_initializer()) + + def call(self, x): + bz, sl = shape_list(x)[:2] + + x = tf.reshape(x, [-1, self.nx]) + x = tf.matmul(x, self.weight) + self.bias + + x = tf.reshape(x, [bz, sl, self.nf]) + + return x + + +class TFSharedEmbeddings(keras.layers.Layer): + r""" + Construct shared token embeddings. + + The weights of the embedding layer is usually shared with the weights of the linear decoder when doing language + modeling. + + Args: + vocab_size (`int`): + The size of the vocabulary, e.g., the number of unique tokens. + hidden_size (`int`): + The size of the embedding vectors. + initializer_range (`float`, *optional*): + The standard deviation to use when initializing the weights. If no value is provided, it will default to + \\(1/\sqrt{hidden\_size}\\). + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments passed along to the `__init__` of `keras.layers.Layer`. + """ + + # TODO (joao): flagged for delection due to embeddings refactor + + def __init__(self, vocab_size: int, hidden_size: int, initializer_range: Optional[float] = None, **kwargs): + super().__init__(**kwargs) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.initializer_range = hidden_size**-0.5 if initializer_range is None else initializer_range + warnings.warn( + "`TFSharedEmbeddings` is scheduled for deletion in v4.32, use `keras.layers.Embedding` instead.", + DeprecationWarning, + ) + + def build(self, input_shape): + """ + Build shared token embedding layer Shared weights logic adapted from + https://github.com/tensorflow/models/blob/a009f4fb9d2fc4949e32192a944688925ef78659/official/transformer/v2/embedding_layer.py#L24 + """ + self.weight = self.add_weight( + "weight", shape=[self.vocab_size, self.hidden_size], initializer=get_initializer(self.initializer_range) + ) + super().build(input_shape) + + def get_config(self): + config = { + "vocab_size": self.vocab_size, + "hidden_size": self.hidden_size, + "initializer_range": self.initializer_range, + } + base_config = super().get_config() + + return dict(list(base_config.items()) + list(config.items())) + + def call(self, inputs: tf.Tensor, mode: str = "embedding") -> tf.Tensor: + """ + Get token embeddings of inputs or decode final hidden state. + + Args: + inputs (`tf.Tensor`): + In embedding mode, should be an int64 tensor with shape `[batch_size, length]`. + + In linear mode, should be a float tensor with shape `[batch_size, length, hidden_size]`. + mode (`str`, defaults to `"embedding"`): + A valid value is either `"embedding"` or `"linear"`, the first one indicates that the layer should be + used as an embedding layer, the second one that the layer should be used as a linear decoder. + + Returns: + `tf.Tensor`: In embedding mode, the output is a float32 embedding tensor, with shape `[batch_size, length, + embedding_size]`. + + In linear mode, the output is a float32 with shape `[batch_size, length, vocab_size]`. + + Raises: + ValueError: if `mode` is not valid. + + Shared weights logic is adapted from + [here](https://github.com/tensorflow/models/blob/a009f4fb9d2fc4949e32192a944688925ef78659/official/transformer/v2/embedding_layer.py#L24). + """ + if mode == "embedding": + return self._embedding(inputs) + elif mode == "linear": + return self._linear(inputs) + else: + raise ValueError(f"mode {mode} is not valid.") + + def _embedding(self, input_ids): + """Applies embedding based on inputs tensor.""" + return tf.gather(self.weight, input_ids) + + def _linear(self, inputs): + """ + Computes logits by running inputs through a linear layer. + + Args: + inputs: A float32 tensor with shape [..., hidden_size] + + Returns: + float32 tensor with shape [..., vocab_size]. + """ + first_dims = shape_list(inputs)[:-1] + x = tf.reshape(inputs, [-1, self.hidden_size]) + logits = tf.matmul(x, self.weight, transpose_b=True) + + return tf.reshape(logits, first_dims + [self.vocab_size]) + + +class TFSequenceSummary(keras.layers.Layer): + """ + Compute a single vector summary of a sequence hidden states. + + Args: + config ([`PretrainedConfig`]): + The config used by the model. Relevant arguments in the config class of the model are (refer to the actual + config class of your model for the default values it uses): + + - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are: + + - `"last"` -- Take the last token hidden state (like XLNet) + - `"first"` -- Take the first token hidden state (like Bert) + - `"mean"` -- Take the mean of all tokens hidden states + - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2) + - `"attn"` -- Not implemented now, use multi-head attention + + - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction. + - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes + (otherwise to `config.hidden_size`). + - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output, + another string or `None` will add no activation. + - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation. + - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation. + + initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation to use to initialize the weights. + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments passed along to the `__init__` of `keras.layers.Layer`. + """ + + def __init__(self, config: PretrainedConfig, initializer_range: float = 0.02, **kwargs): + super().__init__(**kwargs) + + self.summary_type = config.summary_type if hasattr(config, "summary_use_proj") else "last" + if self.summary_type == "attn": + # We should use a standard multi-head attention module with absolute positional embedding for that. + # Cf. https://github.com/zihangdai/xlnet/blob/master/modeling.py#L253-L276 + # We can probably just use the multi-head attention module of PyTorch >=1.1.0 + raise NotImplementedError + + self.has_summary = hasattr(config, "summary_use_proj") and config.summary_use_proj + if self.has_summary: + if hasattr(config, "summary_proj_to_labels") and config.summary_proj_to_labels and config.num_labels > 0: + num_classes = config.num_labels + else: + num_classes = config.hidden_size + self.summary = keras.layers.Dense( + num_classes, kernel_initializer=get_initializer(initializer_range), name="summary" + ) + + self.has_activation = False + activation_string = getattr(config, "summary_activation", None) + if activation_string is not None: + self.has_activation = True + self.activation = get_tf_activation(activation_string) + + self.has_first_dropout = hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0 + if self.has_first_dropout: + self.first_dropout = keras.layers.Dropout(config.summary_first_dropout) + + self.has_last_dropout = hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0 + if self.has_last_dropout: + self.last_dropout = keras.layers.Dropout(config.summary_last_dropout) + self.hidden_size = config.hidden_size + + def call(self, inputs, cls_index=None, training=False): + if not isinstance(inputs, (dict, tuple, list)): + hidden_states = inputs + elif isinstance(inputs, (tuple, list)): + hidden_states = inputs[0] + cls_index = inputs[1] if len(inputs) > 1 else None + assert len(inputs) <= 2, "Too many inputs." + else: + hidden_states = inputs.get("hidden_states") + cls_index = inputs.get("cls_index", None) + + if self.summary_type == "last": + output = hidden_states[:, -1] + elif self.summary_type == "first": + output = hidden_states[:, 0] + elif self.summary_type == "mean": + output = tf.reduce_mean(hidden_states, axis=1) + elif self.summary_type == "cls_index": + hidden_shape = shape_list(hidden_states) # e.g. [batch, num choices, seq length, hidden dims] + if cls_index is None: + cls_index = tf.fill( + hidden_shape[:-2], hidden_shape[-2] - 1 + ) # A tensor full of shape [batch] or [batch, num choices] full of sequence length + cls_shape = shape_list(cls_index) + if len(cls_shape) <= len(hidden_shape) - 2: + cls_index = tf.expand_dims(cls_index, axis=-1) + # else: + # cls_index = cls_index[..., tf.newaxis] + # cls_index = cls_index.expand((-1,) * (cls_index.dim()-1) + (hidden_states.size(-1),)) + # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states + output = tf.gather(hidden_states, cls_index, batch_dims=len(hidden_shape) - 2) + output = tf.squeeze( + output, axis=len(hidden_shape) - 2 + ) # shape of output: (batch, num choices, hidden_size) + elif self.summary_type == "attn": + raise NotImplementedError + + if self.has_first_dropout: + output = self.first_dropout(output, training=training) + + if self.has_summary: + output = self.summary(output) + + if self.has_activation: + output = self.activation(output) + + if self.has_last_dropout: + output = self.last_dropout(output, training=training) + + return output + + def build(self, input_shape): + if self.built: + return + self.built = True + if getattr(self, "summary", None) is not None: + with tf.name_scope("summary"): + self.summary.build(self.hidden_size) + + +def get_initializer(initializer_range: float = 0.02) -> keras.initializers.TruncatedNormal: + """ + Creates a `keras.initializers.TruncatedNormal` with the given range. + + Args: + initializer_range (*float*, defaults to 0.02): Standard deviation of the initializer range. + + Returns: + `keras.initializers.TruncatedNormal`: The truncated normal initializer. + """ + return keras.initializers.TruncatedNormal(stddev=initializer_range) diff --git a/lib/python3.12/site-packages/transformers/modeling_utils.py b/lib/python3.12/site-packages/transformers/modeling_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f51a0c1f8d0eeaf5f0c86a6a8b8610e29c21564f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/modeling_utils.py @@ -0,0 +1,5866 @@ +# coding=utf-8 +# Copyright 2018 The Google AI Language Team Authors, Facebook AI Research authors and The HuggingFace Inc. team. +# Copyright (c) 2018, NVIDIA CORPORATION. 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. +import collections +import copy +import functools +import gc +import importlib.metadata +import inspect +import itertools +import json +import math +import os +import re +import shutil +import tempfile +import warnings +from collections import defaultdict +from collections.abc import MutableMapping +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from functools import partial, wraps +from threading import Thread +from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, TypeVar, Union +from zipfile import is_zipfile + +import torch +import torch.distributed.tensor +from huggingface_hub import split_torch_state_dict_into_shards +from packaging import version +from torch import Tensor, nn +from torch.distributions import constraints +from torch.nn import CrossEntropyLoss, Identity +from torch.utils.checkpoint import checkpoint + +from transformers.utils import is_torchao_available + + +if is_torchao_available(): + from torchao.quantization import Int4WeightOnlyConfig + +from .activations import get_activation +from .configuration_utils import PretrainedConfig +from .dynamic_module_utils import custom_object_save +from .generation import CompileConfig, GenerationConfig, GenerationMixin +from .integrations import PeftAdapterMixin, deepspeed_config, is_deepspeed_zero3_enabled +from .integrations.accelerate import find_tied_parameters, init_empty_weights +from .integrations.deepspeed import _load_state_dict_into_zero3_model, is_deepspeed_available +from .integrations.flash_attention import flash_attention_forward +from .integrations.flex_attention import flex_attention_forward +from .integrations.sdpa_attention import sdpa_attention_forward +from .integrations.tensor_parallel import ( + SUPPORTED_TP_STYLES, + shard_and_distribute_module, +) +from .loss.loss_utils import LOSS_MAPPING +from .pytorch_utils import ( # noqa: F401 + Conv1D, + apply_chunking_to_forward, + find_pruneable_heads_and_indices, + id_tensor_storage, + prune_conv1d_layer, + prune_layer, + prune_linear_layer, +) +from .quantizers import AutoHfQuantizer, HfQuantizer +from .quantizers.quantizers_utils import get_module_from_name +from .safetensors_conversion import auto_conversion +from .utils import ( + ADAPTER_SAFE_WEIGHTS_NAME, + ADAPTER_WEIGHTS_NAME, + CONFIG_NAME, + DUMMY_INPUTS, + FLAX_WEIGHTS_NAME, + SAFE_WEIGHTS_INDEX_NAME, + SAFE_WEIGHTS_NAME, + TF2_WEIGHTS_NAME, + TF_WEIGHTS_NAME, + WEIGHTS_INDEX_NAME, + WEIGHTS_NAME, + ContextManagers, + ModelOutput, + PushToHubMixin, + cached_file, + copy_func, + download_url, + extract_commit_hash, + has_file, + is_accelerate_available, + is_bitsandbytes_available, + is_flash_attn_2_available, + is_offline_mode, + is_optimum_available, + is_peft_available, + is_remote_url, + is_safetensors_available, + is_torch_flex_attn_available, + is_torch_greater_or_equal, + is_torch_mlu_available, + is_torch_npu_available, + is_torch_sdpa_available, + is_torch_xla_available, + logging, + replace_return_docstrings, + strtobool, +) +from .utils.hub import create_and_tag_model_card, get_checkpoint_shard_files +from .utils.import_utils import ( + ENV_VARS_TRUE_VALUES, + is_sagemaker_mp_enabled, + is_torch_fx_proxy, + is_torchdynamo_compiling, +) +from .utils.quantization_config import BitsAndBytesConfig, QuantizationMethod + + +XLA_USE_BF16 = os.environ.get("XLA_USE_BF16", "0").upper() +XLA_DOWNCAST_BF16 = os.environ.get("XLA_DOWNCAST_BF16", "0").upper() + + +if is_accelerate_available(): + from accelerate import dispatch_model, infer_auto_device_map + from accelerate.hooks import add_hook_to_module + from accelerate.utils import ( + check_tied_parameters_on_same_device, + extract_model_from_parallel, + get_balanced_memory, + get_max_memory, + load_offloaded_weights, + offload_weight, + save_offload_index, + ) + + accelerate_version = version.parse(importlib.metadata.version("accelerate")) + if accelerate_version >= version.parse("0.31"): + from accelerate.utils.modeling import get_state_dict_from_offload + +if is_safetensors_available(): + from safetensors import safe_open + from safetensors.torch import load_file as safe_load_file + from safetensors.torch import save_file as safe_save_file + + +if is_deepspeed_available(): + import deepspeed + +logger = logging.get_logger(__name__) + + +_init_weights = True +_is_quantized = False +_is_ds_init_called = False +_torch_distributed_available = torch.distributed.is_available() + + +def is_fsdp_enabled(): + return ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and strtobool(os.environ.get("ACCELERATE_USE_FSDP", "False")) == 1 + and strtobool(os.environ.get("FSDP_CPU_RAM_EFFICIENT_LOADING", "False")) == 1 + ) + + +def is_local_dist_rank_0(): + return ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and int(os.environ.get("LOCAL_RANK", -1)) == 0 + ) + + +if is_sagemaker_mp_enabled(): + import smdistributed.modelparallel.torch as smp + from smdistributed.modelparallel import __version__ as SMP_VERSION + + IS_SAGEMAKER_MP_POST_1_10 = version.parse(SMP_VERSION) >= version.parse("1.10") +else: + IS_SAGEMAKER_MP_POST_1_10 = False + +if is_peft_available(): + from .utils import find_adapter_config_file + + +SpecificPreTrainedModelType = TypeVar("SpecificPreTrainedModelType", bound="PreTrainedModel") + +TORCH_INIT_FUNCTIONS = { + "uniform_": nn.init.uniform_, + "normal_": nn.init.normal_, + "trunc_normal_": nn.init.trunc_normal_, + "constant_": nn.init.constant_, + "xavier_uniform_": nn.init.xavier_uniform_, + "xavier_normal_": nn.init.xavier_normal_, + "kaiming_uniform_": nn.init.kaiming_uniform_, + "kaiming_normal_": nn.init.kaiming_normal_, + "uniform": nn.init.uniform, + "normal": nn.init.normal, + "xavier_uniform": nn.init.xavier_uniform, + "xavier_normal": nn.init.xavier_normal, + "kaiming_uniform": nn.init.kaiming_uniform, + "kaiming_normal": nn.init.kaiming_normal, +} + + +@contextmanager +def no_init_weights(): + """ + Context manager to globally disable weight initialization to speed up loading large models. + """ + global _init_weights + old_init_weights = _init_weights + + _init_weights = False + + def _skip_init(*args, **kwargs): + pass + + # Save the original initialization functions + for name, init_func in TORCH_INIT_FUNCTIONS.items(): + setattr(torch.nn.init, name, _skip_init) + + try: + yield + finally: + _init_weights = old_init_weights + # Restore the original initialization functions + for name, init_func in TORCH_INIT_FUNCTIONS.items(): + setattr(torch.nn.init, name, init_func) + + +@contextmanager +def set_quantized_state(): + global _is_quantized + _is_quantized = True + try: + yield + finally: + _is_quantized = False + + +# Skip recursive calls to deepspeed.zero.Init to avoid pinning errors. +# This issue occurs with ZeRO stage 3 when using NVMe offloading. +# For more details, refer to issue #34429. +@contextmanager +def set_zero3_state(): + global _is_ds_init_called + _is_ds_init_called = True + try: + yield + finally: + _is_ds_init_called = False + + +def restore_default_torch_dtype(func): + """ + Decorator to restore the default torch dtype + at the end of the function. Serves + as a backup in case calling the function raises + an error after the function has changed the default dtype but before it could restore it. + """ + + @wraps(func) + def _wrapper(*args, **kwargs): + old_dtype = torch.get_default_dtype() + try: + return func(*args, **kwargs) + finally: + torch.set_default_dtype(old_dtype) + + return _wrapper + + +def get_parameter_device(parameter: Union[nn.Module, "ModuleUtilsMixin"]): + try: + return next(parameter.parameters()).device + except StopIteration: + # For nn.DataParallel compatibility in PyTorch 1.5 + + def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]: + tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] + return tuples + + gen = parameter._named_members(get_members_fn=find_tensor_attributes) + first_tuple = next(gen) + return first_tuple[1].device + + +def get_first_parameter_dtype(parameter: Union[nn.Module, "ModuleUtilsMixin"]): + """ + Returns the first parameter dtype (can be non-floating) or asserts if none were found. + """ + try: + return next(parameter.parameters()).dtype + except StopIteration: + # For nn.DataParallel compatibility in PyTorch > 1.5 + + def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]: + tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] + return tuples + + gen = parameter._named_members(get_members_fn=find_tensor_attributes) + first_tuple = next(gen) + return first_tuple[1].dtype + + +def get_parameter_dtype(parameter: Union[nn.Module, "ModuleUtilsMixin"]): + """ + Returns the first found floating dtype in parameters if there is one, otherwise returns the last dtype it found. + """ + last_dtype = None + for t in parameter.parameters(): + last_dtype = t.dtype + if t.is_floating_point(): + # Adding fix for https://github.com/pytorch/xla/issues/4152 + # Fixes issue where the model code passes a value that is out of range for XLA_USE_BF16=1 + # and XLA_DOWNCAST_BF16=1 so the conversion would cast it to -inf + # NOTE: `is_torch_xla_available()` is checked last as it induces a graph break in torch dynamo + if XLA_USE_BF16 in ENV_VARS_TRUE_VALUES and is_torch_xla_available(): + return torch.bfloat16 + if XLA_DOWNCAST_BF16 in ENV_VARS_TRUE_VALUES and is_torch_xla_available(): + if t.dtype == torch.float: + return torch.bfloat16 + if t.dtype == torch.double: + return torch.float32 + return t.dtype + + if last_dtype is not None: + # if no floating dtype was found return whatever the first dtype is + return last_dtype + + # For nn.DataParallel compatibility in PyTorch > 1.5 + def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]: + tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)] + return tuples + + gen = parameter._named_members(get_members_fn=find_tensor_attributes) + last_tuple = None + for tuple in gen: + last_tuple = tuple + if tuple[1].is_floating_point(): + return tuple[1].dtype + + if last_tuple is not None: + # fallback to the last dtype + return last_tuple[1].dtype + + # fallback to buffer dtype + for t in parameter.buffers(): + last_dtype = t.dtype + if t.is_floating_point(): + return t.dtype + return last_dtype + + +def get_state_dict_float_dtype(state_dict): + """ + Returns the first found floating dtype in `state_dict` or asserts if none were found. + """ + for t in state_dict.values(): + if t.is_floating_point(): + return t.dtype + + raise ValueError("couldn't find any floating point dtypes in state_dict") + + +def get_state_dict_dtype(state_dict): + """ + Returns the first found floating dtype in `state_dict` if there is one, otherwise returns the first dtype. + """ + for t in state_dict.values(): + if t.is_floating_point(): + return t.dtype + + # if no floating dtype was found return whatever the first dtype is + else: + return next(state_dict.values()).dtype + + +def load_sharded_checkpoint(model, folder, strict=True, prefer_safe=True): + """ + This is the same as + [`torch.nn.Module.load_state_dict`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=load_state_dict#torch.nn.Module.load_state_dict) + but for a sharded checkpoint. + + This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being + loaded in the model. + + Args: + model (`torch.nn.Module`): The model in which to load the checkpoint. + folder (`str` or `os.PathLike`): A path to a folder containing the sharded checkpoint. + strict (`bool`, *optional*, defaults to `True`): + Whether to strictly enforce that the keys in the model state dict match the keys in the sharded checkpoint. + prefer_safe (`bool`, *optional*, defaults to `False`): + If both safetensors and PyTorch save files are present in checkpoint and `prefer_safe` is True, the + safetensors files will be loaded. Otherwise, PyTorch files are always loaded when possible. + + Returns: + `NamedTuple`: A named tuple with `missing_keys` and `unexpected_keys` fields + - `missing_keys` is a list of str containing the missing keys + - `unexpected_keys` is a list of str containing the unexpected keys + """ + # Load the index + index_file = os.path.join(folder, WEIGHTS_INDEX_NAME) + safe_index_file = os.path.join(folder, SAFE_WEIGHTS_INDEX_NAME) + + index_present = os.path.isfile(index_file) + safe_index_present = os.path.isfile(safe_index_file) + + if not index_present and not (safe_index_present and is_safetensors_available()): + filenames = ( + (WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME) if is_safetensors_available() else (WEIGHTS_INDEX_NAME,) + ) + raise ValueError(f"Can't find a checkpoint index ({' or '.join(filenames)}) in {folder}.") + + load_safe = False + if safe_index_present: + if prefer_safe: + if is_safetensors_available(): + load_safe = True # load safe due to preference + else: + logger.warning( + f"Cannot load sharded checkpoint at {folder} safely since safetensors is not installed!" + ) + elif not index_present: + load_safe = True # load safe since we have no other choice + + load_index = safe_index_file if load_safe else index_file + + with open(load_index, "r", encoding="utf-8") as f: + index = json.load(f) + + shard_files = list(set(index["weight_map"].values())) + + # If strict=True, error before loading any of the state dicts. + loaded_keys = index["weight_map"].keys() + model_keys = model.state_dict().keys() + missing_keys = [key for key in model_keys if key not in loaded_keys] + unexpected_keys = [key for key in loaded_keys if key not in model_keys] + if strict and (len(missing_keys) > 0 or len(unexpected_keys) > 0): + error_message = f"Error(s) in loading state_dict for {model.__class__.__name__}" + if len(missing_keys) > 0: + str_missing_keys = ",".join([f'"{k}"' for k in missing_keys]) + error_message += f"\nMissing key(s): {str_missing_keys}." + if len(unexpected_keys) > 0: + str_unexpected_keys = ",".join([f'"{k}"' for k in unexpected_keys]) + error_message += f"\nMissing key(s): {str_unexpected_keys}." + raise RuntimeError(error_message) + + loader = safe_load_file if load_safe else partial(torch.load, map_location="cpu", weights_only=True) + + for shard_file in shard_files: + state_dict = loader(os.path.join(folder, shard_file)) + model.load_state_dict(state_dict, strict=False) + + # Make sure memory is freed before we load the next state dict. + del state_dict + gc.collect() + + # Return the same thing as PyTorch load_state_dict function. + return torch.nn.modules.module._IncompatibleKeys(missing_keys, unexpected_keys) + + +str_to_torch_dtype = { + "BOOL": torch.bool, + "U8": torch.uint8, + "I8": torch.int8, + "I16": torch.int16, + "F16": torch.float16, + "BF16": torch.bfloat16, + "I32": torch.int32, + "F32": torch.float32, + "F64": torch.float64, + "I64": torch.int64, + "F8_E4M3": torch.float8_e4m3fn, +} + +if is_torch_greater_or_equal("2.1.0"): + str_to_torch_dtype["F8_E4M3"] = torch.float8_e4m3fn + +if is_torch_greater_or_equal("2.3.0"): + str_to_torch_dtype["U16"] = torch.uint16 + str_to_torch_dtype["U32"] = torch.uint32 + str_to_torch_dtype["U64"] = torch.uint64 + +if is_torch_greater_or_equal("2.1.0"): + str_to_torch_dtype["F8_E4M3"] = torch.float8_e4m3fn + str_to_torch_dtype["F8_E5M2"] = torch.float8_e5m2 + + +def load_state_dict( + checkpoint_file: Union[str, os.PathLike], + is_quantized: bool = False, + map_location: Optional[Union[str, torch.device]] = "cpu", + weights_only: bool = True, +): + """ + Reads a `safetensor` or a `.bin` checkpoint file. We load the checkpoint on "cpu" by default. + """ + if checkpoint_file.endswith(".safetensors") and is_safetensors_available(): + with safe_open(checkpoint_file, framework="pt") as f: + metadata = f.metadata() + + if metadata is not None and metadata.get("format") not in ["pt", "tf", "flax", "mlx"]: + raise OSError( + f"The safetensors archive passed at {checkpoint_file} does not contain the valid metadata. Make sure " + "you save your model with the `save_pretrained` method." + ) + state_dict = {} + for k in f.keys(): + k_dtype = f.get_slice(k).get_dtype() + if k_dtype in str_to_torch_dtype: + dtype = str_to_torch_dtype[k_dtype] + else: + raise ValueError(f"Cannot load safetensors of unknown dtype {k_dtype}") + if map_location == "meta": + state_dict[k] = torch.empty(size=f.get_slice(k).get_shape(), dtype=dtype, device="meta") + else: + state_dict[k] = f.get_tensor(k) + return state_dict + + try: + if map_location is None: + if ( + ( + is_deepspeed_zero3_enabled() + and torch.distributed.is_initialized() + and torch.distributed.get_rank() > 0 + ) + or (is_fsdp_enabled() and not is_local_dist_rank_0()) + ) and not is_quantized: + map_location = "meta" + else: + map_location = "cpu" + extra_args = {} + # mmap can only be used with files serialized with zipfile-based format. + if ( + isinstance(checkpoint_file, str) + and map_location != "meta" + and version.parse(torch.__version__) >= version.parse("2.1.0") + and is_zipfile(checkpoint_file) + ): + extra_args = {"mmap": True} + return torch.load( + checkpoint_file, + map_location=map_location, + weights_only=weights_only, + **extra_args, + ) + except Exception as e: + try: + with open(checkpoint_file) as f: + if f.read(7) == "version": + raise OSError( + "You seem to have cloned a repository without having git-lfs installed. Please install " + "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder " + "you cloned." + ) + else: + raise ValueError( + f"Unable to locate the file {checkpoint_file} which is necessary to load this pretrained " + "model. Make sure you have saved the model properly." + ) from e + except (UnicodeDecodeError, ValueError): + raise OSError( + f"Unable to load weights from pytorch checkpoint file for '{checkpoint_file}' " + f"at '{checkpoint_file}'. " + "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True." + ) + + +def set_initialized_submodules(model, state_dict_keys): + """ + Sets the `_is_hf_initialized` flag in all submodules of a given model when all its weights are in the loaded state + dict. + """ + state_dict_keys = set(state_dict_keys) + not_initialized_submodules = {} + for module_name, module in model.named_modules(): + if module_name == "": + # When checking if the root module is loaded there's no need to prepend module_name. + module_keys = set(module.state_dict()) + else: + module_keys = {f"{module_name}.{k}" for k in module.state_dict()} + if module_keys.issubset(state_dict_keys): + module._is_hf_initialized = True + else: + not_initialized_submodules[module_name] = module + return not_initialized_submodules + + +def _end_ptr(tensor: torch.Tensor) -> int: + # extract the end of the pointer if the tensor is a slice of a bigger tensor + if tensor.nelement(): + stop = tensor.view(-1)[-1].data_ptr() + tensor.element_size() + else: + stop = tensor.data_ptr() + return stop + + +def _get_tied_weight_keys(module: nn.Module, prefix=""): + tied_weight_keys = [] + if getattr(module, "_tied_weights_keys", None) is not None: + names = [f"{prefix}.{k}" if prefix else k for k in module._tied_weights_keys] + tied_weight_keys.extend(names) + if getattr(module, "_dynamic_tied_weights_keys", None) is not None: + names = [f"{prefix}.{k}" if prefix else k for k in module._dynamic_tied_weights_keys] + tied_weight_keys.extend(names) + for name, submodule in module.named_children(): + local_prefix = f"{prefix}.{name}" if prefix else name + tied_weight_keys.extend(_get_tied_weight_keys(submodule, prefix=local_prefix)) + return tied_weight_keys + + +def _find_disjoint(tensors: List[Set[str]], state_dict: Dict[str, torch.Tensor]) -> Tuple[List[Set[str]], List[str]]: + filtered_tensors = [] + for shared in tensors: + if len(shared) < 2: + filtered_tensors.append(shared) + continue + + areas = [] + for name in shared: + tensor = state_dict[name] + areas.append((tensor.data_ptr(), _end_ptr(tensor), name)) + areas.sort() + + _, last_stop, last_name = areas[0] + filtered_tensors.append({last_name}) + for start, stop, name in areas[1:]: + if start >= last_stop: + filtered_tensors.append({name}) + else: + filtered_tensors[-1].add(name) + last_stop = stop + disjoint_tensors = [] + shared_tensors = [] + for tensors in filtered_tensors: + if len(tensors) == 1: + disjoint_tensors.append(tensors.pop()) + else: + shared_tensors.append(tensors) + return shared_tensors, disjoint_tensors + + +def _find_identical(tensors: List[Set[str]], state_dict: Dict[str, torch.Tensor]) -> Tuple[List[Set[str]], Set[str]]: + shared_tensors = [] + identical = [] + for shared in tensors: + if len(shared) < 2: + continue + + areas = collections.defaultdict(set) + for name in shared: + tensor = state_dict[name] + area = (tensor.device, tensor.data_ptr(), _end_ptr(tensor)) + areas[area].add(name) + if len(areas) == 1: + identical.append(shared) + else: + shared_tensors.append(shared) + return shared_tensors, identical + + +def _infer_parameter_dtype( + model: "PreTrainedModel", + param_name: str, + empty_param: torch.Tensor, + keep_in_fp32_regex: Optional[re.Pattern] = None, + hf_quantizer: Optional[HfQuantizer] = None, +) -> Union[bool, Optional[torch.dtype]]: + try: + old_param = model.get_parameter_or_buffer(param_name) + except Exception as e: + if hf_quantizer is not None and hf_quantizer.quantization_config.quant_method == QuantizationMethod.HQQ: + return True, None + else: + raise e + is_torch_e4m3fn_available = hasattr(torch, "float8_e4m3fn") + # We convert floating dtypes to the `dtype` passed except for float8_e4m3fn type. We also want to keep the buffers/params + # in int/uint/bool and not cast them. + casting_dtype = None + is_param_float8_e4m3fn = is_torch_e4m3fn_available and empty_param.dtype == torch.float8_e4m3fn + if empty_param.dtype.is_floating_point and not is_param_float8_e4m3fn: + # First fp32 if part of the exception list + if keep_in_fp32_regex is not None and keep_in_fp32_regex.search(param_name): + casting_dtype = torch.float32 + # Then dtype that was instantiated in the meta model -- note that this respects subconfigs dtypes + elif hf_quantizer is not None: + casting_dtype = model.config._pre_quantization_dtype + else: + casting_dtype = old_param.dtype + return old_param is not None and old_param.is_contiguous(), casting_dtype + + +def _load_parameter_into_model(model: "PreTrainedModel", param_name: str, tensor: torch.Tensor): + """Cast a single parameter `param_name` into the `model`, with value `tensor`.""" + module, param_type = get_module_from_name(model, param_name) + # This will check potential shape mismatch if skipped before + module.load_state_dict({param_type: tensor}, strict=False, assign=True) + + +@torch.no_grad() +def _load_state_dict_into_meta_model( + model: "PreTrainedModel", + state_dict: Dict, + shard_file: str, + expected_keys: List[str], + reverse_renaming_mapping: Dict[str, str], + device_map: Optional[Dict] = None, + disk_offload_folder: Optional[str] = None, + disk_offload_index: Optional[Dict] = None, + cpu_offload_folder: Optional[str] = None, + cpu_offload_index: Optional[Dict] = None, + hf_quantizer: Optional[HfQuantizer] = None, + is_safetensors: bool = False, + keep_in_fp32_regex: Optional[re.Pattern] = None, + unexpected_keys: Optional[List[str]] = None, # passing `unexpected` for cleanup from quantization items + device_mesh: Optional["torch.distributed.device_mesh.DeviceMesh"] = None, +) -> Tuple[Optional[Dict], Optional[Dict]]: + """Load parameters from `meta_state_dict` into the model. The parameters of the `meta_state_dict` are on the meta + device in order to easily infer the shapes and dtypes that they will have. Then proper parameters are then loaded + from `shard_file`, which is the actual state dict file on disk. + This function takes care of correctly casting dtypes, devices, and sharding tensors in case of tensor parallelism. + """ + tensor_device = "cpu" + if device_map is not None and device_map.get("", None) is not None: + if device_map[""] not in ("cpu", torch.device("cpu")): + tensor_device = device_map[""].index if isinstance(device_map[""], torch.device) else device_map[""] + if device_map is not None: + device_map_regex = "|".join([re.escape(k) for k in sorted(device_map.keys(), reverse=True)]) + + is_quantized = hf_quantizer is not None + is_hqq_or_bnb = is_quantized and hf_quantizer.quantization_config.quant_method in [ + QuantizationMethod.HQQ, + QuantizationMethod.BITS_AND_BYTES, + ] + is_meta_state_dict = shard_file.endswith(".safetensors") and not is_hqq_or_bnb + file_pointer = None + if is_meta_state_dict: + file_pointer = safe_open(shard_file, framework="pt", device=tensor_device) + + for param_name, empty_param in state_dict.items(): + if param_name not in expected_keys: + continue + + # we need to use serialized_param_name as file pointer is untouched + if is_meta_state_dict: + # This is the name of the parameter as it appears on disk file + serialized_param_name = reverse_renaming_mapping[param_name] + param = file_pointer.get_slice(serialized_param_name) + else: + param = empty_param.to(tensor_device) # It is actually not empty! + + to_contiguous, casting_dtype = _infer_parameter_dtype( + model, + param_name, + empty_param, + keep_in_fp32_regex, + hf_quantizer, + ) + + if device_mesh is not None: # In this case, the param is already on the correct device! + shard_and_distribute_module( + model, + param, + empty_param, + param_name, + casting_dtype, + to_contiguous, + int(os.environ["RANK"]), # the rank + device_mesh, + ) + else: + param = param[...] + if casting_dtype is not None: + param = param.to(casting_dtype) + if to_contiguous: + param = param.contiguous() + + if device_map is None: + param_device = "cpu" + else: + module_layer = re.search(device_map_regex, param_name) + if not module_layer: + raise ValueError(f"{param_name} doesn't have any device set.") + else: + param_device = device_map[module_layer.group()] + + if param_device == "disk": + if not is_safetensors: + disk_offload_index = offload_weight(param, param_name, disk_offload_folder, disk_offload_index) + elif param_device == "cpu" and cpu_offload_index is not None: + cpu_offload_index = offload_weight(param, param_name, cpu_offload_folder, cpu_offload_index) + elif ( + not is_quantized + or (not hf_quantizer.requires_parameters_quantization) + or ( + not hf_quantizer.check_quantized_param( + model, + param, + param_name, + state_dict, + param_device=param_device, + device_map=device_map, + ) + ) + ): + if is_fsdp_enabled(): + param_device = "cpu" if is_local_dist_rank_0() else "meta" + + _load_parameter_into_model(model, param_name, param.to(param_device)) + + else: + hf_quantizer.create_quantized_param( + model, param, param_name, param_device, state_dict, unexpected_keys + ) + # For quantized modules with FSDP/DeepSpeed Stage 3, we need to quantize the parameter on the GPU + # and then cast it to CPU to avoid excessive memory usage on each GPU + # in comparison to the sharded model across GPUs. + if is_fsdp_enabled() or is_deepspeed_zero3_enabled(): + module, param_type = get_module_from_name(model, param_name) + value = getattr(module, param_type) + param_to = "cpu" + if is_fsdp_enabled() and not is_local_dist_rank_0(): + param_to = "meta" + val_kwargs = {} + if hasattr(module, "weight") and module.weight.__class__.__name__ == "Int8Params": + val_kwargs["requires_grad"] = False + value = type(value)(value.data.to(param_to), **val_kwargs, **value.__dict__) + setattr(module, param_type, value) + + if file_pointer is not None: + file_pointer.__exit__(None, None, None) + + return disk_offload_index, cpu_offload_index + + +def _add_variant(weights_name: str, variant: Optional[str] = None) -> str: + if variant is not None: + path, name = weights_name.rsplit(".", 1) + weights_name = f"{path}.{variant}.{name}" + return weights_name + + +def _get_resolved_checkpoint_files( + pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], + subfolder: str, + variant: Optional[str], + gguf_file: Optional[str], + from_tf: bool, + from_flax: bool, + use_safetensors: bool, + cache_dir: str, + force_download: bool, + proxies: Optional[Dict[str, str]], + local_files_only: bool, + token: Optional[Union[str, bool]], + user_agent: dict, + revision: str, + commit_hash: Optional[str], +) -> Tuple[Optional[List[str]], Optional[Dict]]: + """Get all the checkpoint filenames based on `pretrained_model_name_or_path`, and optional metadata if the + checkpoints are sharded. + This function will download the data if necesary. + """ + is_sharded = False + + if pretrained_model_name_or_path is not None and gguf_file is None: + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + is_local = os.path.isdir(pretrained_model_name_or_path) + if is_local: + if from_tf and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, TF_WEIGHTS_NAME + ".index") + ): + # Load from a TF 1.0 checkpoint in priority if from_tf + archive_file = os.path.join(pretrained_model_name_or_path, subfolder, TF_WEIGHTS_NAME + ".index") + elif from_tf and os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, TF2_WEIGHTS_NAME)): + # Load from a TF 2.0 checkpoint in priority if from_tf + archive_file = os.path.join(pretrained_model_name_or_path, subfolder, TF2_WEIGHTS_NAME) + elif from_flax and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME) + ): + # Load from a Flax checkpoint in priority if from_flax + archive_file = os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME) + elif use_safetensors is not False and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant)) + ): + # Load from a safetensors checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant) + ) + elif use_safetensors is not False and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)) + ): + # Load from a sharded safetensors checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant) + ) + is_sharded = True + elif not use_safetensors and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant)) + ): + # Load from a PyTorch checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant) + ) + elif not use_safetensors and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant)) + ): + # Load from a sharded PyTorch checkpoint + archive_file = os.path.join( + pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant) + ) + is_sharded = True + # At this stage we don't have a weight file so we will raise an error. + elif not use_safetensors and ( + os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, TF_WEIGHTS_NAME + ".index")) + or os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, TF2_WEIGHTS_NAME)) + ): + raise EnvironmentError( + f"Error no file named {_add_variant(WEIGHTS_NAME, variant)} found in directory" + f" {pretrained_model_name_or_path} but there is a file for TensorFlow weights. Use" + " `from_tf=True` to load this model from those weights." + ) + elif not use_safetensors and os.path.isfile( + os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME) + ): + raise EnvironmentError( + f"Error no file named {_add_variant(WEIGHTS_NAME, variant)} found in directory" + f" {pretrained_model_name_or_path} but there is a file for Flax weights. Use `from_flax=True`" + " to load this model from those weights." + ) + elif use_safetensors: + raise EnvironmentError( + f"Error no file named {_add_variant(SAFE_WEIGHTS_NAME, variant)} found in directory" + f" {pretrained_model_name_or_path}." + ) + else: + raise EnvironmentError( + f"Error no file named {_add_variant(WEIGHTS_NAME, variant)}, {_add_variant(SAFE_WEIGHTS_NAME, variant)}," + f" {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME + '.index'} or {FLAX_WEIGHTS_NAME} found in directory" + f" {pretrained_model_name_or_path}." + ) + elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)): + archive_file = pretrained_model_name_or_path + is_local = True + elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path + ".index")): + if not from_tf: + raise ValueError( + f"We found a TensorFlow checkpoint at {pretrained_model_name_or_path + '.index'}, please set " + "from_tf to True to load from this checkpoint." + ) + archive_file = os.path.join(subfolder, pretrained_model_name_or_path + ".index") + is_local = True + elif is_remote_url(pretrained_model_name_or_path): + filename = pretrained_model_name_or_path + resolved_archive_file = download_url(pretrained_model_name_or_path) + else: + # set correct filename + if from_tf: + filename = TF2_WEIGHTS_NAME + elif from_flax: + filename = FLAX_WEIGHTS_NAME + elif use_safetensors is not False: + filename = _add_variant(SAFE_WEIGHTS_NAME, variant) + else: + filename = _add_variant(WEIGHTS_NAME, variant) + + try: + # Load from URL or cache if already cached + cached_file_kwargs = { + "cache_dir": cache_dir, + "force_download": force_download, + "proxies": proxies, + "local_files_only": local_files_only, + "token": token, + "user_agent": user_agent, + "revision": revision, + "subfolder": subfolder, + "_raise_exceptions_for_gated_repo": False, + "_raise_exceptions_for_missing_entries": False, + "_commit_hash": commit_hash, + } + resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs) + + # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None + # result when internet is up, the repo and revision exist, but the file does not. + if resolved_archive_file is None and filename == _add_variant(SAFE_WEIGHTS_NAME, variant): + # Maybe the checkpoint is sharded, we try to grab the index name in this case. + resolved_archive_file = cached_file( + pretrained_model_name_or_path, + _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant), + **cached_file_kwargs, + ) + if resolved_archive_file is not None: + is_sharded = True + elif use_safetensors: + if revision == "main": + resolved_archive_file, revision, is_sharded = auto_conversion( + pretrained_model_name_or_path, **cached_file_kwargs + ) + cached_file_kwargs["revision"] = revision + if resolved_archive_file is None: + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(SAFE_WEIGHTS_NAME, variant)} or {_add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)} " + "and thus cannot be loaded with `safetensors`. Please make sure that the model has " + "been saved with `safe_serialization=True` or do not set `use_safetensors=True`." + ) + else: + # This repo has no safetensors file of any kind, we switch to PyTorch. + filename = _add_variant(WEIGHTS_NAME, variant) + resolved_archive_file = cached_file( + pretrained_model_name_or_path, filename, **cached_file_kwargs + ) + if resolved_archive_file is None and filename == _add_variant(WEIGHTS_NAME, variant): + # Maybe the checkpoint is sharded, we try to grab the index name in this case. + resolved_archive_file = cached_file( + pretrained_model_name_or_path, + _add_variant(WEIGHTS_INDEX_NAME, variant), + **cached_file_kwargs, + ) + if resolved_archive_file is not None: + is_sharded = True + if not local_files_only and not is_offline_mode(): + if resolved_archive_file is not None: + if filename in [WEIGHTS_NAME, WEIGHTS_INDEX_NAME]: + # If the PyTorch file was found, check if there is a safetensors file on the repository + # If there is no safetensors file on the repositories, start an auto conversion + safe_weights_name = SAFE_WEIGHTS_INDEX_NAME if is_sharded else SAFE_WEIGHTS_NAME + has_file_kwargs = { + "revision": revision, + "proxies": proxies, + "token": token, + "cache_dir": cache_dir, + "local_files_only": local_files_only, + } + cached_file_kwargs = { + "cache_dir": cache_dir, + "force_download": force_download, + "local_files_only": local_files_only, + "user_agent": user_agent, + "subfolder": subfolder, + "_raise_exceptions_for_gated_repo": False, + "_raise_exceptions_for_missing_entries": False, + "_commit_hash": commit_hash, + **has_file_kwargs, + } + if not has_file(pretrained_model_name_or_path, safe_weights_name, **has_file_kwargs): + Thread( + target=auto_conversion, + args=(pretrained_model_name_or_path,), + kwargs={"ignore_errors_during_conversion": True, **cached_file_kwargs}, + name="Thread-auto_conversion", + ).start() + else: + # Otherwise, no PyTorch file was found, maybe there is a TF or Flax model file. + # We try those to give a helpful error message. + has_file_kwargs = { + "revision": revision, + "proxies": proxies, + "token": token, + "cache_dir": cache_dir, + "local_files_only": local_files_only, + } + if has_file(pretrained_model_name_or_path, TF2_WEIGHTS_NAME, **has_file_kwargs): + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(WEIGHTS_NAME, variant)} but there is a file for TensorFlow weights." + " Use `from_tf=True` to load this model from those weights." + ) + elif has_file(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME, **has_file_kwargs): + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(WEIGHTS_NAME, variant)} but there is a file for Flax weights. Use" + " `from_flax=True` to load this model from those weights." + ) + elif variant is not None and has_file( + pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs + ): + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(WEIGHTS_NAME, variant)} but there is a file without the variant" + f" {variant}. Use `variant=None` to load this model from those weights." + ) + else: + raise EnvironmentError( + f"{pretrained_model_name_or_path} does not appear to have a file named" + f" {_add_variant(WEIGHTS_NAME, variant)}, {_add_variant(SAFE_WEIGHTS_NAME, variant)}," + f" {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or {FLAX_WEIGHTS_NAME}." + ) + + except EnvironmentError: + # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted + # to the original exception. + raise + except Exception as e: + # For any other exception, we throw a generic error. + raise EnvironmentError( + f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it" + " from 'https://huggingface.co/models', make sure you don't have a local directory with the" + f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a" + f" directory containing a file named {_add_variant(WEIGHTS_NAME, variant)}," + f" {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or {FLAX_WEIGHTS_NAME}." + ) from e + + if is_local: + logger.info(f"loading weights file {archive_file}") + resolved_archive_file = archive_file + else: + logger.info(f"loading weights file {filename} from cache at {resolved_archive_file}") + + elif gguf_file: + # Case 1: the GGUF file is present locally + if os.path.isfile(gguf_file): + resolved_archive_file = gguf_file + # Case 2: The GGUF path is a location on the Hub + # Load from URL or cache if already cached + else: + cached_file_kwargs = { + "cache_dir": cache_dir, + "force_download": force_download, + "proxies": proxies, + "local_files_only": local_files_only, + "token": token, + "user_agent": user_agent, + "revision": revision, + "subfolder": subfolder, + "_raise_exceptions_for_gated_repo": False, + "_raise_exceptions_for_missing_entries": False, + "_commit_hash": commit_hash, + } + + resolved_archive_file = cached_file(pretrained_model_name_or_path, gguf_file, **cached_file_kwargs) + + # We now download and resolve all checkpoint files if the checkpoint is sharded + sharded_metadata = None + if is_sharded: + checkpoint_files, sharded_metadata = get_checkpoint_shard_files( + pretrained_model_name_or_path, + resolved_archive_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _commit_hash=commit_hash, + ) + else: + checkpoint_files = [resolved_archive_file] if pretrained_model_name_or_path is not None else None + + return checkpoint_files, sharded_metadata + + +def _get_torch_dtype( + cls, + torch_dtype: Optional[Union[str, torch.dtype, Dict]], + checkpoint_files: Optional[List[str]], + config: PretrainedConfig, + sharded_metadata: Optional[Dict], + state_dict: Optional[Dict], + weights_only: bool, +) -> Tuple[PretrainedConfig, Optional[torch.dtype], Optional[torch.dtype]]: + """Find the correct `torch_dtype` to use based on provided arguments. Also update the `config` based on the + inferred dtype. We do the following: + 1. If torch_dtype is not None, we use that dtype + 2. If torch_dtype is "auto", we auto-detect dtype from the loaded state_dict, by checking its first + weights entry that is of a floating type - we assume all floating dtype weights are of the same dtype + we also may have config.torch_dtype available, but we won't rely on it till v5 + """ + dtype_orig = None + is_sharded = sharded_metadata is not None + + if torch_dtype is not None: + if isinstance(torch_dtype, str): + if torch_dtype == "auto": + if hasattr(config, "torch_dtype") and config.torch_dtype is not None: + torch_dtype = config.torch_dtype + logger.info(f"Will use torch_dtype={torch_dtype} as defined in model's config object") + else: + if is_sharded and "dtype" in sharded_metadata: + torch_dtype = sharded_metadata["dtype"] + elif state_dict is not None: + torch_dtype = get_state_dict_dtype(state_dict) + else: + state_dict = load_state_dict( + checkpoint_files[0], map_location="meta", weights_only=weights_only + ) + torch_dtype = get_state_dict_dtype(state_dict) + logger.info( + "Since the `torch_dtype` attribute can't be found in model's config object, " + "will use torch_dtype={torch_dtype} as derived from model's weights" + ) + elif hasattr(torch, torch_dtype): + torch_dtype = getattr(torch, torch_dtype) + config.torch_dtype = torch_dtype + for sub_config_key in config.sub_configs.keys(): + sub_config = getattr(config, sub_config_key) + sub_config.torch_dtype = torch_dtype + elif isinstance(torch_dtype, torch.dtype): + config.torch_dtype = torch_dtype + for sub_config_key in config.sub_configs.keys(): + sub_config = getattr(config, sub_config_key) + sub_config.torch_dtype = torch_dtype + elif isinstance(torch_dtype, dict): + for key, curr_dtype in torch_dtype.items(): + if hasattr(config, key): + value = getattr(config, key) + curr_dtype = curr_dtype if not isinstance(curr_dtype, str) else getattr(torch, curr_dtype) + value.torch_dtype = curr_dtype + # main torch dtype for modules that aren't part of any sub-config + torch_dtype = torch_dtype.get("") + torch_dtype = torch_dtype if not isinstance(torch_dtype, str) else getattr(torch, torch_dtype) + config.torch_dtype = torch_dtype + if torch_dtype is None: + torch_dtype = torch.float32 + else: + raise ValueError( + f"`torch_dtype` can be one of: `torch.dtype`, `'auto'`, a string of a valid `torch.dtype` or a `dict` with valid `torch_dtype` " + f"for each sub-config in composite configs, but received {torch_dtype}" + ) + + dtype_orig = cls._set_default_torch_dtype(torch_dtype) + else: + # set fp32 as the default dtype for BC + default_dtype = torch.get_default_dtype() + config.torch_dtype = default_dtype + for key in config.sub_configs.keys(): + value = getattr(config, key) + value.torch_dtype = default_dtype + + return config, torch_dtype, dtype_orig + + +def _get_device_map( + model: "PreTrainedModel", + device_map: Optional[Union[str, Dict]], + max_memory: Optional[Dict], + hf_quantizer: Optional[HfQuantizer], + torch_dtype: Optional[torch.dtype], + keep_in_fp32_regex: Optional[re.Pattern], +) -> Dict: + """Compute the final `device_map` to use if we passed a value in ['auto', 'balanced', 'balanced_low_0', 'sequential']. + Otherwise, we check for any device inconsistencies in the device_map. + """ + if isinstance(device_map, str): + special_dtypes = {} + if hf_quantizer is not None: + special_dtypes.update(hf_quantizer.get_special_dtypes_update(model, torch_dtype)) + if keep_in_fp32_regex is not None: + special_dtypes.update( + {name: torch.float32 for name, _ in model.named_parameters() if keep_in_fp32_regex.search(name)} + ) + + target_dtype = torch_dtype + + if hf_quantizer is not None: + target_dtype = hf_quantizer.adjust_target_dtype(target_dtype) + + no_split_modules = model._get_no_split_modules(device_map) + device_map_kwargs = {"no_split_module_classes": no_split_modules} + + if "special_dtypes" in inspect.signature(infer_auto_device_map).parameters: + device_map_kwargs["special_dtypes"] = special_dtypes + elif len(special_dtypes) > 0: + logger.warning( + "This model has some weights that should be kept in higher precision, you need to upgrade " + "`accelerate` to properly deal with them (`pip install --upgrade accelerate`)." + ) + + if device_map != "sequential": + max_memory = get_balanced_memory( + model, + dtype=target_dtype, + low_zero=(device_map == "balanced_low_0"), + max_memory=max_memory, + **device_map_kwargs, + ) + else: + max_memory = get_max_memory(max_memory) + if hf_quantizer is not None: + max_memory = hf_quantizer.adjust_max_memory(max_memory) + device_map_kwargs["max_memory"] = max_memory + + device_map = infer_auto_device_map(model, dtype=target_dtype, **device_map_kwargs) + + if hf_quantizer is not None: + hf_quantizer.validate_environment(device_map=device_map) + + elif device_map is not None: + tied_params = find_tied_parameters(model) + # check if we don't have tied param in different devices + check_tied_parameters_on_same_device(tied_params, device_map) + + return device_map + + +def _find_missing_and_unexpected_keys( + cls, + model: "PreTrainedModel", + original_checkpoint_keys: List[str], + checkpoint_keys: List[str], + loading_base_model_from_task_state_dict: bool, + hf_quantizer: Optional[HfQuantizer], + device_map: Dict, +) -> Tuple[List[str], List[str]]: + """Find missing keys (keys that are part of the model parameters but were NOT found in the loaded state dict keys) and unexpected keys + (keys found in the loaded state dict keys, but that are NOT part of the model parameters) + """ + prefix = model.base_model_prefix + + # Compute expected keys, i.e. keys that the FULL model (not model_to_load) expects + expected_keys = list(model.state_dict().keys()) + if hf_quantizer is not None: + expected_keys = hf_quantizer.update_expected_keys(model, expected_keys, checkpoint_keys) + + # Adjust prefix of the keys to make them match loaded keys before removing them + missing_keys = sorted(set(expected_keys) - set(checkpoint_keys)) + unexpected_keys = set(checkpoint_keys) - set(expected_keys) + # If a module has the same name under the base and task specific model, we have to re-add it to unexpected keys + if loading_base_model_from_task_state_dict: + task_specific_keys = [k for k in original_checkpoint_keys if not k.startswith(f"{prefix}.")] + unexpected_keys.update(task_specific_keys) + + # Remove nonpersistent buffers from unexpected keys: they are not in the expected keys (model state dict), but + # may be in the loaded keys. Note that removing all buffers does the job, as they were part of the expected keys anyway + model_buffers = {n for n, _ in model.named_buffers()} + unexpected_keys = sorted(unexpected_keys - model_buffers) + + # Old checkpoints may have keys for rotary_emb.inv_freq for each layer, however we moved this buffer to the main model + # (so the buffer name has changed). Remove them in such a case + has_inv_freq_buffers = any(buffer.endswith("rotary_emb.inv_freq") for buffer in model_buffers) + if has_inv_freq_buffers: + unexpected_keys = [k for k in unexpected_keys if "rotary_emb.inv_freq" not in k] + + tied_params = find_tied_parameters(model) + for group in tied_params: + missing_in_group = [k for k in missing_keys if k in group] + if len(missing_in_group) > 0 and len(missing_in_group) < len(group): + missing_keys = [k for k in missing_keys if k not in missing_in_group] + + if hf_quantizer is not None: + missing_keys = hf_quantizer.update_missing_keys(model, missing_keys, prefix) + unexpected_keys = hf_quantizer.update_unexpected_keys(model, unexpected_keys, prefix) + + # Model-specific exceptions for missing and unexpected keys (e.g. if the modeling change over time, or any other reason...) + if cls._keys_to_ignore_on_load_missing is not None: + for pattern in cls._keys_to_ignore_on_load_missing: + missing_keys = [k for k in missing_keys if re.search(pattern, k) is None] + + if cls._keys_to_ignore_on_load_unexpected is not None: + for pattern in cls._keys_to_ignore_on_load_unexpected: + unexpected_keys = [k for k in unexpected_keys if re.search(pattern, k) is None] + + return missing_keys, unexpected_keys + + +def _find_mismatched_keys( + model: "PreTrainedModel", + state_dict: Optional[Dict], + checkpoint_files: Optional[List[str]], + ignore_mismatched_sizes: bool, + keys_to_rename_mapping: Dict[str, str], + is_quantized: bool, + weights_only: bool, +) -> Tuple[List[str], List[Tuple[int, int]]]: + """ + Find potential shape mismatch between the different state dicts and the model parameters, but only if `ignore_mismatched_sizes` + is True. Otherwise, return immediately and any shape mismatch that may exist will be raised later on. This avoids checking + every parameter in advance, as shape mismatch are extremely rare in practice. If we want to ignore them however, we do + need to check in advance as we need to know which parameters we need to move back from meta to cpu, and initialize + correctly. Indeed, as our model initialization takes place at the module level, and not the weight level, in the + case of a sharded checkpoint we cannot correctly initialize the weights according to `model._init_weights()` if we perform + this check on each state dict at loading time (after the first loaded checkpoint, there are no way to initialize only the + mismatched weights if any, without overwriting the previously loaded weights as well because all the module will be + initialized, not only the weights that are mismatched). + """ + + # An error will be raised later on anyway if there is a mismatch - this avoids running the rest of this function + # if there are no mismatch (which is almost always the case) + if not ignore_mismatched_sizes: + return [], [] + + if state_dict is not None: + checkpoint_files = [""] + + model_state_dict = model.state_dict() + mismatched_keys = [] + mismatched_shapes = [] + for shard_file in checkpoint_files: + # If shard_file is "", we use the existing state_dict instead of loading it + if shard_file != "": + state_dict = load_state_dict( + shard_file, is_quantized=is_quantized, map_location="meta", weights_only=weights_only + ) + + # Fix the key names + new_state_dict = {keys_to_rename_mapping[k]: v for k, v in state_dict.items() if k in keys_to_rename_mapping} + + for key in new_state_dict.keys(): + if key in model_state_dict and new_state_dict[key].shape != model_state_dict[key].shape: + # This skips size mismatches for 4-bit weights. Two 4-bit values share an 8-bit container, causing size differences. + # Without matching with module type or paramter type it seems like a practical way to detect valid 4bit weights. + if not ( + new_state_dict[key].shape[-1] == 1 + and new_state_dict[key].numel() * 2 == model_state_dict[key].numel() + ): + mismatched_keys.append(key) + mismatched_shapes.append((new_state_dict[key].shape, model_state_dict[key].shape)) + + return mismatched_keys, mismatched_shapes + + +class PipelineParallel(Enum): + inputs: 0 + outputs: 1 + + +class ModuleUtilsMixin: + """ + A few utilities for `torch.nn.Modules`, to be used as a mixin. + """ + + @staticmethod + def _hook_rss_memory_pre_forward(module, *args, **kwargs): + try: + import psutil + except ImportError: + raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.") + + process = psutil.Process(os.getpid()) + mem = process.memory_info() + module.mem_rss_pre_forward = mem.rss + return None + + @staticmethod + def _hook_rss_memory_post_forward(module, *args, **kwargs): + try: + import psutil + except ImportError: + raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.") + + process = psutil.Process(os.getpid()) + mem = process.memory_info() + module.mem_rss_post_forward = mem.rss + mem_rss_diff = module.mem_rss_post_forward - module.mem_rss_pre_forward + module.mem_rss_diff = mem_rss_diff + (module.mem_rss_diff if hasattr(module, "mem_rss_diff") else 0) + return None + + def add_memory_hooks(self): + """ + Add a memory hook before and after each sub-module forward pass to record increase in memory consumption. + + Increase in memory consumption is stored in a `mem_rss_diff` attribute for each module and can be reset to zero + with `model.reset_memory_hooks_state()`. + """ + for module in self.modules(): + module.register_forward_pre_hook(self._hook_rss_memory_pre_forward) + module.register_forward_hook(self._hook_rss_memory_post_forward) + self.reset_memory_hooks_state() + + def reset_memory_hooks_state(self): + """ + Reset the `mem_rss_diff` attribute of each module (see [`~modeling_utils.ModuleUtilsMixin.add_memory_hooks`]). + """ + for module in self.modules(): + module.mem_rss_diff = 0 + module.mem_rss_post_forward = 0 + module.mem_rss_pre_forward = 0 + + @property + def device(self) -> torch.device: + """ + `torch.device`: The device on which the module is (assuming that all the module parameters are on the same + device). + """ + return get_parameter_device(self) + + @property + def dtype(self) -> torch.dtype: + """ + `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype). + """ + return get_parameter_dtype(self) + + def invert_attention_mask(self, encoder_attention_mask: Tensor) -> Tensor: + """ + Invert an attention mask (e.g., switches 0. and 1.). + + Args: + encoder_attention_mask (`torch.Tensor`): An attention mask. + + Returns: + `torch.Tensor`: The inverted attention mask. + """ + if encoder_attention_mask.dim() == 3: + encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :] + if encoder_attention_mask.dim() == 2: + encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :] + # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition + # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow + # /transformer/transformer_layers.py#L270 + # encoder_extended_attention_mask = (encoder_extended_attention_mask == + # encoder_extended_attention_mask.transpose(-1, -2)) + encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility + encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * torch.finfo(self.dtype).min + + return encoder_extended_attention_mask + + @staticmethod + def create_extended_attention_mask_for_decoder(input_shape, attention_mask, device=None): + if device is not None: + warnings.warn( + "The `device` argument is deprecated and will be removed in v5 of Transformers.", FutureWarning + ) + else: + device = attention_mask.device + batch_size, seq_length = input_shape + seq_ids = torch.arange(seq_length, device=device) + causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None] + # in case past_key_values are used we need to add a prefix ones mask to the causal mask + # causal and attention masks must have same type with pytorch version < 1.3 + causal_mask = causal_mask.to(attention_mask.dtype) + + if causal_mask.shape[1] < attention_mask.shape[1]: + prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1] + causal_mask = torch.cat( + [ + torch.ones((batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype), + causal_mask, + ], + axis=-1, + ) + + extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :] + return extended_attention_mask + + def get_extended_attention_mask( + self, attention_mask: Tensor, input_shape: Tuple[int], device: torch.device = None, dtype: torch.float = None + ) -> Tensor: + """ + Makes broadcastable attention and causal masks so that future and masked tokens are ignored. + + Arguments: + attention_mask (`torch.Tensor`): + Mask with ones indicating tokens to attend to, zeros for tokens to ignore. + input_shape (`Tuple[int]`): + The shape of the input to the model. + + Returns: + `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`. + """ + if dtype is None: + dtype = self.dtype + + if not (attention_mask.dim() == 2 and self.config.is_decoder): + # show warning only if it won't be shown in `create_extended_attention_mask_for_decoder` + if device is not None: + warnings.warn( + "The `device` argument is deprecated and will be removed in v5 of Transformers.", FutureWarning + ) + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + if attention_mask.dim() == 3: + extended_attention_mask = attention_mask[:, None, :, :] + elif attention_mask.dim() == 2: + # Provided a padding mask of dimensions [batch_size, seq_length] + # - if the model is a decoder, apply a causal mask in addition to the padding mask + # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.config.is_decoder: + extended_attention_mask = ModuleUtilsMixin.create_extended_attention_mask_for_decoder( + input_shape, attention_mask, device + ) + else: + extended_attention_mask = attention_mask[:, None, None, :] + else: + raise ValueError( + f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})" + ) + + # Since attention_mask is 1.0 for positions we want to attend and 0.0 for + # masked positions, this operation will create a tensor which is 0.0 for + # positions we want to attend and the dtype's smallest value for masked positions. + # Since we are adding it to the raw scores before the softmax, this is + # effectively the same as removing these entirely. + extended_attention_mask = extended_attention_mask.to(dtype=dtype) # fp16 compatibility + extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min + return extended_attention_mask + + def get_head_mask( + self, head_mask: Optional[Tensor], num_hidden_layers: int, is_attention_chunked: bool = False + ) -> Tensor: + """ + Prepare the head mask if needed. + + Args: + head_mask (`torch.Tensor` with shape `[num_heads]` or `[num_hidden_layers x num_heads]`, *optional*): + The mask indicating if we should keep the heads or not (1.0 for keep, 0.0 for discard). + num_hidden_layers (`int`): + The number of hidden layers in the model. + is_attention_chunked (`bool`, *optional*, defaults to `False`): + Whether or not the attentions scores are computed by chunks or not. + + Returns: + `torch.Tensor` with shape `[num_hidden_layers x batch x num_heads x seq_length x seq_length]` or list with + `[None]` for each layer. + """ + if head_mask is not None: + head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers) + if is_attention_chunked is True: + head_mask = head_mask.unsqueeze(-1) + else: + head_mask = [None] * num_hidden_layers + + return head_mask + + def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers): + """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]""" + if head_mask.dim() == 1: + head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) + head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1) + elif head_mask.dim() == 2: + head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) # We can specify head_mask for each layer + assert head_mask.dim() == 5, f"head_mask.dim != 5, instead {head_mask.dim()}" + head_mask = head_mask.to(dtype=self.dtype) # switch to float if need + fp16 compatibility + return head_mask + + def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int: + """ + Get number of (optionally, trainable or non-embeddings) parameters in the module. + + Args: + only_trainable (`bool`, *optional*, defaults to `False`): + Whether or not to return only the number of trainable parameters + + exclude_embeddings (`bool`, *optional*, defaults to `False`): + Whether or not to return only the number of non-embeddings parameters + + Returns: + `int`: The number of parameters. + """ + + if exclude_embeddings: + embedding_param_names = [ + f"{name}.weight" for name, module_type in self.named_modules() if isinstance(module_type, nn.Embedding) + ] + total_parameters = [ + parameter for name, parameter in self.named_parameters() if name not in embedding_param_names + ] + else: + total_parameters = list(self.parameters()) + + total_numel = [] + is_loaded_in_4bit = getattr(self, "is_loaded_in_4bit", False) + + if is_loaded_in_4bit: + if is_bitsandbytes_available(): + import bitsandbytes as bnb + else: + raise ValueError( + "bitsandbytes is not installed but it seems that the model has been loaded in 4bit precision, something went wrong" + " make sure to install bitsandbytes with `pip install bitsandbytes`. You also need a GPU. " + ) + + for param in total_parameters: + if param.requires_grad or not only_trainable: + # For 4bit models, we need to multiply the number of parameters by 2 as half of the parameters are + # used for the 4bit quantization (uint8 tensors are stored) + if is_loaded_in_4bit and isinstance(param, bnb.nn.Params4bit): + if hasattr(param, "element_size"): + num_bytes = param.element_size() + elif hasattr(param, "quant_storage"): + num_bytes = param.quant_storage.itemsize + else: + num_bytes = 1 + total_numel.append(param.numel() * 2 * num_bytes) + else: + total_numel.append(param.numel()) + + return sum(total_numel) + + def estimate_tokens(self, input_dict: Dict[str, Union[torch.Tensor, Any]]) -> int: + """ + Helper function to estimate the total number of tokens from the model inputs. + + Args: + inputs (`dict`): The model inputs. + + Returns: + `int`: The total number of tokens. + """ + if not hasattr(self, "warnings_issued"): + self.warnings_issued = {} + if self.main_input_name in input_dict: + return input_dict[self.main_input_name].numel() + elif "estimate_tokens" not in self.warnings_issued: + logger.warning( + "Could not estimate the number of tokens of the input, floating-point operations will not be computed" + ) + self.warnings_issued["estimate_tokens"] = True + return 0 + + def floating_point_ops( + self, input_dict: Dict[str, Union[torch.Tensor, Any]], exclude_embeddings: bool = True + ) -> int: + """ + Get number of (optionally, non-embeddings) floating-point operations for the forward and backward passes of a + batch with this transformer model. Default approximation neglects the quadratic dependency on the number of + tokens (valid if `12 * d_model << sequence_length`) as laid out in [this + paper](https://arxiv.org/pdf/2001.08361.pdf) section 2.1. Should be overridden for transformers with parameter + re-use e.g. Albert or Universal Transformers, or if doing long-range modeling with very high sequence lengths. + + Args: + batch_size (`int`): + The batch size for the forward pass. + + sequence_length (`int`): + The number of tokens in each line of the batch. + + exclude_embeddings (`bool`, *optional*, defaults to `True`): + Whether or not to count embedding and softmax operations. + + Returns: + `int`: The number of floating-point operations. + """ + + return 6 * self.estimate_tokens(input_dict) * self.num_parameters(exclude_embeddings=exclude_embeddings) + + +# TODO (joao): remove `GenerationMixin` inheritance in v4.50 +class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMixin, PeftAdapterMixin): + r""" + Base class for all models. + + [`PreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading, + downloading and saving models as well as a few methods common to all models to: + + - resize the input embeddings, + - prune heads in the self-attention heads. + + Class attributes (overridden by derived classes): + + - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class + for this model architecture. + - **load_tf_weights** (`Callable`) -- A python *method* for loading a TensorFlow checkpoint in a PyTorch model, + taking as arguments: + + - **model** ([`PreTrainedModel`]) -- An instance of the model on which to load the TensorFlow checkpoint. + - **config** ([`PreTrainedConfig`]) -- An instance of the configuration associated to the model. + - **path** (`str`) -- A path to the TensorFlow checkpoint. + + - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived + classes of the same architecture adding modules on top of the base model. + - **is_parallelizable** (`bool`) -- A flag indicating whether this model supports model parallelization. + - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP + models, `pixel_values` for vision models and `input_values` for speech models). + """ + + config_class = None + base_model_prefix = "" + main_input_name = "input_ids" + model_tags = None + + _auto_class = None + _no_split_modules = None + _skip_keys_device_placement = None + _keep_in_fp32_modules = None + + # a list of `re` patterns of `state_dict` keys that should be removed from the list of missing + # keys we find (keys inside the model but not in the checkpoint) and avoid unnecessary warnings. + _keys_to_ignore_on_load_missing = None + # a list of `re` patterns of `state_dict` keys that should be removed from the list of + # unexpected keys we find (keys inside the checkpoint but not the model) and avoid unnecessary + # warnings. + _keys_to_ignore_on_load_unexpected = None + # a list of `state_dict` keys to ignore when saving the model (useful for keys that aren't + # trained, but which are either deterministic or tied variables) + _keys_to_ignore_on_save = None + # a list of `state_dict` keys that are potentially tied to another key in the state_dict. + _tied_weights_keys = None + + is_parallelizable = False + supports_gradient_checkpointing = False + _is_stateful = False + + # Flash Attention 2 support + _supports_flash_attn_2 = False + + # SDPA support + _supports_sdpa = False + + # Flex Attention support + _supports_flex_attn = False + + # Has support for a `Cache` instance as `past_key_values`? Does it support a `StaticCache`? + _supports_cache_class = False + _supports_static_cache = False + + # Has support for a `QuantoQuantizedCache` instance as `past_key_values` + _supports_quantized_cache = False + + # A tensor parallel plan to be applied to the model when TP is enabled. For + # top-level models, this attribute is currently defined in respective model + # code. For base models, this attribute comes from + # `config.base_model_tp_plan` during `__init__`. + # It should identify the layers exactly: if you want to TP model.language_model.layers.fc1 + # by passing `tp_plan` to the init, it should be {"model.language_model.layers.fc1":"colwise"} + # for example. + _tp_plan = None + + # A pipeline parallel plan specifying the layers which may not be present + # on all ranks when PP is enabled. For top-level models, this attribute is + # currently defined in respective model code. For base models, this + # attribute comes from `config.base_model_pp_plan` during `post_init`. + # + # The variable names for the inputs and outputs of the specified layers can + # be indexed using the `PipelineParallel` enum as follows: + # - `_pp_plan["layers"][PipelineParallel.inputs]` + # - `_pp_plan["layers"][PipelineParallel.outputs]` + _pp_plan = None + + # This flag signal that the model can be used as an efficient backend in TGI and vLLM + # In practice, it means that they support attention interface functions, fully pass the kwargs + # through all modules up to the Attention layer, can slice logits with Tensor, and have a default TP plan + _supports_attention_backend = False + + @property + def dummy_inputs(self) -> Dict[str, torch.Tensor]: + """ + `Dict[str, torch.Tensor]`: Dummy inputs to do a forward pass in the network. + """ + return {"input_ids": torch.tensor(DUMMY_INPUTS)} + + @property + def framework(self) -> str: + """ + :str: Identifies that this is a PyTorch model. + """ + return "pt" + + def __init__(self, config: PretrainedConfig, *inputs, **kwargs): + super().__init__() + if not isinstance(config, PretrainedConfig): + raise ValueError( + f"Parameter config in `{self.__class__.__name__}(config)` should be an instance of class " + "`PretrainedConfig`. To create a model from a pretrained model use " + f"`model = {self.__class__.__name__}.from_pretrained(PRETRAINED_MODEL_NAME)`" + ) + if not getattr(config, "_attn_implementation_autoset", False): + # config usually has a `torch_dtype` but we need the next line for the `no_super_init` tests + dtype = config.torch_dtype if hasattr(config, "torch_dtype") else torch.get_default_dtype() + config = self._autoset_attn_implementation(config, torch_dtype=dtype, check_device_map=False) + self.config = config + + # for initialization of the loss + loss_type = self.__class__.__name__ + if loss_type not in LOSS_MAPPING: + loss_groups = f"({'|'.join(LOSS_MAPPING)})" + loss_type = re.findall(loss_groups, self.__class__.__name__) + if len(loss_type) > 0: + loss_type = loss_type[0] + else: + loss_type = None + self.loss_type = loss_type + + self.name_or_path = config.name_or_path + self.warnings_issued = {} + self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None + # Overwrite the class attribute to make it an instance attribute, so models like + # `InstructBlipForConditionalGeneration` can dynamically update it without modifying the class attribute + # when a different component (e.g. language_model) is used. + self._keep_in_fp32_modules = copy.copy(self.__class__._keep_in_fp32_modules) + + self._no_split_modules = self._no_split_modules or [] + + def post_init(self): + """ + A method executed at the end of each Transformer model initialization, to execute code that needs the model's + modules properly initialized (such as weight initialization). + """ + self.init_weights() + self._backward_compatibility_gradient_checkpointing() + + # Make sure the modules correctly exist if the flag is active + if self._keep_in_fp32_modules is not None: + all_parameters = {name for name, _ in self.named_parameters() if len(name) > 0} + unique_module_names = set() + # Get all unique module names in the module graph, without the prefixes + for param in all_parameters: + unique_module_names.update( + [name for name in param.split(".") if not name.isnumeric() and name not in ["weight", "bias"]] + ) + # Check that every module in the keep_in_fp32 list is part of the module graph + for module in self._keep_in_fp32_modules: + if module not in unique_module_names: + raise ValueError( + f"{module} was specified in the `_keep_in_fp32_modules` list, but is not part of the modules in" + f" {self.__class__.__name__}" + ) + + # If current model is a base model, attach `base_model_tp_plan` and `base_model_pp_plan` from config + self._pp_plan = self.config.base_model_pp_plan.copy() if self.config.base_model_pp_plan is not None else None + self._tp_plan = self.config.base_model_tp_plan.copy() if self.config.base_model_tp_plan is not None else {} + for name, module in self.named_children(): + if plan := getattr(module, "_tp_plan", None): + self._tp_plan.update({f"{name}.{k}": v for k, v in plan.copy().items()}) + + if self._tp_plan is not None and is_torch_greater_or_equal("2.3"): + for _, v in self._tp_plan.items(): + if v not in SUPPORTED_TP_STYLES: + raise ValueError( + f"Unsupported tensor parallel style {v}. Supported styles are {SUPPORTED_TP_STYLES}" + ) + + def dequantize(self): + """ + Potentially dequantize the model in case it has been quantized by a quantization method that support + dequantization. + """ + hf_quantizer = getattr(self, "hf_quantizer", None) + + if hf_quantizer is None: + raise ValueError("You need to first quantize your model in order to dequantize it") + + return hf_quantizer.dequantize(self) + + def _backward_compatibility_gradient_checkpointing(self): + if self.supports_gradient_checkpointing and getattr(self.config, "gradient_checkpointing", False): + self.gradient_checkpointing_enable() + # Remove the attribute now that is has been consumed, so it's no saved in the config. + delattr(self.config, "gradient_checkpointing") + + def add_model_tags(self, tags: Union[List[str], str]) -> None: + r""" + Add custom tags into the model that gets pushed to the Hugging Face Hub. Will + not overwrite existing tags in the model. + + Args: + tags (`Union[List[str], str]`): + The desired tags to inject in the model + + Examples: + + ```python + from transformers import AutoModel + + model = AutoModel.from_pretrained("google-bert/bert-base-cased") + + model.add_model_tags(["custom", "custom-bert"]) + + # Push the model to your namespace with the name "my-custom-bert". + model.push_to_hub("my-custom-bert") + ``` + """ + if isinstance(tags, str): + tags = [tags] + + if self.model_tags is None: + self.model_tags = [] + + for tag in tags: + if tag not in self.model_tags: + self.model_tags.append(tag) + + @classmethod + @restore_default_torch_dtype + def _from_config(cls, config, **kwargs): + """ + All context managers that the model should be initialized under go here. + + Args: + torch_dtype (`torch.dtype`, *optional*): + Override the default `torch.dtype` and load the model under this dtype. + """ + # when we init a model from within another model (e.g. VLMs) and dispatch on FA2 + # a warning is raised that dtype should be fp16. Since we never pass dtype from within + # modeling code, we can try to infer it here same way as done in `from_pretrained` + torch_dtype = kwargs.pop("torch_dtype", config.torch_dtype) + if isinstance(torch_dtype, str): + torch_dtype = getattr(torch, torch_dtype) + + use_flash_attention_2 = kwargs.pop("use_flash_attention_2", False) + + # override default dtype if needed + dtype_orig = None + if torch_dtype is not None: + dtype_orig = cls._set_default_torch_dtype(torch_dtype) + + config = copy.deepcopy(config) # We do not want to modify the config inplace in _from_config. + + if config._attn_implementation_internal is not None: + # In this case, the config has been created with the attn_implementation set by the user, which we + # should respect. + attn_implementation = config._attn_implementation_internal + else: + attn_implementation = None + + config._attn_implementation = kwargs.pop("attn_implementation", attn_implementation) + if not getattr(config, "_attn_implementation_autoset", False): + config = cls._autoset_attn_implementation( + config, + use_flash_attention_2=use_flash_attention_2, + check_device_map=False, + torch_dtype=torch_dtype, + ) + + if is_deepspeed_zero3_enabled() and not _is_quantized and not _is_ds_init_called: + logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model") + # this immediately partitions the model across all gpus, to avoid the overhead in time + # and memory copying it on CPU or each GPU first + init_contexts = [deepspeed.zero.Init(config_dict_or_path=deepspeed_config()), set_zero3_state()] + with ContextManagers(init_contexts): + model = cls(config, **kwargs) + + else: + model = cls(config, **kwargs) + + # restore default dtype if it was modified + if dtype_orig is not None: + torch.set_default_dtype(dtype_orig) + + return model + + @classmethod + def _autoset_attn_implementation( + cls, + config, + use_flash_attention_2: bool = False, + torch_dtype: Optional[torch.dtype] = None, + device_map: Optional[Union[str, Dict[str, int]]] = None, + check_device_map: bool = True, + ): + """ + Automatically checks and dispatches to a default attention implementation. In order of priority: + 1. An implementation specified in `config._attn_implementation` (due for example to the argument attn_implementation="sdpa" in from_pretrained). + 2. DEPRECATED: if use_flash_attention_2 is set to `True` and `flash_attn` is available, flash attention. (`LlamaFlashAttention` for example) + 3. SDPA implementation, if available and supported by the model type. (`LlamaSdpaAttention` for example) + 4. The default model's implementation otherwise (`LlamaAttention` for example) . + """ + # Here we use config._attn_implementation_internal to check whether the attention implementation was explicitly set by the user. + # The property `PretrainedConfig._attn_implementation` is never `None`, for backward compatibility (always fall back on "eager"). + # The `hasattr` here is used as some Transformers tests for some reason do not call PretrainedConfig __init__ (e.g. test_no_super_init_config_and_model) + requested_attn_implementation = None + if hasattr(config, "_attn_implementation_internal") and config._attn_implementation_internal is not None: + if config._attn_implementation != "flash_attention_2" and use_flash_attention_2: + raise ValueError( + f'Both attn_implementation="{config._attn_implementation}" and `use_flash_attention_2=True` were used when loading the model, which are not compatible.' + ' We recommend to just use `attn_implementation="flash_attention_2"` when loading the model.' + ) + + if ( + not isinstance(config._attn_implementation, dict) + and config._attn_implementation not in ["eager"] + ALL_ATTENTION_FUNCTIONS.valid_keys() + ): + message = f'Specified `attn_implementation="{config._attn_implementation}"` is not supported. The only possible arguments are `attn_implementation="eager"` (manual attention implementation)' + if cls._supports_flash_attn_2: + message += ', `"attn_implementation=flash_attention_2"` (implementation using flash attention 2)' + if cls._supports_sdpa: + message += ', `"attn_implementation=sdpa"` (implementation using torch.nn.functional.scaled_dot_product_attention)' + if cls._supports_flex_attn: + message += ( + ', `"attn_implementation=flex_attention"` (implementation using torch\'s flex_attention)' + ) + raise ValueError(message + ".") + + # If a config is passed with a preset attn_implementation, we skip the automatic dispatch and use the user-provided config, with hard checks that the requested attention implementation is available. + requested_attn_implementation = config._attn_implementation_internal + + # Composite models consisting of several PretrainedModels have to specify attention impl as a dict + # where keys are sub-config names. But most people will specify one `str` which means that should dispatch it + # for all sub-models. + # Below we check if a config is composite and manually prepare a dict of attn impl if not already passed as a dict. + # Later each sub-module will dispatch with its own attn impl, by calling `XXXModel._from_config(config.text_config)` + # If any of sub-modules doesn't support requested attn, an error will be raised. See https://github.com/huggingface/transformers/pull/32238 + for key in config.sub_configs.keys(): + sub_config = getattr(config, key) + curr_attn_implementation = ( + requested_attn_implementation + if not isinstance(requested_attn_implementation, dict) + else requested_attn_implementation.get(key, None) + ) + # For models with backbone sub-config might be not initialized + if sub_config is not None: + sub_config._attn_implementation_internal = curr_attn_implementation + + if use_flash_attention_2: + logger.warning_once( + 'The model was loaded with use_flash_attention_2=True, which is deprecated and may be removed in a future release. Please use `attn_implementation="flash_attention_2"` instead.' + ) + config._attn_implementation = "flash_attention_2" + + if config._attn_implementation == "flash_attention_2": + cls._check_and_enable_flash_attn_2( + config, + torch_dtype=torch_dtype, + device_map=device_map, + hard_check_only=False, + check_device_map=check_device_map, + ) + elif requested_attn_implementation == "flex_attention": + config = cls._check_and_enable_flex_attn(config, hard_check_only=True) + elif requested_attn_implementation in [None, "sdpa"] and not is_torch_xla_available(): + # use_flash_attention_2 takes priority over SDPA, hence SDPA treated in this elif. + config = cls._check_and_enable_sdpa( + config, + hard_check_only=False if requested_attn_implementation is None else True, + ) + + if ( + torch.version.hip is not None + and config._attn_implementation == "sdpa" + and torch.cuda.device_count() > 1 + and version.parse(torch.__version__) < version.parse("2.4.1") + ): + logger.warning_once( + "Using the `SDPA` attention implementation on multi-gpu setup with ROCM may lead to performance issues due to the FA backend. Disabling it to use alternative backends." + ) + torch.backends.cuda.enable_flash_sdp(False) + elif requested_attn_implementation in ALL_ATTENTION_FUNCTIONS.valid_keys(): + config._attn_implementation = requested_attn_implementation + elif isinstance(requested_attn_implementation, dict): + config._attn_implementation = None + else: + config._attn_implementation = "eager" + + config._attn_implementation_autoset = True + return config + + @classmethod + def _set_default_torch_dtype(cls, dtype: torch.dtype) -> torch.dtype: + """ + Change the default dtype and return the previous one. This is needed when wanting to instantiate the model + under specific dtype. + + Args: + dtype (`torch.dtype`): + a floating dtype to set to. + + Returns: + `torch.dtype`: the original `dtype` that can be used to restore `torch.set_default_dtype(dtype)` if it was + modified. If it wasn't, returns `None`. + + Note `set_default_dtype` currently only works with floating-point types and asserts if for example, + `torch.int64` is passed. So if a non-float `dtype` is passed this functions will throw an exception. + """ + if not dtype.is_floating_point: + raise ValueError( + f"Can't instantiate {cls.__name__} model under dtype={dtype} since it is not a floating point dtype" + ) + + logger.info(f"Instantiating {cls.__name__} model under default dtype {dtype}.") + dtype_orig = torch.get_default_dtype() + torch.set_default_dtype(dtype) + return dtype_orig + + @property + def base_model(self) -> nn.Module: + """ + `torch.nn.Module`: The main body of the model. + """ + return getattr(self, self.base_model_prefix, self) + + @classmethod + def can_generate(cls) -> bool: + """ + Returns whether this model can generate sequences with `.generate()` from the `GenerationMixin`. + + Under the hood, on classes where this function returns True, some generation-specific changes are triggered: + for instance, the model instance will have a populated `generation_config` attribute. + + Returns: + `bool`: Whether this model can generate sequences with `.generate()`. + """ + # Directly inherits `GenerationMixin` -> can generate + if "GenerationMixin" in str(cls.__bases__): + return True + # The class inherits from a class that can generate (recursive check) -> can generate + for base in cls.__bases__: + if not hasattr(base, "can_generate"): + continue + if "PreTrainedModel" not in str(base) and base.can_generate(): + return True + # BC: Detects whether `prepare_inputs_for_generation` has been overwritten in the model. Prior to v4.45, this + # was how we detected whether a model could generate. + if "GenerationMixin" not in str(cls.prepare_inputs_for_generation): + logger.warning_once( + f"{cls.__name__} has generative capabilities, as `prepare_inputs_for_generation` is explicitly " + "overwritten. However, it doesn't directly inherit from `GenerationMixin`. From 👉v4.50👈 onwards, " + "`PreTrainedModel` will NOT inherit from `GenerationMixin`, and this model will lose the ability " + "to call `generate` and other related functions." + "\n - If you're using `trust_remote_code=True`, you can get rid of this warning by loading the " + "model with an auto class. See https://huggingface.co/docs/transformers/en/model_doc/auto#auto-classes" + "\n - If you are the owner of the model architecture code, please modify your model class such that " + "it inherits from `GenerationMixin` (after `PreTrainedModel`, otherwise you'll get an exception)." + "\n - If you are not the owner of the model architecture class, please contact the model code owner " + "to update it." + ) + return True + # Otherwise, can't generate + return False + + @classmethod + def _check_and_enable_flash_attn_2( + cls, + config, + torch_dtype: Optional[torch.dtype] = None, + device_map: Optional[Union[str, Dict[str, int]]] = None, + check_device_map: bool = True, + hard_check_only: bool = False, + ) -> PretrainedConfig: + """ + Checks the availability of Flash Attention 2 and compatibility with the current model. + + If all checks pass and `hard_check_only` is False, the method will set the config attribute `attn_implementation` to "flash_attention_2" so that the model can initialize the correct attention module. + """ + if not cls._supports_flash_attn_2: + raise ValueError( + f"{cls.__name__} does not support Flash Attention 2.0 yet. Please request to add support where" + f" the model is hosted, on its model hub page: https://huggingface.co/{config._name_or_path}/discussions/new" + " or in the Transformers GitHub repo: https://github.com/huggingface/transformers/issues/new" + ) + + if not is_flash_attn_2_available(): + preface = "FlashAttention2 has been toggled on, but it cannot be used due to the following error:" + install_message = "Please refer to the documentation of https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2 to install Flash Attention 2." + + if importlib.util.find_spec("flash_attn") is None: + # package `flash-attn` can not be installed on Ascend NPU, ignore related validation logic and early exit. + if is_torch_npu_available(): + if not hard_check_only: + config._attn_implementation = "flash_attention_2" + + logger.info("Detect using FlashAttention2 on Ascend NPU.") + return config + else: + raise ImportError(f"{preface} the package flash_attn seems to be not installed. {install_message}") + + flash_attention_version = version.parse(importlib.metadata.version("flash_attn")) + if torch.version.cuda: + if flash_attention_version < version.parse("2.1.0"): + raise ImportError( + f"{preface} you need flash_attn package version to be greater or equal than 2.1.0. Detected version {flash_attention_version}. {install_message}" + ) + elif not torch.cuda.is_available(): + raise ValueError( + f"{preface} Flash Attention 2 is not available on CPU. Please make sure torch can access a CUDA device." + ) + else: + raise ImportError(f"{preface} Flash Attention 2 is not available. {install_message}") + elif torch.version.hip: + if flash_attention_version < version.parse("2.0.4"): + raise ImportError( + f"{preface} you need flash_attn package version to be greater or equal than 2.0.4. Make sure to have that version installed - detected version {flash_attention_version}. {install_message}" + ) + else: + raise ImportError(f"{preface} Flash Attention 2 is not available. {install_message}") + + _is_bettertransformer = getattr(cls, "use_bettertransformer", False) + + if _is_bettertransformer: + raise ValueError( + "Flash Attention 2 and BetterTransformer API are not compatible. Please make sure to disable BetterTransformers by doing model.reverse_bettertransformer()" + ) + + if torch_dtype is None: + logger.warning_once( + "You are attempting to use Flash Attention 2.0 without specifying a torch dtype. This might lead to unexpected behaviour" + ) + elif torch_dtype is not None and torch_dtype not in [torch.float16, torch.bfloat16]: + logger.warning_once( + "Flash Attention 2.0 only supports torch.float16 and torch.bfloat16 dtypes, but" + f" the current dype in {cls.__name__} is {torch_dtype}. You should run training or inference using Automatic Mixed-Precision via the `with torch.autocast(device_type='torch_device'):` decorator," + ' or load the model with the `torch_dtype` argument. Example: `model = AutoModel.from_pretrained("openai/whisper-tiny", attn_implementation="flash_attention_2", torch_dtype=torch.float16)`' + ) + + # The check `torch.empty(0).device.type != "cuda"` is needed as the model may be initialized after `torch.set_default_device` has been called, + # or the model may be initialized under the context manager `with torch.device("cuda"):`. + if check_device_map and device_map is None and torch.empty(0).device.type not in ["cuda", "mlu"]: + if torch.cuda.is_available(): + logger.warning_once( + "You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU" + " after initializing it on CPU with `model.to('cuda')`." + ) + elif is_torch_mlu_available(): + logger.warning_once( + "You are attempting to use Flash Attention 2.0 with a model not initialized on MLU. Make sure to move the model to MLU" + " after initializing it on CPU with `model.to('mlu')`." + ) + else: + raise ValueError( + "You are attempting to use Flash Attention 2.0 with a model not initialized on GPU and with no GPU available. " + "This is not supported yet. Please make sure to have access to a GPU and either initialise the model on a GPU by passing a device_map " + "or initialising the model on CPU and then moving it to GPU." + ) + elif ( + check_device_map + and device_map is not None + and isinstance(device_map, dict) + and ("cpu" in device_map.values() or "disk" in device_map.values()) + ): + raise ValueError( + "You are attempting to use Flash Attention 2.0 with a model dispatched on CPU or disk. This is not supported. Please make sure to " + "initialise the model on a GPU by passing a device_map that contains only GPU devices as keys." + ) + if not hard_check_only: + config._attn_implementation = "flash_attention_2" + return config + + @classmethod + def _check_and_enable_sdpa(cls, config, hard_check_only: bool = False) -> PretrainedConfig: + """ + Checks the availability of SDPA for a given model. + + If all checks pass and `hard_check_only` is False, the method will set the config attribute `_attn_implementation` to "sdpa" so that the model can initialize the correct attention module. + """ + if hard_check_only: + if not cls._supports_sdpa: + raise ValueError( + f"{cls.__name__} does not support an attention implementation through torch.nn.functional.scaled_dot_product_attention yet." + " Please request the support for this architecture: https://github.com/huggingface/transformers/issues/28005. If you believe" + ' this error is a bug, please open an issue in Transformers GitHub repository and load your model with the argument `attn_implementation="eager"` meanwhile. Example: `model = AutoModel.from_pretrained("openai/whisper-tiny", attn_implementation="eager")`' + ) + if not is_torch_sdpa_available(): + raise ImportError( + "PyTorch SDPA requirements in Transformers are not met. Please install torch>=2.1.1." + ) + + if not is_torch_sdpa_available() or not cls._supports_sdpa: + return config + + _is_bettertransformer = getattr(cls, "use_bettertransformer", False) + if _is_bettertransformer: + return config + + if not hard_check_only: + config._attn_implementation = "sdpa" + return config + + @classmethod + def _check_and_enable_flex_attn(cls, config, hard_check_only: bool = False) -> PretrainedConfig: + """ + Checks the availability of Flex Attention for a given model. + + If all checks pass and `hard_check_only` is False, the method will set the config attribute `_attn_implementation` to "flex_attention" so that the model can initialize the correct attention module. + """ + if hard_check_only: + if not cls._supports_flex_attn: + raise ValueError( + f"{cls.__name__} does not support an attention implementation through torch's flex_attention." + " Please request the support for this architecture: https://github.com/huggingface/transformers/issues/34809." + " If you believe this error is a bug, please open an issue in Transformers GitHub repository" + ' and load your model with the argument `attn_implementation="eager"` meanwhile.' + ' Example: `model = AutoModel.from_pretrained("openai/whisper-tiny", attn_implementation="eager")`' + ) + if not is_torch_flex_attn_available(): + raise ImportError( + "PyTorch Flex Attention requirements in Transformers are not met. Please install torch>=2.5.0." + ) + + if not is_torch_flex_attn_available() or not cls._supports_flex_attn: + return config + + if not hard_check_only: + config._attn_implementation = "flex_attention" + + return config + + def enable_input_require_grads(self): + """ + Enables the gradients for the input embeddings. This is useful for fine-tuning adapter weights while keeping + the model weights fixed. + """ + + def make_inputs_require_grads(module, input, output): + output.requires_grad_(True) + + self._require_grads_hook = self.get_input_embeddings().register_forward_hook(make_inputs_require_grads) + + def disable_input_require_grads(self): + """ + Removes the `_require_grads_hook`. + """ + self._require_grads_hook.remove() + + def get_input_embeddings(self) -> nn.Module: + """ + Returns the model's input embeddings. + + Returns: + `nn.Module`: A torch module mapping vocabulary to hidden states. + """ + base_model = getattr(self, self.base_model_prefix, self) + if base_model is not self: + return base_model.get_input_embeddings() + else: + raise NotImplementedError + + def set_input_embeddings(self, value: nn.Module): + """ + Set model's input embeddings. + + Args: + value (`nn.Module`): A module mapping vocabulary to hidden states. + """ + base_model = getattr(self, self.base_model_prefix, self) + if base_model is not self: + base_model.set_input_embeddings(value) + else: + raise NotImplementedError + + def get_output_embeddings(self) -> nn.Module: + """ + Returns the model's output embeddings. + + Returns: + `nn.Module`: A torch module mapping hidden states to vocabulary. + """ + return None # Overwrite for models with output embeddings + + def _init_weights(self, module): + """ + Initialize the weights. This method should be overridden by derived class and is + the only initialization method that will be called when loading a checkpoint + using `from_pretrained`. Any attempt to initialize outside of this function + will be useless as the torch.nn.init function are all replaced with skip. + """ + pass + + def _initialize_weights(self, module): + """ + Initialize the weights if they are not already initialized. + """ + if getattr(module, "_is_hf_initialized", False): + return + self._init_weights(module) + module._is_hf_initialized = True + + def tie_weights(self): + """ + Tie the weights between the input embeddings and the output embeddings. + + If the `torchscript` flag is set in the configuration, can't handle parameter sharing so we are cloning the + weights instead. + """ + if getattr(self.config.get_text_config(decoder=True), "tie_word_embeddings", True): + output_embeddings = self.get_output_embeddings() + if output_embeddings is not None: + self._tie_or_clone_weights(output_embeddings, self.get_input_embeddings()) + + if getattr(self.config, "is_encoder_decoder", False) and getattr(self.config, "tie_encoder_decoder", False): + if hasattr(self, self.base_model_prefix): + self = getattr(self, self.base_model_prefix) + tied_weights = self._tie_encoder_decoder_weights( + self.encoder, self.decoder, self.base_model_prefix, "encoder" + ) + # Setting a dynamic variable instead of `_tied_weights_keys` because it's a class + # attributed not an instance member, therefore modifying it will modify the entire class + # Leading to issues on subsequent calls by different tests or subsequent calls. + self._dynamic_tied_weights_keys = tied_weights + + for module in self.modules(): + if hasattr(module, "_tie_weights"): + module._tie_weights() + + @staticmethod + def _tie_encoder_decoder_weights( + encoder: nn.Module, decoder: nn.Module, base_model_prefix: str, base_encoder_name: str + ): + uninitialized_encoder_weights: List[str] = [] + tied_weights: List[str] = [] + if decoder.__class__ != encoder.__class__: + logger.info( + f"{decoder.__class__} and {encoder.__class__} are not equal. In this case make sure that all encoder" + " weights are correctly initialized." + ) + + def tie_encoder_to_decoder_recursively( + decoder_pointer: nn.Module, + encoder_pointer: nn.Module, + module_name: str, + base_encoder_name: str, + uninitialized_encoder_weights: List[str], + depth=0, + total_decoder_name="", + total_encoder_name="", + ): + assert isinstance(decoder_pointer, nn.Module) and isinstance(encoder_pointer, nn.Module), ( + f"{decoder_pointer} and {encoder_pointer} have to be of type nn.Module" + ) + if hasattr(decoder_pointer, "weight"): + assert hasattr(encoder_pointer, "weight") + encoder_pointer.weight = decoder_pointer.weight + tied_weights.append(f"{base_encoder_name}{total_encoder_name}.weight") + if hasattr(decoder_pointer, "bias"): + assert hasattr(encoder_pointer, "bias") + tied_weights.append(f"{base_encoder_name}{total_encoder_name}.bias") + encoder_pointer.bias = decoder_pointer.bias + return + + encoder_modules = encoder_pointer._modules + decoder_modules = decoder_pointer._modules + if len(decoder_modules) > 0: + assert len(encoder_modules) > 0, ( + f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}" + ) + + all_encoder_weights = {module_name + "/" + sub_name for sub_name in encoder_modules.keys()} + encoder_layer_pos = 0 + for name, module in decoder_modules.items(): + if name.isdigit(): + encoder_name = str(int(name) + encoder_layer_pos) + decoder_name = name + if not isinstance(decoder_modules[decoder_name], type(encoder_modules[encoder_name])) and len( + encoder_modules + ) != len(decoder_modules): + # this can happen if the name corresponds to the position in a list module list of layers + # in this case the decoder has added a cross-attention that the encoder does not have + # thus skip this step and subtract one layer pos from encoder + encoder_layer_pos -= 1 + continue + elif name not in encoder_modules: + continue + elif depth > 500: + raise ValueError( + "Max depth of recursive function `tie_encoder_to_decoder` reached. It seems that there is" + " a circular dependency between two or more `nn.Modules` of your model." + ) + else: + decoder_name = encoder_name = name + tie_encoder_to_decoder_recursively( + decoder_modules[decoder_name], + encoder_modules[encoder_name], + module_name + "/" + name, + base_encoder_name, + uninitialized_encoder_weights, + depth=depth + 1, + total_encoder_name=f"{total_encoder_name}.{encoder_name}", + total_decoder_name=f"{total_decoder_name}.{decoder_name}", + ) + all_encoder_weights.remove(module_name + "/" + encoder_name) + + uninitialized_encoder_weights += list(all_encoder_weights) + + # tie weights recursively + tie_encoder_to_decoder_recursively( + decoder, encoder, base_model_prefix, base_encoder_name, uninitialized_encoder_weights + ) + + if len(uninitialized_encoder_weights) > 0: + logger.warning( + f"The following encoder weights were not tied to the decoder {uninitialized_encoder_weights}" + ) + return tied_weights + + def _tie_or_clone_weights(self, output_embeddings, input_embeddings): + """Tie or clone module weights depending of whether we are using TorchScript or not""" + if self.config.torchscript: + output_embeddings.weight = nn.Parameter(input_embeddings.weight.clone()) + else: + output_embeddings.weight = input_embeddings.weight + + if getattr(output_embeddings, "bias", None) is not None: + output_embeddings.bias.data = nn.functional.pad( + output_embeddings.bias.data, + ( + 0, + output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0], + ), + "constant", + 0, + ) + if hasattr(output_embeddings, "out_features") and hasattr(input_embeddings, "num_embeddings"): + output_embeddings.out_features = input_embeddings.num_embeddings + + def _get_no_split_modules(self, device_map: str): + """ + Get the modules of the model that should not be spit when using device_map. We iterate through the modules to + get the underlying `_no_split_modules`. + + Args: + device_map (`str`): + The device map value. Options are ["auto", "balanced", "balanced_low_0", "sequential"] + + Returns: + `List[str]`: List of modules that should not be split + """ + _no_split_modules = set() + modules_to_check = [self] + while len(modules_to_check) > 0: + module = modules_to_check.pop(-1) + # if the module does not appear in _no_split_modules, we also check the children + if module.__class__.__name__ not in _no_split_modules: + if isinstance(module, PreTrainedModel): + if module._no_split_modules is None: + raise ValueError( + f"{module.__class__.__name__} does not support `device_map='{device_map}'`. To implement support, the model " + "class needs to implement the `_no_split_modules` attribute." + ) + else: + _no_split_modules = _no_split_modules | set(module._no_split_modules) + modules_to_check += list(module.children()) + return list(_no_split_modules) + + def resize_token_embeddings( + self, + new_num_tokens: Optional[int] = None, + pad_to_multiple_of: Optional[int] = None, + mean_resizing: bool = True, + ) -> nn.Embedding: + """ + Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`. + + Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method. + + Arguments: + new_num_tokens (`int`, *optional*): + The new number of tokens in the embedding matrix. Increasing the size will add newly initialized + vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just + returns a pointer to the input tokens `torch.nn.Embedding` module of the model without doing anything. + pad_to_multiple_of (`int`, *optional*): + If set will pad the embedding matrix to a multiple of the provided value.If `new_num_tokens` is set to + `None` will just pad the embedding to a multiple of `pad_to_multiple_of`. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. For more + details about this, or help on choosing the correct value for resizing, refer to this guide: + https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc + mean_resizing (`bool`): + Whether to initialize the added embeddings from a multivariate normal distribution that has old embeddings' mean and + covariance or to initialize them with a normal distribution that has a mean of zero and std equals `config.initializer_range`. + + Setting `mean_resizing` to `True` is useful when increasing the size of the embeddings of causal language models, + where the generated tokens' probabilities won't be affected by the added embeddings because initializing the new embeddings with the + old embeddings' mean will reduce the kl-divergence between the next token probability before and after adding the new embeddings. + Refer to this article for more information: https://nlp.stanford.edu/~johnhew/vocab-expansion.html + + Return: + `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model. + """ + model_embeds = self._resize_token_embeddings(new_num_tokens, pad_to_multiple_of, mean_resizing) + if new_num_tokens is None and pad_to_multiple_of is None: + return model_embeds + + # Since we are basically reusing the same old embeddings with new weight values, gathering is required + is_quantized = hasattr(self, "hf_quantizer") and self.hf_quantizer is not None + if is_deepspeed_zero3_enabled() and not is_quantized: + with deepspeed.zero.GatheredParameters(model_embeds.weight, modifier_rank=None): + vocab_size = model_embeds.weight.shape[0] + else: + vocab_size = model_embeds.weight.shape[0] + + # Update base model and current model config. + self.config.get_text_config().vocab_size = vocab_size + self.vocab_size = vocab_size + + # Tie weights again if needed + self.tie_weights() + + return model_embeds + + def _resize_token_embeddings(self, new_num_tokens, pad_to_multiple_of=None, mean_resizing=True): + old_embeddings = self.get_input_embeddings() + new_embeddings = self._get_resized_embeddings( + old_embeddings, new_num_tokens, pad_to_multiple_of, mean_resizing + ) + if hasattr(old_embeddings, "_hf_hook"): + hook = old_embeddings._hf_hook + add_hook_to_module(new_embeddings, hook) + old_embeddings_requires_grad = old_embeddings.weight.requires_grad + new_embeddings.requires_grad_(old_embeddings_requires_grad) + self.set_input_embeddings(new_embeddings) + is_quantized = hasattr(self, "hf_quantizer") and self.hf_quantizer is not None + + # Update new_num_tokens with the actual size of new_embeddings + if pad_to_multiple_of is not None: + if is_deepspeed_zero3_enabled() and not is_quantized: + with deepspeed.zero.GatheredParameters(new_embeddings.weight, modifier_rank=None): + new_num_tokens = new_embeddings.weight.shape[0] + else: + new_num_tokens = new_embeddings.weight.shape[0] + + # if word embeddings are not tied, make sure that lm head is resized as well + if ( + self.get_output_embeddings() is not None + and not self.config.get_text_config(decoder=True).tie_word_embeddings + ): + old_lm_head = self.get_output_embeddings() + if isinstance(old_lm_head, torch.nn.Embedding): + new_lm_head = self._get_resized_embeddings(old_lm_head, new_num_tokens, mean_resizing=mean_resizing) + else: + new_lm_head = self._get_resized_lm_head(old_lm_head, new_num_tokens, mean_resizing=mean_resizing) + if hasattr(old_lm_head, "_hf_hook"): + hook = old_lm_head._hf_hook + add_hook_to_module(new_lm_head, hook) + old_lm_head_requires_grad = old_lm_head.weight.requires_grad + new_lm_head.requires_grad_(old_lm_head_requires_grad) + self.set_output_embeddings(new_lm_head) + + return self.get_input_embeddings() + + def _get_resized_embeddings( + self, + old_embeddings: nn.Embedding, + new_num_tokens: Optional[int] = None, + pad_to_multiple_of: Optional[int] = None, + mean_resizing: bool = True, + ) -> nn.Embedding: + """ + Build a resized Embedding Module from a provided token Embedding Module. Increasing the size will add newly + initialized vectors at the end. Reducing the size will remove vectors from the end + + Args: + old_embeddings (`torch.nn.Embedding`): + Old embeddings to be resized. + new_num_tokens (`int`, *optional*): + New number of tokens in the embedding matrix. + + Increasing the size will add newly initialized vectors at the end. Reducing the size will remove + vectors from the end. If not provided or `None`, just returns a pointer to the input tokens + `torch.nn.Embedding` module of the model without doing anything. + pad_to_multiple_of (`int`, *optional*): + If set will pad the embedding matrix to a multiple of the provided value. If `new_num_tokens` is set to + `None` will just pad the embedding to a multiple of `pad_to_multiple_of`. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. For more + details about this, or help on choosing the correct value for resizing, refer to this guide: + https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc + mean_resizing (`bool`): + Whether to initialize the added embeddings from a multivariate normal distribution that has old embeddings' mean and + covariance or to initialize them with a normal distribution that has a mean of zero and std equals `config.initializer_range`. + + Setting `mean_resizing` to `True` is useful when increasing the size of the embeddings of causal language models, + where the generated tokens' probabilities will not be affected by the added embeddings because initializing the new embeddings with the + old embeddings' mean will reduce the kl-divergence between the next token probability before and after adding the new embeddings. + Refer to this article for more information: https://nlp.stanford.edu/~johnhew/vocab-expansion.html + + + Return: + `torch.nn.Embedding`: Pointer to the resized Embedding Module or the old Embedding Module if + `new_num_tokens` is `None` + """ + + if pad_to_multiple_of is not None: + if not isinstance(pad_to_multiple_of, int): + raise ValueError( + f"Asking to pad the embedding matrix to a multiple of `{pad_to_multiple_of}`, which is not and integer. Please make sure to pass an integer" + ) + if new_num_tokens is None: + new_num_tokens = old_embeddings.weight.shape[0] + new_num_tokens = ((new_num_tokens + pad_to_multiple_of - 1) // pad_to_multiple_of) * pad_to_multiple_of + else: + logger.info( + "You are resizing the embedding layer without providing a `pad_to_multiple_of` parameter. This means that the new embedding" + f" dimension will be {new_num_tokens}. This might induce some performance reduction as *Tensor Cores* will not be available." + " For more details about this, or help on choosing the correct value for resizing, refer to this guide:" + " https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc" + ) + + if new_num_tokens is None: + return old_embeddings + + is_quantized = hasattr(self, "hf_quantizer") and self.hf_quantizer is not None + if is_deepspeed_zero3_enabled() and not is_quantized: + with deepspeed.zero.GatheredParameters(old_embeddings.weight, modifier_rank=None): + old_num_tokens, old_embedding_dim = old_embeddings.weight.size() + else: + old_num_tokens, old_embedding_dim = old_embeddings.weight.size() + + if old_num_tokens == new_num_tokens and not is_deepspeed_zero3_enabled(): + return old_embeddings + + if not isinstance(old_embeddings, nn.Embedding): + raise TypeError( + f"Old embeddings are of type {type(old_embeddings)}, which is not an instance of {nn.Embedding}. You" + " should either use a different resize function or make sure that `old_embeddings` are an instance of" + f" {nn.Embedding}." + ) + + # Build new embeddings + + # When using DeepSpeed ZeRO-3, we shouldn't create new embeddings with DeepSpeed init + # because the shape of the new embedding layer is used across various modeling files + # as well as to update config vocab size. Shape will be 0 when using DeepSpeed init leading + # to errors when training. + new_embeddings = nn.Embedding( + new_num_tokens, + old_embedding_dim, + device=old_embeddings.weight.device, + dtype=old_embeddings.weight.dtype, + ) + + if new_num_tokens > old_num_tokens and not mean_resizing: + # initialize new embeddings (in particular added tokens) with a mean of 0 and std equals `config.initializer_range`. + self._init_weights(new_embeddings) + + elif new_num_tokens > old_num_tokens and mean_resizing: + # initialize new embeddings (in particular added tokens). The new embeddings will be initialized + # from a multivariate normal distribution that has old embeddings' mean and covariance. + # as described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html + logger.warning_once( + "The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. " + "As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. " + "To disable this, use `mean_resizing=False`" + ) + + added_num_tokens = new_num_tokens - old_num_tokens + if is_deepspeed_zero3_enabled() and not is_quantized: + with deepspeed.zero.GatheredParameters([old_embeddings.weight], modifier_rank=None): + self._init_added_embeddings_weights_with_mean( + old_embeddings, new_embeddings, old_embedding_dim, old_num_tokens, added_num_tokens + ) + else: + self._init_added_embeddings_weights_with_mean( + old_embeddings, new_embeddings, old_embedding_dim, old_num_tokens, added_num_tokens + ) + + # Copy token embeddings from the previous weights + + # numbers of tokens to copy + n = min(old_num_tokens, new_num_tokens) + + if is_deepspeed_zero3_enabled() and not is_quantized: + params = [old_embeddings.weight, new_embeddings.weight] + with deepspeed.zero.GatheredParameters(params, modifier_rank=0): + new_embeddings.weight.data[:n, :] = old_embeddings.weight.data[:n, :] + else: + new_embeddings.weight.data[:n, :] = old_embeddings.weight.data[:n, :] + + # Replace weights in old_embeddings and return to maintain the same embedding type. + # This ensures correct functionality when a Custom Embedding class is passed as input. + # The input and output embedding types remain consistent. (c.f. https://github.com/huggingface/transformers/pull/31979) + if is_deepspeed_zero3_enabled() and not is_quantized: + params = [old_embeddings.weight, new_embeddings.weight] + with deepspeed.zero.GatheredParameters(params, modifier_rank=0): + old_embeddings.weight = new_embeddings.weight + old_embeddings.num_embeddings = new_embeddings.weight.data.shape[0] + + # If the new number of tokens is smaller than the original `padding_idx`, the `padding_idx` + # will be set to `None` in the resized embeddings. + if old_embeddings.padding_idx is not None and (new_num_tokens - 1) < old_embeddings.padding_idx: + old_embeddings.padding_idx = None + else: + old_embeddings.weight.data = new_embeddings.weight.data + old_embeddings.num_embeddings = new_embeddings.weight.data.shape[0] + if old_embeddings.padding_idx is not None and (new_num_tokens - 1) < old_embeddings.padding_idx: + old_embeddings.padding_idx = None + + return old_embeddings + + def _get_resized_lm_head( + self, + old_lm_head: nn.Linear, + new_num_tokens: Optional[int] = None, + transposed: Optional[bool] = False, + mean_resizing: bool = True, + ) -> nn.Linear: + """ + Build a resized Linear Module from a provided old Linear Module. Increasing the size will add newly initialized + vectors at the end. Reducing the size will remove vectors from the end + + Args: + old_lm_head (`torch.nn.Linear`): + Old lm head liner layer to be resized. + new_num_tokens (`int`, *optional*): + New number of tokens in the linear matrix. + + Increasing the size will add newly initialized vectors at the end. Reducing the size will remove + vectors from the end. If not provided or `None`, just returns a pointer to the input tokens + `torch.nn.Linear` module of the model without doing anything. transposed (`bool`, *optional*, defaults + to `False`): Whether `old_lm_head` is transposed or not. If True `old_lm_head.size()` is `lm_head_dim, + vocab_size` else `vocab_size, lm_head_dim`. + mean_resizing (`bool`): + Whether to initialize the added embeddings from a multivariate normal distribution that has old embeddings' mean and + covariance or to initialize them with a normal distribution that has a mean of zero and std equals `config.initializer_range`. + + Setting `mean_resizing` to `True` is useful when increasing the size of the embeddings of causal language models, + where the generated tokens' probabilities will not be affected by the added embeddings because initializing the new embeddings with the + old embeddings' mean will reduce the kl-divergence between the next token probability before and after adding the new embeddings. + Refer to this article for more information: https://nlp.stanford.edu/~johnhew/vocab-expansion.html + + Return: + `torch.nn.Linear`: Pointer to the resized Linear Module or the old Linear Module if `new_num_tokens` is + `None` + """ + if new_num_tokens is None: + return old_lm_head + + is_quantized = hasattr(self, "hf_quantizer") and self.hf_quantizer is not None + if is_deepspeed_zero3_enabled() and not is_quantized: + with deepspeed.zero.GatheredParameters(old_lm_head.weight, modifier_rank=None): + old_num_tokens, old_lm_head_dim = ( + old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size() + ) + else: + old_num_tokens, old_lm_head_dim = ( + old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size() + ) + + if old_num_tokens == new_num_tokens and not is_deepspeed_zero3_enabled(): + return old_lm_head + + if not isinstance(old_lm_head, nn.Linear): + raise TypeError( + f"Old language model head is of type {type(old_lm_head)}, which is not an instance of {nn.Linear}. You" + " should either use a different resize function or make sure that `old_lm_head` are an instance of" + f" {nn.Linear}." + ) + + # Build new lm head + new_lm_head_shape = (old_lm_head_dim, new_num_tokens) if not transposed else (new_num_tokens, old_lm_head_dim) + has_new_lm_head_bias = old_lm_head.bias is not None + + # When using DeepSpeed ZeRO-3, we shouldn't create new embeddings with DeepSpeed init + # because the shape of the new embedding layer is used across various modeling files + # as well as to update config vocab size. Shape will be 0 when using DeepSpeed init leading + # to errors when training. + new_lm_head = nn.Linear( + *new_lm_head_shape, + bias=has_new_lm_head_bias, + device=old_lm_head.weight.device, + dtype=old_lm_head.weight.dtype, + ) + + if new_num_tokens > old_num_tokens and not mean_resizing: + # initialize new embeddings (in particular added tokens) with a mean of 0 and std equals `config.initializer_range`. + self._init_weights(new_lm_head) + + elif new_num_tokens > old_num_tokens and mean_resizing: + # initialize new lm_head weights (in particular added tokens). The new lm_head weights + # will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. + # as described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html + logger.warning_once( + "The new lm_head weights will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. " + "As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. " + "To disable this, use `mean_resizing=False`" + ) + + added_num_tokens = new_num_tokens - old_num_tokens + if is_deepspeed_zero3_enabled() and not is_quantized: + params = [old_lm_head.weight] + if has_new_lm_head_bias: + params += [old_lm_head.bias] + with deepspeed.zero.GatheredParameters(params, modifier_rank=None): + self._init_added_lm_head_weights_with_mean( + old_lm_head, new_lm_head, old_lm_head_dim, old_num_tokens, added_num_tokens, transposed + ) + if has_new_lm_head_bias: + self._init_added_lm_head_bias_with_mean(old_lm_head, new_lm_head, added_num_tokens) + + else: + self._init_added_lm_head_weights_with_mean( + old_lm_head, new_lm_head, old_lm_head_dim, old_num_tokens, added_num_tokens, transposed + ) + if has_new_lm_head_bias: + self._init_added_lm_head_bias_with_mean(old_lm_head, new_lm_head, added_num_tokens) + + num_tokens_to_copy = min(old_num_tokens, new_num_tokens) + + if is_deepspeed_zero3_enabled() and not is_quantized: + params = [old_lm_head.weight, old_lm_head.bias, new_lm_head.weight, new_lm_head.bias] + with deepspeed.zero.GatheredParameters(params, modifier_rank=0): + self._copy_lm_head_original_to_resized( + new_lm_head, old_lm_head, num_tokens_to_copy, transposed, has_new_lm_head_bias + ) + else: + self._copy_lm_head_original_to_resized( + new_lm_head, old_lm_head, num_tokens_to_copy, transposed, has_new_lm_head_bias + ) + + return new_lm_head + + def _init_added_embeddings_weights_with_mean( + self, old_embeddings, new_embeddings, old_embedding_dim, old_num_tokens, added_num_tokens + ): + old_embeddings_weight = old_embeddings.weight.data.to(torch.float32) + mean_embeddings = torch.mean(old_embeddings_weight, axis=0) + old_centered_embeddings = old_embeddings_weight - mean_embeddings + covariance = old_centered_embeddings.T @ old_centered_embeddings / old_num_tokens + + # Check if the covariance is positive definite. + epsilon = 1e-9 + is_covariance_psd = constraints.positive_definite.check(epsilon * covariance).all() + if is_covariance_psd: + # If covariances is positive definite, a distribution can be created. and we can sample new weights from it. + distribution = torch.distributions.multivariate_normal.MultivariateNormal( + mean_embeddings, covariance_matrix=epsilon * covariance + ) + new_embeddings.weight.data[-1 * added_num_tokens :, :] = distribution.sample( + sample_shape=(added_num_tokens,) + ).to(old_embeddings.weight.dtype) + else: + # Otherwise, just initialize with the mean. because distribution will not be created. + new_embeddings.weight.data[-1 * added_num_tokens :, :] = ( + mean_embeddings[None, :].repeat(added_num_tokens, 1).to(old_embeddings.weight.dtype) + ) + + def _init_added_lm_head_weights_with_mean( + self, + old_lm_head, + new_lm_head, + old_lm_head_dim, + old_num_tokens, + added_num_tokens, + transposed=False, + ): + if transposed: + # Transpose to the desired shape for the function. + new_lm_head.weight.data = new_lm_head.weight.data.T + old_lm_head.weight.data = old_lm_head.weight.data.T + + # The same initialization logic as Embeddings. + self._init_added_embeddings_weights_with_mean( + old_lm_head, new_lm_head, old_lm_head_dim, old_num_tokens, added_num_tokens + ) + + if transposed: + # Transpose again to the correct shape. + new_lm_head.weight.data = new_lm_head.weight.data.T + old_lm_head.weight.data = old_lm_head.weight.data.T + + def _init_added_lm_head_bias_with_mean(self, old_lm_head, new_lm_head, added_num_tokens): + bias_mean = torch.mean(old_lm_head.bias.data, axis=0, dtype=torch.float32) + bias_std = torch.std(old_lm_head.bias.data, axis=0).to(torch.float32) + new_lm_head.bias.data[-1 * added_num_tokens :].normal_(mean=bias_mean, std=1e-9 * bias_std) + + def _copy_lm_head_original_to_resized( + self, new_lm_head, old_lm_head, num_tokens_to_copy, transposed, has_new_lm_head_bias + ): + # Copy old lm head weights to new lm head + if not transposed: + new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[:num_tokens_to_copy, :] + else: + new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[:, :num_tokens_to_copy] + + # Copy bias weights to new lm head + if has_new_lm_head_bias: + new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy] + + def resize_position_embeddings(self, new_num_position_embeddings: int): + raise NotImplementedError( + f"`resize_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should " + f"overwrite this method in the class {self.__class__} in `modeling_{self.__class__.__module__}.py`" + ) + + def get_position_embeddings(self) -> Union[nn.Embedding, Tuple[nn.Embedding]]: + raise NotImplementedError( + f"`get_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should " + f"overwrite this method in the class {self.__class__} in `modeling_{self.__class__.__module__}.py`" + ) + + def init_weights(self): + """ + If needed prunes and maybe initializes weights. If using a custom `PreTrainedModel`, you need to implement any + initialization logic in `_init_weights`. + """ + # Prune heads if needed + if self.config.pruned_heads: + self.prune_heads(self.config.pruned_heads) + + if _init_weights: + # Initialize weights + self.apply(self._initialize_weights) + + # Tie weights should be skipped when not initializing all weights + # since from_pretrained(...) calls tie weights anyways + self.tie_weights() + + def prune_heads(self, heads_to_prune: Dict[int, List[int]]): + """ + Prunes heads of the base model. + + Arguments: + heads_to_prune (`Dict[int, List[int]]`): + Dictionary with keys being selected layer indices (`int`) and associated values being the list of heads + to prune in said layer (list of `int`). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on + layer 1 and heads 2 and 3 on layer 2. + """ + # save new sets of pruned heads as union of previously stored pruned heads and newly pruned heads + for layer, heads in heads_to_prune.items(): + union_heads = set(self.config.pruned_heads.get(layer, [])) | set(heads) + self.config.pruned_heads[layer] = list(union_heads) # Unfortunately we have to store it as list for JSON + + self.base_model._prune_heads(heads_to_prune) + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + """ + Activates gradient checkpointing for the current model. + + Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint + activations". + + We pass the `__call__` method of the modules instead of `forward` because `__call__` attaches all the hooks of + the module. https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690/2 + + Args: + gradient_checkpointing_kwargs (dict, *optional*): + Additional keyword arguments passed along to the `torch.utils.checkpoint.checkpoint` function. + """ + if not self.supports_gradient_checkpointing: + raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.") + + if gradient_checkpointing_kwargs is None: + gradient_checkpointing_kwargs = {"use_reentrant": True} + + gradient_checkpointing_func = functools.partial(checkpoint, **gradient_checkpointing_kwargs) + + # For old GC format (transformers < 4.35.0) for models that live on the Hub + # we will fall back to the overwritten `_set_gradient_checkpointing` method + _is_using_old_format = "value" in inspect.signature(self._set_gradient_checkpointing).parameters + + if not _is_using_old_format: + self._set_gradient_checkpointing(enable=True, gradient_checkpointing_func=gradient_checkpointing_func) + else: + self.apply(partial(self._set_gradient_checkpointing, value=True)) + logger.warning( + "You are using an old version of the checkpointing format that is deprecated (We will also silently ignore `gradient_checkpointing_kwargs` in case you passed it)." + "Please update to the new format on your modeling file. To use the new format, you need to completely remove the definition of the method `_set_gradient_checkpointing` in your model." + ) + + if getattr(self, "_hf_peft_config_loaded", False): + # When using PEFT + gradient checkpointing + Trainer we need to make sure the input has requires_grad=True + # we do it also on PEFT: https://github.com/huggingface/peft/blob/85013987aa82aa1af3da1236b6902556ce3e483e/src/peft/peft_model.py#L334 + # When training with PEFT, only LoRA layers will have requires grad set to True, but the output of frozen layers need to propagate + # the gradients to make sure the gradient flows. + self.enable_input_require_grads() + + def _set_gradient_checkpointing(self, enable: bool = True, gradient_checkpointing_func: Callable = checkpoint): + is_gradient_checkpointing_set = False + + # Apply it on the top-level module in case the top-level modules supports it + # for example, LongT5Stack inherits from `PreTrainedModel`. + if hasattr(self, "gradient_checkpointing"): + self._gradient_checkpointing_func = gradient_checkpointing_func + self.gradient_checkpointing = enable + is_gradient_checkpointing_set = True + + for module in self.modules(): + if hasattr(module, "gradient_checkpointing"): + module._gradient_checkpointing_func = gradient_checkpointing_func + module.gradient_checkpointing = enable + is_gradient_checkpointing_set = True + + if not is_gradient_checkpointing_set: + raise ValueError( + f"{self.__class__.__name__} is not compatible with gradient checkpointing. Make sure all the architecture support it by setting a boolean attribute" + " `gradient_checkpointing` to modules of the model that uses checkpointing." + ) + + def gradient_checkpointing_disable(self): + """ + Deactivates gradient checkpointing for the current model. + + Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint + activations". + """ + if self.supports_gradient_checkpointing: + # For old GC format (transformers < 4.35.0) for models that live on the Hub + # we will fall back to the overwritten `_set_gradient_checkpointing` method + _is_using_old_format = "value" in inspect.signature(self._set_gradient_checkpointing).parameters + if not _is_using_old_format: + self._set_gradient_checkpointing(enable=False) + else: + logger.warning( + "You are using an old version of the checkpointing format that is deprecated (We will also silently ignore `gradient_checkpointing_kwargs` in case you passed it)." + "Please update to the new format on your modeling file. To use the new format, you need to completely remove the definition of the method `_set_gradient_checkpointing` in your model." + ) + self.apply(partial(self._set_gradient_checkpointing, value=False)) + + if getattr(self, "_hf_peft_config_loaded", False): + self.disable_input_require_grads() + + @property + def is_gradient_checkpointing(self) -> bool: + """ + Whether gradient checkpointing is activated for this model or not. + + Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint + activations". + """ + return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules()) + + def save_pretrained( + self, + save_directory: Union[str, os.PathLike], + is_main_process: bool = True, + state_dict: Optional[dict] = None, + save_function: Callable = torch.save, + push_to_hub: bool = False, + max_shard_size: Union[int, str] = "5GB", + safe_serialization: bool = True, + variant: Optional[str] = None, + token: Optional[Union[str, bool]] = None, + save_peft_format: bool = True, + **kwargs, + ): + """ + Save a model and its configuration file to a directory, so that it can be re-loaded using the + [`~PreTrainedModel.from_pretrained`] class method. + + Arguments: + save_directory (`str` or `os.PathLike`): + Directory to which to save. Will be created if it doesn't exist. + is_main_process (`bool`, *optional*, defaults to `True`): + Whether the process calling this is the main process or not. Useful when in distributed training like + TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on + the main process to avoid race conditions. + state_dict (nested dictionary of `torch.Tensor`): + The state dictionary of the model to save. Will default to `self.state_dict()`, but can be used to only + save parts of the model or if special precautions need to be taken when recovering the state dictionary + of a model (like when using model parallelism). + save_function (`Callable`): + The function to use to save the state dictionary. Useful on distributed training like TPUs when one + need to replace `torch.save` by another method. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + max_shard_size (`int` or `str`, *optional*, defaults to `"5GB"`): + The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size + lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`). + We default it to 5GB in order for models to be able to run easily on free-tier google colab instances + without CPU OOM issues. + + + + If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard + which will be bigger than `max_shard_size`. + + + + safe_serialization (`bool`, *optional*, defaults to `True`): + Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). + variant (`str`, *optional*): + If specified, weights are saved in the format pytorch_model..bin. + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + save_peft_format (`bool`, *optional*, defaults to `True`): + For backward compatibility with PEFT library, in case adapter weights are attached to the model, all + keys of the state dict of adapters needs to be pre-pended with `base_model.model`. Advanced users can + disable this behaviours by setting `save_peft_format` to `False`. + kwargs (`Dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + use_auth_token = kwargs.pop("use_auth_token", None) + ignore_metadata_errors = kwargs.pop("ignore_metadata_errors", False) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if token is not None: + kwargs["token"] = token + + _hf_peft_config_loaded = getattr(self, "_hf_peft_config_loaded", False) + + hf_quantizer = getattr(self, "hf_quantizer", None) + quantization_serializable = ( + hf_quantizer is not None + and isinstance(hf_quantizer, HfQuantizer) + and hf_quantizer.is_serializable(safe_serialization=safe_serialization) + ) + + if hf_quantizer is not None and not _hf_peft_config_loaded and not quantization_serializable: + raise ValueError( + f"The model is quantized with {hf_quantizer.quantization_config.quant_method} and is not serializable - check out the warnings from" + " the logger on the traceback to understand the reason why the quantized model is not serializable." + ) + + if "save_config" in kwargs: + warnings.warn( + "`save_config` is deprecated and will be removed in v5 of Transformers. Use `is_main_process` instead." + ) + is_main_process = kwargs.pop("save_config") + if safe_serialization and not is_safetensors_available(): + raise ImportError("`safe_serialization` requires the `safetensors library: `pip install safetensors`.") + + if os.path.isfile(save_directory): + logger.error(f"Provided path ({save_directory}) should be a directory, not a file") + return + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = self._create_repo(repo_id, **kwargs) + files_timestamps = self._get_files_timestamps(save_directory) + + # Only save the model itself if we are using distributed training + model_to_save = unwrap_model(self) + + # save the string version of dtype to the config, e.g. convert torch.float32 => "float32" + # we currently don't use this setting automatically, but may start to use with v5 + dtype = get_parameter_dtype(model_to_save) + model_to_save.config.torch_dtype = str(dtype).split(".")[1] + + # Attach architecture to the config + model_to_save.config.architectures = [model_to_save.__class__.__name__] + + # Unset attn implementation so it can be set to another one when loading back + model_to_save.config._attn_implementation_autoset = False + + # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + custom_object_save(self, save_directory, config=self.config) + + # Save the config + if is_main_process: + if not _hf_peft_config_loaded: + # If the model config has set attributes that should be in the generation config, move them there. + misplaced_generation_parameters = model_to_save.config._get_non_default_generation_parameters() + if self.can_generate() and len(misplaced_generation_parameters) > 0: + warnings.warn( + "Moving the following attributes in the config to the generation config: " + f"{misplaced_generation_parameters}. You are seeing this warning because you've set " + "generation parameters in the model config, as opposed to in the generation config.", + UserWarning, + ) + for param_name, param_value in misplaced_generation_parameters.items(): + setattr(model_to_save.generation_config, param_name, param_value) + setattr(model_to_save.config, param_name, None) + + model_to_save.config.save_pretrained(save_directory) + if self.can_generate(): + model_to_save.generation_config.save_pretrained(save_directory) + + if _hf_peft_config_loaded: + logger.info( + "Detected adapters on the model, saving the model in the PEFT format, only adapter weights will be saved." + ) + state_dict = model_to_save.get_adapter_state_dict(state_dict=state_dict) + + if save_peft_format: + logger.info( + "To match the expected format of the PEFT library, all keys of the state dict of adapters will be pre-pended with `base_model.model`." + ) + peft_state_dict = {} + for key, value in state_dict.items(): + peft_state_dict[f"base_model.model.{key}"] = value + state_dict = peft_state_dict + + active_adapter = self.active_adapters() + + if len(active_adapter) > 1: + raise ValueError( + "Multiple active adapters detected, saving multiple active adapters is not supported yet. You can save adapters separately one by one " + "by iteratively calling `model.set_adapter(adapter_name)` then `model.save_pretrained(...)`" + ) + active_adapter = active_adapter[0] + + current_peft_config = self.peft_config[active_adapter] + current_peft_config.save_pretrained(save_directory) + + # for offloaded modules + module_map = {} + + # Save the model + if state_dict is None: + # if any model parameters are offloaded, make module map + if ( + hasattr(self, "hf_device_map") + and len(set(self.hf_device_map.values())) > 1 + and ("cpu" in self.hf_device_map.values() or "disk" in self.hf_device_map.values()) + ): + warnings.warn( + "Attempting to save a model with offloaded modules. Ensure that unallocated cpu memory exceeds the `shard_size` (5GB default)" + ) + for name, module in model_to_save.named_modules(): + if name == "": + continue + module_state_dict = module.state_dict() + + for key in module_state_dict: + module_map[name + f".{key}"] = module + state_dict = model_to_save.state_dict() + + # Translate state_dict from smp to hf if saving with smp >= 1.10 + if IS_SAGEMAKER_MP_POST_1_10: + for smp_to_hf, _ in smp.state.module_manager.translate_functions: + state_dict = smp_to_hf(state_dict) + + # Handle the case where some state_dict keys shouldn't be saved + if self._keys_to_ignore_on_save is not None: + for ignore_key in self._keys_to_ignore_on_save: + if ignore_key in state_dict.keys(): + del state_dict[ignore_key] + + # Rename state_dict keys before saving to file. Do nothing unless overridden in a particular model. + # (initially introduced with TimmWrapperModel to remove prefix and make checkpoints compatible with timm) + state_dict = self._fix_state_dict_keys_on_save(state_dict) + + if safe_serialization: + # Safetensors does not allow tensor aliasing. + # We're going to remove aliases before saving + ptrs = collections.defaultdict(list) + for name, tensor in state_dict.items(): + # Sometimes in the state_dict we have non-tensor objects. + # e.g. in bitsandbytes we have some `str` objects in the state_dict + if isinstance(tensor, torch.Tensor): + ptrs[id_tensor_storage(tensor)].append(name) + else: + # In the non-tensor case, fall back to the pointer of the object itself + ptrs[id(tensor)].append(name) + + # These are all the pointers of shared tensors + if hasattr(self, "hf_device_map"): + # if the model has offloaded parameters, we must check using find_tied_parameters() + tied_params = find_tied_parameters(self) + if tied_params: + tied_names = tied_params[0] + shared_ptrs = { + ptr: names for ptr, names in ptrs.items() if any(name in tied_names for name in names) + } + else: + shared_ptrs = {} + else: + shared_ptrs = {ptr: names for ptr, names in ptrs.items() if len(names) > 1} + + # Recursively descend to find tied weight keys + _tied_weights_keys = _get_tied_weight_keys(self) + error_names = [] + to_delete_names = set() + for names in shared_ptrs.values(): + # Removing the keys which are declared as known duplicates on + # load. This allows to make sure the name which is kept is consistent. + if _tied_weights_keys is not None: + found = 0 + for name in sorted(names): + matches_pattern = any(re.search(pat, name) for pat in _tied_weights_keys) + if matches_pattern and name in state_dict: + found += 1 + if found < len(names): + to_delete_names.add(name) + # We are entering a place where the weights and the transformers configuration do NOT match. + shared_names, disjoint_names = _find_disjoint(shared_ptrs.values(), state_dict) + # Those are actually tensor sharing but disjoint from each other, we can safely clone them + # Reloaded won't have the same property, but it shouldn't matter in any meaningful way. + for name in disjoint_names: + state_dict[name] = state_dict[name].clone() + + # When not all duplicates have been cleaned, still remove those keys, but put a clear warning. + # If the link between tensors was done at runtime then `from_pretrained` will not get + # the key back leading to random tensor. A proper warning will be shown + # during reload (if applicable), but since the file is not necessarily compatible with + # the config, better show a proper warning. + shared_names, identical_names = _find_identical(shared_names, state_dict) + # delete tensors that have identical storage + for inames in identical_names: + known = inames.intersection(to_delete_names) + for name in known: + del state_dict[name] + unknown = inames.difference(to_delete_names) + if len(unknown) > 1: + error_names.append(unknown) + + if shared_names: + error_names.append(set(shared_names)) + + if len(error_names) > 0: + raise RuntimeError( + f"The weights trying to be saved contained shared tensors {error_names} that are mismatching the transformers base configuration. Try saving using `safe_serialization=False` or remove this tensor sharing.", + ) + + # Shard the model if it is too big. + if not _hf_peft_config_loaded: + weights_name = SAFE_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME + weights_name = _add_variant(weights_name, variant) + else: + weights_name = ADAPTER_SAFE_WEIGHTS_NAME if safe_serialization else ADAPTER_WEIGHTS_NAME + + filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors") + state_dict_split = split_torch_state_dict_into_shards( + state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size + ) + # Save index if sharded + index = None + if state_dict_split.is_sharded: + index = { + "metadata": state_dict_split.metadata, + "weight_map": state_dict_split.tensor_to_filename, + } + + # Clean the folder from a previous save + for filename in os.listdir(save_directory): + full_filename = os.path.join(save_directory, filename) + # If we have a shard file that is not going to be replaced, we delete it, but only from the main process + # in distributed settings to avoid race conditions. + weights_no_suffix = weights_name.replace(".bin", "").replace(".safetensors", "") + + # make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005 + filename_no_suffix = filename.replace(".bin", "").replace(".safetensors", "") + reg = re.compile(r"(.*?)-\d{5}-of-\d{5}") + + if ( + filename.startswith(weights_no_suffix) + and os.path.isfile(full_filename) + and filename not in state_dict_split.filename_to_tensors.keys() + and is_main_process + and reg.fullmatch(filename_no_suffix) is not None + ): + os.remove(full_filename) + # Save the model + filename_to_tensors = state_dict_split.filename_to_tensors.items() + if module_map: + filename_to_tensors = logging.tqdm(filename_to_tensors, desc="Saving checkpoint shards") + for shard_file, tensors in filename_to_tensors: + shard = {} + for tensor in tensors: + shard[tensor] = state_dict[tensor].contiguous() + # delete reference, see https://github.com/huggingface/transformers/pull/34890 + del state_dict[tensor] + + # remake shard with onloaded parameters if necessary + if module_map: + if accelerate_version < version.parse("0.31"): + raise ImportError( + f"You need accelerate version to be greater or equal than 0.31 to save models with offloaded parameters. Detected version {accelerate_version}. " + f"Please upgrade accelerate with `pip install -U accelerate`" + ) + # init state_dict for this shard + shard_state_dict = dict.fromkeys(shard, "") + for module_name in shard: + # skip to collect this weight again + if shard_state_dict.get(module_name) != "": + continue + module = module_map[module_name] + # update state dict with onloaded parameters + shard_state_dict = get_state_dict_from_offload(module, module_name, shard_state_dict) + + # assign shard to be the completed state dict + shard = shard_state_dict + del shard_state_dict + gc.collect() + + if safe_serialization: + # At some point we will need to deal better with save_function (used for TPU and other distributed + # joyfulness), but for now this enough. + safe_save_file(shard, os.path.join(save_directory, shard_file), metadata={"format": "pt"}) + else: + save_function(shard, os.path.join(save_directory, shard_file)) + + del state_dict + + if index is None: + path_to_weights = os.path.join(save_directory, weights_name) + logger.info(f"Model weights saved in {path_to_weights}") + else: + save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME + save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant)) + # Save the index as well + with open(save_index_file, "w", encoding="utf-8") as f: + content = json.dumps(index, indent=2, sort_keys=True) + "\n" + f.write(content) + logger.info( + f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be " + f"split in {len(state_dict_split.filename_to_tensors)} checkpoint shards. You can find where each parameters has been saved in the " + f"index located at {save_index_file}." + ) + + if push_to_hub: + # Eventually create an empty model card + model_card = create_and_tag_model_card( + repo_id, self.model_tags, token=token, ignore_metadata_errors=ignore_metadata_errors + ) + + # Update model card if needed: + model_card.save(os.path.join(save_directory, "README.md")) + + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=token, + ) + + @wraps(PushToHubMixin.push_to_hub) + def push_to_hub(self, *args, **kwargs): + tags = self.model_tags if self.model_tags is not None else [] + + tags_kwargs = kwargs.get("tags", []) + if isinstance(tags_kwargs, str): + tags_kwargs = [tags_kwargs] + + for tag in tags_kwargs: + if tag not in tags: + tags.append(tag) + + if tags: + kwargs["tags"] = tags + return super().push_to_hub(*args, **kwargs) + + def get_memory_footprint(self, return_buffers=True): + r""" + Get the memory footprint of a model. This will return the memory footprint of the current model in bytes. + Useful to benchmark the memory footprint of the current model and design some tests. Solution inspired from the + PyTorch discussions: https://discuss.pytorch.org/t/gpu-memory-that-model-uses/56822/2 + + Arguments: + return_buffers (`bool`, *optional*, defaults to `True`): + Whether to return the size of the buffer tensors in the computation of the memory footprint. Buffers + are tensors that do not require gradients and not registered as parameters. E.g. mean and std in batch + norm layers. Please see: https://discuss.pytorch.org/t/what-pytorch-means-by-buffers/120266/2 + """ + mem = sum([param.nelement() * param.element_size() for param in self.parameters()]) + if return_buffers: + mem_bufs = sum([buf.nelement() * buf.element_size() for buf in self.buffers()]) + mem = mem + mem_bufs + return mem + + @wraps(torch.nn.Module.cuda) + def cuda(self, *args, **kwargs): + if getattr(self, "quantization_method", None) == QuantizationMethod.HQQ: + raise ValueError("`.cuda` is not supported for HQQ-quantized models.") + # Checks if the model has been loaded in 4-bit or 8-bit with BNB + if getattr(self, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES: + if getattr(self, "is_loaded_in_8bit", False): + raise ValueError( + "Calling `cuda()` is not supported for `8-bit` quantized models. " + " Please use the model as it is, since the model has already been set to the correct devices." + ) + elif version.parse(importlib.metadata.version("bitsandbytes")) < version.parse("0.43.2"): + raise ValueError( + "Calling `cuda()` is not supported for `4-bit` quantized models with the installed version of bitsandbytes. " + f"The current device is `{self.device}`. If you intended to move the model, please install bitsandbytes >= 0.43.2." + ) + else: + return super().cuda(*args, **kwargs) + + @wraps(torch.nn.Module.to) + def to(self, *args, **kwargs): + # For BNB/GPTQ models, we prevent users from casting the model to another dtype to restrict unwanted behaviours. + # the correct API should be to load the model with the desired dtype directly through `from_pretrained`. + dtype_present_in_args = "dtype" in kwargs + + if not dtype_present_in_args: + for arg in args: + if isinstance(arg, torch.dtype): + dtype_present_in_args = True + break + + if getattr(self, "quantization_method", None) == QuantizationMethod.HQQ: + raise ValueError("`.to` is not supported for HQQ-quantized models.") + + if dtype_present_in_args and getattr(self, "quantization_method", None) == QuantizationMethod.QUARK: + raise ValueError("Casting a Quark quantized model to a new `dtype` is not supported.") + + # Checks if the model has been loaded in 4-bit or 8-bit with BNB + if getattr(self, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES: + if dtype_present_in_args: + raise ValueError( + "You cannot cast a bitsandbytes model in a new `dtype`. Make sure to load the model using `from_pretrained` using the" + " desired `dtype` by passing the correct `torch_dtype` argument." + ) + + if getattr(self, "is_loaded_in_8bit", False): + raise ValueError( + "`.to` is not supported for `8-bit` bitsandbytes models. Please use the model as it is, since the" + " model has already been set to the correct devices and casted to the correct `dtype`." + ) + elif version.parse(importlib.metadata.version("bitsandbytes")) < version.parse("0.43.2"): + raise ValueError( + "Calling `to()` is not supported for `4-bit` quantized models with the installed version of bitsandbytes. " + f"The current device is `{self.device}`. If you intended to move the model, please install bitsandbytes >= 0.43.2." + ) + elif getattr(self, "quantization_method", None) == QuantizationMethod.GPTQ: + if dtype_present_in_args: + raise ValueError( + "You cannot cast a GPTQ model in a new `dtype`. Make sure to load the model using `from_pretrained` using the desired" + " `dtype` by passing the correct `torch_dtype` argument." + ) + return super().to(*args, **kwargs) + + def half(self, *args): + # Checks if the model is quantized + if getattr(self, "is_quantized", False): + raise ValueError( + "`.half()` is not supported for quantized model. Please use the model as it is, since the" + " model has already been casted to the correct `dtype`." + ) + else: + return super().half(*args) + + def float(self, *args): + # Checks if the model is quantized + if getattr(self, "is_quantized", False): + raise ValueError( + "`.float()` is not supported for quantized model. Please use the model as it is, since the" + " model has already been casted to the correct `dtype`." + ) + else: + return super().float(*args) + + @classmethod + def get_init_context(cls, is_quantized: bool, _is_ds_init_called: bool): + if is_deepspeed_zero3_enabled(): + init_contexts = [no_init_weights()] + # We cannot initialize the model on meta device with deepspeed when not quantized + if not is_quantized and not _is_ds_init_called: + logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model") + init_contexts.extend([deepspeed.zero.Init(config_dict_or_path=deepspeed_config()), set_zero3_state()]) + elif is_quantized: + init_contexts.extend([init_empty_weights(), set_quantized_state()]) + else: + init_contexts = [no_init_weights(), init_empty_weights()] + + return init_contexts + + @classmethod + @restore_default_torch_dtype + def from_pretrained( + cls: Type[SpecificPreTrainedModelType], + pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], + *model_args, + config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None, + cache_dir: Optional[Union[str, os.PathLike]] = None, + ignore_mismatched_sizes: bool = False, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + use_safetensors: Optional[bool] = None, + weights_only: bool = True, + **kwargs, + ) -> SpecificPreTrainedModelType: + r""" + Instantiate a pretrained pytorch model from a pre-trained model configuration. + + The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train + the model, you should first set it back in training mode with `model.train()`. + + The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come + pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning + task. + + The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those + weights are discarded. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*): + Can be either: + + - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co. + - A path to a *directory* containing model weights saved using + [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`. + - A path or url to a *tensorflow index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In + this case, `from_tf` should be set to `True` and a configuration object should be provided as + `config` argument. This loading path is slower than converting the TensorFlow checkpoint in a + PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards. + - A path or url to a model folder containing a *flax checkpoint file* in *.msgpack* format (e.g, + `./flax_model/` containing `flax_model.msgpack`). In this case, `from_flax` should be set to + `True`. + - `None` if you are both providing the configuration and state dictionary (resp. with keyword + arguments `config` and `state_dict`). + model_args (sequence of positional arguments, *optional*): + All remaining positional arguments will be passed to the underlying model's `__init__` method. + config (`Union[PretrainedConfig, str, os.PathLike]`, *optional*): + Can be either: + + - an instance of a class derived from [`PretrainedConfig`], + - a string or path valid as input to [`~PretrainedConfig.from_pretrained`]. + + Configuration for the model to use instead of an automatically loaded configuration. Configuration can + be automatically loaded when: + + - The model is a model provided by the library (loaded with the *model id* string of a pretrained + model). + - The model was saved using [`~PreTrainedModel.save_pretrained`] and is reloaded by supplying the + save directory. + - The model is loaded by supplying a local directory as `pretrained_model_name_or_path` and a + configuration JSON file named *config.json* is found in the directory. + state_dict (`Dict[str, torch.Tensor]`, *optional*): + A state dictionary to use instead of a state dictionary loaded from saved weights file. + + This option can be used if you want to create a model from a pretrained configuration but load your own + weights. In this case though, you should check if using [`~PreTrainedModel.save_pretrained`] and + [`~PreTrainedModel.from_pretrained`] is not a simpler option. + cache_dir (`Union[str, 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. + from_tf (`bool`, *optional*, defaults to `False`): + Load the model weights from a TensorFlow checkpoint save file (see docstring of + `pretrained_model_name_or_path` argument). + from_flax (`bool`, *optional*, defaults to `False`): + Load the model weights from a Flax checkpoint save file (see docstring of + `pretrained_model_name_or_path` argument). + ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`): + Whether or not to raise an error if some of the weights from the checkpoint do not have the same size + as the weights of the model (if for instance, you are instantiating a model with 10 labels from a + checkpoint with 3 labels). + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download of the model weights and configuration files, overriding the + cached versions if they exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + output_loading_info(`bool`, *optional*, defaults to `False`): + Whether ot not to also return a dictionary containing missing keys, unexpected keys and error messages. + local_files_only(`bool`, *optional*, defaults to `False`): + Whether or not to only look at local files (i.e., do not try to download the model). + token (`str` or `bool`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use + the token generated when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`. + + + attn_implementation (`str`, *optional*): + The attention implementation to use in the model (if relevant). Can be any of `"eager"` (manual implementation of the attention), `"sdpa"` (using [`F.scaled_dot_product_attention`](https://pytorch.org/docs/master/generated/torch.nn.functional.scaled_dot_product_attention.html)), or `"flash_attention_2"` (using [Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention)). By default, if available, SDPA will be used for torch>=2.1.1. The default is otherwise the manual `"eager"` implementation. + + > Parameters for big model inference + + torch_dtype (`str` or `torch.dtype`, *optional*): + Override the default `torch.dtype` and load the model under a specific `dtype`. The different options + are: + + 1. `torch.float16` or `torch.bfloat16` or `torch.float`: load in a specified + `dtype`, ignoring the model's `config.torch_dtype` if one exists. If not specified + - the model will get loaded in `torch.float` (fp32). + + 2. `"auto"` - A `torch_dtype` entry in the `config.json` file of the model will be + attempted to be used. If this entry isn't found then next check the `dtype` of the first weight in + the checkpoint that's of a floating point type and use that as `dtype`. This will load the model + using the `dtype` it was saved in at the end of the training. It can't be used as an indicator of how + the model was trained. Since it could be trained in one of half precision dtypes, but saved in fp32. + + 3. A string that is a valid `torch.dtype`. E.g. "float32" loads the model in `torch.float32`, "float16" loads in `torch.float16` etc. + + + + For some models the `dtype` they were trained in is unknown - you may try to check the model's paper or + reach out to the authors and ask them to add this information to the model's card and to insert the + `torch_dtype` entry in `config.json` on the hub. + + + + device_map (`str` or `Dict[str, Union[int, str, torch.device]]` or `int` or `torch.device`, *optional*): + A map that specifies where each submodule should go. It doesn't need to be refined to each + parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the + same device. If we only pass the device (*e.g.*, `"cpu"`, `"cuda:1"`, `"mps"`, or a GPU ordinal rank + like `1`) on which the model will be allocated, the device map will map the entire model to this + device. Passing `device_map = 0` means put the whole model on GPU 0. + + To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For + more information about each option see [designing a device + map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map). + max_memory (`Dict`, *optional*): + A dictionary device identifier to maximum memory if using `device_map`. Will default to the maximum memory available for each + GPU and the available CPU RAM if unset. + tp_plan (`str`, *optional*): + A torch tensor parallel plan, see [here](https://pytorch.org/tutorials/intermediate/TP_tutorial.html). Currently, it only accepts + `tp_plan="auto"` to use predefined plan based on the model. Note that if you use it, you should launch your script accordingly with + `torchrun [args] script.py`. This will be much faster than using a `device_map`, but has limitations. + offload_folder (`str` or `os.PathLike`, *optional*): + If the `device_map` contains any value `"disk"`, the folder where we will offload weights. + offload_state_dict (`bool`, *optional*): + If `True`, will temporarily offload the CPU state dict to the hard drive to avoid getting out of CPU + RAM if the weight of the CPU state dict + the biggest shard of the checkpoint does not fit. Defaults to + `True` when there is some disk offload. + offload_buffers (`bool`, *optional*): + Whether or not to offload the buffers with the model parameters. + quantization_config (`Union[QuantizationConfigMixin,Dict]`, *optional*): + A dictionary of configuration parameters or a QuantizationConfigMixin object for quantization (e.g + bitsandbytes, gptq). There may be other quantization-related kwargs, including `load_in_4bit` and + `load_in_8bit`, which are parsed by QuantizationConfigParser. Supported only for bitsandbytes + quantizations and not preferred. consider inserting all such arguments into quantization_config + instead. + 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. + variant (`str`, *optional*): + If specified load weights from `variant` filename, *e.g.* pytorch_model..bin. `variant` is + ignored when using `from_tf` or `from_flax`. + use_safetensors (`bool`, *optional*, defaults to `None`): + Whether or not to use `safetensors` checkpoints. Defaults to `None`. If not specified and `safetensors` + is not installed, it will be set to `False`. + weights_only (`bool`, *optional*, defaults to `True`): + Indicates whether unpickler should be restricted to loading only tensors, primitive types, + dictionaries and any types added via torch.serialization.add_safe_globals(). + When set to False, we can load wrapper tensor subclass weights. + key_mapping (`Dict[str, str], *optional*): + A potential mapping of the weight names if using a model on the Hub which is compatible to a Transformers + architecture, but was not converted accordingly. + kwargs (remaining dictionary of keyword arguments, *optional*): + Can be used to update the configuration object (after it being loaded) and initiate the model (e.g., + `output_attentions=True`). Behaves differently depending on whether a `config` is provided or + automatically loaded: + + - If a configuration is provided with `config`, `**kwargs` will be directly passed to the + underlying model's `__init__` method (we assume all relevant updates to the configuration have + already been done) + - If a configuration is not provided, `kwargs` will be first passed to the configuration class + initialization function ([`~PretrainedConfig.from_pretrained`]). Each key of `kwargs` that + corresponds to a configuration attribute will be used to override said attribute with the + supplied `kwargs` value. Remaining keys that do not correspond to any configuration attribute + will be passed to the underlying model's `__init__` function. + + + + Activate the special ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to + use this method in a firewalled environment. + + + + Examples: + + ```python + >>> from transformers import BertConfig, BertModel + + >>> # Download model and configuration from huggingface.co and cache. + >>> model = BertModel.from_pretrained("google-bert/bert-base-uncased") + >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable). + >>> model = BertModel.from_pretrained("./test/saved_model/") + >>> # Update configuration during loading. + >>> model = BertModel.from_pretrained("google-bert/bert-base-uncased", output_attentions=True) + >>> assert model.config.output_attentions == True + >>> # Loading from a TF checkpoint file instead of a PyTorch model (slower, for example purposes, not runnable). + >>> config = BertConfig.from_json_file("./tf_model/my_tf_model_config.json") + >>> model = BertModel.from_pretrained("./tf_model/my_tf_checkpoint.ckpt.index", from_tf=True, config=config) + >>> # Loading from a Flax checkpoint file instead of a PyTorch model (slower) + >>> model = BertModel.from_pretrained("google-bert/bert-base-uncased", from_flax=True) + ``` + """ + state_dict = kwargs.pop("state_dict", None) + from_tf = kwargs.pop("from_tf", False) + from_flax = kwargs.pop("from_flax", False) + proxies = kwargs.pop("proxies", None) + output_loading_info = kwargs.pop("output_loading_info", False) + use_auth_token = kwargs.pop("use_auth_token", None) + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + torch_dtype = kwargs.pop("torch_dtype", None) + device_map = kwargs.pop("device_map", None) + max_memory = kwargs.pop("max_memory", None) + offload_folder = kwargs.pop("offload_folder", None) + offload_state_dict = kwargs.pop("offload_state_dict", False) + offload_buffers = kwargs.pop("offload_buffers", False) + load_in_8bit = kwargs.pop("load_in_8bit", False) + load_in_4bit = kwargs.pop("load_in_4bit", False) + quantization_config = kwargs.pop("quantization_config", None) + subfolder = kwargs.pop("subfolder", "") + commit_hash = kwargs.pop("_commit_hash", None) + variant = kwargs.pop("variant", None) + adapter_kwargs = kwargs.pop("adapter_kwargs", {}) + adapter_name = kwargs.pop("adapter_name", "default") + use_flash_attention_2 = kwargs.pop("use_flash_attention_2", False) + generation_config = kwargs.pop("generation_config", None) + gguf_file = kwargs.pop("gguf_file", None) + tp_plan = kwargs.pop("tp_plan", None) + key_mapping = kwargs.pop("key_mapping", None) + # Not used anymore -- remove them from the kwargs + _ = kwargs.pop("resume_download", None) + _ = kwargs.pop("trust_remote_code", None) + _ = kwargs.pop("mirror", None) + _ = kwargs.pop("_fast_init", True) + _ = kwargs.pop("low_cpu_mem_usage", None) + + if state_dict is not None and (pretrained_model_name_or_path is not None or gguf_file is not None): + raise ValueError( + "`state_dict` cannot be passed together with a model name or a `gguf_file`. Use one of the two loading strategies." + ) + + if tp_plan is not None and tp_plan != "auto": + # TODO: we can relax this check when we support taking tp_plan from a json file, for example. + raise ValueError(f"tp_plan supports 'auto' only for now but got {tp_plan}.") + if tp_plan is not None and device_map is not None: + raise ValueError( + "`tp_plan` and `device_map` are mutually exclusive. Choose either one for parallelization." + ) + + # If torchrun was used, make sure to TP by default. This way people don't need to change tp or device map + if device_map == "auto" and tp_plan is None and int(os.environ.get("WORLD_SIZE", 0)): + tp_plan = "auto" # device_map = "auto" in torchrun equivalent to TP plan = AUTO! + device_map = None + + # We need to correctly dispatch the model on the current process device. The easiest way for this is to use a simple + # `device_map` pointing to the correct device + device_mesh = None + if tp_plan is not None: + if not is_torch_greater_or_equal("2.5"): + raise EnvironmentError("tensor parallel is only supported for `torch>=2.5`.") + + # Detect the accelerator on the machine. If no accelerator is available, it returns CPU. + device_type = torch._C._get_accelerator().type + + if not torch.distributed.is_initialized(): + try: + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + if device_type == "cuda": + torch.distributed.init_process_group( + "nccl", rank=rank, world_size=world_size, init_method="env://" + ) + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + elif device_type == "cpu": + cpu_backend = "ccl" if int(os.environ.get("CCL_WORKER_COUNT", 0)) else "gloo" + torch.distributed.init_process_group(cpu_backend, rank=rank, world_size=world_size) + + except Exception as e: + raise EnvironmentError( + "We tried to initialize torch.distributed for you, but it failed, make" + "sure you init torch distributed in your script to use `tp_plan='auto'`" + ) from e + + # Get device with index assuming equal number of devices per host + index = None if device_type == "cpu" else torch.cuda.current_device() + tp_device = torch.device(device_type, index) + + if index is not None and index > 0: + import sys + + sys.stdout = open(os.devnull, "w") + sys.stderr = open(os.devnull, "w") + # This is the easiest way to dispatch to the current process device + device_map = tp_device + # Assuming sharding the model onto the world + world_size = torch.distributed.get_world_size() + device_mesh = torch.distributed.init_device_mesh(tp_device.type, (world_size,)) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if token is not None and adapter_kwargs is not None and "token" not in adapter_kwargs: + adapter_kwargs["token"] = token + + if use_safetensors is None and not is_safetensors_available(): + use_safetensors = False + + if gguf_file is not None and not is_accelerate_available(): + raise ValueError("accelerate is required when loading a GGUF file `pip install accelerate`.") + + if commit_hash is None: + if not isinstance(config, PretrainedConfig): + # We make a call to the config file first (which may be absent) to get the commit hash as soon as possible + resolved_config_file = cached_file( + pretrained_model_name_or_path, + CONFIG_NAME, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder=subfolder, + _raise_exceptions_for_gated_repo=False, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + ) + commit_hash = extract_commit_hash(resolved_config_file, commit_hash) + else: + commit_hash = getattr(config, "_commit_hash", None) + + if is_peft_available(): + _adapter_model_path = adapter_kwargs.pop("_adapter_model_path", None) + + if _adapter_model_path is None: + _adapter_model_path = find_adapter_config_file( + pretrained_model_name_or_path, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + _commit_hash=commit_hash, + **adapter_kwargs, + ) + if _adapter_model_path is not None and os.path.isfile(_adapter_model_path): + with open(_adapter_model_path, "r", encoding="utf-8") as f: + _adapter_model_path = pretrained_model_name_or_path + pretrained_model_name_or_path = json.load(f)["base_model_name_or_path"] + else: + _adapter_model_path = None + + # change device_map into a map if we passed an int, a str or a torch.device + if isinstance(device_map, torch.device): + device_map = {"": device_map} + elif isinstance(device_map, str) and device_map not in ["auto", "balanced", "balanced_low_0", "sequential"]: + try: + device_map = {"": torch.device(device_map)} + except RuntimeError: + raise ValueError( + "When passing device_map as a string, the value needs to be a device name (e.g. cpu, cuda:0) or " + f"'auto', 'balanced', 'balanced_low_0', 'sequential' but found {device_map}." + ) + elif isinstance(device_map, int): + if device_map < 0: + raise ValueError( + "You can't pass device_map as a negative int. If you want to put the model on the cpu, pass device_map = 'cpu' " + ) + else: + device_map = {"": device_map} + + if device_map is not None: + if is_deepspeed_zero3_enabled(): + raise ValueError("DeepSpeed Zero-3 is not compatible with passing a `device_map`.") + if not is_accelerate_available(): + raise ValueError( + "Using a `device_map` or `tp_plan` requires `accelerate`. You can install it with `pip install accelerate`" + ) + + # handling bnb config from kwargs, remove after `load_in_{4/8}bit` deprecation. + if load_in_4bit or load_in_8bit: + if quantization_config is not None: + raise ValueError( + "You can't pass `load_in_4bit`or `load_in_8bit` as a kwarg when passing " + "`quantization_config` argument at the same time." + ) + + # preparing BitsAndBytesConfig from kwargs + config_dict = {k: v for k, v in kwargs.items() if k in inspect.signature(BitsAndBytesConfig).parameters} + config_dict = {**config_dict, "load_in_4bit": load_in_4bit, "load_in_8bit": load_in_8bit} + quantization_config, kwargs = BitsAndBytesConfig.from_dict( + config_dict=config_dict, return_unused_kwargs=True, **kwargs + ) + logger.warning( + "The `load_in_4bit` and `load_in_8bit` arguments are deprecated and will be removed in the future versions. " + "Please, pass a `BitsAndBytesConfig` object in `quantization_config` argument instead." + ) + + from_pt = not (from_tf | from_flax) + + user_agent = {"file_type": "model", "framework": "pytorch", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + # Load config if we don't provide a configuration + if not isinstance(config, PretrainedConfig): + config_path = config if config is not None else pretrained_model_name_or_path + config, model_kwargs = cls.config_class.from_pretrained( + config_path, + cache_dir=cache_dir, + return_unused_kwargs=True, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder=subfolder, + gguf_file=gguf_file, + _from_auto=from_auto_class, + _from_pipeline=from_pipeline, + **kwargs, + ) + if "gguf_file" in model_kwargs: + model_kwargs.pop("gguf_file") + else: + # In case one passes a config to `from_pretrained` + "attn_implementation" + # override the `_attn_implementation` attribute to `attn_implementation` of the kwargs + # Please see: https://github.com/huggingface/transformers/issues/28038 + + # Overwrite `config._attn_implementation` by the one from the kwargs --> in auto-factory + # we pop attn_implementation from the kwargs but this handles the case where users + # passes manually the config to `from_pretrained`. + config = copy.deepcopy(config) + + kwarg_attn_imp = kwargs.pop("attn_implementation", None) + if kwarg_attn_imp is not None: + config._attn_implementation = kwarg_attn_imp + + model_kwargs = kwargs + + pre_quantized = hasattr(config, "quantization_config") + if pre_quantized and not AutoHfQuantizer.supports_quant_method(config.quantization_config): + pre_quantized = False + + if pre_quantized or quantization_config is not None: + if pre_quantized: + config.quantization_config = AutoHfQuantizer.merge_quantization_configs( + config.quantization_config, quantization_config + ) + else: + config.quantization_config = quantization_config + + hf_quantizer = AutoHfQuantizer.from_config( + config.quantization_config, + pre_quantized=pre_quantized, + ) + else: + hf_quantizer = None + + if hf_quantizer is not None: + hf_quantizer.validate_environment( + torch_dtype=torch_dtype, + from_tf=from_tf, + from_flax=from_flax, + device_map=device_map, + weights_only=weights_only, + ) + torch_dtype = hf_quantizer.update_torch_dtype(torch_dtype) + device_map = hf_quantizer.update_device_map(device_map) + config = hf_quantizer.update_tp_plan(config) + + # In order to ensure popular quantization methods are supported. Can be disable with `disable_telemetry` + if hasattr(hf_quantizer.quantization_config.quant_method, "value"): + user_agent["quant"] = hf_quantizer.quantization_config.quant_method.value + else: + user_agent["quant"] = hf_quantizer.quantization_config.quant_method + + if gguf_file is not None and hf_quantizer is not None: + raise ValueError( + "You cannot combine Quantization and loading a model from a GGUF file, try again by making sure you did not passed a `quantization_config` or that you did not load a quantized model from the Hub." + ) + + if ( + gguf_file + and device_map is not None + and ((isinstance(device_map, dict) and "disk" in device_map.values()) or "disk" in device_map) + ): + raise RuntimeError( + "One or more modules is configured to be mapped to disk. Disk offload is not supported for models " + "loaded from GGUF files." + ) + + checkpoint_files, sharded_metadata = _get_resolved_checkpoint_files( + pretrained_model_name_or_path=pretrained_model_name_or_path, + subfolder=subfolder, + variant=variant, + gguf_file=gguf_file, + from_tf=from_tf, + from_flax=from_flax, + use_safetensors=use_safetensors, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + commit_hash=commit_hash, + ) + + is_sharded = sharded_metadata is not None + is_quantized = hf_quantizer is not None + is_from_file = pretrained_model_name_or_path is not None or gguf_file is not None + + if ( + is_safetensors_available() + and is_from_file + and not is_sharded + and checkpoint_files[0].endswith(".safetensors") + ): + with safe_open(checkpoint_files[0], framework="pt") as f: + metadata = f.metadata() + + if metadata is None: + # Assume it's a pytorch checkpoint (introduced for timm checkpoints) + pass + elif metadata.get("format") == "pt": + pass + elif metadata.get("format") == "tf": + from_tf = True + logger.info("A TensorFlow safetensors file is being loaded in a PyTorch model.") + elif metadata.get("format") == "flax": + from_flax = True + logger.info("A Flax safetensors file is being loaded in a PyTorch model.") + elif metadata.get("format") == "mlx": + # This is a mlx file, we assume weights are compatible with pt + pass + else: + raise ValueError( + f"Incompatible safetensors file. File metadata is not ['pt', 'tf', 'flax', 'mlx'] but {metadata.get('format')}" + ) + + from_pt = not (from_tf | from_flax) + + if from_pt: + if gguf_file: + from .modeling_gguf_pytorch_utils import load_gguf_checkpoint + + # we need a dummy model to get the state_dict - for this reason, we keep the state_dict as if it was + # passed directly as a kwarg from now on + with torch.device("meta"): + dummy_model = cls(config) + state_dict = load_gguf_checkpoint(checkpoint_files[0], return_tensors=True, model_to_load=dummy_model)[ + "tensors" + ] + + # Find the correct dtype based on current state + config, torch_dtype, dtype_orig = _get_torch_dtype( + cls, torch_dtype, checkpoint_files, config, sharded_metadata, state_dict, weights_only + ) + + config.name_or_path = pretrained_model_name_or_path + + # Instantiate model. + model_init_context = cls.get_init_context(is_quantized, _is_ds_init_called) + + config = copy.deepcopy(config) # We do not want to modify the config inplace in from_pretrained. + if not getattr(config, "_attn_implementation_autoset", False): + config = cls._autoset_attn_implementation( + config, use_flash_attention_2=use_flash_attention_2, torch_dtype=torch_dtype, device_map=device_map + ) + + with ContextManagers(model_init_context): + # Let's make sure we don't run the init function of buffer modules + model = cls(config, *model_args, **model_kwargs) + + # Make sure to tie the weights correctly + model.tie_weights() + + # Last check for tp + if device_mesh is not None and not model.supports_tp_plan: + if config.base_model_tp_plan is None and config.get_text_config().base_model_tp_plan is None: + raise NotImplementedError("This model does not have a tensor parallel plan.") + + # make sure we use the model's config since the __init__ call might have copied it + config = model.config + + # Find fp32 modules if needed + keep_in_fp32_regex = None + # The _keep_in_fp32_modules flag is only used to avoid bf16 -> fp16 casting precision issues. It was introduced + # in case of force loading a model that should stay bf16 in fp16 (which includes a few quantizers as this is a pre-processing + # step for e.g. bitsandbytes). See https://github.com/huggingface/transformers/issues/20287 for details. + if model._keep_in_fp32_modules is not None and ( + torch_dtype == torch.float16 or getattr(hf_quantizer, "use_keep_in_fp32_modules", False) + ): + # We need to match exact layers, so we add either `.` on each side, or start/end of string + keep_in_fp32_regex = re.compile( + "|".join([rf"((^|\.){module}($|\.))" for module in model._keep_in_fp32_modules]) + ) + + if hf_quantizer is not None: + hf_quantizer.preprocess_model( + model=model, device_map=device_map, keep_in_fp32_modules=model._keep_in_fp32_modules, config=config + ) + # We store the original dtype for quantized models as we cannot easily retrieve it + # once the weights have been quantized + # Note that once you have loaded a quantized model, you can't change its dtype so this will + # remain a single source of truth + config._pre_quantization_dtype = torch_dtype if torch_dtype is not None else torch.get_default_dtype() + + # Prepare the full device map + if device_map is not None: + device_map = _get_device_map(model, device_map, max_memory, hf_quantizer, torch_dtype, keep_in_fp32_regex) + + # Finalize model weight initialization + if from_tf: + model, loading_info = cls._load_from_tf(model, config, checkpoint_files) + elif from_flax: + model = cls._load_from_flax(model, checkpoint_files) + elif from_pt: + # restore default dtype + if dtype_orig is not None: + torch.set_default_dtype(dtype_orig) + + ( + model, + missing_keys, + unexpected_keys, + mismatched_keys, + offload_index, + error_msgs, + ) = cls._load_pretrained_model( + model, + state_dict, + checkpoint_files, + pretrained_model_name_or_path, + ignore_mismatched_sizes=ignore_mismatched_sizes, + sharded_metadata=sharded_metadata, + device_map=device_map, + disk_offload_folder=offload_folder, + offload_state_dict=offload_state_dict, + dtype=torch_dtype, + hf_quantizer=hf_quantizer, + keep_in_fp32_regex=keep_in_fp32_regex, + device_mesh=device_mesh, + key_mapping=key_mapping, + weights_only=weights_only, + ) + + # make sure token embedding weights are still tied if needed + model.tie_weights() + + # Set model in evaluation mode to deactivate DropOut modules by default + model.eval() + + # If it is a model with generation capabilities, attempt to load the generation config + if model.can_generate() and generation_config is not None: + logger.info("The user-defined `generation_config` will be used to override the default generation config.") + model.generation_config = model.generation_config.from_dict(generation_config.to_dict()) + elif model.can_generate() and pretrained_model_name_or_path is not None: + try: + model.generation_config = GenerationConfig.from_pretrained( + pretrained_model_name_or_path, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder=subfolder, + _from_auto=from_auto_class, + _from_pipeline=from_pipeline, + **kwargs, + ) + except OSError: + logger.info( + "Generation config file not found, using a generation config created from the model config." + ) + pass + + # Dispatch model with hooks on all devices if necessary (not needed with a tp_plan, so we skip it as it slightly + # harm performances) + if device_map is not None and device_mesh is None: + device_map_kwargs = { + "device_map": device_map, + "offload_dir": offload_folder, + "offload_index": offload_index, + "offload_buffers": offload_buffers, + } + if "skip_keys" in inspect.signature(dispatch_model).parameters: + device_map_kwargs["skip_keys"] = model._skip_keys_device_placement + # For HQQ method we force-set the hooks for single GPU envs + if ( + "force_hooks" in inspect.signature(dispatch_model).parameters + and hf_quantizer is not None + and hf_quantizer.quantization_config.quant_method == QuantizationMethod.HQQ + ): + device_map_kwargs["force_hooks"] = True + if ( + hf_quantizer is not None + and hf_quantizer.quantization_config.quant_method == QuantizationMethod.FBGEMM_FP8 + and isinstance(device_map, dict) + and ("cpu" in device_map.values() or "disk" in device_map.values()) + ): + device_map_kwargs["offload_buffers"] = True + + if not is_fsdp_enabled() and not is_deepspeed_zero3_enabled(): + dispatch_model(model, **device_map_kwargs) + + if hf_quantizer is not None: + hf_quantizer.postprocess_model(model, config=config) + model.hf_quantizer = hf_quantizer + + if _adapter_model_path is not None: + model.load_adapter( + _adapter_model_path, + adapter_name=adapter_name, + token=token, + adapter_kwargs=adapter_kwargs, + ) + + if output_loading_info: + if from_pt: + loading_info = { + "missing_keys": missing_keys, + "unexpected_keys": unexpected_keys, + "mismatched_keys": mismatched_keys, + "error_msgs": error_msgs, + } + elif from_flax: + loading_info = None + return model, loading_info + + return model + + @staticmethod + def _fix_state_dict_key_on_load(key: str) -> Tuple[str, bool]: + """Replace legacy parameter names with their modern equivalents. E.g. beta -> bias, gamma -> weight.""" + # Rename LayerNorm beta & gamma params for some early models ported from Tensorflow (e.g. Bert) + # This rename is logged. + if key.endswith("LayerNorm.beta"): + return key.replace("LayerNorm.beta", "LayerNorm.bias"), True + if key.endswith("LayerNorm.gamma"): + return key.replace("LayerNorm.gamma", "LayerNorm.weight"), True + + # Rename weight norm parametrizations to match changes across torch versions. + # Impacts a number of speech/wav2vec models. e.g. Hubert, Wav2Vec2, and others. + # This rename is not logged. + if hasattr(nn.utils.parametrizations, "weight_norm"): + if key.endswith("weight_g"): + return key.replace("weight_g", "parametrizations.weight.original0"), True + if key.endswith("weight_v"): + return key.replace("weight_v", "parametrizations.weight.original1"), True + else: + if key.endswith("parametrizations.weight.original0"): + return key.replace("parametrizations.weight.original0", "weight_g"), True + if key.endswith("parametrizations.weight.original1"): + return key.replace("parametrizations.weight.original1", "weight_v"), True + + return key, False + + def _get_key_renaming_mapping( + self, + checkpoint_keys: List[str], + key_mapping: Optional[Dict[str, str]] = None, + loading_base_model_from_task_state_dict: bool = False, + loading_task_model_from_base_state_dict: bool = False, + ): + """ + Compute a mapping between the serialized keys on disk `checkpoint_keys`, and the keys that the model + that we are loading expects. This is the single entry point for key renaming that will be used during + loading. + Log if any parameters have been renamed. + """ + prefix = self.base_model_prefix + _prefix = f"{prefix}." + + renamed_keys = {} + key_renaming_mapping = {} + for key in checkpoint_keys: + # Class specific rename + new_key, has_changed = self._fix_state_dict_key_on_load(key) + + # Optionally map the key according to `key_mapping` + if key_mapping is not None: + for pattern, replacement in key_mapping.items(): + new_key, n_replace = re.subn(pattern, replacement, new_key) + # Early exit of the loop + if n_replace > 0: + has_changed = True + break + + # In this case, we need to add the prefix to the keys, to match them to the expected keys + if loading_task_model_from_base_state_dict: + new_key = ".".join([prefix, new_key]) + # In this case we need to remove the prefix from the key to match them to the expected keys, and use + # only the keys starting with the prefix + elif loading_base_model_from_task_state_dict: + if not new_key.startswith(_prefix): + continue + new_key = new_key[len(_prefix) :] + + key_renaming_mapping[key] = new_key + + # track gamma/beta rename for logging + if has_changed: + if key.endswith("LayerNorm.gamma"): + renamed_keys["LayerNorm.gamma"] = (key, new_key) + elif key.endswith("LayerNorm.beta"): + renamed_keys["LayerNorm.beta"] = (key, new_key) + + if renamed_keys: + warning_msg = f"A pretrained model of type `{self.__class__.__name__}` " + warning_msg += "contains parameters that have been renamed internally (a few are listed below but more are present in the model):\n" + for old_key, new_key in renamed_keys.values(): + warning_msg += f"* `{old_key}` -> `{new_key}`\n" + warning_msg += "If you are using a model from the Hub, consider submitting a PR to adjust these weights and help future users." + logger.info_once(warning_msg) + + return key_renaming_mapping + + @staticmethod + def _fix_state_dict_key_on_save(key) -> Tuple[str, bool]: + """ + Similar to `_fix_state_dict_key_on_load` allows to define hook for state dict key renaming on model save. + Do nothing by default, but can be overridden in particular models. + """ + return key, False + + def _fix_state_dict_keys_on_save(self, state_dict): + """ + Similar to `_fix_state_dict_keys_on_load` allows to define hook for state dict key renaming on model save. + Apply `_fix_state_dict_key_on_save` to all keys in `state_dict`. + """ + return {self._fix_state_dict_key_on_save(key)[0]: value for key, value in state_dict.items()} + + @classmethod + def _load_pretrained_model( + cls, + model: "PreTrainedModel", + state_dict: Optional[Dict], + checkpoint_files: Optional[List[str]], + pretrained_model_name_or_path: Optional[str], + ignore_mismatched_sizes: bool = False, + sharded_metadata: Optional[Dict] = None, + device_map: Optional[Dict] = None, + disk_offload_folder: Optional[str] = None, + offload_state_dict: Optional[bool] = None, + dtype: Optional[torch.dtype] = None, + hf_quantizer: Optional[HfQuantizer] = None, + keep_in_fp32_regex: Optional[re.Pattern] = None, + device_mesh: Optional["torch.distributed.device_mesh.DeviceMesh"] = None, + key_mapping: Optional[Dict[str, str]] = None, + weights_only: bool = True, + ): + # Useful flags + is_quantized = hf_quantizer is not None + is_hqq = is_quantized and hf_quantizer.quantization_config.quant_method == QuantizationMethod.HQQ + is_hqq_or_bnb = is_quantized and hf_quantizer.quantization_config.quant_method in [ + QuantizationMethod.HQQ, + QuantizationMethod.BITS_AND_BYTES, + ] + + # Get all the keys of the state dicts that we have to initialize the model + if sharded_metadata is not None: + original_checkpoint_keys = sharded_metadata["all_checkpoint_keys"] + elif state_dict is not None: + original_checkpoint_keys = list(state_dict.keys()) + else: + original_checkpoint_keys = list( + load_state_dict(checkpoint_files[0], map_location="meta", weights_only=weights_only).keys() + ) + + # Check if we are in a special state, i.e. loading from a state dict coming from a different architecture + prefix = model.base_model_prefix + _prefix = f"{prefix}." + has_prefix_module = any(s.startswith(prefix) for s in original_checkpoint_keys) if len(prefix) > 0 else False + expects_prefix_module = hasattr(model, prefix) if len(prefix) > 0 else False + loading_task_model_from_base_state_dict = not has_prefix_module and expects_prefix_module + loading_base_model_from_task_state_dict = has_prefix_module and not expects_prefix_module + + # Find the key names that the model expects from the serialized keys + key_renaming_mapping = model._get_key_renaming_mapping( + original_checkpoint_keys, + key_mapping, + loading_base_model_from_task_state_dict, + loading_task_model_from_base_state_dict, + ) + checkpoint_keys = list(key_renaming_mapping.values()) + + # Find missing and unexpected keys from the state dict + missing_keys, unexpected_keys = _find_missing_and_unexpected_keys( + cls, + model, + original_checkpoint_keys, + checkpoint_keys, + loading_base_model_from_task_state_dict, + hf_quantizer, + device_map, + ) + # Find all the keys with shape mismatch (if we ignore the mismatch, the weights need to be newly initialized the + # same way as missing keys) + mismatched_keys, mismatched_shapes = _find_mismatched_keys( + model, + state_dict, + checkpoint_files, + ignore_mismatched_sizes, + key_renaming_mapping, + is_quantized, + weights_only, + ) + + # We need to update both the mapping and the list of checkpoint keys to remove the mismatched ones + key_renaming_mapping = {k: v for k, v in key_renaming_mapping.items() if v not in mismatched_keys} + checkpoint_keys = list(key_renaming_mapping.values()) + + # Move missing (and potentially mismatched) keys back to cpu from meta device (because they won't be moved when + # loading the weights as they are not in the loaded state dict) + model._move_missing_keys_from_meta_to_cpu(missing_keys + mismatched_keys, unexpected_keys, dtype, hf_quantizer) + + # correctly initialize the missing (and potentially mismatched) keys + model._initialize_missing_keys(checkpoint_keys, ignore_mismatched_sizes, is_quantized) + + # Set some modules to fp32 if needed + if keep_in_fp32_regex is not None: + for name, param in model.named_parameters(): + if keep_in_fp32_regex.search(name): + # param = param.to(torch.float32) does not work here as only in the local scope. + param.data = param.data.to(torch.float32) + + # Make sure we are able to load base models as well as derived models (specific task models, with heads) + model_to_load = model + # In this case, we load a ForTaskModel with keys from a BaseModel -> only load keys to the BaseModel + if loading_task_model_from_base_state_dict: + model_to_load = getattr(model, prefix) + # Here we need to remove the prefix we added to correctly find missing/unexpected keys, as we will load + # in the submodule + key_renaming_mapping = {k: v[len(_prefix) :] for k, v in key_renaming_mapping.items()} + checkpoint_keys = list(key_renaming_mapping.values()) + # We need to update the device map as well + if device_map is not None: + device_map = {k[len(_prefix) :] if k.startswith(_prefix) else k: v for k, v in device_map.items()} + # small sanity check: the base model should not contain task-specific head keys + task_specific_expected_keys = [s for s in model.state_dict().keys() if not s.startswith(_prefix)] + base_model_expected_keys = list(model_to_load.state_dict().keys()) + if any( + key in task_specific_expected_keys and key not in base_model_expected_keys for key in checkpoint_keys + ): + raise ValueError( + "The state dictionary of the model you are trying to load is corrupted. Are you sure it was " + "properly saved?" + ) + + # Get reverse key mapping + reverse_key_renaming_mapping = {v: k for k, v in key_renaming_mapping.items()} + + is_offloaded_safetensors = False + # This offload index if for params explicitly on the "disk" in the device_map + disk_offload_index = None + disk_only_shard_files = [] + # Prepare parameters offloading if needed + if device_map is not None and "disk" in device_map.values(): + if offload_state_dict is None: + offload_state_dict = True + if disk_offload_folder is not None: + os.makedirs(disk_offload_folder, exist_ok=True) + is_offloaded_safetensors = checkpoint_files is not None and checkpoint_files[0].endswith(".safetensors") + if disk_offload_folder is None and not is_offloaded_safetensors: + raise ValueError( + "The current `device_map` had weights offloaded to the disk. Please provide an `offload_folder`" + " for them. Alternatively, make sure you have `safetensors` installed if the model you are using" + " offers the weights in this format." + ) + if is_offloaded_safetensors: + param_device_map = expand_device_map(device_map, checkpoint_keys) + str_dtype = str(dtype).replace("torch.", "") if dtype is not None else "float32" + if sharded_metadata is None: + weight_map = dict.fromkeys(checkpoint_keys, checkpoint_files[0]) + else: + folder = os.path.sep.join(checkpoint_files[0].split(os.path.sep)[:-1]) + # Fix the weight map keys according to the key mapping + weight_map = { + key_renaming_mapping[k]: v + for k, v in sharded_metadata["weight_map"].items() + if k in key_renaming_mapping + } + weight_map = {k: os.path.join(folder, v) for k, v in weight_map.items()} + # Find potential checkpoints containing only offloaded weights + disk_only_shard_files = get_disk_only_shard_files(device_map, weight_map) + disk_offload_index = { + name: { + "safetensors_file": file, + "weight_name": reverse_key_renaming_mapping[name], + "dtype": str_dtype, + } + for name, file in weight_map.items() + if param_device_map[name] == "disk" + } + else: + disk_offload_index = {} + + # This offload index if for params that are supposed to be on the "cpu", either with or without a device_map + # It allows to load parameters one-by-one from the state dict, avoiding a memory peak of 2 x state_dict_size, + # i.e. 1x to load it, and 1x to copy it to model + cpu_offload_folder = None + cpu_offload_index = None + if offload_state_dict: + cpu_offload_folder = tempfile.mkdtemp() + cpu_offload_index = {} + + # For nice tqdm bars + if checkpoint_files is not None and len(checkpoint_files) > 1: + checkpoint_files = logging.tqdm(checkpoint_files, desc="Loading checkpoint shards") + # To be able to iterate, even if we don't use it if the state_dict is already provided + elif state_dict is not None: + checkpoint_files = [""] + + # Compute expected model keys + expected_keys = list(model_to_load.state_dict().keys()) + if hf_quantizer is not None: + expected_keys = hf_quantizer.update_expected_keys(model_to_load, expected_keys, checkpoint_keys) + + # Warmup cuda to load the weights much faster on devices + if device_map is not None and not is_hqq: + expanded_device_map = expand_device_map(device_map, expected_keys) + caching_allocator_warmup(model_to_load, expanded_device_map, factor=2 if hf_quantizer is None else 4) + + error_msgs = [] + # Iterate on all the shards to load the weights + for shard_file in checkpoint_files: + # Skip the load for shards that only contain disk-offloaded weights + if shard_file in disk_only_shard_files: + continue + + map_location = "cpu" + if ( + shard_file.endswith(".safetensors") + and not is_hqq_or_bnb + and not (is_deepspeed_zero3_enabled() and not is_quantized) + ): + map_location = "meta" + elif ( + device_map is not None + and hf_quantizer is not None + and hf_quantizer.quantization_config.quant_method == QuantizationMethod.TORCHAO + and ( + hf_quantizer.quantization_config.quant_type in ["int4_weight_only", "autoquant"] + or isinstance(hf_quantizer.quantization_config.quant_type, Int4WeightOnlyConfig) + ) + ): + map_location = torch.device([d for d in device_map.values() if d not in ["cpu", "disk"]][0]) + + # If shard_file is "", we use the existing state_dict instead of loading it + if shard_file != "": + state_dict = load_state_dict( + shard_file, is_quantized=is_quantized, map_location=map_location, weights_only=weights_only + ) + + # Fix the key names + state_dict = {key_renaming_mapping[k]: v for k, v in state_dict.items() if k in key_renaming_mapping} + + if is_deepspeed_zero3_enabled() and not is_quantized: + error_msgs += _load_state_dict_into_zero3_model(model_to_load, state_dict) + # Skip it with fsdp on ranks other than 0 + elif not (is_fsdp_enabled() and not is_local_dist_rank_0() and not is_quantized): + disk_offload_index, cpu_offload_index = _load_state_dict_into_meta_model( + model_to_load, + state_dict, + shard_file, + expected_keys, + reverse_key_renaming_mapping, + device_map=device_map, + disk_offload_folder=disk_offload_folder, + disk_offload_index=disk_offload_index, + cpu_offload_folder=cpu_offload_folder, + cpu_offload_index=cpu_offload_index, + hf_quantizer=hf_quantizer, + is_safetensors=is_offloaded_safetensors, + keep_in_fp32_regex=keep_in_fp32_regex, + unexpected_keys=unexpected_keys, + device_mesh=device_mesh, + ) + + # force memory release if loading multiple shards, to avoid having 2 state dicts in memory in next loop + del state_dict + + # Adjust offloaded weights name and save if needed + if disk_offload_index is not None and len(disk_offload_index) > 0: + if loading_task_model_from_base_state_dict: + # We need to add the prefix of the base model + prefix = cls.base_model_prefix + if not is_offloaded_safetensors: + for weight_name in disk_offload_index: + shutil.move( + os.path.join(disk_offload_folder, f"{weight_name}.dat"), + os.path.join(disk_offload_folder, f"{prefix}.{weight_name}.dat"), + ) + disk_offload_index = {f"{prefix}.{key}": value for key, value in disk_offload_index.items()} + if not is_offloaded_safetensors: + save_offload_index(disk_offload_index, disk_offload_folder) + disk_offload_index = None + + # one-at-a-time param loading for the cpu offloaded params + if offload_state_dict: + # Load back temporarily offloaded state dict + load_offloaded_weights(model_to_load, cpu_offload_index, cpu_offload_folder) + shutil.rmtree(cpu_offload_folder) + + if hf_quantizer is not None: + missing_keys = hf_quantizer.update_missing_keys_after_loading(model_to_load, missing_keys, prefix) + + # Post-processing for tensor parallelism + if device_mesh is not None: + # When using TP, the device map is a single device for all parameters + tp_device = list(device_map.values())[0] + # This is needed for the RotaryEmbedding, which was not initialized on the correct device as it is + # not part of the state_dict (persistent=False) + for buffer in model.buffers(): + if buffer.device != tp_device: + buffer.data = buffer.to(tp_device) + + # In this case, the top-most task module weights were not moved to device and parallelized as they + # were not part of the loaded weights: do it now + if loading_task_model_from_base_state_dict: + parameters_to_initialize = { + name: param for name, param in model.named_parameters() if not name.startswith(prefix) + } + for name, param in parameters_to_initialize.items(): + # First move data to correct + to_contiguous, casting_dtype = _infer_parameter_dtype(model, name, param, keep_in_fp32_regex) + shard_and_distribute_module( + model, + param.to(tp_device), + param, + name, + casting_dtype, + to_contiguous, + os.environ["RANK"], + device_mesh, + ) + + # All potential warnings/infos + if len(error_msgs) > 0: + error_msg = "\n\t".join(error_msgs) + if "size mismatch" in error_msg: + error_msg += ( + "\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method." + ) + raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}") + if len(unexpected_keys) > 0: + archs = [] if model.config.architectures is None else model.config.architectures + warner = logger.warning if model.__class__.__name__ in archs else logger.info + warner( + f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when" + f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are" + f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or" + " with another architecture (e.g. initializing a BertForSequenceClassification model from a" + " BertForPreTraining model).\n- This IS NOT expected if you are initializing" + f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical" + " (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)." + ) + else: + logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n") + if len(missing_keys) > 0: + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably" + " TRAIN this model on a down-stream task to be able to use it for predictions and inference." + ) + elif len(mismatched_keys) == 0: + logger.info( + f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at" + f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint" + f" was trained on, you can already use {model.__class__.__name__} for predictions without further" + " training." + ) + if len(mismatched_keys) > 0: + mismatched_warning = "\n".join( + [ + f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated" + for key, (shape1, shape2) in zip(mismatched_keys, mismatched_shapes) + ] + ) + logger.warning( + f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at" + f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not" + f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able" + " to use it for predictions and inference." + ) + + return model, missing_keys, unexpected_keys, mismatched_keys, disk_offload_index, error_msgs + + @classmethod + def _load_from_tf(cls, model, config, checkpoint_files): + if checkpoint_files[0].endswith(".index"): + # Load from a TensorFlow 1.X checkpoint - provided by original authors + model = cls.load_tf_weights(model, config, checkpoint_files[0][:-6]) # Remove the '.index' + loading_info = None + else: + # Load from our TensorFlow 2.0 checkpoints + try: + from .modeling_tf_pytorch_utils import load_tf2_checkpoint_in_pytorch_model + + model, loading_info = load_tf2_checkpoint_in_pytorch_model( + model, checkpoint_files[0], allow_missing_keys=True, output_loading_info=True + ) + except ImportError: + logger.error( + "Loading a TensorFlow model in PyTorch, requires both PyTorch and TensorFlow to be installed." + " Please see https://pytorch.org/ and https://www.tensorflow.org/install/ for installation" + " instructions." + ) + raise + return model, loading_info + + @classmethod + def _load_from_flax(cls, model, checkpoint_files): + try: + from .modeling_flax_pytorch_utils import load_flax_checkpoint_in_pytorch_model + + model = load_flax_checkpoint_in_pytorch_model(model, checkpoint_files[0]) + except ImportError: + logger.error( + "Loading a Flax model in PyTorch, requires both PyTorch and Flax to be installed. Please see" + " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for" + " installation instructions." + ) + raise + return model + + def retrieve_modules_from_names(self, names, add_prefix=False, remove_prefix=False): + module_keys = {".".join(key.split(".")[:-1]) for key in names} + + # torch.nn.ParameterList is a special case where two parameter keywords + # are appended to the module name, *e.g.* bert.special_embeddings.0 + module_keys = module_keys.union( + {".".join(key.split(".")[:-2]) for key in names if len(key) > 0 and key[-1].isdigit()} + ) + + retrieved_modules = [] + # retrieve all modules that has at least one missing weight name + for name, module in self.named_modules(): + if remove_prefix: + _prefix = f"{self.base_model_prefix}." + name = name[len(_prefix) :] if name.startswith(_prefix) else name + elif add_prefix: + name = ".".join([self.base_model_prefix, name]) if len(name) > 0 else self.base_model_prefix + + if name in module_keys: + retrieved_modules.append(module) + + return retrieved_modules + + @classmethod + def register_for_auto_class(cls, auto_class="AutoModel"): + """ + Register this class with a given auto class. This should only be used for custom models as the ones in the + library are already mapped with an auto class. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"AutoModel"`): + The auto class to register this new model with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + def to_bettertransformer(self) -> "PreTrainedModel": + """ + Converts the model to use [PyTorch's native attention + implementation](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html), integrated to + Transformers through [Optimum library](https://huggingface.co/docs/optimum/bettertransformer/overview). Only a + subset of all Transformers models are supported. + + PyTorch's attention fastpath allows to speed up inference through kernel fusions and the use of [nested + tensors](https://pytorch.org/docs/stable/nested.html). Detailed benchmarks can be found in [this blog + post](https://medium.com/pytorch/bettertransformer-out-of-the-box-performance-for-huggingface-transformers-3fbe27d50ab2). + + Returns: + [`PreTrainedModel`]: The model converted to BetterTransformer. + """ + if not is_optimum_available(): + raise ImportError("The package `optimum` is required to use Better Transformer.") + + from optimum.version import __version__ as optimum_version + + if version.parse(optimum_version) < version.parse("1.7.0"): + raise ImportError( + f"Please install optimum>=1.7.0 to use Better Transformer. The version {optimum_version} was found." + ) + + from optimum.bettertransformer import BetterTransformer + + return BetterTransformer.transform(self) + + def reverse_bettertransformer(self): + """ + Reverts the transformation from [`~PreTrainedModel.to_bettertransformer`] so that the original modeling is + used, for example in order to save the model. + + Returns: + [`PreTrainedModel`]: The model converted back to the original modeling. + """ + if not is_optimum_available(): + raise ImportError("The package `optimum` is required to use Better Transformer.") + + from optimum.version import __version__ as optimum_version + + if version.parse(optimum_version) < version.parse("1.7.0"): + raise ImportError( + f"Please install optimum>=1.7.0 to use Better Transformer. The version {optimum_version} was found." + ) + + from optimum.bettertransformer import BetterTransformer + + return BetterTransformer.reverse(self) + + def warn_if_padding_and_no_attention_mask(self, input_ids, attention_mask): + """ + Shows a one-time warning if the input_ids appear to contain padding and no attention mask was given. + """ + + # Skip the check during tracing. + if is_torch_fx_proxy(input_ids) or torch.jit.is_tracing() or is_torchdynamo_compiling(): + return + + if (attention_mask is not None) or (self.config.pad_token_id is None): + return + + # Check only the first and last input IDs to reduce overhead. + if self.config.pad_token_id in input_ids[:, [-1, 0]]: + warn_string = ( + "We strongly recommend passing in an `attention_mask` since your input_ids may be padded. See " + "https://huggingface.co/docs/transformers/troubleshooting" + "#incorrect-output-when-padding-tokens-arent-masked." + ) + + # If the pad token is equal to either BOS, EOS, or SEP, we do not know whether the user should use an + # attention_mask or not. In this case, we should still show a warning because this is a rare case. + if ( + (self.config.bos_token_id is not None and self.config.bos_token_id == self.config.pad_token_id) + or (self.config.eos_token_id is not None and self.config.eos_token_id == self.config.pad_token_id) + or (self.config.sep_token_id is not None and self.config.sep_token_id == self.config.pad_token_id) + ): + warn_string += ( + f"\nYou may ignore this warning if your `pad_token_id` ({self.config.pad_token_id}) is identical " + f"to the `bos_token_id` ({self.config.bos_token_id}), `eos_token_id` ({self.config.eos_token_id}), " + f"or the `sep_token_id` ({self.config.sep_token_id}), and your input is not padded." + ) + + logger.warning_once(warn_string) + + @property + def supports_tp_plan(self): + """ + Returns whether the model has a tensor parallelism plan. + """ + if self._tp_plan is not None: + return True + # Check if base model has a TP plan + if getattr(self.base_model, "_tp_plan", None) is not None: + return True + return False + + @property + def supports_pp_plan(self): + if self._pp_plan is not None: + return True + # Check if base model has PP plan + if getattr(self.base_model, "_pp_plan", None) is not None: + return True + return False + + @property + def loss_function(self): + if hasattr(self, "_loss_function"): + return self._loss_function + + loss_type = getattr(self, "loss_type", None) + + if loss_type is None or loss_type not in LOSS_MAPPING: + logger.warning_once( + f"`loss_type={loss_type}` was set in the config but it is unrecognised." + f"Using the default loss: `ForCausalLMLoss`." + ) + loss_type = "ForCausalLM" + return LOSS_MAPPING[loss_type] + + @loss_function.setter + def loss_function(self, value): + self._loss_function = value + + def get_compiled_call(self, compile_config: CompileConfig): + """Return a `torch.compile`'d version of `self.__call__`. This is useful to dynamically choose between + non-compiled/compiled `forward` during inference, especially to switch between prefill (where we don't + want to use compiled version to avoid recomputing the graph with new shapes) and iterative decoding + (where we want the speed-ups of compiled version with static shapes).""" + # Only reset it if not present or different from previous config + if "llama4" in self.config.model_type: # TODO try to enable for FULL COMPILE HYBRID CACHE SUPPORT + return self.__call__ + default_config = getattr(self.generation_config, "compile_config", CompileConfig()) + if ( + not hasattr(self, "_compiled_call") + or getattr(self, "_last_compile_config", default_config) != compile_config + ): + self._last_compile_config = compile_config + self._compiled_call = torch.compile(self.__call__, **compile_config.to_dict()) + return self._compiled_call + + @classmethod + def is_backend_compatible(cls): + return cls._supports_attention_backend + + def _move_missing_keys_from_meta_to_cpu( + self, + missing_keys: List[str], + unexpected_keys: List[str], + dtype: Optional[torch.dtype], + hf_quantizer: Optional[HfQuantizer], + ) -> "PreTrainedModel": + """Move the missing keys (keys that are part of the model parameters, but were NOT found in the loaded state dicts) back + from meta device to cpu. + """ + is_quantized = hf_quantizer is not None + + # In this case we need to move everything back + if is_fsdp_enabled() and not is_local_dist_rank_0() and not is_quantized: + # We only do it for the parameters, as the buffers are not initialized on the meta device by default + for key, param in self.named_parameters(): + value = torch.empty_like(param, dtype=dtype, device="cpu") + _load_parameter_into_model(self, key, value) + return + + model_state_dict = self.state_dict() + for key in missing_keys: + param = model_state_dict[key] + # Buffers are not initialized on the meta device, so we still need this check to avoid overwriting them + if param.device == torch.device("meta"): + value = torch.empty_like(param, dtype=dtype, device="cpu") + if ( + not is_quantized + or (getattr(hf_quantizer, "requires_parameters_quantization", False)) + or not hf_quantizer.check_quantized_param(self, param_value=value, param_name=key, state_dict={}) + ): + _load_parameter_into_model(self, key, value) + else: + hf_quantizer.create_quantized_param(self, value, key, "cpu", model_state_dict, unexpected_keys) + + def _initialize_missing_keys( + self, + loaded_keys: List[str], + ignore_mismatched_sizes: bool, + is_quantized: bool, + ) -> "PreTrainedModel": + """Initialize the missing keys (keys that are part of the model parameters, but were NOT found in the loaded state dicts), according to + `_initialize_weights`. Indeed, since the corresponding weights are missing from the state dict, they will not be replaced and need to + be initialized correctly (i.e. weight initialization distribution). + Also take care of setting the `_is_hf_initialized` flag for keys that are not missing. + """ + if not ignore_mismatched_sizes: + not_initialized_submodules = set_initialized_submodules(self, loaded_keys) + # If we're about to tie the output embeds to the input embeds we don't need to init them + if ( + hasattr(self.config.get_text_config(decoder=True), "tie_word_embeddings") + and self.config.get_text_config(decoder=True).tie_word_embeddings + ): + output_embeddings = self.get_output_embeddings() + if output_embeddings is not None: + # Still need to initialize if there is a bias term since biases are not tied. + if not hasattr(output_embeddings, "bias") or output_embeddings.bias is None: + output_embeddings._is_hf_initialized = True + else: + not_initialized_submodules = dict(self.named_modules()) + # This will only initialize submodules that are not marked as initialized by the line above. + if is_deepspeed_zero3_enabled() and not is_quantized: + not_initialized_parameters = list( + set( + itertools.chain.from_iterable( + submodule.parameters(recurse=False) for submodule in not_initialized_submodules.values() + ) + ) + ) + with deepspeed.zero.GatheredParameters(not_initialized_parameters, modifier_rank=0): + self.apply(self._initialize_weights) + else: + self.apply(self._initialize_weights) + + def get_parameter_or_buffer(self, target: str): + """ + Return the parameter or buffer given by `target` if it exists, otherwise throw an error. This combines + `get_parameter()` and `get_buffer()` in a single handy function. Note that it only work if `target` is a + leaf of the model. + """ + try: + return self.get_parameter(target) + except AttributeError: + pass + try: + return self.get_buffer(target) + except AttributeError: + pass + raise AttributeError(f"`{target}` is neither a parameter nor a buffer.") + + +PreTrainedModel.push_to_hub = copy_func(PreTrainedModel.push_to_hub) +if PreTrainedModel.push_to_hub.__doc__ is not None: + PreTrainedModel.push_to_hub.__doc__ = PreTrainedModel.push_to_hub.__doc__.format( + object="model", object_class="AutoModel", object_files="model file" + ) + + +class PoolerStartLogits(nn.Module): + """ + Compute SQuAD start logits from sequence hidden states. + + Args: + config ([`PretrainedConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model. + """ + + def __init__(self, config: PretrainedConfig): + super().__init__() + self.dense = nn.Linear(config.hidden_size, 1) + + def forward( + self, hidden_states: torch.FloatTensor, p_mask: Optional[torch.FloatTensor] = None + ) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + + Returns: + `torch.FloatTensor`: The start logits for SQuAD. + """ + x = self.dense(hidden_states).squeeze(-1) + + if p_mask is not None: + if get_parameter_dtype(self) == torch.float16: + x = x * (1 - p_mask) - 65500 * p_mask + else: + x = x * (1 - p_mask) - 1e30 * p_mask + + return x + + +class PoolerEndLogits(nn.Module): + """ + Compute SQuAD end logits from sequence hidden states. + + Args: + config ([`PretrainedConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps` + to use. + """ + + def __init__(self, config: PretrainedConfig): + super().__init__() + self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size) + self.activation = nn.Tanh() + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dense_1 = nn.Linear(config.hidden_size, 1) + + def forward( + self, + hidden_states: torch.FloatTensor, + start_states: Optional[torch.FloatTensor] = None, + start_positions: Optional[torch.LongTensor] = None, + p_mask: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*): + The hidden states of the first tokens for the labeled span. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The position of the first token for the labeled span. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + + + + One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides + `start_states`. + + + + Returns: + `torch.FloatTensor`: The end logits for SQuAD. + """ + assert start_states is not None or start_positions is not None, ( + "One of start_states, start_positions should be not None" + ) + if start_positions is not None: + slen, hsz = hidden_states.shape[-2:] + start_positions = start_positions[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + start_states = hidden_states.gather(-2, start_positions) # shape (bsz, 1, hsz) + start_states = start_states.expand(-1, slen, -1) # shape (bsz, slen, hsz) + + x = self.dense_0(torch.cat([hidden_states, start_states], dim=-1)) + x = self.activation(x) + x = self.LayerNorm(x) + x = self.dense_1(x).squeeze(-1) + + if p_mask is not None: + if get_parameter_dtype(self) == torch.float16: + x = x * (1 - p_mask) - 65500 * p_mask + else: + x = x * (1 - p_mask) - 1e30 * p_mask + + return x + + +class PoolerAnswerClass(nn.Module): + """ + Compute SQuAD 2.0 answer class from classification and start tokens hidden states. + + Args: + config ([`PretrainedConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model. + """ + + def __init__(self, config): + super().__init__() + self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size) + self.activation = nn.Tanh() + self.dense_1 = nn.Linear(config.hidden_size, 1, bias=False) + + def forward( + self, + hidden_states: torch.FloatTensor, + start_states: Optional[torch.FloatTensor] = None, + start_positions: Optional[torch.LongTensor] = None, + cls_index: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + The final hidden states of the model. + start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*): + The hidden states of the first tokens for the labeled span. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + The position of the first token for the labeled span. + cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Position of the CLS token for each sentence in the batch. If `None`, takes the last token. + + + + One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides + `start_states`. + + + + Returns: + `torch.FloatTensor`: The SQuAD 2.0 answer class. + """ + # No dependency on end_feature so that we can obtain one single `cls_logits` for each sample. + hsz = hidden_states.shape[-1] + assert start_states is not None or start_positions is not None, ( + "One of start_states, start_positions should be not None" + ) + if start_positions is not None: + start_positions = start_positions[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + start_states = hidden_states.gather(-2, start_positions).squeeze(-2) # shape (bsz, hsz) + + if cls_index is not None: + cls_index = cls_index[:, None, None].expand(-1, -1, hsz) # shape (bsz, 1, hsz) + cls_token_state = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, hsz) + else: + cls_token_state = hidden_states[:, -1, :] # shape (bsz, hsz) + + x = self.dense_0(torch.cat([start_states, cls_token_state], dim=-1)) + x = self.activation(x) + x = self.dense_1(x).squeeze(-1) + + return x + + +@dataclass +class SquadHeadOutput(ModelOutput): + """ + Base class for outputs of question answering models using a [`~modeling_utils.SQuADHead`]. + + Args: + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned if both `start_positions` and `end_positions` are provided): + Classification loss as the sum of start token, end token (and is_impossible if provided) classification + losses. + start_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the top config.start_n_top start token possibilities (beam-search). + start_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Indices for the top config.start_n_top start token possibilities (beam-search). + end_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the top `config.start_n_top * config.end_n_top` end token possibilities + (beam-search). + end_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Indices for the top `config.start_n_top * config.end_n_top` end token possibilities (beam-search). + cls_logits (`torch.FloatTensor` of shape `(batch_size,)`, *optional*, returned if `start_positions` or `end_positions` is not provided): + Log probabilities for the `is_impossible` label of the answers. + + """ + + loss: Optional[torch.FloatTensor] = None + start_top_log_probs: Optional[torch.FloatTensor] = None + start_top_index: Optional[torch.LongTensor] = None + end_top_log_probs: Optional[torch.FloatTensor] = None + end_top_index: Optional[torch.LongTensor] = None + cls_logits: Optional[torch.FloatTensor] = None + + +class SQuADHead(nn.Module): + r""" + A SQuAD head inspired by XLNet. + + Args: + config ([`PretrainedConfig`]): + The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps` + to use. + """ + + def __init__(self, config): + super().__init__() + self.start_n_top = config.start_n_top + self.end_n_top = config.end_n_top + + self.start_logits = PoolerStartLogits(config) + self.end_logits = PoolerEndLogits(config) + self.answer_class = PoolerAnswerClass(config) + + @replace_return_docstrings(output_type=SquadHeadOutput, config_class=PretrainedConfig) + def forward( + self, + hidden_states: torch.FloatTensor, + start_positions: Optional[torch.LongTensor] = None, + end_positions: Optional[torch.LongTensor] = None, + cls_index: Optional[torch.LongTensor] = None, + is_impossible: Optional[torch.LongTensor] = None, + p_mask: Optional[torch.FloatTensor] = None, + return_dict: bool = False, + ) -> Union[SquadHeadOutput, Tuple[torch.FloatTensor]]: + """ + Args: + hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`): + Final hidden states of the model on the sequence tokens. + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Positions of the first token for the labeled span. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Positions of the last token for the labeled span. + cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Position of the CLS token for each sentence in the batch. If `None`, takes the last token. + is_impossible (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Whether the question has a possible answer in the paragraph or not. + p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*): + Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token + should be masked. + return_dict (`bool`, *optional*, defaults to `False`): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + + Returns: + """ + start_logits = self.start_logits(hidden_states, p_mask=p_mask) + + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, let's remove the dimension added by batch splitting + for x in (start_positions, end_positions, cls_index, is_impossible): + if x is not None and x.dim() > 1: + x.squeeze_(-1) + + # during training, compute the end logits based on the ground truth of the start position + end_logits = self.end_logits(hidden_states, start_positions=start_positions, p_mask=p_mask) + + loss_fct = CrossEntropyLoss() + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if cls_index is not None and is_impossible is not None: + # Predict answerability from the representation of CLS and START + cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index) + loss_fct_cls = nn.BCEWithLogitsLoss() + cls_loss = loss_fct_cls(cls_logits, is_impossible) + + # note(zhiliny): by default multiply the loss by 0.5 so that the scale is comparable to start_loss and end_loss + total_loss += cls_loss * 0.5 + + return SquadHeadOutput(loss=total_loss) if return_dict else (total_loss,) + + else: + # during inference, compute the end logits based on beam search + bsz, slen, hsz = hidden_states.size() + start_log_probs = nn.functional.softmax(start_logits, dim=-1) # shape (bsz, slen) + + start_top_log_probs, start_top_index = torch.topk( + start_log_probs, self.start_n_top, dim=-1 + ) # shape (bsz, start_n_top) + start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz) # shape (bsz, start_n_top, hsz) + start_states = torch.gather(hidden_states, -2, start_top_index_exp) # shape (bsz, start_n_top, hsz) + start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1) # shape (bsz, slen, start_n_top, hsz) + + hidden_states_expanded = hidden_states.unsqueeze(2).expand_as( + start_states + ) # shape (bsz, slen, start_n_top, hsz) + p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None + end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask) + end_log_probs = nn.functional.softmax(end_logits, dim=1) # shape (bsz, slen, start_n_top) + + end_top_log_probs, end_top_index = torch.topk( + end_log_probs, self.end_n_top, dim=1 + ) # shape (bsz, end_n_top, start_n_top) + end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top) + end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top) + + start_states = torch.einsum("blh,bl->bh", hidden_states, start_log_probs) + cls_logits = self.answer_class(hidden_states, start_states=start_states, cls_index=cls_index) + + if not return_dict: + return (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits) + else: + return SquadHeadOutput( + start_top_log_probs=start_top_log_probs, + start_top_index=start_top_index, + end_top_log_probs=end_top_log_probs, + end_top_index=end_top_index, + cls_logits=cls_logits, + ) + + +class SequenceSummary(nn.Module): + r""" + Compute a single vector summary of a sequence hidden states. + + Args: + config ([`PretrainedConfig`]): + The config used by the model. Relevant arguments in the config class of the model are (refer to the actual + config class of your model for the default values it uses): + + - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are: + + - `"last"` -- Take the last token hidden state (like XLNet) + - `"first"` -- Take the first token hidden state (like Bert) + - `"mean"` -- Take the mean of all tokens hidden states + - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2) + - `"attn"` -- Not implemented now, use multi-head attention + + - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction. + - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes + (otherwise to `config.hidden_size`). + - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output, + another string or `None` will add no activation. + - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation. + - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation. + """ + + def __init__(self, config: PretrainedConfig): + super().__init__() + + self.summary_type = getattr(config, "summary_type", "last") + if self.summary_type == "attn": + # We should use a standard multi-head attention module with absolute positional embedding for that. + # Cf. https://github.com/zihangdai/xlnet/blob/master/modeling.py#L253-L276 + # We can probably just use the multi-head attention module of PyTorch >=1.1.0 + raise NotImplementedError + + self.summary = Identity() + if hasattr(config, "summary_use_proj") and config.summary_use_proj: + if hasattr(config, "summary_proj_to_labels") and config.summary_proj_to_labels and config.num_labels > 0: + num_classes = config.num_labels + else: + num_classes = config.hidden_size + self.summary = nn.Linear(config.hidden_size, num_classes) + + activation_string = getattr(config, "summary_activation", None) + self.activation: Callable = get_activation(activation_string) if activation_string else Identity() + + self.first_dropout = Identity() + if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0: + self.first_dropout = nn.Dropout(config.summary_first_dropout) + + self.last_dropout = Identity() + if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0: + self.last_dropout = nn.Dropout(config.summary_last_dropout) + + def forward( + self, hidden_states: torch.FloatTensor, cls_index: Optional[torch.LongTensor] = None + ) -> torch.FloatTensor: + """ + Compute a single vector summary of a sequence hidden states. + + Args: + hidden_states (`torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`): + The hidden states of the last layer. + cls_index (`torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*): + Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token. + + Returns: + `torch.FloatTensor`: The summary of the sequence hidden states. + """ + if self.summary_type == "last": + output = hidden_states[:, -1] + elif self.summary_type == "first": + output = hidden_states[:, 0] + elif self.summary_type == "mean": + output = hidden_states.mean(dim=1) + elif self.summary_type == "cls_index": + if cls_index is None: + cls_index = torch.full_like( + hidden_states[..., :1, :], + hidden_states.shape[-2] - 1, + dtype=torch.long, + ) + else: + cls_index = cls_index.unsqueeze(-1).unsqueeze(-1) + cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),)) + # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states + output = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, XX, hidden_size) + elif self.summary_type == "attn": + raise NotImplementedError + + output = self.first_dropout(output) + output = self.summary(output) + output = self.activation(output) + output = self.last_dropout(output) + + return output + + +def unwrap_model(model: nn.Module, recursive: bool = False) -> nn.Module: + """ + Recursively unwraps a model from potential containers (as used in distributed training). + + Args: + model (`torch.nn.Module`): The model to unwrap. + recursive (`bool`, *optional*, defaults to `False`): + Whether to recursively extract all cases of `module.module` from `model` as well as unwrap child sublayers + recursively, not just the top-level distributed containers. + """ + # Use accelerate implementation if available (should always be the case when using torch) + # This is for pytorch, as we also have to handle things like dynamo + if is_accelerate_available(): + kwargs = {} + if recursive: + if not is_accelerate_available("0.29.0"): + raise RuntimeError( + "Setting `recursive=True` to `unwrap_model` requires `accelerate` v0.29.0. Please upgrade your version of accelerate" + ) + else: + kwargs["recursive"] = recursive + return extract_model_from_parallel(model, **kwargs) + else: + # since there could be multiple levels of wrapping, unwrap recursively + if hasattr(model, "module"): + return unwrap_model(model.module) + else: + return model + + +def expand_device_map(device_map, param_names): + """ + Expand a device map to return the correspondence parameter name to device. + """ + new_device_map = {} + for module, device in device_map.items(): + new_device_map.update( + {p: device for p in param_names if p == module or p.startswith(f"{module}.") or module == ""} + ) + return new_device_map + + +def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: Dict, factor=2): + """This function warm-ups the caching allocator based on the size of the model tensors that will reside on each + device. It allows to have one large call to Malloc, instead of recursively calling it later when loading + the model, which is actually the loading speed botteneck. + Calling this function allows to cut the model loading time by a very large margin. + + A few facts related to loading speed (taking into account the use of this function): + - When loading a model the first time, it is usually slower than the subsequent times, because the OS is very likely + to cache the different state dicts (if enough ressources/RAM are available) + - Trying to force the OS to cache the files in advance (by e.g. accessing a small portion of them) is really hard, + and not a good idea in general as this is low level OS optimizations that depend on ressource usage anyway + - As of 18/03/2025, loading a Llama 70B model with TP takes ~1 min without file cache, and ~13s with full file cache. + The baseline, i.e. only loading the tensor shards on device and adjusting dtype (i.e. copying them) is ~5s with full cache. + These numbers are reported for TP on 4 H100 GPUs. + - It is useless to pre-allocate more than the model size in this function (i.e. using an `allocation_factor` > 1) as + cudaMalloc is not a bottleneck at all anymore + - Loading speed bottleneck is now almost only tensor copy (i.e. changing the dtype) and moving the tensors to the devices. + However, we cannot really improve on those aspects obviously, as the data needs to be moved/copied in the end. + """ + # Remove disk and cpu devices, and cast to proper torch.device + accelerator_device_map = { + param: torch.device(device) for param, device in expanded_device_map.items() if device not in ["cpu", "disk"] + } + if not len(accelerator_device_map): + return + + tp_plan_regex = ( + re.compile("|".join([re.escape(plan) for plan in model._tp_plan])) + if _torch_distributed_available and torch.distributed.is_initialized() + else None + ) + total_byte_count = defaultdict(lambda: 0) + for param_name, device in accelerator_device_map.items(): + param = model.get_parameter_or_buffer(param_name) + # The dtype of different parameters may be different with composite models or `keep_in_fp32_modules` + param_byte_count = math.prod(param.shape) * param.element_size() + + if tp_plan_regex is not None: + generic_name = re.sub(r"\.\d+\.", ".*.", param_name) + param_byte_count //= torch.distributed.get_world_size() if tp_plan_regex.search(generic_name) else 1 + + total_byte_count[device] += param_byte_count + + # This will kick off the caching allocator to avoid having to Malloc afterwards + for device, byte_count in total_byte_count.items(): + if device.type == "cuda": + index = device.index if device.index is not None else torch.cuda.current_device() + device_memory = torch.cuda.mem_get_info(index)[0] + # Allow up to 95% of max device memory + byte_count = min(byte_count, int(0.95 * device_memory)) + # Allocate memory + _ = torch.empty(byte_count // factor, dtype=torch.float16, device=device, requires_grad=False) + + +def get_disk_only_shard_files(device_map, weight_map): + """ + Returns the list of shard files containing only weights offloaded to disk. + """ + files_content = collections.defaultdict(list) + for weight_name, filename in weight_map.items(): + while len(weight_name) > 0 and weight_name not in device_map: + weight_name = ".".join(weight_name.split(".")[:-1]) + files_content[filename].append(device_map[weight_name]) + + return [fname for fname, devices in files_content.items() if set(devices) == {"disk"}] + + +class AttentionInterface(MutableMapping): + """ + Dict-like object keeping track of allowed attention functions. You can easily add a new attention function + with a call to `register()`. If a model needs to locally overwrite an existing attention function, say `sdpa`, + it needs to declare a new instance of this class inside the `modeling_.py`, and declare it on that instance. + """ + + # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if + # a new instance is created (in order to locally override a given function) + _global_mapping = { + "flash_attention_2": flash_attention_forward, + "flex_attention": flex_attention_forward, + "sdpa": sdpa_attention_forward, + } + + def __init__(self): + self._local_mapping = {} + + def __getitem__(self, key): + # First check if instance has a local override + if key in self._local_mapping: + return self._local_mapping[key] + return self._global_mapping[key] + + def __setitem__(self, key, value): + # Allow local update of the default functions without impacting other instances + self._local_mapping.update({key: value}) + + def __delitem__(self, key): + del self._local_mapping[key] + + def __iter__(self): + # Ensure we use all keys, with the overwritten ones on top + return iter({**self._global_mapping, **self._local_mapping}) + + def __len__(self): + return len(self._global_mapping.keys() | self._local_mapping.keys()) + + @classmethod + def register(cls, key: str, value: Callable): + cls._global_mapping.update({key: value}) + + def valid_keys(self) -> List[str]: + return list(self.keys()) + + +# Global AttentionInterface shared by all models which do not need to overwrite any of the existing ones +ALL_ATTENTION_FUNCTIONS: AttentionInterface = AttentionInterface() diff --git a/lib/python3.12/site-packages/transformers/onnx/__init__.py b/lib/python3.12/site-packages/transformers/onnx/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..33350c83a2c161ee228677e8f6fc4b495e9c05bb --- /dev/null +++ b/lib/python3.12/site-packages/transformers/onnx/__init__.py @@ -0,0 +1,49 @@ +# Copyright 2020 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. + +from typing import TYPE_CHECKING + +from ..utils import _LazyModule + + +_import_structure = { + "config": [ + "EXTERNAL_DATA_FORMAT_SIZE_LIMIT", + "OnnxConfig", + "OnnxConfigWithPast", + "OnnxSeq2SeqConfigWithPast", + "PatchingSpec", + ], + "convert": ["export", "validate_model_outputs"], + "features": ["FeaturesManager"], + "utils": ["ParameterFormat", "compute_serialized_parameters_size"], +} + + +if TYPE_CHECKING: + from .config import ( + EXTERNAL_DATA_FORMAT_SIZE_LIMIT, + OnnxConfig, + OnnxConfigWithPast, + OnnxSeq2SeqConfigWithPast, + PatchingSpec, + ) + from .convert import export, validate_model_outputs + from .features import FeaturesManager + from .utils import ParameterFormat, compute_serialized_parameters_size + +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/lib/python3.12/site-packages/transformers/onnx/__main__.py b/lib/python3.12/site-packages/transformers/onnx/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..e3dc6dfb78aa16fa32ea4a30eb4933a0653c7264 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/onnx/__main__.py @@ -0,0 +1,242 @@ +# Copyright 2021 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. +import subprocess +import sys +import warnings +from argparse import ArgumentParser +from pathlib import Path + +from packaging import version + +from .. import AutoFeatureExtractor, AutoImageProcessor, AutoProcessor, AutoTokenizer +from ..utils import logging +from ..utils.import_utils import is_optimum_available +from .convert import export, validate_model_outputs +from .features import FeaturesManager +from .utils import get_preprocessor + + +MIN_OPTIMUM_VERSION = "1.5.0" + +ENCODER_DECODER_MODELS = ["vision-encoder-decoder"] + + +def export_with_optimum(args): + if is_optimum_available(): + from optimum.version import __version__ as optimum_version + + parsed_optimum_version = version.parse(optimum_version) + if parsed_optimum_version < version.parse(MIN_OPTIMUM_VERSION): + raise RuntimeError( + f"transformers.onnx requires optimum >= {MIN_OPTIMUM_VERSION} but {optimum_version} is installed. You " + "can upgrade optimum by running: pip install -U optimum[exporters]" + ) + else: + raise RuntimeError( + "transformers.onnx requires optimum to run, you can install the library by running: pip install " + "optimum[exporters]" + ) + cmd_line = [ + sys.executable, + "-m", + "optimum.exporters.onnx", + f"--model {args.model}", + f"--task {args.feature}", + f"--framework {args.framework}" if args.framework is not None else "", + f"{args.output}", + ] + proc = subprocess.Popen(cmd_line, stdout=subprocess.PIPE) + proc.wait() + + logger.info( + "The export was done by optimum.exporters.onnx. We recommend using to use this package directly in future, as " + "transformers.onnx is deprecated, and will be removed in v5. You can find more information here: " + "https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model." + ) + + +def export_with_transformers(args): + args.output = args.output if args.output.is_file() else args.output.joinpath("model.onnx") + if not args.output.parent.exists(): + args.output.parent.mkdir(parents=True) + + # Allocate the model + model = FeaturesManager.get_model_from_feature( + args.feature, args.model, framework=args.framework, cache_dir=args.cache_dir + ) + + model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=args.feature) + onnx_config = model_onnx_config(model.config) + + if model_kind in ENCODER_DECODER_MODELS: + encoder_model = model.get_encoder() + decoder_model = model.get_decoder() + + encoder_onnx_config = onnx_config.get_encoder_config(encoder_model.config) + decoder_onnx_config = onnx_config.get_decoder_config( + encoder_model.config, decoder_model.config, feature=args.feature + ) + + if args.opset is None: + args.opset = max(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset) + + if args.opset < min(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset): + raise ValueError( + f"Opset {args.opset} is not sufficient to export {model_kind}. At least " + f" {min(encoder_onnx_config.default_onnx_opset, decoder_onnx_config.default_onnx_opset)} is required." + ) + + preprocessor = AutoFeatureExtractor.from_pretrained(args.model) + + onnx_inputs, onnx_outputs = export( + preprocessor, + encoder_model, + encoder_onnx_config, + args.opset, + args.output.parent.joinpath("encoder_model.onnx"), + ) + + validate_model_outputs( + encoder_onnx_config, + preprocessor, + encoder_model, + args.output.parent.joinpath("encoder_model.onnx"), + onnx_outputs, + args.atol if args.atol else encoder_onnx_config.atol_for_validation, + ) + + preprocessor = AutoTokenizer.from_pretrained(args.model) + + onnx_inputs, onnx_outputs = export( + preprocessor, + decoder_model, + decoder_onnx_config, + args.opset, + args.output.parent.joinpath("decoder_model.onnx"), + ) + + validate_model_outputs( + decoder_onnx_config, + preprocessor, + decoder_model, + args.output.parent.joinpath("decoder_model.onnx"), + onnx_outputs, + args.atol if args.atol else decoder_onnx_config.atol_for_validation, + ) + logger.info( + f"All good, model saved at: {args.output.parent.joinpath('encoder_model.onnx').as_posix()}," + f" {args.output.parent.joinpath('decoder_model.onnx').as_posix()}" + ) + + else: + # Instantiate the appropriate preprocessor + if args.preprocessor == "auto": + preprocessor = get_preprocessor(args.model) + elif args.preprocessor == "tokenizer": + preprocessor = AutoTokenizer.from_pretrained(args.model) + elif args.preprocessor == "image_processor": + preprocessor = AutoImageProcessor.from_pretrained(args.model) + elif args.preprocessor == "feature_extractor": + preprocessor = AutoFeatureExtractor.from_pretrained(args.model) + elif args.preprocessor == "processor": + preprocessor = AutoProcessor.from_pretrained(args.model) + else: + raise ValueError(f"Unknown preprocessor type '{args.preprocessor}'") + + # Ensure the requested opset is sufficient + if args.opset is None: + args.opset = onnx_config.default_onnx_opset + + if args.opset < onnx_config.default_onnx_opset: + raise ValueError( + f"Opset {args.opset} is not sufficient to export {model_kind}. " + f"At least {onnx_config.default_onnx_opset} is required." + ) + + onnx_inputs, onnx_outputs = export( + preprocessor, + model, + onnx_config, + args.opset, + args.output, + ) + + if args.atol is None: + args.atol = onnx_config.atol_for_validation + + validate_model_outputs(onnx_config, preprocessor, model, args.output, onnx_outputs, args.atol) + logger.info(f"All good, model saved at: {args.output.as_posix()}") + warnings.warn( + "The export was done by transformers.onnx which is deprecated and will be removed in v5. We recommend" + " using optimum.exporters.onnx in future. You can find more information here:" + " https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model.", + FutureWarning, + ) + + +def main(): + parser = ArgumentParser("Hugging Face Transformers ONNX exporter") + parser.add_argument( + "-m", "--model", type=str, required=True, help="Model ID on huggingface.co or path on disk to load model from." + ) + parser.add_argument( + "--feature", + default="default", + help="The type of features to export the model with.", + ) + parser.add_argument("--opset", type=int, default=None, help="ONNX opset version to export the model with.") + parser.add_argument( + "--atol", type=float, default=None, help="Absolute difference tolerance when validating the model." + ) + parser.add_argument( + "--framework", + type=str, + choices=["pt", "tf"], + default=None, + help=( + "The framework to use for the ONNX export." + " If not provided, will attempt to use the local checkpoint's original framework" + " or what is available in the environment." + ), + ) + parser.add_argument("output", type=Path, help="Path indicating where to store generated ONNX model.") + parser.add_argument("--cache_dir", type=str, default=None, help="Path indicating where to store cache.") + parser.add_argument( + "--preprocessor", + type=str, + choices=["auto", "tokenizer", "feature_extractor", "image_processor", "processor"], + default="auto", + help="Which type of preprocessor to use. 'auto' tries to automatically detect it.", + ) + parser.add_argument( + "--export_with_transformers", + action="store_true", + help=( + "Whether to use transformers.onnx instead of optimum.exporters.onnx to perform the ONNX export. It can be " + "useful when exporting a model supported in transformers but not in optimum, otherwise it is not " + "recommended." + ), + ) + + args = parser.parse_args() + if args.export_with_transformers or not is_optimum_available(): + export_with_transformers(args) + else: + export_with_optimum(args) + + +if __name__ == "__main__": + logger = logging.get_logger("transformers.onnx") # pylint: disable=invalid-name + logger.setLevel(logging.INFO) + main() diff --git a/lib/python3.12/site-packages/transformers/onnx/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/onnx/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f09352dc99aba572d52810b3439e3a32e0b8811 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/onnx/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/onnx/__pycache__/__main__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/onnx/__pycache__/__main__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7c82a1447a1af86f8892624cd056b14588c84e0 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/onnx/__pycache__/__main__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/onnx/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/transformers/onnx/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1aa644c35cd2eacf8c353fb8af007c88fe833c8 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/onnx/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/onnx/__pycache__/convert.cpython-312.pyc b/lib/python3.12/site-packages/transformers/onnx/__pycache__/convert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e367cca288a70659656d89e8b8d7208d01b073a Binary files /dev/null and b/lib/python3.12/site-packages/transformers/onnx/__pycache__/convert.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/onnx/__pycache__/features.cpython-312.pyc b/lib/python3.12/site-packages/transformers/onnx/__pycache__/features.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..91e3ca3b6b7c41260fa3fec96fbf07b83b467023 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/onnx/__pycache__/features.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/onnx/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/onnx/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52e963312e4ca62a0a59629655440fa64b7a7803 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/onnx/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/onnx/config.py b/lib/python3.12/site-packages/transformers/onnx/config.py new file mode 100644 index 0000000000000000000000000000000000000000..460ee9329977bdb2aadc54aae72e0cd93792aab0 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/onnx/config.py @@ -0,0 +1,741 @@ +# Copyright 2021 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. +import copy +import dataclasses +import warnings +from abc import ABC, abstractmethod +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Union + +import numpy as np +from packaging import version + +from ..utils import TensorType, is_torch_available, is_vision_available, logging +from .utils import ParameterFormat, compute_effective_axis_dimension, compute_serialized_parameters_size + + +if TYPE_CHECKING: + from ..configuration_utils import PretrainedConfig + from ..feature_extraction_utils import FeatureExtractionMixin + from ..image_processing_utils import ImageProcessingMixin + from ..tokenization_utils_base import PreTrainedTokenizerBase + + +if is_vision_available(): + from PIL import Image + +logger = logging.get_logger(__name__) + + +DEFAULT_ONNX_OPSET = 11 + +# 2 Gb +EXTERNAL_DATA_FORMAT_SIZE_LIMIT = 2 * 1024 * 1024 * 1024 + + +@dataclasses.dataclass +class PatchingSpec: + """ + Data class that holds patching specifications. + + Args: + o: Module / object where the op to patch is located + name: Name of the op to monkey patch + custom_op: Custom op that patches the original op + orig_op: Original op that is being patched + op_wrapper: Wrapper (optional) that wraps both the original and custom ops. + It is useful for ops that are class or static methods for instance. + """ + + o: Any + name: str + custom_op: Callable + orig_op: Optional[Callable] = None + op_wrapper: Optional[Callable] = None + + +class OnnxConfig(ABC): + """ + Base class for ONNX exportable model describing metadata on how to export the model through the ONNX format. + """ + + default_fixed_batch = 2 + default_fixed_sequence = 8 + default_fixed_num_choices = 4 + torch_onnx_minimum_version = version.parse("1.8") + _tasks_to_common_outputs = { + "causal-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}), + "default": OrderedDict({"last_hidden_state": {0: "batch", 1: "sequence"}}), + "image-classification": OrderedDict({"logits": {0: "batch", 1: "sequence"}}), + "image-segmentation": OrderedDict( + { + "logits": {0: "batch", 1: "sequence"}, + "pred_boxes": {0: "batch", 1: "sequence"}, + "pred_masks": {0: "batch", 1: "sequence"}, + } + ), + "masked-im": OrderedDict({"logits": {0: "batch", 1: "sequence"}}), + "masked-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}), + "multiple-choice": OrderedDict({"logits": {0: "batch"}}), + "object-detection": OrderedDict( + { + "logits": {0: "batch", 1: "sequence"}, + "pred_boxes": {0: "batch", 1: "sequence"}, + } + ), + "question-answering": OrderedDict( + { + "start_logits": {0: "batch", 1: "sequence"}, + "end_logits": {0: "batch", 1: "sequence"}, + } + ), + "semantic-segmentation": OrderedDict({"logits": {0: "batch", 1: "num_labels", 2: "height", 3: "width"}}), + "seq2seq-lm": OrderedDict({"logits": {0: "batch", 1: "decoder_sequence"}}), + "sequence-classification": OrderedDict({"logits": {0: "batch"}}), + "token-classification": OrderedDict({"logits": {0: "batch", 1: "sequence"}}), + "vision2seq-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}), + "speech2seq-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}), + } + + def __init__(self, config: "PretrainedConfig", task: str = "default", patching_specs: List[PatchingSpec] = None): + self._config = config + + if task not in self._tasks_to_common_outputs: + raise ValueError( + f"{task} is not a supported task, supported tasks: {self._tasks_to_common_outputs.keys()}" + ) + self.task = task + + self._patching_specs = [] + for spec in patching_specs if patching_specs is not None else []: + final_spec = spec + if spec.orig_op is None: + final_spec = dataclasses.replace(spec, orig_op=getattr(spec.o, spec.name)) + self._patching_specs.append(final_spec) + + @classmethod + def from_model_config(cls, config: "PretrainedConfig", task: str = "default") -> "OnnxConfig": + """ + Instantiate a OnnxConfig for a specific model + + Args: + config: The model's configuration to use when exporting to ONNX + + Returns: + OnnxConfig for this model + """ + return cls(config, task=task) + + @property + @abstractmethod + def inputs(self) -> Mapping[str, Mapping[int, str]]: + """ + Mapping containing the axis definition of the input tensors to provide to the model + + Returns: + For each input: its name associated to the axes symbolic name and the axis position within the tensor + """ + raise NotImplementedError() + + @property + def outputs(self) -> Mapping[str, Mapping[int, str]]: + """ + Mapping containing the axis definition of the output tensors to provide to the model + + Returns: + For each output: its name associated to the axes symbolic name and the axis position within the tensor + """ + common_outputs = self._tasks_to_common_outputs[self.task] + return copy.deepcopy(common_outputs) + + @property + def values_override(self) -> Optional[Mapping[str, Any]]: + """ + Dictionary of keys to override in the model's config before exporting + + Returns: + Dictionary with the keys (and their corresponding values) to override + """ + if hasattr(self._config, "use_cache"): + return {"use_cache": False} + + return None + + @property + def default_batch_size(self) -> int: + """ + The default batch size to use if no other indication + + Returns: + Integer > 0 + """ + # Using 2 avoid ONNX making assumption about single sample batch + return OnnxConfig.default_fixed_batch + + @property + def default_sequence_length(self) -> int: + """ + The default sequence length to use if no other indication + + Returns: + Integer > 0 + """ + return OnnxConfig.default_fixed_sequence + + @property + def default_num_choices(self) -> int: + """ + The default number of choices to use if no other indication + + Returns: + Integer > 0 + """ + return OnnxConfig.default_fixed_num_choices + + @property + def default_onnx_opset(self) -> int: + """ + Which onnx opset to use when exporting the model + + Returns: + Integer ONNX Opset version + """ + return DEFAULT_ONNX_OPSET + + @property + def atol_for_validation(self) -> float: + """ + What absolute tolerance value to use during model conversion validation. + + Returns: + Float absolute tolerance value. + """ + return 1e-5 + + @property + def is_torch_support_available(self) -> bool: + """ + The minimum PyTorch version required to export the model. + + Returns: + `bool`: Whether the installed version of PyTorch is compatible with the model. + """ + if is_torch_available(): + from transformers.utils import get_torch_version + + return version.parse(get_torch_version()) >= self.torch_onnx_minimum_version + else: + return False + + @staticmethod + def use_external_data_format(num_parameters: int) -> bool: + """ + Flag indicating if the model requires using external data format + + Args: + num_parameters: Number of parameter on the model + + Returns: + True if model.num_parameters() * size_of(float32) >= 2Gb False otherwise + """ + + return ( + compute_serialized_parameters_size(num_parameters, ParameterFormat.Float) + >= EXTERNAL_DATA_FORMAT_SIZE_LIMIT + ) + + def _generate_dummy_images( + self, batch_size: int = 2, num_channels: int = 3, image_height: int = 40, image_width: int = 40 + ): + images = [] + for _ in range(batch_size): + data = np.random.rand(image_height, image_width, num_channels) * 255 + images.append(Image.fromarray(data.astype("uint8")).convert("RGB")) + return images + + def _generate_dummy_audio( + self, batch_size: int = 2, sampling_rate: int = 22050, time_duration: float = 5.0, frequency: int = 220 + ): + audio_data = [] + for _ in range(batch_size): + # time variable + t = np.linspace(0, time_duration, int(time_duration * sampling_rate), endpoint=False) + + # generate pure sine wave at `frequency` Hz + audio_data.append(0.5 * np.sin(2 * np.pi * frequency * t)) + + return audio_data + + def generate_dummy_inputs( + self, + preprocessor: Union["PreTrainedTokenizerBase", "FeatureExtractionMixin", "ImageProcessingMixin"], + batch_size: int = -1, + seq_length: int = -1, + num_choices: int = -1, + is_pair: bool = False, + framework: Optional[TensorType] = None, + num_channels: int = 3, + image_width: int = 40, + image_height: int = 40, + sampling_rate: int = 22050, + time_duration: float = 5.0, + frequency: int = 220, + tokenizer: Optional["PreTrainedTokenizerBase"] = None, + ) -> Mapping[str, Any]: + """ + Generate inputs to provide to the ONNX exporter for the specific framework + + Args: + preprocessor: ([`PreTrainedTokenizerBase`], [`FeatureExtractionMixin`], or [`ImageProcessingMixin`]): + The preprocessor associated with this model configuration. + batch_size (`int`, *optional*, defaults to -1): + The batch size to export the model for (-1 means dynamic axis). + num_choices (`int`, *optional*, defaults to -1): + The number of candidate answers provided for multiple choice task (-1 means dynamic axis). + seq_length (`int`, *optional*, defaults to -1): + The sequence length to export the model for (-1 means dynamic axis). + is_pair (`bool`, *optional*, defaults to `False`): + Indicate if the input is a pair (sentence 1, sentence 2) + framework (`TensorType`, *optional*, defaults to `None`): + The framework (PyTorch or TensorFlow) that the tokenizer will generate tensors for. + num_channels (`int`, *optional*, defaults to 3): + The number of channels of the generated images. + image_width (`int`, *optional*, defaults to 40): + The width of the generated images. + image_height (`int`, *optional*, defaults to 40): + The height of the generated images. + sampling_rate (`int`, *optional* defaults to 22050) + The sampling rate for audio data generation. + time_duration (`float`, *optional* defaults to 5.0) + Total seconds of sampling for audio data generation. + frequency (`int`, *optional* defaults to 220) + The desired natural frequency of generated audio. + + Returns: + Mapping[str, Tensor] holding the kwargs to provide to the model's forward function + """ + from ..feature_extraction_utils import FeatureExtractionMixin + from ..image_processing_utils import ImageProcessingMixin + from ..tokenization_utils_base import PreTrainedTokenizerBase + + if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None: + raise ValueError("You cannot provide both a tokenizer and a preprocessor to generate dummy inputs.") + if tokenizer is not None: + warnings.warn( + "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use" + " `preprocessor` instead.", + FutureWarning, + ) + logger.warning("Overwriting the `preprocessor` argument with `tokenizer` to generate dummy inputs.") + preprocessor = tokenizer + if isinstance(preprocessor, PreTrainedTokenizerBase): + # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX + batch_size = compute_effective_axis_dimension( + batch_size, fixed_dimension=OnnxConfig.default_fixed_batch, num_token_to_add=0 + ) + # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX + token_to_add = preprocessor.num_special_tokens_to_add(is_pair) + seq_length = compute_effective_axis_dimension( + seq_length, fixed_dimension=OnnxConfig.default_fixed_sequence, num_token_to_add=token_to_add + ) + # Generate dummy inputs according to compute batch and sequence + input_token = ( + preprocessor.unk_token + if (preprocessor.unk_token is not None and len(preprocessor.unk_token) > 0) + else "0" + ) + dummy_input = [" ".join([input_token]) * seq_length] * batch_size + if self.task == "multiple-choice": + # If dynamic axis (-1) we forward with a fixed dimension of 4 candidate answers to avoid optimizations + # made by ONNX + num_choices = compute_effective_axis_dimension( + num_choices, fixed_dimension=OnnxConfig.default_fixed_num_choices, num_token_to_add=0 + ) + dummy_input = dummy_input * num_choices + # The shape of the tokenized inputs values is [batch_size * num_choices, seq_length] + tokenized_input = preprocessor(dummy_input, text_pair=dummy_input) + # Unflatten the tokenized inputs values expanding it to the shape [batch_size, num_choices, seq_length] + for k, v in tokenized_input.items(): + tokenized_input[k] = [v[i : i + num_choices] for i in range(0, len(v), num_choices)] + return dict(tokenized_input.convert_to_tensors(tensor_type=framework)) + return dict(preprocessor(dummy_input, return_tensors=framework)) + elif isinstance(preprocessor, ImageProcessingMixin): + if preprocessor.model_input_names[0] != "pixel_values": + raise ValueError( + f"The `preprocessor` is an image processor ({preprocessor.__class__.__name__}) and expects" + f' `model_input_names[0]` to be "pixel_values", but got {preprocessor.model_input_names[0]}' + ) + # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX + batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch) + dummy_input = self._generate_dummy_images(batch_size, num_channels, image_height, image_width) + return dict(preprocessor(images=dummy_input, return_tensors=framework)) + elif isinstance(preprocessor, FeatureExtractionMixin) and preprocessor.model_input_names[0] == "pixel_values": + # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX + batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch) + dummy_input = self._generate_dummy_images(batch_size, num_channels, image_height, image_width) + return dict(preprocessor(images=dummy_input, return_tensors=framework)) + elif ( + isinstance(preprocessor, FeatureExtractionMixin) and preprocessor.model_input_names[0] == "input_features" + ): + # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX + batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch) + dummy_input = self._generate_dummy_audio(batch_size, sampling_rate, time_duration, frequency) + return dict(preprocessor(dummy_input, return_tensors=framework)) + else: + raise ValueError( + "Unable to generate dummy inputs for the model. Please provide a tokenizer or a preprocessor." + ) + + def generate_dummy_inputs_onnxruntime(self, reference_model_inputs: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Generate inputs for ONNX Runtime using the reference model inputs. Override this to run inference with seq2seq + models which have the encoder and decoder exported as separate ONNX files. + + Args: + reference_model_inputs ([`Mapping[str, Tensor]`): + Reference inputs for the model. + + Returns: + `Mapping[str, Tensor]`: The mapping holding the kwargs to provide to the model's forward function + """ + return reference_model_inputs + + def patch_ops(self): + for spec in self._patching_specs: + custom_op = spec.custom_op if spec.op_wrapper is None else spec.op_wrapper(spec.custom_op) + setattr(spec.o, spec.name, custom_op) + + def restore_ops(self): + for spec in self._patching_specs: + orig_op = spec.orig_op if spec.op_wrapper is None else spec.op_wrapper(spec.orig_op) + setattr(spec.o, spec.name, orig_op) + + @classmethod + def flatten_output_collection_property(cls, name: str, field: Iterable[Any]) -> Dict[str, Any]: + """ + Flatten any potential nested structure expanding the name of the field with the index of the element within the + structure. + + Args: + name: The name of the nested structure + field: The structure to, potentially, be flattened + + Returns: + (Dict[str, Any]): Outputs with flattened structure and key mapping this new structure. + + """ + from itertools import chain + + return {f"{name}.{idx}": item for idx, item in enumerate(chain.from_iterable(field))} + + +class OnnxConfigWithPast(OnnxConfig, ABC): + def __init__( + self, + config: "PretrainedConfig", + task: str = "default", + patching_specs: Optional[list[PatchingSpec]] = None, + use_past: bool = False, + ): + super().__init__(config, task=task, patching_specs=patching_specs) + self.use_past = use_past + + @classmethod + def with_past(cls, config: "PretrainedConfig", task: str = "default") -> "OnnxConfigWithPast": + """ + Instantiate a OnnxConfig with `use_past` attribute set to True + + Args: + config: The underlying model's config to use when exporting to ONNX + + Returns: + OnnxConfig with `.use_past = True` + """ + return cls(config, task=task, use_past=True) + + @property + def outputs(self) -> Mapping[str, Mapping[int, str]]: + common_outputs = super().outputs + if self.use_past: + self.fill_with_past_key_values_(common_outputs, direction="outputs") + + return common_outputs + + @property + def values_override(self) -> Optional[Mapping[str, Any]]: + if hasattr(self._config, "use_cache"): + return {"use_cache": self.use_past} + + return None + + @property + def num_layers(self) -> int: + """ + The number of layers attribute retrieved from the model config. Override this for model configs where the + number of layers attribute is not called `num_layers`. + """ + if not hasattr(self._config, "num_layers"): + raise AttributeError( + "could not find the number of layers attribute in the model configuration, override the num_layers" + " property of the model OnnxConfig to solve this" + ) + return self._config.num_layers + + @property + def num_attention_heads(self) -> int: + """ + The number of attention heads attribute retrieved from the model config. Override this for model configs where + the number of attention heads attribute is not called `num_attention_heads`. + """ + if not hasattr(self._config, "num_attention_heads"): + raise AttributeError( + "could not find the number of attention heads attribute in the model configuration, override the" + " num_attention_heads property of the model OnnxConfig to solve this" + ) + return self._config.num_attention_heads + + def generate_dummy_inputs( + self, + tokenizer: "PreTrainedTokenizerBase", + batch_size: int = -1, + seq_length: int = -1, + is_pair: bool = False, + framework: Optional[TensorType] = None, + ) -> Mapping[str, Any]: + # TODO: should we set seq_length = 1 when self.use_past = True? + common_inputs = super().generate_dummy_inputs( + tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework + ) + + if self.use_past: + if not is_torch_available(): + raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.") + else: + import torch + + batch, seqlen = common_inputs["input_ids"].shape + # Not using the same length for past_key_values + past_key_values_length = seqlen + 2 + shape = ( + batch, + self.num_attention_heads, + past_key_values_length, + self._config.hidden_size // self.num_attention_heads, + ) + + if "attention_mask" in common_inputs: + mask_dtype = common_inputs["attention_mask"].dtype + common_inputs["attention_mask"] = torch.cat( + [common_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], + dim=1, + ) + + common_inputs["past_key_values"] = [] + for _ in range(self.num_layers): + common_inputs["past_key_values"].append((torch.zeros(shape), torch.zeros(shape))) + + return common_inputs + + def fill_with_past_key_values_( + self, inputs_or_outputs: Mapping[str, Mapping[int, str]], direction: str, inverted_values_shape: bool = False + ): + """ + Fill the input_or_outputs mapping with past_key_values dynamic axes considering. + + Args: + inputs_or_outputs: The mapping to fill. + direction: either "inputs" or "outputs", it specifies whether input_or_outputs is the input mapping or the + output mapping, this is important for axes naming. + inverted_values_shape: + If `True`, store values on dynamic axis 1, else on axis 2. + + """ + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + name = "past_key_values" if direction == "inputs" else "present" + for i in range(self.num_layers): + inputs_or_outputs[f"{name}.{i}.key"] = {0: "batch", 2: "past_sequence + sequence"} + if inverted_values_shape: + inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch", 1: "past_sequence + sequence"} + else: + inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch", 2: "past_sequence + sequence"} + + def _flatten_past_key_values_(self, flattened_output, name, idx, t): + flattened_output[f"{name}.{idx}.key"] = t[0] + flattened_output[f"{name}.{idx}.value"] = t[1] + + def flatten_output_collection_property(self, name: str, field: Iterable[Any]) -> Dict[str, Any]: + flattened_output = {} + if name in ["present", "past_key_values"]: + for idx, t in enumerate(field): + self._flatten_past_key_values_(flattened_output, name, idx, t) + else: + flattened_output = super().flatten_output_collection_property(name, field) + + return flattened_output + + +class OnnxSeq2SeqConfigWithPast(OnnxConfigWithPast): + @property + def outputs(self) -> Mapping[str, Mapping[int, str]]: + common_outputs = super(OnnxConfigWithPast, self).outputs + # Renaming the outputs axes properly. + for name, axes_names in common_outputs.items(): + sequence_name = "encoder_sequence" if "encoder" in name else "decoder_sequence" + for axis_idx, name in axes_names.items(): + if "sequence" in name: + axes_names[axis_idx] = sequence_name + # We reset the value as the order in common_outputs (OrderedDict) is lost otherwise + else: + axes_names[axis_idx] = name + if self.use_past: + self.fill_with_past_key_values_(common_outputs, direction="outputs") + + return common_outputs + + @property + def num_layers(self) -> Tuple[int]: + try: + num_layers = super().num_layers + num_layers = (num_layers, num_layers) + except AttributeError: + if hasattr(self._config, "encoder_layers") and hasattr(self._config, "decoder_layers"): + num_layers = (self._config.encoder_layers, self._config.decoder_layers) + else: + raise AttributeError( + "could not find the number of encoder and decoder layers attributes in the model configuration," + " override the num_layers property of the model OnnxConfig to solve this" + ) + + return num_layers + + @property + def num_attention_heads(self) -> Tuple[int]: + try: + num_attention_heads = super().num_attention_heads + num_attention_heads = (num_attention_heads, num_attention_heads) + except AttributeError: + if hasattr(self._config, "encoder_attention_heads") and hasattr(self._config, "decoder_attention_heads"): + num_attention_heads = (self._config.encoder_attention_heads, self._config.decoder_attention_heads) + else: + raise AttributeError( + "could not find the number of attention heads for the encoder and the decoder attributes in the" + " model configuration, override the num_attention_heads property of the model OnnxConfig to solve" + " this" + ) + return num_attention_heads + + def generate_dummy_inputs( + self, + tokenizer: Optional["PreTrainedTokenizerBase"], + batch_size: int = -1, + seq_length: int = -1, + is_pair: bool = False, + framework: Optional[TensorType] = None, + ) -> Mapping[str, Any]: + encoder_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs( + tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework + ) + + # Generate decoder inputs + decoder_seq_length = seq_length if not self.use_past else 1 + decoder_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs( + tokenizer, batch_size=batch_size, seq_length=decoder_seq_length, is_pair=is_pair, framework=framework + ) + decoder_inputs = {f"decoder_{name}": tensor for name, tensor in decoder_inputs.items()} + common_inputs = dict(**encoder_inputs, **decoder_inputs) + + if self.use_past: + if not is_torch_available(): + raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.") + else: + import torch + batch = common_inputs["input_ids"].shape[0] + encoder_seq_length = common_inputs["input_ids"].shape[1] + decoder_seq_length = common_inputs["decoder_input_ids"].shape[1] + num_encoder_attention_heads, num_decoder_attention_heads = self.num_attention_heads + encoder_shape = ( + batch, + num_encoder_attention_heads, + encoder_seq_length, + self._config.hidden_size // num_encoder_attention_heads, + ) + decoder_shape = ( + batch, + num_decoder_attention_heads, + # Not using the same length for past_key_values + decoder_seq_length + 3, + self._config.hidden_size // num_decoder_attention_heads, + ) + + common_inputs["past_key_values"] = [] + # If the number of encoder and decoder layers are present in the model configuration, both are considered + num_encoder_layers, num_decoder_layers = self.num_layers + min_num_layers = min(num_encoder_layers, num_decoder_layers) + max_num_layers = max(num_encoder_layers, num_decoder_layers) - min_num_layers + remaining_side_name = "encoder" if num_encoder_layers > num_decoder_layers else "decoder" + + for _ in range(min_num_layers): + # For encoder-decoder models, past_key_values contains pre-computed values for both the encoder and the + # decoder layers, hence a tuple of 4 tensors instead of 2 + common_inputs["past_key_values"].append( + ( + torch.zeros(decoder_shape), + torch.zeros(decoder_shape), + torch.zeros(encoder_shape), + torch.zeros(encoder_shape), + ) + ) + + # TODO: test this. + shape = encoder_shape if remaining_side_name == "encoder" else decoder_shape + for _ in range(min_num_layers, max_num_layers): + common_inputs["past_key_values"].append((torch.zeros(shape), torch.zeros(shape))) + + return common_inputs + + def fill_with_past_key_values_(self, inputs_or_outputs: Mapping[str, Mapping[int, str]], direction: str): + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + name = "past_key_values" if direction == "inputs" else "present" + + # If the number of encoder and decoder layers are present in the model configuration, both are considered + num_encoder_layers, num_decoder_layers = self.num_layers + min_num_layers = min(num_encoder_layers, num_decoder_layers) + max_num_layers = max(num_encoder_layers, num_decoder_layers) - min_num_layers + remaining_side_name = "encoder" if num_encoder_layers > num_decoder_layers else "decoder" + + encoder_sequence = "past_encoder_sequence" + decoder_sequence = "past_decoder_sequence" if direction == "inputs" else "past_decoder_sequence + sequence" + + for i in range(min_num_layers): + inputs_or_outputs[f"{name}.{i}.decoder.key"] = {0: "batch", 2: decoder_sequence} + inputs_or_outputs[f"{name}.{i}.decoder.value"] = {0: "batch", 2: decoder_sequence} + inputs_or_outputs[f"{name}.{i}.encoder.key"] = {0: "batch", 2: encoder_sequence} + inputs_or_outputs[f"{name}.{i}.encoder.value"] = {0: "batch", 2: encoder_sequence} + + for i in range(min_num_layers, max_num_layers): + if remaining_side_name == "encoder": + axes_info = {0: "batch", 2: encoder_sequence} + else: + axes_info = {0: "batch", 2: decoder_sequence} + inputs_or_outputs[f"{name}.{i}.{remaining_side_name}.key"] = axes_info + + def _flatten_past_key_values_(self, flattened_output, name, idx, t): + flattened_output[f"{name}.{idx}.decoder.key"] = t[0] + flattened_output[f"{name}.{idx}.decoder.value"] = t[1] + flattened_output[f"{name}.{idx}.encoder.key"] = t[2] + flattened_output[f"{name}.{idx}.encoder.value"] = t[3] diff --git a/lib/python3.12/site-packages/transformers/onnx/convert.py b/lib/python3.12/site-packages/transformers/onnx/convert.py new file mode 100644 index 0000000000000000000000000000000000000000..58bc51f8e801cfaeb9423342f330c269fa5c2318 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/onnx/convert.py @@ -0,0 +1,460 @@ +# Copyright 2021 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. + +import warnings +from inspect import signature +from itertools import chain +from pathlib import Path +from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union + +import numpy as np +from packaging.version import Version, parse + +from ..tokenization_utils_base import PreTrainedTokenizerBase +from ..utils import ( + TensorType, + is_tf_available, + is_torch_available, + logging, +) +from .config import OnnxConfig + + +if is_torch_available(): + from ..modeling_utils import PreTrainedModel + +if is_tf_available(): + from ..modeling_tf_utils import TFPreTrainedModel + +if TYPE_CHECKING: + from ..feature_extraction_utils import FeatureExtractionMixin + from ..processing_utils import ProcessorMixin + from ..tokenization_utils import PreTrainedTokenizer + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# This is the minimal required version to support some ONNX Runtime features +ORT_QUANTIZE_MINIMUM_VERSION = parse("1.4.0") + + +def check_onnxruntime_requirements(minimum_version: Version): + """ + Check onnxruntime is installed and if the installed version match is recent enough + + Raises: + ImportError: If onnxruntime is not installed or too old version is found + """ + try: + import onnxruntime + + # Parse the version of the installed onnxruntime + ort_version = parse(onnxruntime.__version__) + + # We require 1.4.0 minimum + if ort_version < ORT_QUANTIZE_MINIMUM_VERSION: + raise ImportError( + f"We found an older version of onnxruntime ({onnxruntime.__version__}) " + f"but we require onnxruntime to be >= {minimum_version} to enable all the conversions options.\n" + "Please update onnxruntime by running `pip install --upgrade onnxruntime`" + ) + + except ImportError: + raise ImportError( + "onnxruntime doesn't seem to be currently installed. " + "Please install the onnxruntime by running `pip install onnxruntime`" + " and relaunch the conversion." + ) + + +def export_pytorch( + preprocessor: Union["PreTrainedTokenizer", "FeatureExtractionMixin", "ProcessorMixin"], + model: "PreTrainedModel", + config: OnnxConfig, + opset: int, + output: Path, + tokenizer: Optional["PreTrainedTokenizer"] = None, + device: str = "cpu", +) -> Tuple[List[str], List[str]]: + """ + Export a PyTorch model to an ONNX Intermediate Representation (IR) + + Args: + preprocessor: ([`PreTrainedTokenizer`], [`FeatureExtractionMixin`] or [`ProcessorMixin`]): + The preprocessor used for encoding the data. + model ([`PreTrainedModel`]): + The model to export. + config ([`~onnx.config.OnnxConfig`]): + The ONNX configuration associated with the exported model. + opset (`int`): + The version of the ONNX operator set to use. + output (`Path`): + Directory to store the exported ONNX model. + device (`str`, *optional*, defaults to `cpu`): + The device on which the ONNX model will be exported. Either `cpu` or `cuda`. + + Returns: + `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from + the ONNX configuration. + """ + + if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None: + raise ValueError("You cannot provide both a tokenizer and a preprocessor to export the model.") + if tokenizer is not None: + warnings.warn( + "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use" + " `preprocessor` instead.", + FutureWarning, + ) + logger.info("Overwriting the `preprocessor` argument with `tokenizer` to generate dummy inputs.") + preprocessor = tokenizer + + if issubclass(type(model), PreTrainedModel): + import torch + from torch.onnx import export as onnx_export + + logger.info(f"Using framework PyTorch: {torch.__version__}") + with torch.no_grad(): + model.config.return_dict = True + model.eval() + + # Check if we need to override certain configuration item + if config.values_override is not None: + logger.info(f"Overriding {len(config.values_override)} configuration item(s)") + for override_config_key, override_config_value in config.values_override.items(): + logger.info(f"\t- {override_config_key} -> {override_config_value}") + setattr(model.config, override_config_key, override_config_value) + + # Ensure inputs match + # TODO: Check when exporting QA we provide "is_pair=True" + model_inputs = config.generate_dummy_inputs(preprocessor, framework=TensorType.PYTORCH) + device = torch.device(device) + if device.type == "cuda" and torch.cuda.is_available(): + model.to(device) + model_inputs_device = {} + for k, v in model_inputs.items(): + if isinstance(v, Tuple): + model_inputs_device[k] = tuple( + x.to(device) if isinstance(x, torch.Tensor) else None for x in v + ) + elif isinstance(v, List): + model_inputs_device[k] = [ + tuple(x.to(device) if isinstance(x, torch.Tensor) else None for x in t) for t in v + ] + else: + model_inputs_device[k] = v.to(device) + + model_inputs = model_inputs_device + + inputs_match, matched_inputs = ensure_model_and_config_inputs_match(model, model_inputs.keys()) + onnx_outputs = list(config.outputs.keys()) + + if not inputs_match: + raise ValueError("Model and config inputs doesn't match") + + config.patch_ops() + + onnx_export( + model, + (model_inputs,), + f=output.as_posix(), + input_names=list(config.inputs.keys()), + output_names=onnx_outputs, + dynamic_axes=dict(chain(config.inputs.items(), config.outputs.items())), + do_constant_folding=True, + opset_version=opset, + ) + + config.restore_ops() + + return matched_inputs, onnx_outputs + + +def export_tensorflow( + preprocessor: Union["PreTrainedTokenizer", "FeatureExtractionMixin"], + model: "TFPreTrainedModel", + config: OnnxConfig, + opset: int, + output: Path, + tokenizer: Optional["PreTrainedTokenizer"] = None, +) -> Tuple[List[str], List[str]]: + """ + Export a TensorFlow model to an ONNX Intermediate Representation (IR) + + Args: + preprocessor: ([`PreTrainedTokenizer`] or [`FeatureExtractionMixin`]): + The preprocessor used for encoding the data. + model ([`TFPreTrainedModel`]): + The model to export. + config ([`~onnx.config.OnnxConfig`]): + The ONNX configuration associated with the exported model. + opset (`int`): + The version of the ONNX operator set to use. + output (`Path`): + Directory to store the exported ONNX model. + + Returns: + `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from + the ONNX configuration. + """ + import onnx + import tensorflow as tf + import tf2onnx + + if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None: + raise ValueError("You cannot provide both a tokenizer and preprocessor to export the model.") + if tokenizer is not None: + warnings.warn( + "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use" + " `preprocessor` instead.", + FutureWarning, + ) + logger.info("Overwriting the `preprocessor` argument with `tokenizer` to generate dummy inputs.") + preprocessor = tokenizer + + model.config.return_dict = True + + # Check if we need to override certain configuration item + if config.values_override is not None: + logger.info(f"Overriding {len(config.values_override)} configuration item(s)") + for override_config_key, override_config_value in config.values_override.items(): + logger.info(f"\t- {override_config_key} -> {override_config_value}") + setattr(model.config, override_config_key, override_config_value) + + # Ensure inputs match + model_inputs = config.generate_dummy_inputs(preprocessor, framework=TensorType.TENSORFLOW) + inputs_match, matched_inputs = ensure_model_and_config_inputs_match(model, model_inputs.keys()) + onnx_outputs = list(config.outputs.keys()) + + input_signature = [ + tf.TensorSpec([None] * tensor.ndim, dtype=tensor.dtype, name=key) for key, tensor in model_inputs.items() + ] + onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature, opset=opset) + onnx.save(onnx_model, output.as_posix()) + config.restore_ops() + + return matched_inputs, onnx_outputs + + +def export( + preprocessor: Union["PreTrainedTokenizer", "FeatureExtractionMixin", "ProcessorMixin"], + model: Union["PreTrainedModel", "TFPreTrainedModel"], + config: OnnxConfig, + opset: int, + output: Path, + tokenizer: Optional["PreTrainedTokenizer"] = None, + device: str = "cpu", +) -> Tuple[List[str], List[str]]: + """ + Export a Pytorch or TensorFlow model to an ONNX Intermediate Representation (IR) + + Args: + preprocessor: ([`PreTrainedTokenizer`], [`FeatureExtractionMixin`] or [`ProcessorMixin`]): + The preprocessor used for encoding the data. + model ([`PreTrainedModel`] or [`TFPreTrainedModel`]): + The model to export. + config ([`~onnx.config.OnnxConfig`]): + The ONNX configuration associated with the exported model. + opset (`int`): + The version of the ONNX operator set to use. + output (`Path`): + Directory to store the exported ONNX model. + device (`str`, *optional*, defaults to `cpu`): + The device on which the ONNX model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for + export on CUDA devices. + + Returns: + `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from + the ONNX configuration. + """ + if not (is_torch_available() or is_tf_available()): + raise ImportError( + "Cannot convert because neither PyTorch nor TensorFlow are not installed. " + "Please install torch or tensorflow first." + ) + + if is_tf_available() and isinstance(model, TFPreTrainedModel) and device == "cuda": + raise RuntimeError("`tf2onnx` does not support export on CUDA device.") + + if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None: + raise ValueError("You cannot provide both a tokenizer and a preprocessor to export the model.") + if tokenizer is not None: + warnings.warn( + "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use" + " `preprocessor` instead.", + FutureWarning, + ) + logger.info("Overwriting the `preprocessor` argument with `tokenizer` to generate dummy inputs.") + preprocessor = tokenizer + + if is_torch_available(): + from ..utils import get_torch_version + + if not config.is_torch_support_available: + logger.warning( + f"Unsupported PyTorch version for this model. Minimum required is {config.torch_onnx_minimum_version}," + f" got: {get_torch_version()}" + ) + + if is_torch_available() and issubclass(type(model), PreTrainedModel): + return export_pytorch(preprocessor, model, config, opset, output, tokenizer=tokenizer, device=device) + elif is_tf_available() and issubclass(type(model), TFPreTrainedModel): + return export_tensorflow(preprocessor, model, config, opset, output, tokenizer=tokenizer) + + +def validate_model_outputs( + config: OnnxConfig, + preprocessor: Union["PreTrainedTokenizer", "FeatureExtractionMixin", "ProcessorMixin"], + reference_model: Union["PreTrainedModel", "TFPreTrainedModel"], + onnx_model: Path, + onnx_named_outputs: List[str], + atol: float, + tokenizer: Optional["PreTrainedTokenizer"] = None, +): + from onnxruntime import InferenceSession, SessionOptions + + logger.info("Validating ONNX model...") + + if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None: + raise ValueError("You cannot provide both a tokenizer and a preprocessor to validate the model outputs.") + if tokenizer is not None: + warnings.warn( + "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use" + " `preprocessor` instead.", + FutureWarning, + ) + logger.info("Overwriting the `preprocessor` argument with `tokenizer` to generate dummy inputs.") + preprocessor = tokenizer + + # generate inputs with a different batch_size and seq_len that was used for conversion to properly test + # dynamic input shapes. + if is_torch_available() and issubclass(type(reference_model), PreTrainedModel): + reference_model_inputs = config.generate_dummy_inputs( + preprocessor, + batch_size=config.default_fixed_batch + 1, + seq_length=config.default_fixed_sequence + 1, + framework=TensorType.PYTORCH, + ) + else: + reference_model_inputs = config.generate_dummy_inputs( + preprocessor, + batch_size=config.default_fixed_batch + 1, + seq_length=config.default_fixed_sequence + 1, + framework=TensorType.TENSORFLOW, + ) + + # Create ONNX Runtime session + options = SessionOptions() + session = InferenceSession(onnx_model.as_posix(), options, providers=["CPUExecutionProvider"]) + + # Compute outputs from the reference model + if is_torch_available() and issubclass(type(reference_model), PreTrainedModel): + reference_model.to("cpu") + ref_outputs = reference_model(**reference_model_inputs) + ref_outputs_dict = {} + + # We flatten potential collection of outputs (i.e. past_keys) to a flat structure + for name, value in ref_outputs.items(): + # Overwriting the output name as "present" since it is the name used for the ONNX outputs + # ("past_key_values" being taken for the ONNX inputs) + if name == "past_key_values": + name = "present" + if isinstance(value, (list, tuple)): + value = config.flatten_output_collection_property(name, value) + ref_outputs_dict.update(value) + else: + ref_outputs_dict[name] = value + + # Create onnxruntime inputs from the reference model inputs + reference_model_inputs_onnxruntime = config.generate_dummy_inputs_onnxruntime(reference_model_inputs) + + # We flatten potential collection of inputs (i.e. past_keys) + onnx_inputs = {} + for name, value in reference_model_inputs_onnxruntime.items(): + if isinstance(value, (list, tuple)): + value = config.flatten_output_collection_property(name, value) + onnx_inputs.update({tensor_name: pt_tensor.numpy() for tensor_name, pt_tensor in value.items()}) + else: + onnx_inputs[name] = value.numpy() + + # Compute outputs from the ONNX model + onnx_outputs = session.run(onnx_named_outputs, onnx_inputs) + + # Check we have a subset of the keys into onnx_outputs against ref_outputs + ref_outputs_set, onnx_outputs_set = set(ref_outputs_dict.keys()), set(onnx_named_outputs) + if not onnx_outputs_set.issubset(ref_outputs_set): + logger.info( + f"\t-[x] ONNX model output names {onnx_outputs_set} do not match reference model {ref_outputs_set}" + ) + + raise ValueError( + "Outputs doesn't match between reference model and ONNX exported model: " + f"{onnx_outputs_set.difference(ref_outputs_set)}" + ) + else: + logger.info(f"\t-[✓] ONNX model output names match reference model ({onnx_outputs_set})") + + # Check the shape and values match + for name, ort_value in zip(onnx_named_outputs, onnx_outputs): + if is_torch_available() and issubclass(type(reference_model), PreTrainedModel): + ref_value = ref_outputs_dict[name].detach().numpy() + else: + ref_value = ref_outputs_dict[name].numpy() + logger.info(f'\t- Validating ONNX Model output "{name}":') + + # Shape + if not ort_value.shape == ref_value.shape: + logger.info(f"\t\t-[x] shape {ort_value.shape} doesn't match {ref_value.shape}") + raise ValueError( + "Outputs shape doesn't match between reference model and ONNX exported model: " + f"Got {ref_value.shape} (reference) and {ort_value.shape} (ONNX)" + ) + else: + logger.info(f"\t\t-[✓] {ort_value.shape} matches {ref_value.shape}") + + # Values + if not np.allclose(ref_value, ort_value, atol=atol): + bad_indices = np.logical_not(np.isclose(ref_value, ort_value, atol=atol)) + logger.info(f"\t\t-[x] values not close enough (atol: {atol})") + raise ValueError( + "Outputs values doesn't match between reference model and ONNX exported model: " + f"Got max absolute difference of: {np.amax(np.abs(ref_value - ort_value))} for " + f"{ref_value[bad_indices]} vs {ort_value[bad_indices]}" + ) + else: + logger.info(f"\t\t-[✓] all values close (atol: {atol})") + + +def ensure_model_and_config_inputs_match( + model: Union["PreTrainedModel", "TFPreTrainedModel"], model_inputs: Iterable[str] +) -> Tuple[bool, List[str]]: + """ + + :param model_inputs: :param config_inputs: :return: + """ + if is_torch_available() and issubclass(type(model), PreTrainedModel): + forward_parameters = signature(model.forward).parameters + else: + forward_parameters = signature(model.call).parameters + model_inputs_set = set(model_inputs) + + # We are fine if config_inputs has more keys than model_inputs + forward_inputs_set = set(forward_parameters.keys()) + is_ok = model_inputs_set.issubset(forward_inputs_set) + + # Make sure the input order match (VERY IMPORTANT !!!!) + matching_inputs = forward_inputs_set.intersection(model_inputs_set) + ordered_inputs = [parameter for parameter in forward_parameters.keys() if parameter in matching_inputs] + return is_ok, ordered_inputs diff --git a/lib/python3.12/site-packages/transformers/onnx/features.py b/lib/python3.12/site-packages/transformers/onnx/features.py new file mode 100644 index 0000000000000000000000000000000000000000..da9ca2355fef785f76f727a66b36bfb95a24ee51 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/onnx/features.py @@ -0,0 +1,749 @@ +import os +from functools import partial, reduce +from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple, Type, Union + +import transformers + +from .. import PretrainedConfig, is_tf_available, is_torch_available +from ..utils import TF2_WEIGHTS_NAME, WEIGHTS_NAME, logging +from .config import OnnxConfig + + +if TYPE_CHECKING: + from transformers import PreTrainedModel, TFPreTrainedModel + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +if is_torch_available(): + from transformers.models.auto import ( + AutoModel, + AutoModelForCausalLM, + AutoModelForImageClassification, + AutoModelForImageSegmentation, + AutoModelForMaskedImageModeling, + AutoModelForMaskedLM, + AutoModelForMultipleChoice, + AutoModelForObjectDetection, + AutoModelForQuestionAnswering, + AutoModelForSemanticSegmentation, + AutoModelForSeq2SeqLM, + AutoModelForSequenceClassification, + AutoModelForSpeechSeq2Seq, + AutoModelForTokenClassification, + AutoModelForVision2Seq, + ) +if is_tf_available(): + from transformers.models.auto import ( + TFAutoModel, + TFAutoModelForCausalLM, + TFAutoModelForMaskedLM, + TFAutoModelForMultipleChoice, + TFAutoModelForQuestionAnswering, + TFAutoModelForSemanticSegmentation, + TFAutoModelForSeq2SeqLM, + TFAutoModelForSequenceClassification, + TFAutoModelForTokenClassification, + ) +if not is_torch_available() and not is_tf_available(): + logger.warning( + "The ONNX export features are only supported for PyTorch or TensorFlow. You will not be able to export models" + " without one of these libraries installed." + ) + + +def supported_features_mapping( + *supported_features: str, onnx_config_cls: Optional[str] = None +) -> Dict[str, Callable[[PretrainedConfig], OnnxConfig]]: + """ + Generate the mapping between supported the features and their corresponding OnnxConfig for a given model. + + Args: + *supported_features: The names of the supported features. + onnx_config_cls: The OnnxConfig full name corresponding to the model. + + Returns: + The dictionary mapping a feature to an OnnxConfig constructor. + """ + if onnx_config_cls is None: + raise ValueError("A OnnxConfig class must be provided") + + config_cls = transformers + for attr_name in onnx_config_cls.split("."): + config_cls = getattr(config_cls, attr_name) + mapping = {} + for feature in supported_features: + if "-with-past" in feature: + task = feature.replace("-with-past", "") + mapping[feature] = partial(config_cls.with_past, task=task) + else: + mapping[feature] = partial(config_cls.from_model_config, task=feature) + + return mapping + + +class FeaturesManager: + _TASKS_TO_AUTOMODELS = {} + _TASKS_TO_TF_AUTOMODELS = {} + if is_torch_available(): + _TASKS_TO_AUTOMODELS = { + "default": AutoModel, + "masked-lm": AutoModelForMaskedLM, + "causal-lm": AutoModelForCausalLM, + "seq2seq-lm": AutoModelForSeq2SeqLM, + "sequence-classification": AutoModelForSequenceClassification, + "token-classification": AutoModelForTokenClassification, + "multiple-choice": AutoModelForMultipleChoice, + "object-detection": AutoModelForObjectDetection, + "question-answering": AutoModelForQuestionAnswering, + "image-classification": AutoModelForImageClassification, + "image-segmentation": AutoModelForImageSegmentation, + "masked-im": AutoModelForMaskedImageModeling, + "semantic-segmentation": AutoModelForSemanticSegmentation, + "vision2seq-lm": AutoModelForVision2Seq, + "speech2seq-lm": AutoModelForSpeechSeq2Seq, + } + if is_tf_available(): + _TASKS_TO_TF_AUTOMODELS = { + "default": TFAutoModel, + "masked-lm": TFAutoModelForMaskedLM, + "causal-lm": TFAutoModelForCausalLM, + "seq2seq-lm": TFAutoModelForSeq2SeqLM, + "sequence-classification": TFAutoModelForSequenceClassification, + "token-classification": TFAutoModelForTokenClassification, + "multiple-choice": TFAutoModelForMultipleChoice, + "question-answering": TFAutoModelForQuestionAnswering, + "semantic-segmentation": TFAutoModelForSemanticSegmentation, + } + + # Set of model topologies we support associated to the features supported by each topology and the factory + _SUPPORTED_MODEL_TYPE = { + "albert": supported_features_mapping( + "default", + "masked-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.albert.AlbertOnnxConfig", + ), + "bart": supported_features_mapping( + "default", + "default-with-past", + "causal-lm", + "causal-lm-with-past", + "seq2seq-lm", + "seq2seq-lm-with-past", + "sequence-classification", + "question-answering", + onnx_config_cls="models.bart.BartOnnxConfig", + ), + # BEiT cannot be used with the masked image modeling autoclass, so this feature is excluded here + "beit": supported_features_mapping( + "default", "image-classification", onnx_config_cls="models.beit.BeitOnnxConfig" + ), + "bert": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.bert.BertOnnxConfig", + ), + "big-bird": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.big_bird.BigBirdOnnxConfig", + ), + "bigbird-pegasus": supported_features_mapping( + "default", + "default-with-past", + "causal-lm", + "causal-lm-with-past", + "seq2seq-lm", + "seq2seq-lm-with-past", + "sequence-classification", + "question-answering", + onnx_config_cls="models.bigbird_pegasus.BigBirdPegasusOnnxConfig", + ), + "blenderbot": supported_features_mapping( + "default", + "default-with-past", + "causal-lm", + "causal-lm-with-past", + "seq2seq-lm", + "seq2seq-lm-with-past", + onnx_config_cls="models.blenderbot.BlenderbotOnnxConfig", + ), + "blenderbot-small": supported_features_mapping( + "default", + "default-with-past", + "causal-lm", + "causal-lm-with-past", + "seq2seq-lm", + "seq2seq-lm-with-past", + onnx_config_cls="models.blenderbot_small.BlenderbotSmallOnnxConfig", + ), + "bloom": supported_features_mapping( + "default", + "default-with-past", + "causal-lm", + "causal-lm-with-past", + "sequence-classification", + "token-classification", + onnx_config_cls="models.bloom.BloomOnnxConfig", + ), + "camembert": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.camembert.CamembertOnnxConfig", + ), + "clip": supported_features_mapping( + "default", + onnx_config_cls="models.clip.CLIPOnnxConfig", + ), + "codegen": supported_features_mapping( + "default", + "causal-lm", + onnx_config_cls="models.codegen.CodeGenOnnxConfig", + ), + "convbert": supported_features_mapping( + "default", + "masked-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.convbert.ConvBertOnnxConfig", + ), + "convnext": supported_features_mapping( + "default", + "image-classification", + onnx_config_cls="models.convnext.ConvNextOnnxConfig", + ), + "data2vec-text": supported_features_mapping( + "default", + "masked-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.data2vec.Data2VecTextOnnxConfig", + ), + "data2vec-vision": supported_features_mapping( + "default", + "image-classification", + # ONNX doesn't support `adaptive_avg_pool2d` yet + # "semantic-segmentation", + onnx_config_cls="models.data2vec.Data2VecVisionOnnxConfig", + ), + "deberta": supported_features_mapping( + "default", + "masked-lm", + "sequence-classification", + "token-classification", + "question-answering", + onnx_config_cls="models.deberta.DebertaOnnxConfig", + ), + "deberta-v2": supported_features_mapping( + "default", + "masked-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.deberta_v2.DebertaV2OnnxConfig", + ), + "deit": supported_features_mapping( + "default", "image-classification", onnx_config_cls="models.deit.DeiTOnnxConfig" + ), + "detr": supported_features_mapping( + "default", + "object-detection", + "image-segmentation", + onnx_config_cls="models.detr.DetrOnnxConfig", + ), + "distilbert": supported_features_mapping( + "default", + "masked-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.distilbert.DistilBertOnnxConfig", + ), + "electra": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.electra.ElectraOnnxConfig", + ), + "flaubert": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.flaubert.FlaubertOnnxConfig", + ), + "gpt2": supported_features_mapping( + "default", + "default-with-past", + "causal-lm", + "causal-lm-with-past", + "sequence-classification", + "token-classification", + onnx_config_cls="models.gpt2.GPT2OnnxConfig", + ), + "gptj": supported_features_mapping( + "default", + "default-with-past", + "causal-lm", + "causal-lm-with-past", + "question-answering", + "sequence-classification", + onnx_config_cls="models.gptj.GPTJOnnxConfig", + ), + "gpt-neo": supported_features_mapping( + "default", + "default-with-past", + "causal-lm", + "causal-lm-with-past", + "sequence-classification", + onnx_config_cls="models.gpt_neo.GPTNeoOnnxConfig", + ), + "groupvit": supported_features_mapping( + "default", + onnx_config_cls="models.groupvit.GroupViTOnnxConfig", + ), + "ibert": supported_features_mapping( + "default", + "masked-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.ibert.IBertOnnxConfig", + ), + "imagegpt": supported_features_mapping( + "default", "image-classification", onnx_config_cls="models.imagegpt.ImageGPTOnnxConfig" + ), + "layoutlm": supported_features_mapping( + "default", + "masked-lm", + "sequence-classification", + "token-classification", + onnx_config_cls="models.layoutlm.LayoutLMOnnxConfig", + ), + "layoutlmv3": supported_features_mapping( + "default", + "question-answering", + "sequence-classification", + "token-classification", + onnx_config_cls="models.layoutlmv3.LayoutLMv3OnnxConfig", + ), + "levit": supported_features_mapping( + "default", "image-classification", onnx_config_cls="models.levit.LevitOnnxConfig" + ), + "longt5": supported_features_mapping( + "default", + "default-with-past", + "seq2seq-lm", + "seq2seq-lm-with-past", + onnx_config_cls="models.longt5.LongT5OnnxConfig", + ), + "longformer": supported_features_mapping( + "default", + "masked-lm", + "multiple-choice", + "question-answering", + "sequence-classification", + "token-classification", + onnx_config_cls="models.longformer.LongformerOnnxConfig", + ), + "marian": supported_features_mapping( + "default", + "default-with-past", + "seq2seq-lm", + "seq2seq-lm-with-past", + "causal-lm", + "causal-lm-with-past", + onnx_config_cls="models.marian.MarianOnnxConfig", + ), + "mbart": supported_features_mapping( + "default", + "default-with-past", + "causal-lm", + "causal-lm-with-past", + "seq2seq-lm", + "seq2seq-lm-with-past", + "sequence-classification", + "question-answering", + onnx_config_cls="models.mbart.MBartOnnxConfig", + ), + "mobilebert": supported_features_mapping( + "default", + "masked-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.mobilebert.MobileBertOnnxConfig", + ), + "mobilenet-v1": supported_features_mapping( + "default", + "image-classification", + onnx_config_cls="models.mobilenet_v1.MobileNetV1OnnxConfig", + ), + "mobilenet-v2": supported_features_mapping( + "default", + "image-classification", + onnx_config_cls="models.mobilenet_v2.MobileNetV2OnnxConfig", + ), + "mobilevit": supported_features_mapping( + "default", + "image-classification", + onnx_config_cls="models.mobilevit.MobileViTOnnxConfig", + ), + "mt5": supported_features_mapping( + "default", + "default-with-past", + "seq2seq-lm", + "seq2seq-lm-with-past", + onnx_config_cls="models.mt5.MT5OnnxConfig", + ), + "m2m-100": supported_features_mapping( + "default", + "default-with-past", + "seq2seq-lm", + "seq2seq-lm-with-past", + onnx_config_cls="models.m2m_100.M2M100OnnxConfig", + ), + "owlvit": supported_features_mapping( + "default", + onnx_config_cls="models.owlvit.OwlViTOnnxConfig", + ), + "perceiver": supported_features_mapping( + "image-classification", + "masked-lm", + "sequence-classification", + onnx_config_cls="models.perceiver.PerceiverOnnxConfig", + ), + "poolformer": supported_features_mapping( + "default", "image-classification", onnx_config_cls="models.poolformer.PoolFormerOnnxConfig" + ), + "rembert": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.rembert.RemBertOnnxConfig", + ), + "resnet": supported_features_mapping( + "default", + "image-classification", + onnx_config_cls="models.resnet.ResNetOnnxConfig", + ), + "roberta": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.roberta.RobertaOnnxConfig", + ), + "roformer": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "token-classification", + "multiple-choice", + "question-answering", + "token-classification", + onnx_config_cls="models.roformer.RoFormerOnnxConfig", + ), + "segformer": supported_features_mapping( + "default", + "image-classification", + "semantic-segmentation", + onnx_config_cls="models.segformer.SegformerOnnxConfig", + ), + "squeezebert": supported_features_mapping( + "default", + "masked-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.squeezebert.SqueezeBertOnnxConfig", + ), + "swin": supported_features_mapping( + "default", "image-classification", onnx_config_cls="models.swin.SwinOnnxConfig" + ), + "t5": supported_features_mapping( + "default", + "default-with-past", + "seq2seq-lm", + "seq2seq-lm-with-past", + onnx_config_cls="models.t5.T5OnnxConfig", + ), + "vision-encoder-decoder": supported_features_mapping( + "vision2seq-lm", onnx_config_cls="models.vision_encoder_decoder.VisionEncoderDecoderOnnxConfig" + ), + "vit": supported_features_mapping( + "default", "image-classification", onnx_config_cls="models.vit.ViTOnnxConfig" + ), + "whisper": supported_features_mapping( + "default", + "default-with-past", + "speech2seq-lm", + "speech2seq-lm-with-past", + onnx_config_cls="models.whisper.WhisperOnnxConfig", + ), + "xlm": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.xlm.XLMOnnxConfig", + ), + "xlm-roberta": supported_features_mapping( + "default", + "masked-lm", + "causal-lm", + "sequence-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx_config_cls="models.xlm_roberta.XLMRobertaOnnxConfig", + ), + "yolos": supported_features_mapping( + "default", + "object-detection", + onnx_config_cls="models.yolos.YolosOnnxConfig", + ), + } + + AVAILABLE_FEATURES = sorted(reduce(lambda s1, s2: s1 | s2, (v.keys() for v in _SUPPORTED_MODEL_TYPE.values()))) + + @staticmethod + def get_supported_features_for_model_type( + model_type: str, model_name: Optional[str] = None + ) -> Dict[str, Callable[[PretrainedConfig], OnnxConfig]]: + """ + Tries to retrieve the feature -> OnnxConfig constructor map from the model type. + + Args: + model_type (`str`): + The model type to retrieve the supported features for. + model_name (`str`, *optional*): + The name attribute of the model object, only used for the exception message. + + Returns: + The dictionary mapping each feature to a corresponding OnnxConfig constructor. + """ + model_type = model_type.lower() + if model_type not in FeaturesManager._SUPPORTED_MODEL_TYPE: + model_type_and_model_name = f"{model_type} ({model_name})" if model_name else model_type + raise KeyError( + f"{model_type_and_model_name} is not supported yet. " + f"Only {list(FeaturesManager._SUPPORTED_MODEL_TYPE.keys())} are supported. " + f"If you want to support {model_type} please propose a PR or open up an issue." + ) + return FeaturesManager._SUPPORTED_MODEL_TYPE[model_type] + + @staticmethod + def feature_to_task(feature: str) -> str: + return feature.replace("-with-past", "") + + @staticmethod + def _validate_framework_choice(framework: str): + """ + Validates if the framework requested for the export is both correct and available, otherwise throws an + exception. + """ + if framework not in ["pt", "tf"]: + raise ValueError( + f"Only two frameworks are supported for ONNX export: pt or tf, but {framework} was provided." + ) + elif framework == "pt" and not is_torch_available(): + raise RuntimeError("Cannot export model to ONNX using PyTorch because no PyTorch package was found.") + elif framework == "tf" and not is_tf_available(): + raise RuntimeError("Cannot export model to ONNX using TensorFlow because no TensorFlow package was found.") + + @staticmethod + def get_model_class_for_feature(feature: str, framework: str = "pt") -> Type: + """ + Attempts to retrieve an AutoModel class from a feature name. + + Args: + feature (`str`): + The feature required. + framework (`str`, *optional*, defaults to `"pt"`): + The framework to use for the export. + + Returns: + The AutoModel class corresponding to the feature. + """ + task = FeaturesManager.feature_to_task(feature) + FeaturesManager._validate_framework_choice(framework) + if framework == "pt": + task_to_automodel = FeaturesManager._TASKS_TO_AUTOMODELS + else: + task_to_automodel = FeaturesManager._TASKS_TO_TF_AUTOMODELS + if task not in task_to_automodel: + raise KeyError( + f"Unknown task: {feature}. Possible values are {list(FeaturesManager._TASKS_TO_AUTOMODELS.values())}" + ) + + return task_to_automodel[task] + + @staticmethod + def determine_framework(model: str, framework: Optional[str] = None) -> str: + """ + Determines the framework to use for the export. + + The priority is in the following order: + 1. User input via `framework`. + 2. If local checkpoint is provided, use the same framework as the checkpoint. + 3. Available framework in environment, with priority given to PyTorch + + Args: + model (`str`): + The name of the model to export. + framework (`str`, *optional*, defaults to `None`): + The framework to use for the export. See above for priority if none provided. + + Returns: + The framework to use for the export. + + """ + if framework is not None: + return framework + + framework_map = {"pt": "PyTorch", "tf": "TensorFlow"} + exporter_map = {"pt": "torch", "tf": "tf2onnx"} + + if os.path.isdir(model): + if os.path.isfile(os.path.join(model, WEIGHTS_NAME)): + framework = "pt" + elif os.path.isfile(os.path.join(model, TF2_WEIGHTS_NAME)): + framework = "tf" + else: + raise FileNotFoundError( + "Cannot determine framework from given checkpoint location." + f" There should be a {WEIGHTS_NAME} for PyTorch" + f" or {TF2_WEIGHTS_NAME} for TensorFlow." + ) + logger.info(f"Local {framework_map[framework]} model found.") + else: + if is_torch_available(): + framework = "pt" + elif is_tf_available(): + framework = "tf" + else: + raise EnvironmentError("Neither PyTorch nor TensorFlow found in environment. Cannot export to ONNX.") + + logger.info(f"Framework not requested. Using {exporter_map[framework]} to export to ONNX.") + + return framework + + @staticmethod + def get_model_from_feature( + feature: str, model: str, framework: Optional[str] = None, cache_dir: Optional[str] = None + ) -> Union["PreTrainedModel", "TFPreTrainedModel"]: + """ + Attempts to retrieve a model from a model's name and the feature to be enabled. + + Args: + feature (`str`): + The feature required. + model (`str`): + The name of the model to export. + framework (`str`, *optional*, defaults to `None`): + The framework to use for the export. See `FeaturesManager.determine_framework` for the priority should + none be provided. + + Returns: + The instance of the model. + + """ + framework = FeaturesManager.determine_framework(model, framework) + model_class = FeaturesManager.get_model_class_for_feature(feature, framework) + try: + model = model_class.from_pretrained(model, cache_dir=cache_dir) + except OSError: + if framework == "pt": + logger.info("Loading TensorFlow model in PyTorch before exporting to ONNX.") + model = model_class.from_pretrained(model, from_tf=True, cache_dir=cache_dir) + else: + logger.info("Loading PyTorch model in TensorFlow before exporting to ONNX.") + model = model_class.from_pretrained(model, from_pt=True, cache_dir=cache_dir) + return model + + @staticmethod + def check_supported_model_or_raise( + model: Union["PreTrainedModel", "TFPreTrainedModel"], feature: str = "default" + ) -> Tuple[str, Callable]: + """ + Check whether or not the model has the requested features. + + Args: + model: The model to export. + feature: The name of the feature to check if it is available. + + Returns: + (str) The type of the model (OnnxConfig) The OnnxConfig instance holding the model export properties. + + """ + model_type = model.config.model_type.replace("_", "-") + model_name = getattr(model, "name", "") + model_features = FeaturesManager.get_supported_features_for_model_type(model_type, model_name=model_name) + if feature not in model_features: + raise ValueError( + f"{model.config.model_type} doesn't support feature {feature}. Supported values are: {model_features}" + ) + + return model.config.model_type, FeaturesManager._SUPPORTED_MODEL_TYPE[model_type][feature] + + def get_config(model_type: str, feature: str) -> OnnxConfig: + """ + Gets the OnnxConfig for a model_type and feature combination. + + Args: + model_type (`str`): + The model type to retrieve the config for. + feature (`str`): + The feature to retrieve the config for. + + Returns: + `OnnxConfig`: config for the combination + """ + return FeaturesManager._SUPPORTED_MODEL_TYPE[model_type][feature] diff --git a/lib/python3.12/site-packages/transformers/onnx/utils.py b/lib/python3.12/site-packages/transformers/onnx/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9672b0a96af88ffa2c7e791d1f4d7c818174247f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/onnx/utils.py @@ -0,0 +1,109 @@ +# Copyright 2021 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. + +from ctypes import c_float, sizeof +from enum import Enum +from typing import TYPE_CHECKING, Optional, Union + + +if TYPE_CHECKING: + from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer # tests_ignore + + +class ParameterFormat(Enum): + Float = c_float + + @property + def size(self) -> int: + """ + Number of byte required for this data type + + Returns: + Integer > 0 + """ + return sizeof(self.value) + + +def compute_effective_axis_dimension(dimension: int, fixed_dimension: int, num_token_to_add: int = 0) -> int: + """ + + Args: + dimension: + fixed_dimension: + num_token_to_add: + + Returns: + + """ + # < 0 is possible if using a dynamic axis + if dimension <= 0: + dimension = fixed_dimension + + dimension -= num_token_to_add + return dimension + + +def compute_serialized_parameters_size(num_parameters: int, dtype: ParameterFormat) -> int: + """ + Compute the size taken by all the parameters in the given the storage format when serializing the model + + Args: + num_parameters: Number of parameters to be saved + dtype: The data format each parameter will be saved + + Returns: + Size (in byte) taken to save all the parameters + """ + return num_parameters * dtype.size + + +def get_preprocessor(model_name: str) -> Optional[Union["AutoTokenizer", "AutoFeatureExtractor", "AutoProcessor"]]: + """ + Gets a preprocessor (tokenizer, feature extractor or processor) that is available for `model_name`. + + Args: + model_name (`str`): Name of the model for which a preprocessor are loaded. + + Returns: + `Optional[Union[AutoTokenizer, AutoFeatureExtractor, AutoProcessor]]`: + If a processor is found, it is returned. Otherwise, if a tokenizer or a feature extractor exists, it is + returned. If both a tokenizer and a feature extractor exist, an error is raised. The function returns + `None` if no preprocessor is found. + """ + # Avoid circular imports by only importing this here. + from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer # tests_ignore + + try: + return AutoProcessor.from_pretrained(model_name) + except (ValueError, OSError, KeyError): + tokenizer, feature_extractor = None, None + try: + tokenizer = AutoTokenizer.from_pretrained(model_name) + except (OSError, KeyError): + pass + try: + feature_extractor = AutoFeatureExtractor.from_pretrained(model_name) + except (OSError, KeyError): + pass + + if tokenizer is not None and feature_extractor is not None: + raise ValueError( + f"Couldn't auto-detect preprocessor for {model_name}. Found both a tokenizer and a feature extractor." + ) + elif tokenizer is None and feature_extractor is None: + return None + elif tokenizer is not None: + return tokenizer + else: + return feature_extractor diff --git a/lib/python3.12/site-packages/transformers/optimization.py b/lib/python3.12/site-packages/transformers/optimization.py new file mode 100644 index 0000000000000000000000000000000000000000..ba9d9920c1c99a50aa16d011bad68d012cba7eb7 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/optimization.py @@ -0,0 +1,892 @@ +# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team. +# +# 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. +"""PyTorch optimization for BERT model.""" + +import math +import warnings +from functools import partial +from typing import Optional, Union + +import torch +from torch.optim import Optimizer +from torch.optim.lr_scheduler import LambdaLR, ReduceLROnPlateau + +from .trainer_pt_utils import LayerWiseDummyOptimizer, LayerWiseDummyScheduler +from .trainer_utils import SchedulerType +from .utils import logging +from .utils.versions import require_version + + +logger = logging.get_logger(__name__) + + +def _get_constant_lambda(_=None): + return 1 + + +def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1): + """ + Create a schedule with a constant learning rate, using the learning rate set in optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + return LambdaLR(optimizer, _get_constant_lambda, last_epoch=last_epoch) + + +def get_reduce_on_plateau_schedule(optimizer: Optimizer, **kwargs): + """ + Create a schedule with a constant learning rate that decreases when a metric has stopped improving. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + kwargs (`dict`, *optional*): + Extra parameters to be passed to the scheduler. See `torch.optim.lr_scheduler.ReduceLROnPlateau` + for possible parameters. + + Return: + `torch.optim.lr_scheduler.ReduceLROnPlateau` with the appropriate schedule. + """ + + return ReduceLROnPlateau(optimizer, **kwargs) + + +def _get_constant_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1.0, num_warmup_steps)) + return 1.0 + + +def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1): + """ + Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate + increases linearly between 0 and the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + lr_lambda = partial(_get_constant_schedule_with_warmup_lr_lambda, num_warmup_steps=num_warmup_steps) + return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch) + + +def _get_linear_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))) + + +def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1): + """ + Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after + a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + lr_lambda = partial( + _get_linear_schedule_with_warmup_lr_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def _get_cosine_schedule_with_warmup_lr_lambda( + current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float +): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) + return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) + + +def get_cosine_schedule_with_warmup( + optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1 +): + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the + initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_cycles (`float`, *optional*, defaults to 0.5): + The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 + following a half-cosine). + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + lr_lambda = partial( + _get_cosine_schedule_with_warmup_lr_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + num_cycles=num_cycles, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda( + current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: int +): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) + if progress >= 1.0: + return 0.0 + return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0)))) + + +def get_cosine_with_hard_restarts_schedule_with_warmup( + optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1 +): + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases + linearly between 0 and the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_cycles (`int`, *optional*, defaults to 1): + The number of hard restarts to use. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + lr_lambda = partial( + _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + num_cycles=num_cycles, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def _get_polynomial_decay_schedule_with_warmup_lr_lambda( + current_step: int, + *, + num_warmup_steps: int, + num_training_steps: int, + lr_end: float, + power: float, + lr_init: int, +): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + elif current_step > num_training_steps: + return lr_end / lr_init # as LambdaLR multiplies by lr_init + else: + lr_range = lr_init - lr_end + decay_steps = num_training_steps - num_warmup_steps + pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps + decay = lr_range * pct_remaining**power + lr_end + return decay / lr_init # as LambdaLR multiplies by lr_init + + +def get_polynomial_decay_schedule_with_warmup( + optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1 +): + """ + Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the + optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the + initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + lr_end (`float`, *optional*, defaults to 1e-7): + The end LR. + power (`float`, *optional*, defaults to 1.0): + Power factor. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT + implementation at + https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37 + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + + """ + + lr_init = optimizer.defaults["lr"] + if not (lr_init > lr_end): + raise ValueError(f"lr_end ({lr_end}) must be smaller than initial lr ({lr_init})") + + lr_lambda = partial( + _get_polynomial_decay_schedule_with_warmup_lr_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + lr_end=lr_end, + power=power, + lr_init=lr_init, + ) + 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): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + shift = timescale - num_warmup_steps + decay = 1.0 / math.sqrt((current_step + shift) / timescale) + return decay + + +def get_inverse_sqrt_schedule( + optimizer: Optimizer, num_warmup_steps: int, timescale: Optional[int] = 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 + warmup period which increases lr linearly from 0 to the initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + timescale (`int`, *optional*, defaults to `num_warmup_steps`): + Time scale. + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + # Note: this implementation is adapted from + # https://github.com/google-research/big_vision/blob/f071ce68852d56099437004fd70057597a95f6ef/big_vision/utils.py#L930 + + if timescale is None: + timescale = num_warmup_steps or 10_000 + + lr_lambda = partial(_get_inverse_sqrt_schedule_lr_lambda, num_warmup_steps=num_warmup_steps, timescale=timescale) + return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch) + + +def _get_cosine_schedule_with_warmup_lr_lambda( + current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float, min_lr_rate: float = 0.0 +): + if current_step < num_warmup_steps: + return float(current_step) / float(max(1, num_warmup_steps)) + progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) + factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)) + factor = factor * (1 - min_lr_rate) + min_lr_rate + return max(0, factor) + + +def get_cosine_with_min_lr_schedule_with_warmup( + optimizer: Optimizer, + num_warmup_steps: int, + num_training_steps: int, + num_cycles: float = 0.5, + last_epoch: int = -1, + min_lr: Optional[float] = None, + min_lr_rate: Optional[float] = None, +): + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to min_lr, after a warmup period during which it increases linearly between 0 and the + initial lr set in the optimizer. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_training_steps (`int`): + The total number of training steps. + num_cycles (`float`, *optional*, defaults to 0.5): + The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 + following a half-cosine). + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + min_lr (`float`, *optional*): + The minimum learning rate to reach after the cosine schedule. + min_lr_rate (`float`, *optional*): + The minimum learning rate as a ratio of the initial learning rate. If set, `min_lr` should not be set. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + if min_lr is not None and min_lr_rate is not None: + raise ValueError("Only one of min_lr or min_lr_rate should be set") + elif min_lr is not None: + min_lr_rate = min_lr / optimizer.defaults["lr"] + elif min_lr_rate is None: + raise ValueError("One of min_lr or min_lr_rate should be set through the `lr_scheduler_kwargs`") + + lr_lambda = partial( + _get_cosine_schedule_with_warmup_lr_lambda, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + num_cycles=num_cycles, + min_lr_rate=min_lr_rate, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +def _get_wsd_scheduler_lambda( + current_step: int, + *, + num_warmup_steps: int, + num_stable_steps: int, + num_decay_steps: int, + warmup_type: str, + decay_type: str, + min_lr_ratio: float, + num_cycles: float, +): + if current_step < num_warmup_steps: + progress = float(current_step) / float(max(1, num_warmup_steps)) + if warmup_type == "linear": + factor = progress + elif warmup_type == "cosine": + factor = 0.5 * (1.0 - math.cos(math.pi * progress)) + elif warmup_type == "1-sqrt": + factor = 1.0 - math.sqrt(1.0 - progress) + factor = factor * (1.0 - min_lr_ratio) + min_lr_ratio + return max(0.0, factor) + + if current_step < num_warmup_steps + num_stable_steps: + return 1.0 + + if current_step < num_warmup_steps + num_stable_steps + num_decay_steps: + progress = float(current_step - num_warmup_steps - num_stable_steps) / float(max(1, num_decay_steps)) + if decay_type == "linear": + factor = 1.0 - progress + elif decay_type == "cosine": + factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)) + elif decay_type == "1-sqrt": + factor = 1.0 - math.sqrt(progress) + factor = factor * (1.0 - min_lr_ratio) + min_lr_ratio + return max(0.0, factor) + return min_lr_ratio + + +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, + warmup_type: str = "linear", + decay_type: str = "cosine", + min_lr_ratio: float = 0, + num_cycles: float = 0.5, + last_epoch: int = -1, +): + """ + Create a schedule with a learning rate that has three stages: + 1. warmup: increase from min_lr_ratio times the initial learning rate to the initial learning rate following a warmup_type. + 2. stable: constant learning rate. + 3. decay: decrease from the initial learning rate to min_lr_ratio times the initial learning rate following a decay_type. + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + num_warmup_steps (`int`): + The number of steps for the warmup phase. + num_decay_steps (`int`): + The number of steps for the decay phase. + num_training_steps (`int`, *optional*): + The total number of training steps. This is the sum of the warmup, stable and decay steps. If `num_stable_steps` is not provided, the stable phase will be `num_training_steps - num_warmup_steps - num_decay_steps`. + num_stable_steps (`int`, *optional*): + The number of steps for the stable phase. Please ensure that `num_warmup_steps + num_stable_steps + num_decay_steps` equals `num_training_steps`, otherwise the other steps will default to the minimum learning rate. + warmup_type (`str`, *optional*, defaults to "linear"): + The type of warmup to use. Can be 'linear', 'cosine' or '1-sqrt'. + decay_type (`str`, *optional*, defaults to "cosine"): + The type of decay to use. Can be 'linear', 'cosine' or '1-sqrt'. + min_lr_ratio (`float`, *optional*, defaults to 0): + The minimum learning rate as a ratio of the initial learning rate. + num_cycles (`float`, *optional*, defaults to 0.5): + The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 + following a half-cosine). + last_epoch (`int`, *optional*, defaults to -1): + The index of the last epoch when resuming training. + + Return: + `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + + if num_training_steps is None and num_stable_steps is None: + raise ValueError("Either num_training_steps or num_stable_steps must be specified.") + + if num_training_steps is not None and num_stable_steps is not None: + warnings.warn("Both num_training_steps and num_stable_steps are specified. num_stable_steps will be used.") + + if warmup_type not in ["linear", "cosine", "1-sqrt"]: + raise ValueError(f"Unknown warmup type: {warmup_type}, expected 'linear', 'cosine' or '1-sqrt'") + + if decay_type not in ["linear", "cosine", "1-sqrt"]: + raise ValueError(f"Unknown decay type: {decay_type}, expected 'linear', 'cosine' or '1-sqrt'") + + if num_stable_steps is None: + num_stable_steps = num_training_steps - num_warmup_steps - num_decay_steps + + lr_lambda = partial( + _get_wsd_scheduler_lambda, + num_warmup_steps=num_warmup_steps, + num_stable_steps=num_stable_steps, + num_decay_steps=num_decay_steps, + warmup_type=warmup_type, + decay_type=decay_type, + min_lr_ratio=min_lr_ratio, + num_cycles=num_cycles, + ) + return LambdaLR(optimizer, lr_lambda, last_epoch) + + +TYPE_TO_SCHEDULER_FUNCTION = { + SchedulerType.LINEAR: get_linear_schedule_with_warmup, + SchedulerType.COSINE: get_cosine_schedule_with_warmup, + SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup, + SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup, + SchedulerType.CONSTANT: get_constant_schedule, + SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup, + SchedulerType.INVERSE_SQRT: get_inverse_sqrt_schedule, + SchedulerType.REDUCE_ON_PLATEAU: get_reduce_on_plateau_schedule, + SchedulerType.COSINE_WITH_MIN_LR: get_cosine_with_min_lr_schedule_with_warmup, + SchedulerType.WARMUP_STABLE_DECAY: get_wsd_schedule, +} + + +def get_scheduler( + name: Union[str, SchedulerType], + optimizer: Optimizer, + num_warmup_steps: Optional[int] = None, + num_training_steps: Optional[int] = None, + scheduler_specific_kwargs: Optional[dict] = None, +): + """ + Unified API to get any scheduler from its name. + + Args: + name (`str` or `SchedulerType`): + The name of the scheduler to use. + optimizer (`torch.optim.Optimizer`): + The optimizer that will be used during training. + num_warmup_steps (`int`, *optional*): + The number of warmup steps to do. This is not required by all schedulers (hence the argument being + optional), the function will raise an error if it's unset and the scheduler type requires it. + num_training_steps (`int``, *optional*): + The number of training steps to do. This is not required by all schedulers (hence the argument being + optional), the function will raise an error if it's unset and the scheduler type requires it. + scheduler_specific_kwargs (`dict`, *optional*): + Extra parameters for schedulers such as cosine with restarts. Mismatched scheduler types and scheduler + parameters will cause the scheduler function to raise a TypeError. + """ + name = SchedulerType(name) + schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name] + + # If a `LayerWiseDummyOptimizer` is passed we extract the optimizer dict and + # recursively call `get_scheduler` to get the proper schedulers on each parameter + if optimizer is not None and isinstance(optimizer, LayerWiseDummyOptimizer): + optimizer_dict = optimizer.optimizer_dict + scheduler_dict = {} + + for param in optimizer_dict.keys(): + scheduler_dict[param] = get_scheduler( + name, + optimizer=optimizer_dict[param], + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + ) + + def scheduler_hook(param): + # Since the optimizer hook has been already attached we only need to + # attach the scheduler hook, the gradients have been zeroed here + scheduler_dict[param].step() + + for param in optimizer_dict.keys(): + if param.requires_grad: + param.register_post_accumulate_grad_hook(scheduler_hook) + + return LayerWiseDummyScheduler(optimizer_dict=optimizer_dict, lr=optimizer.defaults["lr"]) + + if name == SchedulerType.CONSTANT: + return schedule_func(optimizer) + + if scheduler_specific_kwargs is None: + scheduler_specific_kwargs = {} + + if name == SchedulerType.REDUCE_ON_PLATEAU: + return schedule_func(optimizer, **scheduler_specific_kwargs) + + # All other schedulers require `num_warmup_steps` + if num_warmup_steps is None: + raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.") + + if name == SchedulerType.CONSTANT_WITH_WARMUP: + return schedule_func(optimizer, num_warmup_steps=num_warmup_steps) + + if name == SchedulerType.INVERSE_SQRT: + return schedule_func(optimizer, num_warmup_steps=num_warmup_steps) + + # wsd scheduler requires either num_training_steps or num_stable_steps + if name == SchedulerType.WARMUP_STABLE_DECAY: + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + **scheduler_specific_kwargs, + ) + + # All other schedulers require `num_training_steps` + if num_training_steps is None: + raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.") + + return schedule_func( + optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=num_training_steps, + **scheduler_specific_kwargs, + ) + + +class Adafactor(Optimizer): + """ + AdaFactor pytorch implementation can be used as a drop in replacement for Adam original fairseq code: + https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py + + Paper: *Adafactor: Adaptive Learning Rates with Sublinear Memory Cost* https://arxiv.org/abs/1804.04235 Note that + this optimizer internally adjusts the learning rate depending on the `scale_parameter`, `relative_step` and + `warmup_init` options. To use a manual (external) learning rate schedule you should set `scale_parameter=False` and + `relative_step=False`. + + Arguments: + params (`Iterable[nn.parameter.Parameter]`): + Iterable of parameters to optimize or dictionaries defining parameter groups. + lr (`float`, *optional*): + The external learning rate. + eps (`Tuple[float, float]`, *optional*, defaults to `(1e-30, 0.001)`): + Regularization constants for square gradient and parameter scale respectively + clip_threshold (`float`, *optional*, defaults to 1.0): + Threshold of root mean square of final gradient update + decay_rate (`float`, *optional*, defaults to -0.8): + Coefficient used to compute running averages of square + beta1 (`float`, *optional*): + Coefficient used for computing running averages of gradient + weight_decay (`float`, *optional*, defaults to 0.0): + Weight decay (L2 penalty) + scale_parameter (`bool`, *optional*, defaults to `True`): + If True, learning rate is scaled by root mean square + relative_step (`bool`, *optional*, defaults to `True`): + If True, time-dependent learning rate is computed instead of external learning rate + warmup_init (`bool`, *optional*, defaults to `False`): + Time-dependent learning rate computation depends on whether warm-up initialization is being used + + This implementation handles low-precision (FP16, bfloat) values, but we have not thoroughly tested. + + Recommended T5 finetuning settings (https://discuss.huggingface.co/t/t5-finetuning-tips/684/3): + + - Training without LR warmup or clip_threshold is not recommended. + + - use scheduled LR warm-up to fixed LR + - use clip_threshold=1.0 (https://arxiv.org/abs/1804.04235) + - Disable relative updates + - Use scale_parameter=False + - Additional optimizer operations like gradient clipping should not be used alongside Adafactor + + Example: + + ```python + Adafactor(model.parameters(), scale_parameter=False, relative_step=False, warmup_init=False, lr=1e-3) + ``` + + Others reported the following combination to work well: + + ```python + Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None) + ``` + + When using `lr=None` with [`Trainer`] you will most likely need to use [`~optimization.AdafactorSchedule`] + scheduler as following: + + ```python + from transformers.optimization import Adafactor, AdafactorSchedule + + optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None) + lr_scheduler = AdafactorSchedule(optimizer) + trainer = Trainer(..., optimizers=(optimizer, lr_scheduler)) + ``` + + Usage: + + ```python + # replace AdamW with Adafactor + optimizer = Adafactor( + model.parameters(), + lr=1e-3, + eps=(1e-30, 1e-3), + clip_threshold=1.0, + decay_rate=-0.8, + beta1=None, + weight_decay=0.0, + relative_step=False, + scale_parameter=False, + warmup_init=False, + ) + ```""" + + def __init__( + self, + params, + lr=None, + eps=(1e-30, 1e-3), + clip_threshold=1.0, + decay_rate=-0.8, + beta1=None, + weight_decay=0.0, + scale_parameter=True, + relative_step=True, + warmup_init=False, + ): + require_version("torch>=1.5.0") # add_ with alpha + if lr is not None and relative_step: + raise ValueError("Cannot combine manual `lr` and `relative_step=True` options") + if warmup_init and not relative_step: + raise ValueError("`warmup_init=True` requires `relative_step=True`") + + defaults = { + "lr": lr, + "eps": eps, + "clip_threshold": clip_threshold, + "decay_rate": decay_rate, + "beta1": beta1, + "weight_decay": weight_decay, + "scale_parameter": scale_parameter, + "relative_step": relative_step, + "warmup_init": warmup_init, + } + super().__init__(params, defaults) + + @staticmethod + def _get_lr(param_group, param_state): + rel_step_sz = param_group["lr"] + if param_group["relative_step"]: + min_step = 1e-6 * param_state["step"] if param_group["warmup_init"] else 1e-2 + rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"])) + param_scale = 1.0 + if param_group["scale_parameter"]: + param_scale = max(param_group["eps"][1], param_state["RMS"]) + return param_scale * rel_step_sz + + @staticmethod + def _get_options(param_group, param_shape): + factored = len(param_shape) >= 2 + use_first_moment = param_group["beta1"] is not None + return factored, use_first_moment + + @staticmethod + def _rms(tensor): + return tensor.norm(2) / (tensor.numel() ** 0.5) + + @staticmethod + def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col): + # copy from fairseq's adafactor implementation: + # https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505 + r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1) + c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt() + return torch.mul(r_factor, c_factor) + + @torch.no_grad() + def step(self, closure=None): + """ + Performs a single optimization step + + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad + if grad.dtype in {torch.float16, torch.bfloat16}: + grad = grad.float() + if grad.is_sparse: + raise RuntimeError("Adafactor does not support sparse gradients.") + + state = self.state[p] + grad_shape = grad.shape + + factored, use_first_moment = self._get_options(group, grad_shape) + # State Initialization + if len(state) == 0: + state["step"] = 0 + + if use_first_moment: + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like(grad) + if factored: + state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad) + state["exp_avg_sq_col"] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad) + else: + state["exp_avg_sq"] = torch.zeros_like(grad) + + state["RMS"] = 0 + else: + if use_first_moment: + state["exp_avg"] = state["exp_avg"].to(grad) + if factored: + state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad) + state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad) + else: + state["exp_avg_sq"] = state["exp_avg_sq"].to(grad) + + p_data_fp32 = p + if p.dtype in {torch.float16, torch.bfloat16}: + p_data_fp32 = p_data_fp32.float() + + state["step"] += 1 + state["RMS"] = self._rms(p_data_fp32) + lr = self._get_lr(group, state) + + beta2t = 1.0 - math.pow(state["step"], group["decay_rate"]) + update = (grad**2) + group["eps"][0] + if factored: + exp_avg_sq_row = state["exp_avg_sq_row"] + exp_avg_sq_col = state["exp_avg_sq_col"] + + exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t)) + exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t)) + + # Approximation of exponential moving average of square of gradient + update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col) + update.mul_(grad) + else: + exp_avg_sq = state["exp_avg_sq"] + + exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t)) + update = exp_avg_sq.rsqrt().mul_(grad) + + update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0)) + update.mul_(lr) + + if use_first_moment: + exp_avg = state["exp_avg"] + exp_avg.mul_(group["beta1"]).add_(update, alpha=(1 - group["beta1"])) + update = exp_avg + + if group["weight_decay"] != 0: + p_data_fp32.add_(p_data_fp32, alpha=(-group["weight_decay"] * lr)) + + p_data_fp32.add_(-update) + + if p.dtype in {torch.float16, torch.bfloat16}: + p.copy_(p_data_fp32) + + return loss + + +class AdafactorSchedule(LambdaLR): + """ + Since [`~optimization.Adafactor`] performs its own scheduling, if the training loop relies on a scheduler (e.g., + for logging), this class creates a proxy object that retrieves the current lr values from the optimizer. + + It returns `initial_lr` during startup and the actual `lr` during stepping. + """ + + def __init__(self, optimizer, initial_lr=0.0): + def lr_lambda(_): + return initial_lr + + for group in optimizer.param_groups: + group["initial_lr"] = initial_lr + super().__init__(optimizer, lr_lambda) + for group in optimizer.param_groups: + del group["initial_lr"] + + def get_lr(self): + opt = self.optimizer + lrs = [ + opt._get_lr(group, opt.state[group["params"][0]]) + for group in opt.param_groups + if group["params"][0].grad is not None + ] + if len(lrs) == 0: + lrs = self.base_lrs # if called before stepping + return lrs + + +def get_adafactor_schedule(optimizer, initial_lr=0.0): + """ + Get a proxy schedule for [`~optimization.Adafactor`] + + Args: + optimizer ([`~torch.optim.Optimizer`]): + The optimizer for which to schedule the learning rate. + initial_lr (`float`, *optional*, defaults to 0.0): + Initial lr + + Return: + [`~optimization.Adafactor`] proxy schedule object. + + + """ + return AdafactorSchedule(optimizer, initial_lr) diff --git a/lib/python3.12/site-packages/transformers/optimization_tf.py b/lib/python3.12/site-packages/transformers/optimization_tf.py new file mode 100644 index 0000000000000000000000000000000000000000..4da4ecc901914bc12a94c8ac3193309b85661865 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/optimization_tf.py @@ -0,0 +1,379 @@ +# Copyright 2019 The TensorFlow Authors, The Hugging Face 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. +# ============================================================================== +"""Functions and classes related to optimization (weight updates).""" + +import re +from typing import Callable, Optional, Union + +import tensorflow as tf + + +try: + from tf_keras.optimizers.legacy import Adam +except (ImportError, ModuleNotFoundError): + from tensorflow.keras.optimizers.legacy import Adam + +from .modeling_tf_utils import keras + + +# This block because Keras loves randomly moving things to different places - this changed somewhere between 2.10 - 2.15 +if hasattr(keras.optimizers.schedules, "learning_rate_schedule"): + schedules = keras.optimizers.schedules.learning_rate_schedule +else: + schedules = keras.optimizers.schedules + + +class WarmUp(schedules.LearningRateSchedule): + """ + Applies a warmup schedule on a given learning rate decay schedule. + + Args: + initial_learning_rate (`float`): + The initial learning rate for the schedule after the warmup (so this will be the learning rate at the end + of the warmup). + decay_schedule_fn (`Callable`): + The schedule function to apply after the warmup for the rest of training. + warmup_steps (`int`): + The number of steps for the warmup part of training. + power (`float`, *optional*, defaults to 1.0): + The power to use for the polynomial warmup (defaults is a linear warmup). + name (`str`, *optional*): + Optional name prefix for the returned tensors during the schedule. + """ + + def __init__( + self, + initial_learning_rate: float, + decay_schedule_fn: Callable, + warmup_steps: int, + power: float = 1.0, + name: Optional[str] = None, + ): + super().__init__() + self.initial_learning_rate = initial_learning_rate + self.warmup_steps = warmup_steps + self.power = power + self.decay_schedule_fn = decay_schedule_fn + self.name = name + + def __call__(self, step): + with tf.name_scope(self.name or "WarmUp") as name: + # Implements polynomial warmup. i.e., if global_step < warmup_steps, the + # learning rate will be `global_step/num_warmup_steps * init_lr`. + global_step_float = tf.cast(step, tf.float32) + warmup_steps_float = tf.cast(self.warmup_steps, tf.float32) + warmup_percent_done = global_step_float / warmup_steps_float + warmup_learning_rate = self.initial_learning_rate * tf.math.pow(warmup_percent_done, self.power) + return tf.cond( + global_step_float < warmup_steps_float, + lambda: warmup_learning_rate, + lambda: self.decay_schedule_fn(step - self.warmup_steps), + name=name, + ) + + def get_config(self): + return { + "initial_learning_rate": self.initial_learning_rate, + "decay_schedule_fn": self.decay_schedule_fn, + "warmup_steps": self.warmup_steps, + "power": self.power, + "name": self.name, + } + + +def create_optimizer( + init_lr: float, + num_train_steps: int, + num_warmup_steps: int, + min_lr_ratio: float = 0.0, + adam_beta1: float = 0.9, + adam_beta2: float = 0.999, + adam_epsilon: float = 1e-8, + adam_clipnorm: Optional[float] = None, + adam_global_clipnorm: Optional[float] = None, + weight_decay_rate: float = 0.0, + power: float = 1.0, + include_in_weight_decay: Optional[list[str]] = None, +): + """ + Creates an optimizer with a learning rate schedule using a warmup phase followed by a linear decay. + + Args: + init_lr (`float`): + The desired learning rate at the end of the warmup phase. + num_train_steps (`int`): + The total number of training steps. + num_warmup_steps (`int`): + The number of warmup steps. + min_lr_ratio (`float`, *optional*, defaults to 0): + The final learning rate at the end of the linear decay will be `init_lr * min_lr_ratio`. + adam_beta1 (`float`, *optional*, defaults to 0.9): + The beta1 to use in Adam. + adam_beta2 (`float`, *optional*, defaults to 0.999): + The beta2 to use in Adam. + adam_epsilon (`float`, *optional*, defaults to 1e-8): + The epsilon to use in Adam. + adam_clipnorm (`float`, *optional*, defaults to `None`): + If not `None`, clip the gradient norm for each weight tensor to this value. + adam_global_clipnorm (`float`, *optional*, defaults to `None`) + If not `None`, clip gradient norm to this value. When using this argument, the norm is computed over all + weight tensors, as if they were concatenated into a single vector. + weight_decay_rate (`float`, *optional*, defaults to 0): + The weight decay to use. + power (`float`, *optional*, defaults to 1.0): + The power to use for PolynomialDecay. + include_in_weight_decay (`List[str]`, *optional*): + List of the parameter names (or re patterns) to apply weight decay to. If none is passed, weight decay is + applied to all parameters except bias and layer norm parameters. + """ + # Implements linear decay of the learning rate. + lr_schedule = schedules.PolynomialDecay( + initial_learning_rate=init_lr, + decay_steps=num_train_steps - num_warmup_steps, + end_learning_rate=init_lr * min_lr_ratio, + power=power, + ) + if num_warmup_steps: + lr_schedule = WarmUp( + initial_learning_rate=init_lr, + decay_schedule_fn=lr_schedule, + warmup_steps=num_warmup_steps, + ) + if weight_decay_rate > 0.0: + optimizer = AdamWeightDecay( + learning_rate=lr_schedule, + weight_decay_rate=weight_decay_rate, + beta_1=adam_beta1, + beta_2=adam_beta2, + epsilon=adam_epsilon, + clipnorm=adam_clipnorm, + global_clipnorm=adam_global_clipnorm, + exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"], + include_in_weight_decay=include_in_weight_decay, + ) + else: + optimizer = keras.optimizers.Adam( + learning_rate=lr_schedule, + beta_1=adam_beta1, + beta_2=adam_beta2, + epsilon=adam_epsilon, + clipnorm=adam_clipnorm, + global_clipnorm=adam_global_clipnorm, + ) + # We return the optimizer and the LR scheduler in order to better track the + # evolution of the LR independently of the optimizer. + return optimizer, lr_schedule + + +class AdamWeightDecay(Adam): + """ + Adam enables L2 weight decay and clip_by_global_norm on gradients. Just adding the square of the weights to the + loss function is *not* the correct way of using L2 regularization/weight decay with Adam, since that will interact + with the m and v parameters in strange ways as shown in [Decoupled Weight Decay + Regularization](https://arxiv.org/abs/1711.05101). + + Instead we want to decay the weights in a manner that doesn't interact with the m/v parameters. This is equivalent + to adding the square of the weights to the loss with plain (non-momentum) SGD. + + Args: + learning_rate (`Union[float, LearningRateSchedule]`, *optional*, defaults to 0.001): + The learning rate to use or a schedule. + beta_1 (`float`, *optional*, defaults to 0.9): + The beta1 parameter in Adam, which is the exponential decay rate for the 1st momentum estimates. + beta_2 (`float`, *optional*, defaults to 0.999): + The beta2 parameter in Adam, which is the exponential decay rate for the 2nd momentum estimates. + epsilon (`float`, *optional*, defaults to 1e-07): + The epsilon parameter in Adam, which is a small constant for numerical stability. + amsgrad (`bool`, *optional*, defaults to `False`): + Whether to apply AMSGrad variant of this algorithm or not, see [On the Convergence of Adam and + Beyond](https://arxiv.org/abs/1904.09237). + weight_decay_rate (`float`, *optional*, defaults to 0.0): + The weight decay to apply. + include_in_weight_decay (`List[str]`, *optional*): + List of the parameter names (or re patterns) to apply weight decay to. If none is passed, weight decay is + applied to all parameters by default (unless they are in `exclude_from_weight_decay`). + exclude_from_weight_decay (`List[str]`, *optional*): + List of the parameter names (or re patterns) to exclude from applying weight decay to. If a + `include_in_weight_decay` is passed, the names in it will supersede this list. + name (`str`, *optional*, defaults to `"AdamWeightDecay"`): + Optional name for the operations created when applying gradients. + kwargs (`Dict[str, Any]`, *optional*): + Keyword arguments. Allowed to be {`clipnorm`, `clipvalue`, `lr`, `decay`}. `clipnorm` is clip gradients by + norm; `clipvalue` is clip gradients by value, `decay` is included for backward compatibility to allow time + inverse decay of learning rate. `lr` is included for backward compatibility, recommended to use + `learning_rate` instead. + """ + + def __init__( + self, + learning_rate: Union[float, schedules.LearningRateSchedule] = 0.001, + beta_1: float = 0.9, + beta_2: float = 0.999, + epsilon: float = 1e-7, + amsgrad: bool = False, + weight_decay_rate: float = 0.0, + include_in_weight_decay: Optional[list[str]] = None, + exclude_from_weight_decay: Optional[list[str]] = None, + name: str = "AdamWeightDecay", + **kwargs, + ): + super().__init__(learning_rate, beta_1, beta_2, epsilon, amsgrad, name, **kwargs) + self.weight_decay_rate = weight_decay_rate + self._include_in_weight_decay = include_in_weight_decay + self._exclude_from_weight_decay = exclude_from_weight_decay + + @classmethod + def from_config(cls, config): + """Creates an optimizer from its config with WarmUp custom object.""" + custom_objects = {"WarmUp": WarmUp} + return super().from_config(config, custom_objects=custom_objects) + + def _prepare_local(self, var_device, var_dtype, apply_state): + super()._prepare_local(var_device, var_dtype, apply_state) + apply_state[(var_device, var_dtype)]["weight_decay_rate"] = tf.constant( + self.weight_decay_rate, name="adam_weight_decay_rate" + ) + + def _decay_weights_op(self, var, learning_rate, apply_state): + do_decay = self._do_use_weight_decay(var.name) + if do_decay: + return var.assign_sub( + learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]["weight_decay_rate"], + use_locking=self._use_locking, + ) + return tf.no_op() + + def apply_gradients(self, grads_and_vars, name=None, **kwargs): + grads, tvars = list(zip(*grads_and_vars)) + return super().apply_gradients(zip(grads, tvars), name=name, **kwargs) + + def _get_lr(self, var_device, var_dtype, apply_state): + """Retrieves the learning rate with the given state.""" + if apply_state is None: + return self._decayed_lr_t[var_dtype], {} + + apply_state = apply_state or {} + coefficients = apply_state.get((var_device, var_dtype)) + if coefficients is None: + coefficients = self._fallback_apply_state(var_device, var_dtype) + apply_state[(var_device, var_dtype)] = coefficients + + return coefficients["lr_t"], {"apply_state": apply_state} + + def _resource_apply_dense(self, grad, var, apply_state=None): + lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state) + decay = self._decay_weights_op(var, lr_t, apply_state) + with tf.control_dependencies([decay]): + return super()._resource_apply_dense(grad, var, **kwargs) + + def _resource_apply_sparse(self, grad, var, indices, apply_state=None): + lr_t, kwargs = self._get_lr(var.device, var.dtype.base_dtype, apply_state) + decay = self._decay_weights_op(var, lr_t, apply_state) + with tf.control_dependencies([decay]): + return super()._resource_apply_sparse(grad, var, indices, **kwargs) + + def get_config(self): + config = super().get_config() + config.update({"weight_decay_rate": self.weight_decay_rate}) + return config + + def _do_use_weight_decay(self, param_name): + """Whether to use L2 weight decay for `param_name`.""" + if self.weight_decay_rate == 0: + return False + + if self._include_in_weight_decay: + for r in self._include_in_weight_decay: + if re.search(r, param_name) is not None: + return True + + if self._exclude_from_weight_decay: + for r in self._exclude_from_weight_decay: + if re.search(r, param_name) is not None: + return False + return True + + +# Extracted from https://github.com/OpenNMT/OpenNMT-tf/blob/master/opennmt/optimizers/utils.py +class GradientAccumulator: + """ + Gradient accumulation utility. When used with a distribution strategy, the accumulator should be called in a + replica context. Gradients will be accumulated locally on each replica and without synchronization. Users should + then call `.gradients`, scale the gradients if required, and pass the result to `apply_gradients`. + """ + + # We use the ON_READ synchronization policy so that no synchronization is + # performed on assignment. To get the value, we call .value() which returns the + # value on the current replica without synchronization. + + def __init__(self): + """Initializes the accumulator.""" + self._gradients = [] + self._accum_steps = None + + @property + def step(self): + """Number of accumulated steps.""" + if self._accum_steps is None: + self._accum_steps = tf.Variable( + tf.constant(0, dtype=tf.int64), + trainable=False, + synchronization=tf.VariableSynchronization.ON_READ, + aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA, + ) + + return self._accum_steps.value() + + @property + def gradients(self): + """The accumulated gradients on the current replica.""" + if not self._gradients: + raise ValueError("The accumulator should be called first to initialize the gradients") + return [gradient.value() if gradient is not None else gradient for gradient in self._gradients] + + def __call__(self, gradients): + """Accumulates `gradients` on the current replica.""" + if not self._gradients: + _ = self.step # Create the step variable. + self._gradients.extend( + [ + tf.Variable( + tf.zeros_like(gradient), + trainable=False, + synchronization=tf.VariableSynchronization.ON_READ, + aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA, + ) + if gradient is not None + else gradient + for gradient in gradients + ] + ) + if len(gradients) != len(self._gradients): + raise ValueError(f"Expected {len(self._gradients)} gradients, but got {len(gradients)}") + + for accum_gradient, gradient in zip(self._gradients, gradients): + if accum_gradient is not None and gradient is not None: + accum_gradient.assign_add(gradient) + + self._accum_steps.assign_add(1) + + def reset(self): + """Resets the accumulated gradients on the current replica.""" + if not self._gradients: + return + self._accum_steps.assign(0) + for gradient in self._gradients: + if gradient is not None: + gradient.assign(tf.zeros_like(gradient)) diff --git a/lib/python3.12/site-packages/transformers/processing_utils.py b/lib/python3.12/site-packages/transformers/processing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b1c40e7ff2d7c08e8b8e741a59f933f58c13fb30 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/processing_utils.py @@ -0,0 +1,1537 @@ +# Copyright 2022 The HuggingFace Inc. team. +# +# 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. +""" +Processing saving/loading class for common processors. +""" + +import copy +import inspect +import json +import os +import sys +import typing +import warnings +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, TypedDict, Union + +import numpy as np +import typing_extensions + +from .audio_utils import load_audio +from .dynamic_module_utils import custom_object_save +from .image_utils import ( + ChannelDimension, + ImageInput, + VideoInput, + is_valid_image, + is_vision_available, + load_image, + load_video, +) + + +if is_vision_available(): + from .image_utils import PILImageResampling + +from .tokenization_utils_base import ( + PaddingStrategy, + PreTokenizedInput, + PreTrainedTokenizerBase, + TextInput, + TruncationStrategy, +) +from .utils import ( + PROCESSOR_NAME, + PushToHubMixin, + TensorType, + add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, + cached_file, + copy_func, + direct_transformers_import, + download_url, + is_offline_mode, + is_remote_url, + logging, +) + + +logger = logging.get_logger(__name__) + +# Dynamically import the Transformers module to grab the attribute classes of the processor from their names. +transformers_module = direct_transformers_import(Path(__file__).parent) + + +AUTO_TO_BASE_CLASS_MAPPING = { + "AutoTokenizer": "PreTrainedTokenizerBase", + "AutoFeatureExtractor": "FeatureExtractionMixin", + "AutoImageProcessor": "ImageProcessingMixin", +} + +if sys.version_info >= (3, 11): + Unpack = typing.Unpack +else: + Unpack = typing_extensions.Unpack + + +class TextKwargs(TypedDict, total=False): + """ + Keyword arguments for text processing. For extended documentation, check out tokenization_utils_base methods and + docstrings associated. + + Attributes: + add_special_tokens (`bool`, *optional*) + Whether or not to add special tokens when encoding the sequences. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*) + Activates and controls padding. + truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*): + Activates and controls truncation. + max_length (`int`, *optional*): + Controls the maximum length to use by one of the truncation/padding parameters. + stride (`int`, *optional*): + If set, the overflowing tokens will contain some tokens from the end of the truncated sequence. + is_split_into_words (`bool`, *optional*): + Whether or not the input is already pre-tokenized. + pad_to_multiple_of (`int`, *optional*): + If set, will pad the sequence to a multiple of the provided value. + return_token_type_ids (`bool`, *optional*): + Whether to return token type IDs. + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. + return_overflowing_tokens (`bool`, *optional*): + Whether or not to return overflowing token sequences. + return_special_tokens_mask (`bool`, *optional*): + Whether or not to return special tokens mask information. + return_offsets_mapping (`bool`, *optional*): + Whether or not to return `(char_start, char_end)` for each token. + return_length (`bool`, *optional*): + Whether or not to return the lengths of the encoded inputs. + verbose (`bool`, *optional*): + Whether or not to print more information and warnings. + padding_side (`str`, *optional*): + The side on which padding will be applied. + """ + + text_pair: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]] + text_target: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] + text_pair_target: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]] + add_special_tokens: Optional[bool] + padding: Union[bool, str, PaddingStrategy] + truncation: Union[bool, str, TruncationStrategy] + max_length: Optional[int] + stride: Optional[int] + is_split_into_words: Optional[bool] + pad_to_multiple_of: Optional[int] + return_token_type_ids: Optional[bool] + return_attention_mask: Optional[bool] + return_overflowing_tokens: Optional[bool] + return_special_tokens_mask: Optional[bool] + return_offsets_mapping: Optional[bool] + return_length: Optional[bool] + verbose: Optional[bool] + padding_side: Optional[str] + + +class ImagesKwargs(TypedDict, total=False): + """ + Keyword arguments for image processing. For extended documentation, check the appropriate ImageProcessor + class methods and docstrings. + + Attributes: + do_resize (`bool`, *optional*): + Whether to resize the image. + size (`Dict[str, int]`, *optional*): + Resize the shorter side of the input to `size["shortest_edge"]`. + size_divisor (`int`, *optional*): + The size by which to make sure both the height and width can be divided. + crop_size (`Dict[str, int]`, *optional*): + Desired output size when applying center-cropping. + resample (`PILImageResampling`, *optional*): + Resampling filter to use if resizing the image. + do_rescale (`bool`, *optional*): + Whether to rescale the image by the specified scale `rescale_factor`. + rescale_factor (`int` or `float`, *optional*): + Scale factor to use if rescaling the image. + do_normalize (`bool`, *optional*): + Whether to normalize the image. + image_mean (`float` or `List[float]`, *optional*): + Mean to use if normalizing the image. + image_std (`float` or `List[float]`, *optional*): + Standard deviation to use if normalizing the image. + do_pad (`bool`, *optional*): + Whether to pad the image to the `(max_height, max_width)` of the images in the batch. + pad_size (`Dict[str, int]`, *optional*): + The size `{"height": int, "width" int}` to pad the images to. + do_center_crop (`bool`, *optional*): + Whether to center crop the image. + data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the output image. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the input image. + device (`str`, *optional*): + The device to use for processing (e.g. "cpu", "cuda"), only relevant for fast image processing. + """ + + do_resize: Optional[bool] + size: Optional[dict[str, int]] + size_divisor: Optional[int] + crop_size: Optional[dict[str, int]] + resample: Optional[Union["PILImageResampling", int]] + do_rescale: Optional[bool] + rescale_factor: Optional[float] + do_normalize: Optional[bool] + image_mean: Optional[Union[float, list[float]]] + image_std: Optional[Union[float, list[float]]] + do_pad: Optional[bool] + pad_size: Optional[dict[str, int]] + do_center_crop: Optional[bool] + data_format: Optional[ChannelDimension] + input_data_format: Optional[Union[str, ChannelDimension]] + device: Optional[str] + + +class VideosKwargs(TypedDict, total=False): + """ + Keyword arguments for video processing. + + Attributes: + do_resize (`bool`): + Whether to resize the image. + size (`Dict[str, int]`, *optional*): + Resize the shorter side of the input to `size["shortest_edge"]`. + size_divisor (`int`, *optional*): + The size by which to make sure both the height and width can be divided. + resample (`PILImageResampling`, *optional*): + Resampling filter to use if resizing the image. + do_rescale (`bool`, *optional*): + Whether to rescale the image by the specified scale `rescale_factor`. + rescale_factor (`int` or `float`, *optional*): + Scale factor to use if rescaling the image. + do_normalize (`bool`, *optional*): + Whether to normalize the image. + image_mean (`float` or `List[float]`, *optional*): + Mean to use if normalizing the image. + image_std (`float` or `List[float]`, *optional*): + Standard deviation to use if normalizing the image. + do_pad (`bool`, *optional*): + Whether to pad the image to the `(max_height, max_width)` of the images in the batch. + do_center_crop (`bool`, *optional*): + Whether to center crop the image. + data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the output image. + input_data_format (`ChannelDimension` or `str`, *optional*): + The channel dimension format for the input image. + """ + + do_resize: Optional[bool] + size: Optional[dict[str, int]] + size_divisor: Optional[int] + resample: Optional["PILImageResampling"] + do_rescale: Optional[bool] + rescale_factor: Optional[float] + do_normalize: Optional[bool] + image_mean: Optional[Union[float, list[float]]] + image_std: Optional[Union[float, list[float]]] + do_pad: Optional[bool] + do_center_crop: Optional[bool] + data_format: Optional[ChannelDimension] + input_data_format: Optional[Union[str, ChannelDimension]] + + +class AudioKwargs(TypedDict, total=False): + """ + Keyword arguments for audio processing. + + Attributes: + sampling_rate (`int`, *optional*): + The sampling rate at which the `raw_speech` input was sampled. + raw_speech (`np.ndarray`, `List[float]`, `List[np.ndarray]`, `List[List[float]]`): + The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float + values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not + stereo, i.e. single float per timestep. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` + max_length (`int`, *optional*): + Maximum length of the returned list and optionally padding length (see above). + truncation (`bool`, *optional*): + Activates truncation to cut input sequences longer than *max_length* to *max_length*. + pad_to_multiple_of (`int`, *optional*): + If set, will pad the sequence to a multiple of the provided value. + return_attention_mask (`bool`, *optional*): + Whether or not [`~ASTFeatureExtractor.__call__`] should return `attention_mask`. + """ + + sampling_rate: Optional[int] + raw_speech: Optional[Union["np.ndarray", list[float], list["np.ndarray"], list[list[float]]]] + padding: Optional[Union[bool, str, PaddingStrategy]] + max_length: Optional[int] + truncation: Optional[bool] + pad_to_multiple_of: Optional[int] + return_attention_mask: Optional[bool] + + +class CommonKwargs(TypedDict, total=False): + return_tensors: Optional[Union[str, TensorType]] + + +class ProcessingKwargs(TextKwargs, ImagesKwargs, VideosKwargs, AudioKwargs, CommonKwargs, total=False): + """ + Base class for kwargs passing to processors. + A model should have its own `ModelProcessorKwargs` class that inherits from `ProcessingKwargs` to provide: + 1) Additional typed keys and that this model requires to process inputs. + 2) Default values for existing keys under a `_defaults` attribute. + New keys have to be defined as follows to ensure type hinting is done correctly. + + ```python + # adding a new image kwarg for this model + class ModelImagesKwargs(ImagesKwargs, total=False): + new_image_kwarg: Optional[bool] + + class ModelProcessorKwargs(ProcessingKwargs, total=False): + images_kwargs: ModelImagesKwargs + _defaults = { + "images_kwargs: { + "new_image_kwarg": False, + } + "text_kwargs": { + "padding": "max_length", + }, + } + + ``` + + For Python 3.8 compatibility, when inheriting from this class and overriding one of the kwargs, + you need to manually update the __annotations__ dictionary. This can be done as follows: + + ```python + class CustomProcessorKwargs(ProcessingKwargs, total=False): + images_kwargs: CustomImagesKwargs + + CustomProcessorKwargs.__annotations__["images_kwargs"] = CustomImagesKwargs # python 3.8 compatibility + ```python + + """ + + common_kwargs: CommonKwargs = { + **CommonKwargs.__annotations__, + } + text_kwargs: TextKwargs = { + **TextKwargs.__annotations__, + } + images_kwargs: ImagesKwargs = { + **ImagesKwargs.__annotations__, + } + videos_kwargs: VideosKwargs = { + **VideosKwargs.__annotations__, + } + audio_kwargs: AudioKwargs = { + **AudioKwargs.__annotations__, + } + + +class TokenizerChatTemplateKwargs(TypedDict, total=False): + """ + Keyword arguments for tokenizer's `apply_chat_template`, when it is called from within a processor. + + tools (`List[Dict]`, *optional*): + A list of tools (callable functions) that will be accessible to the model. If the template does not + support function calling, this argument will have no effect. Each tool should be passed as a JSON Schema, + giving the name, description and argument types for the tool. See our + [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use) + for more information. + documents (`List[Dict[str, str]]`, *optional*): + A list of dicts representing documents that will be accessible to the model if it is performing RAG + (retrieval-augmented generation). If the template does not support RAG, this argument will have no + effect. We recommend that each document should be a dict containing "title" and "text" keys. Please + see the RAG section of the [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#arguments-for-RAG) + for examples of passing documents with chat templates. + add_generation_prompt (bool, *optional*): + If this is set, a prompt with the token(s) that indicate + the start of an assistant message will be appended to the formatted output. This is useful when you want to generate a response from the model. + Note that this argument will be passed to the chat template, and so it must be supported in the + template for this argument to have any effect. + continue_final_message (bool, *optional*): + If this is set, the chat will be formatted so that the final + message in the chat is open-ended, without any EOS tokens. The model will continue this message + rather than starting a new one. This allows you to "prefill" part of + the model's response for it. Cannot be used at the same time as `add_generation_prompt`. + return_assistant_tokens_mask (`bool`, defaults to `False`): + Whether to return a mask of the assistant generated tokens. For tokens generated by the assistant, + the mask will contain 1. For user and system tokens, the mask will contain 0. + This functionality is only available for chat templates that support it via the `{% generation %}` keyword. + """ + + tools: Optional[list[dict]] = None + documents: Optional[list[dict[str, str]]] = None + add_generation_prompt: Optional[bool] = False + continue_final_message: Optional[bool] = False + return_assistant_tokens_mask: Optional[bool] = False + + +class ChatTemplateLoadKwargs(TypedDict, total=False): + """ + Keyword arguments used to load multimodal data in processor chat templates. + + num_frames (`int`, *optional*): + Number of frames to sample uniformly. If not passed, the whole video is loaded. + video_load_backend (`str`, *optional*, defaults to `"pyav"`): + The backend to use when loading the video which will be used only when there are videos in the conversation. + Can be any of ["decord", "pyav", "opencv", "torchvision"]. Defaults to "pyav" because it is the only backend + that supports all types of sources to load from. + video_fps (`int`, *optional*): + Number of frames to sample per second. Should be passed only when `num_frames=None`. + If not specified and `num_frames==None`, all frames are sampled. + sample_indices_fn (`Callable`, *optional*): + A callable function that will return indices at which the video should be sampled. If the video has to be loaded using + by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`. + If not provided, simple uniformt sampling with fps is performed, otherwise `sample_indices_fn` has priority over other args. + The function expects at input the all args along with all kwargs passed to `load_video` and should output valid + indices at which the video should be sampled. For example: + + def sample_indices_fn(num_frames, fps, metadata, **kwargs): + # add you sampling logic here ... + return np.linspace(start_idx, end_idx, num_frames, dtype=int) + """ + + num_frames: Optional[int] = None + video_load_backend: Optional[str] = "pyav" + video_fps: Optional[int] = None + sampling_rate: Optional[int] = 16_000 + sample_indices_fn: Optional[Callable] = None + load_audio_from_video: Optional[bool] = False + + +class ProcessorChatTemplateKwargs(ChatTemplateLoadKwargs, TokenizerChatTemplateKwargs, total=False): + """ + Keyword arguments for processor's `apply_chat_template`. + + tokenize (`bool`, *optional*, defaults to `False`): + Whether to tokenize the output or not. + return_dict (`bool`, defaults to `False`): + Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`. + """ + + tokenize: Optional[bool] = False + return_dict: Optional[bool] = False + + +class AllKwargsForChatTemplate( + TextKwargs, ImagesKwargs, VideosKwargs, AudioKwargs, CommonKwargs, ProcessorChatTemplateKwargs +): ... + + +class ProcessorMixin(PushToHubMixin): + """ + This is a mixin used to provide saving/loading functionality for all processor classes. + """ + + attributes = ["feature_extractor", "tokenizer"] + optional_attributes = ["chat_template"] + optional_call_args: list[str] = [] + # Names need to be attr_class for attr in attributes + feature_extractor_class = None + tokenizer_class = None + _auto_class = None + valid_kwargs: list[str] = [] + + # args have to match the attributes class attribute + def __init__(self, *args, **kwargs): + # First, extract optional attributes from kwargs if present + # Optional attributes can never be positional arguments + for optional_attribute in self.optional_attributes: + setattr(self, optional_attribute, kwargs.pop(optional_attribute, None)) + # Sanitize args and kwargs + for key in kwargs: + if key not in self.attributes: + raise TypeError(f"Unexpected keyword argument {key}.") + for arg, attribute_name in zip(args, self.attributes): + if attribute_name in kwargs: + raise TypeError(f"Got multiple values for argument {attribute_name}.") + else: + kwargs[attribute_name] = arg + + if len(kwargs) != len(self.attributes): + raise ValueError( + f"This processor requires {len(self.attributes)} arguments: {', '.join(self.attributes)}. Got " + f"{len(args)} arguments instead." + ) + + # Check each arg is of the proper class (this will also catch a user initializing in the wrong order) + for attribute_name, arg in kwargs.items(): + class_name = getattr(self, f"{attribute_name}_class") + # Nothing is ever going to be an instance of "AutoXxx", in that case we check the base class. + class_name = AUTO_TO_BASE_CLASS_MAPPING.get(class_name, class_name) + if isinstance(class_name, tuple): + proper_class = tuple(self.get_possibly_dynamic_module(n) for n in class_name if n is not None) + else: + proper_class = self.get_possibly_dynamic_module(class_name) + + if not isinstance(arg, proper_class): + raise TypeError( + f"Received a {type(arg).__name__} for argument {attribute_name}, but a {class_name} was expected." + ) + + setattr(self, attribute_name, arg) + + def to_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. + + Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this processor instance. + """ + output = copy.deepcopy(self.__dict__) + + # Get the kwargs in `__init__`. + sig = inspect.signature(self.__init__) + # Only save the attributes that are presented in the kwargs of `__init__`. + attrs_to_save = sig.parameters + # Don't save attributes like `tokenizer`, `image processor` etc. + attrs_to_save = [x for x in attrs_to_save if x not in self.__class__.attributes] + # extra attributes to be kept + attrs_to_save += ["auto_map"] + + output = {k: v for k, v in output.items() if k in attrs_to_save} + + output["processor_class"] = self.__class__.__name__ + + if "tokenizer" in output: + del output["tokenizer"] + if "image_processor" in output: + del output["image_processor"] + if "feature_extractor" in output: + del output["feature_extractor"] + if "chat_template" in output: + del output["chat_template"] + + # Some attributes have different names but containing objects that are not simple strings + output = { + k: v + for k, v in output.items() + if not (isinstance(v, PushToHubMixin) or v.__class__.__name__ == "BeamSearchDecoderCTC") + } + + return output + + def to_json_string(self) -> str: + """ + Serializes this instance to a JSON string. + + Returns: + `str`: String containing all the attributes that make up this feature_extractor instance in JSON format. + """ + dictionary = self.to_dict() + + return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" + + def to_json_file(self, json_file_path: Union[str, os.PathLike]): + """ + Save this instance to a JSON file. + + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file in which this processor instance's parameters will be saved. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + writer.write(self.to_json_string()) + + def __repr__(self): + attributes_repr = [f"- {name}: {repr(getattr(self, name))}" for name in self.attributes] + attributes_repr = "\n".join(attributes_repr) + return f"{self.__class__.__name__}:\n{attributes_repr}\n\n{self.to_json_string()}" + + def save_pretrained(self, save_directory, push_to_hub: bool = False, **kwargs): + """ + Saves the attributes of this processor (feature extractor, tokenizer...) in the specified directory so that it + can be reloaded using the [`~ProcessorMixin.from_pretrained`] method. + + + + This class method is simply calling [`~feature_extraction_utils.FeatureExtractionMixin.save_pretrained`] and + [`~tokenization_utils_base.PreTrainedTokenizerBase.save_pretrained`]. Please refer to the docstrings of the + methods above for more information. + + + + Args: + save_directory (`str` or `os.PathLike`): + Directory where the feature extractor JSON file and the tokenizer files will be saved (directory will + be created if it does not exist). + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + """ + use_auth_token = kwargs.pop("use_auth_token", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if kwargs.get("token", None) is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + kwargs["token"] = use_auth_token + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = self._create_repo(repo_id, **kwargs) + files_timestamps = self._get_files_timestamps(save_directory) + # If we have a custom config, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + attrs = [getattr(self, attribute_name) for attribute_name in self.attributes] + configs = [(a.init_kwargs if isinstance(a, PreTrainedTokenizerBase) else a) for a in attrs] + configs.append(self) + custom_object_save(self, save_directory, config=configs) + + for attribute_name in self.attributes: + attribute = getattr(self, attribute_name) + # Include the processor class in the attribute config so this processor can then be reloaded with the + # `AutoProcessor` API. + if hasattr(attribute, "_set_processor_class"): + attribute._set_processor_class(self.__class__.__name__) + attribute.save_pretrained(save_directory) + + if self._auto_class is not None: + # We added an attribute to the init_kwargs of the tokenizers, which needs to be cleaned up. + for attribute_name in self.attributes: + attribute = getattr(self, attribute_name) + if isinstance(attribute, PreTrainedTokenizerBase): + del attribute.init_kwargs["auto_map"] + + # If we save using the predefined names, we can load using `from_pretrained` + # plus we save chat_template in its own file + output_processor_file = os.path.join(save_directory, PROCESSOR_NAME) + output_raw_chat_template_file = os.path.join(save_directory, "chat_template.jinja") + output_chat_template_file = os.path.join(save_directory, "chat_template.json") + + processor_dict = self.to_dict() + # Save `chat_template` in its own file. We can't get it from `processor_dict` as we popped it in `to_dict` + # to avoid serializing chat template in json config file. So let's get it from `self` directly + if self.chat_template is not None: + if kwargs.get("save_raw_chat_template", False): + with open(output_raw_chat_template_file, "w", encoding="utf-8") as writer: + writer.write(self.chat_template) + logger.info(f"chat template saved in {output_raw_chat_template_file}") + else: + chat_template_json_string = ( + json.dumps({"chat_template": self.chat_template}, indent=2, sort_keys=True) + "\n" + ) + with open(output_chat_template_file, "w", encoding="utf-8") as writer: + writer.write(chat_template_json_string) + logger.info(f"chat template saved in {output_chat_template_file}") + + # For now, let's not save to `processor_config.json` if the processor doesn't have extra attributes and + # `auto_map` is not specified. + if set(processor_dict.keys()) != {"processor_class"}: + self.to_json_file(output_processor_file) + logger.info(f"processor saved in {output_processor_file}") + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=kwargs.get("token"), + ) + + if set(processor_dict.keys()) == {"processor_class"}: + return [] + return [output_processor_file] + + @classmethod + def get_processor_dict( + cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs + ) -> tuple[dict[str, Any], dict[str, Any]]: + """ + From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a + processor of type [`~processing_utils.ProcessingMixin`] using `from_args_and_dict`. + + Parameters: + pretrained_model_name_or_path (`str` or `os.PathLike`): + The identifier of the pre-trained checkpoint from which we want the dictionary of parameters. + 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: + `Tuple[Dict, Dict]`: The dictionary(ies) that will be used to instantiate the processor object. + """ + cache_dir = kwargs.pop("cache_dir", None) + force_download = kwargs.pop("force_download", False) + resume_download = kwargs.pop("resume_download", None) + proxies = kwargs.pop("proxies", None) + token = kwargs.pop("token", None) + local_files_only = kwargs.pop("local_files_only", False) + revision = kwargs.pop("revision", None) + subfolder = kwargs.pop("subfolder", "") + + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + + user_agent = {"file_type": "processor", "from_auto_class": from_auto_class} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + is_local = os.path.isdir(pretrained_model_name_or_path) + if os.path.isdir(pretrained_model_name_or_path): + processor_file = os.path.join(pretrained_model_name_or_path, PROCESSOR_NAME) + + if os.path.isfile(pretrained_model_name_or_path): + resolved_processor_file = pretrained_model_name_or_path + # cant't load chat-template when given a file as pretrained_model_name_or_path + resolved_chat_template_file = None + resolved_raw_chat_template_file = None + is_local = True + elif is_remote_url(pretrained_model_name_or_path): + processor_file = pretrained_model_name_or_path + resolved_processor_file = download_url(pretrained_model_name_or_path) + # can't load chat-template when given a file url as pretrained_model_name_or_path + resolved_chat_template_file = None + resolved_raw_chat_template_file = None + else: + processor_file = PROCESSOR_NAME + chat_template_file = "chat_template.json" + raw_chat_template_file = "chat_template.jinja" + try: + # Load from local folder or from cache or download from model Hub and cache + resolved_processor_file = cached_file( + pretrained_model_name_or_path, + processor_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _raise_exceptions_for_missing_entries=False, + ) + + # Load chat template from a separate json if exists + # because making it part of processor-config break BC. + # Processors in older version do not accept any kwargs + resolved_chat_template_file = cached_file( + pretrained_model_name_or_path, + chat_template_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _raise_exceptions_for_missing_entries=False, + ) + + resolved_raw_chat_template_file = cached_file( + pretrained_model_name_or_path, + raw_chat_template_file, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _raise_exceptions_for_missing_entries=False, + ) + except OSError: + # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted to + # the original exception. + raise + except Exception: + # For any other exception, we throw a generic error. + raise OSError( + f"Can't load processor for '{pretrained_model_name_or_path}'. If you were trying to load" + " it from 'https://huggingface.co/models', make sure you don't have a local directory with the" + f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a" + f" directory containing a {PROCESSOR_NAME} file" + ) + + # Add chat template as kwarg before returning because most models don't have processor config + if resolved_raw_chat_template_file is not None: + with open(resolved_raw_chat_template_file, encoding="utf-8") as reader: + chat_template = reader.read() + kwargs["chat_template"] = chat_template + elif resolved_chat_template_file is not None: + with open(resolved_chat_template_file, encoding="utf-8") as reader: + text = reader.read() + chat_template = json.loads(text)["chat_template"] + kwargs["chat_template"] = chat_template + + # Existing processors on the Hub created before #27761 being merged don't have `processor_config.json` (if not + # updated afterward), and we need to keep `from_pretrained` work. So here it fallbacks to the empty dict. + # (`cached_file` called using `_raise_exceptions_for_missing_entries=False` to avoid exception) + # However, for models added in the future, we won't get the expected error if this file is missing. + if resolved_processor_file is None: + # In any case we need to pass `chat_template` if it is available + processor_dict = {} + if "chat_template" in kwargs: + processor_dict = {"chat_template": kwargs.pop("chat_template")} + return processor_dict, kwargs + + try: + # Load processor dict + with open(resolved_processor_file, encoding="utf-8") as reader: + text = reader.read() + processor_dict = json.loads(text) + + except json.JSONDecodeError: + raise OSError(f"It looks like the config file at '{resolved_processor_file}' is not a valid JSON file.") + + if is_local: + logger.info(f"loading configuration file {resolved_processor_file}") + else: + logger.info(f"loading configuration file {processor_file} from cache at {resolved_processor_file}") + + if "chat_template" in processor_dict and processor_dict["chat_template"] is not None: + logger.warning_once( + "Chat templates should be in a 'chat_template.jinja' file but found key='chat_template' " + "in the processor's config. Make sure to move your template to its own file." + ) + + if "chat_template" in kwargs: + processor_dict["chat_template"] = kwargs.pop("chat_template") + + if not is_local: + if "auto_map" in processor_dict: + processor_dict["auto_map"] = add_model_info_to_auto_map( + processor_dict["auto_map"], pretrained_model_name_or_path + ) + if "custom_pipelines" in processor_dict: + processor_dict["custom_pipelines"] = add_model_info_to_custom_pipelines( + processor_dict["custom_pipelines"], pretrained_model_name_or_path + ) + + return processor_dict, kwargs + + @classmethod + def from_args_and_dict(cls, args, processor_dict: dict[str, Any], **kwargs): + """ + Instantiates a type of [`~processing_utils.ProcessingMixin`] from a Python dictionary of parameters. + + Args: + processor_dict (`Dict[str, Any]`): + Dictionary that will be used to instantiate the processor object. Such a dictionary can be + retrieved from a pretrained checkpoint by leveraging the + [`~processing_utils.ProcessingMixin.to_dict`] method. + kwargs (`Dict[str, Any]`): + Additional parameters from which to initialize the processor object. + + Returns: + [`~processing_utils.ProcessingMixin`]: The processor object instantiated from those + parameters. + """ + processor_dict = processor_dict.copy() + return_unused_kwargs = kwargs.pop("return_unused_kwargs", False) + + # We have to pop up some unused (but specific) kwargs and then validate that it doesn't contain unused kwargs + # If we don't pop, some specific kwargs will raise a warning + if "processor_class" in processor_dict: + del processor_dict["processor_class"] + + if "auto_map" in processor_dict: + del processor_dict["auto_map"] + + unused_kwargs = cls.validate_init_kwargs(processor_config=processor_dict, valid_kwargs=cls.valid_kwargs) + processor = cls(*args, **processor_dict) + + # Update processor with kwargs if needed + for key in set(kwargs.keys()): + if hasattr(processor, key): + setattr(processor, key, kwargs.pop(key)) + + kwargs.update(unused_kwargs) + logger.info(f"Processor {processor}") + if return_unused_kwargs: + return processor, kwargs + else: + return processor + + def _merge_kwargs( + self, + ModelProcessorKwargs: ProcessingKwargs, + tokenizer_init_kwargs: Optional[dict] = None, + **kwargs, + ) -> dict[str, dict]: + """ + Method to merge dictionaries of kwargs cleanly separated by modality within a Processor instance. + The order of operations is as follows: + 1) kwargs passed as before have highest priority to preserve BC. + ```python + high_priority_kwargs = {"crop_size" = {"height": 222, "width": 222}, "padding" = "max_length"} + processor(..., **high_priority_kwargs) + ``` + 2) kwargs passed as modality-specific kwargs have second priority. This is the recommended API. + ```python + processor(..., text_kwargs={"padding": "max_length"}, images_kwargs={"crop_size": {"height": 222, "width": 222}}}) + ``` + 3) kwargs passed during instantiation of a modality processor have fourth priority. + ```python + tokenizer = tokenizer_class(..., {"padding": "max_length"}) + image_processor = image_processor_class(...) + processor(tokenizer, image_processor) # will pass max_length unless overridden by kwargs at call + ``` + 4) defaults kwargs specified at processor level have lowest priority. + ```python + class MyProcessingKwargs(ProcessingKwargs, CommonKwargs, TextKwargs, ImagesKwargs, total=False): + _defaults = { + "text_kwargs": { + "padding": "max_length", + "max_length": 64, + }, + } + ``` + Args: + ModelProcessorKwargs (`ProcessingKwargs`): + Typed dictionary of kwargs specifically required by the model passed. + tokenizer_init_kwargs (`Dict`, *optional*): + Dictionary of kwargs the tokenizer was instantiated with and need to take precedence over defaults. + + Returns: + output_kwargs (`Dict`): + Dictionary of per-modality kwargs to be passed to each modality-specific processor. + + """ + # Initialize dictionaries + output_kwargs = { + "text_kwargs": {}, + "images_kwargs": {}, + "audio_kwargs": {}, + "videos_kwargs": {}, + "common_kwargs": {}, + } + + default_kwargs = { + "text_kwargs": {}, + "images_kwargs": {}, + "audio_kwargs": {}, + "videos_kwargs": {}, + "common_kwargs": {}, + } + + possible_modality_keywords = {"text", "audio", "videos", "images"} + used_keys = set() + + # get defaults from set model processor kwargs if they exist + for modality in default_kwargs: + default_kwargs[modality] = ModelProcessorKwargs._defaults.get(modality, {}).copy() + # update defaults with arguments from tokenizer init + for modality_key in ModelProcessorKwargs.__annotations__[modality].__annotations__.keys(): + # init with tokenizer init kwargs if necessary + if modality_key in tokenizer_init_kwargs: + value = ( + getattr(self.tokenizer, modality_key) + if hasattr(self.tokenizer, modality_key) + else tokenizer_init_kwargs[modality_key] + ) + default_kwargs[modality][modality_key] = value + # now defaults kwargs are updated with the tokenizers defaults. + # pass defaults to output dictionary + output_kwargs.update(default_kwargs) + + # update modality kwargs with passed kwargs + non_modality_kwargs = set(kwargs) - set(output_kwargs) + for modality in output_kwargs: + for modality_key in ModelProcessorKwargs.__annotations__[modality].__annotations__.keys(): + # check if we received a structured kwarg dict or not to handle it correctly + if modality in kwargs: + kwarg_value = kwargs[modality].pop(modality_key, "__empty__") + # check if this key was passed as a flat kwarg. + if kwarg_value != "__empty__" and modality_key in non_modality_kwargs: + raise ValueError( + f"Keyword argument {modality_key} was passed two times:\n" + f"in a dictionary for {modality} and as a **kwarg." + ) + elif modality_key in kwargs: + # we get a modality_key instead of popping it because modality-specific processors + # can have overlapping kwargs + kwarg_value = kwargs.get(modality_key, "__empty__") + else: + kwarg_value = "__empty__" + if not isinstance(kwarg_value, str) or kwarg_value != "__empty__": + output_kwargs[modality][modality_key] = kwarg_value + used_keys.add(modality_key) + + # Determine if kwargs is a flat dictionary or contains nested dictionaries + if any(key in default_kwargs for key in kwargs): + # kwargs is dictionary-based, and some keys match modality names + for modality, subdict in kwargs.items(): + if modality in default_kwargs: + for subkey, subvalue in subdict.items(): + if subkey not in used_keys: + output_kwargs[modality][subkey] = subvalue + used_keys.add(subkey) + else: + # kwargs is a flat dictionary + for key in kwargs: + if key not in used_keys: + if key in ModelProcessorKwargs.__annotations__["common_kwargs"].__annotations__.keys(): + output_kwargs["common_kwargs"][key] = kwargs[key] + elif key not in possible_modality_keywords: + logger.warning_once( + f"Keyword argument `{key}` is not a valid argument for this processor and will be ignored." + ) + + # all modality-specific kwargs are updated with common kwargs + for modality in output_kwargs: + output_kwargs[modality].update(output_kwargs["common_kwargs"]) + return output_kwargs + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Union[str, os.PathLike], + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + **kwargs, + ): + r""" + Instantiate a processor associated with a pretrained model. + + + + This class method is simply calling the feature extractor + [`~feature_extraction_utils.FeatureExtractionMixin.from_pretrained`], image processor + [`~image_processing_utils.ImageProcessingMixin`] and the tokenizer + [`~tokenization_utils_base.PreTrainedTokenizer.from_pretrained`] methods. Please refer to the docstrings of the + methods above for more information. + + + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + This can be either: + + - a string, the *model id* of a pretrained feature_extractor hosted inside a model repo on + huggingface.co. + - a path to a *directory* containing a feature extractor file saved using the + [`~SequenceFeatureExtractor.save_pretrained`] method, e.g., `./my_model_directory/`. + - a path or url to a saved feature extractor JSON *file*, e.g., + `./my_model_directory/preprocessor_config.json`. + **kwargs + Additional keyword arguments passed along to both + [`~feature_extraction_utils.FeatureExtractionMixin.from_pretrained`] and + [`~tokenization_utils_base.PreTrainedTokenizer.from_pretrained`]. + """ + kwargs["cache_dir"] = cache_dir + kwargs["force_download"] = force_download + kwargs["local_files_only"] = local_files_only + kwargs["revision"] = revision + + use_auth_token = kwargs.pop("use_auth_token", None) + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + if token is not None: + kwargs["token"] = token + + args = cls._get_arguments_from_pretrained(pretrained_model_name_or_path, **kwargs) + processor_dict, kwargs = cls.get_processor_dict(pretrained_model_name_or_path, **kwargs) + processor_dict.update({k: v for k, v in kwargs.items() if k in processor_dict.keys()}) + return cls.from_args_and_dict(args, processor_dict, **kwargs) + + @classmethod + def register_for_auto_class(cls, auto_class="AutoProcessor"): + """ + Register this class with a given auto class. This should only be used for custom feature extractors as the ones + in the library are already mapped with `AutoProcessor`. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"AutoProcessor"`): + The auto class to register this new feature extractor with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + @classmethod + def _get_arguments_from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + """ + Identify and instantiate the subcomponents of Processor classes, like image processors and + tokenizers. This method uses the Processor attributes like `tokenizer_class` to figure out what class those + subcomponents should be. Note that any subcomponents must either be library classes that are accessible in + the `transformers` root, or they must be custom code that has been registered with the relevant autoclass, + via methods like `AutoTokenizer.register()`. If neither of these conditions are fulfilled, this method + will be unable to find the relevant subcomponent class and will raise an error. + """ + args = [] + for attribute_name in cls.attributes: + class_name = getattr(cls, f"{attribute_name}_class") + if isinstance(class_name, tuple): + classes = tuple(cls.get_possibly_dynamic_module(n) if n is not None else None for n in class_name) + if attribute_name == "image_processor": + # TODO: @yoni, change logic in v4.52 (when use_fast set to True by default) + use_fast = kwargs.get("use_fast", None) + if use_fast is None: + logger.warning_once( + "Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. " + "`use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. " + "This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`." + ) + else: + use_fast = kwargs.get("use_fast", True) + if use_fast and classes[1] is not None: + attribute_class = classes[1] + else: + attribute_class = classes[0] + else: + attribute_class = cls.get_possibly_dynamic_module(class_name) + + args.append(attribute_class.from_pretrained(pretrained_model_name_or_path, **kwargs)) + return args + + @staticmethod + def get_possibly_dynamic_module(module_name): + if hasattr(transformers_module, module_name): + return getattr(transformers_module, module_name) + lookup_locations = [ + transformers_module.IMAGE_PROCESSOR_MAPPING, + transformers_module.TOKENIZER_MAPPING, + transformers_module.FEATURE_EXTRACTOR_MAPPING, + ] + for lookup_location in lookup_locations: + for custom_class in lookup_location._extra_content.values(): + if isinstance(custom_class, tuple): + for custom_subclass in custom_class: + if custom_subclass is not None and custom_subclass.__name__ == module_name: + return custom_subclass + elif custom_class is not None and custom_class.__name__ == module_name: + return custom_class + else: + raise ValueError( + f"Could not find module {module_name} in `transformers`. If this is a custom class, " + f"it should be registered using the relevant `AutoClass.register()` function so that " + f"other functions can find it!" + ) + + @property + def model_input_names(self): + first_attribute = getattr(self, self.attributes[0]) + return getattr(first_attribute, "model_input_names", None) + + @staticmethod + def validate_init_kwargs(processor_config, valid_kwargs): + kwargs_from_config = processor_config.keys() + unused_kwargs = {} + unused_keys = set(kwargs_from_config) - set(valid_kwargs) + if unused_keys: + unused_kwargs = {k: processor_config[k] for k in unused_keys} + return unused_kwargs + + def prepare_and_validate_optional_call_args(self, *args): + """ + Matches optional positional arguments to their corresponding names in `optional_call_args` + in the processor class in the order they are passed to the processor call. + + Note that this should only be used in the `__call__` method of the processors with special + arguments. Special arguments are arguments that aren't `text`, `images`, `audio`, nor `videos` + but also aren't passed to the tokenizer, image processor, etc. Examples of such processors are: + - `CLIPSegProcessor` + - `LayoutLMv2Processor` + - `OwlViTProcessor` + + Also note that passing by position to the processor call is now deprecated and will be disallowed + in future versions. We only have this for backward compatibility. + + Example: + Suppose that the processor class has `optional_call_args = ["arg_name_1", "arg_name_2"]`. + And we define the call method as: + ```python + def __call__( + self, + text: str, + images: Optional[ImageInput] = None, + *arg, + audio=None, + videos=None, + ) + ``` + + Then, if we call the processor as: + ```python + images = [...] + processor("What is common in these images?", images, arg_value_1, arg_value_2) + ``` + + Then, this method will return: + ```python + { + "arg_name_1": arg_value_1, + "arg_name_2": arg_value_2, + } + ``` + which we could then pass as kwargs to `self._merge_kwargs` + """ + if len(args): + warnings.warn( + "Passing positional arguments to the processor call is now deprecated and will be disallowed in v4.47. " + "Please pass all arguments as keyword arguments." + ) + if len(args) > len(self.optional_call_args): + raise ValueError( + f"Expected *at most* {len(self.optional_call_args)} optional positional arguments in processor call" + f"which will be matched with {' '.join(self.optional_call_args)} in the order they are passed." + f"However, got {len(args)} positional arguments instead." + "Please pass all arguments as keyword arguments instead (e.g. `processor(arg_name_1=..., arg_name_2=...))`." + ) + return {arg_name: arg_value for arg_value, arg_name in zip(args, self.optional_call_args)} + + def _process_messages_for_chat_template( + self, + conversation: List[List[Dict[str, str]]], + batch_images: List[ImageInput], + batch_videos: List[VideoInput], + batch_video_metadata: List[List[Dict[str, any]]], + **mm_load_kwargs: Unpack[ChatTemplateLoadKwargs], + ): + """ + Used within `apply_chat_template` when a model has a special way to process conversation history. For example, + video models might want to specify in the prompt the duration of video or which frame indices at which timestamps + were sampled. This information cannot be accessed before the video is loaded. + + For most models it is a no-op, and must be overridden by model processors which require special processing. + + Args: + conversation (`List[Dict, str, str]`): + The conversation to process. Always comes in batched format. + batch_images (`List[List[ImageInput]]`): + Batch of images that were loaded from url/path defined in the conversation. The images + are ordered in the same way as in the conversation. Comes in nested list format, one list of `PIL` images + per batch. + batch_videos (`List[List[ImageInput]]`): + Batch of videos that were loaded from url/path defined in the conversation. The videos + are ordered in the samm way as in the conversation. Comes in nested list format, one list of 4D video arrays + per batch. + batch_video_metadata (`List[List[Dict[[str, any]]]]`): + Batch of metadata returned from loading videos. That includes video fps, duration and total number of framer in original video. + Metadata are ordered in the same way as `batch_videos`. Comes in nested list format, one list of 4D video arrays + per batch. + + """ + return conversation + + def apply_chat_template( + self, + conversation: Union[list[dict[str, str]], list[list[dict[str, str]]]], + chat_template: Optional[str] = None, + **kwargs: Unpack[AllKwargsForChatTemplate], + ) -> str: + """ + Similar to the `apply_chat_template` method on tokenizers, this method applies a Jinja template to input + conversations to turn them into a single tokenizable string. + + The input is expected to be in the following format, where each message content is a list consisting of text and + optionally image or video inputs. One can also provide an image, video, URL or local path which will be used to form + `pixel_values` when `return_dict=True`. If not provided, one will get only the formatted text, optionally tokenized text. + + conversation = [ + { + "role": "user", + "content": [ + {"type": "image", "image": "https://www.ilankelman.org/stopsigns/australia.jpg"}, + {"type": "text", "text": "Please describe this image in detail."}, + ], + }, + ] + + Args: + conversation (`Union[List[Dict, [str, str]], List[List[Dict[str, str]]]]`): + The conversation to format. + chat_template (`Optional[str]`, *optional*): + The Jinja template to use for formatting the conversation. If not provided, the tokenizer's + chat template is used. + """ + + if chat_template is None: + if self.chat_template is not None: + chat_template = self.chat_template + else: + raise ValueError( + "No chat template is set for this processor. Please either set the `chat_template` attribute, " + "or provide a chat template as an argument. See " + "https://huggingface.co/docs/transformers/main/en/chat_templating for more information." + ) + + # Fill two sets of kwargs that should be used by tokenizer's `apply_chat_template` + # and for multimodal data loading. Everything else will be used in `__call__` + tokenizer_template_kwargs = {} + for tokenizer_key in TokenizerChatTemplateKwargs.__annotations__.keys(): + default_value = getattr(TokenizerChatTemplateKwargs, tokenizer_key, None) + value = kwargs.pop(tokenizer_key, default_value) + tokenizer_template_kwargs[tokenizer_key] = value + + mm_load_kwargs = {} + for mm_load_key in ChatTemplateLoadKwargs.__annotations__.keys(): + default_value = getattr(ChatTemplateLoadKwargs, mm_load_key, None) + value = kwargs.pop(mm_load_key, default_value) + mm_load_kwargs[mm_load_key] = value + + if isinstance(conversation, (list, tuple)) and ( + isinstance(conversation[0], (list, tuple)) or hasattr(conversation[0], "content") + ): + is_batched = True + conversations = conversation + else: + is_batched = False + conversations = [conversation] + + tokenize = kwargs.pop("tokenize", False) + return_dict = kwargs.pop("return_dict", False) + + if tokenize: + batch_images, batch_videos = [], [] + batch_audios = [] + batch_video_metadata = [] + for conversation in conversations: + images, videos = [], [] + video_metadata = [] + for message in conversation: + visuals = [content for content in message["content"] if content["type"] in ["image", "video"]] + audio_fnames = [ + content[key] + for content in message["content"] + for key in ["audio", "url", "path"] + if key in content and content["type"] == "audio" + ] + image_fnames = [ + vision_info[key] + for vision_info in visuals + for key in ["image", "url", "path", "base64"] + if key in vision_info and vision_info["type"] == "image" + ] + video_fnames = [ + vision_info[key] + for vision_info in visuals + for key in ["video", "url", "path"] + if key in vision_info and vision_info["type"] == "video" + ] + + for fname in image_fnames: + images.append(load_image(fname)) + + # Audio models do not accept nested list of audios (yet!) so we construct a flat input audio list + if not mm_load_kwargs["load_audio_from_video"]: + for fname in audio_fnames: + batch_audios.append(load_audio(fname, sampling_rate=mm_load_kwargs["sampling_rate"])) + else: + for fname in video_fnames: + batch_audios.append(load_audio(fname, sampling_rate=mm_load_kwargs["sampling_rate"])) + + for fname in video_fnames: + if isinstance(fname, (list, tuple)) and isinstance(fname[0], str): + video = [np.array(load_image(image_fname)).T for image_fname in fname] + # create a 4D video because `load_video` always returns a 4D array + video = np.stack(video) + metadata = None + logger.warning( + "When loading the video from list of images, we cannot infer metadata such as `fps` or `duration`. " + "If your model uses this metadata during processing, please load the whole video and let the model sample frames instead." + ) + else: + video, metadata = load_video( + fname, + num_frames=mm_load_kwargs["num_frames"], + fps=mm_load_kwargs["video_fps"], + backend=mm_load_kwargs["video_load_backend"], + sample_indices_fn=mm_load_kwargs["sample_indices_fn"], + ) + videos.append(video) + video_metadata.append(metadata) + + # Currently all processors can accept nested list of batches, but not flat list of visuals + # So we'll make a batched list of images and let the processor handle it + if images: + batch_images.append(images) + if videos: + batch_videos.append(videos) + batch_video_metadata.append(video_metadata) + + # Process conversation with video/image information if needed. Then convert into a prompt using Jinja template + conversations = self._process_messages_for_chat_template( + conversations, + batch_images=batch_images, + batch_videos=batch_videos, + batch_video_metadata=batch_video_metadata, + **mm_load_kwargs, + ) + + prompt = self.tokenizer.apply_chat_template( + conversations, + chat_template=chat_template, + tokenize=False, + return_dict=False, + **tokenizer_template_kwargs, + ) + + if not is_batched: + prompt = prompt[0] + + if tokenize: + # Tokenizer's `apply_chat_template` never adds special tokens when tokenizing + # But processor's `apply_chat_template` didn't have an option to tokenize, so users had to format the prompt + # and pass it to the processor. Users thus never worried about special tokens relying on processor handling + # everything internally. The below line is to keep BC for that and be able to work with model that have + # special tokens in the template (consistent with tokenizers). We dont want to raise warning, it will flood command line + # without actionable solution for users + single_prompt = prompt[0] if is_batched else prompt + if self.tokenizer.bos_token is not None and single_prompt.startswith(self.tokenizer.bos_token): + kwargs["add_special_tokens"] = False + + out = self( + text=prompt, + images=batch_images if batch_images else None, + videos=batch_videos if batch_videos else None, + audio=batch_audios if batch_audios else None, + **kwargs, + ) + if return_dict: + return out + else: + return out["input_ids"] + return prompt + + def post_process_image_text_to_text(self, generated_outputs, skip_special_tokens=True, **kwargs): + """ + Post-process the output of a vlm to decode the text. + + Args: + generated_outputs (`torch.Tensor` or `np.ndarray`): + The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)` + or `(sequence_length,)`. + skip_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method. + **kwargs: + Additional arguments to be passed to the tokenizer's `batch_decode method`. + + Returns: + `List[str]`: The decoded text. + """ + return self.tokenizer.batch_decode(generated_outputs, skip_special_tokens=skip_special_tokens, **kwargs) + + +def _validate_images_text_input_order(images, text): + """ + For backward compatibility: reverse the order of `images` and `text` inputs if they are swapped. + This method should only be called for processors where `images` and `text` have been swapped for uniformization purposes. + Note that this method assumes that two `None` inputs are valid inputs. If this is not the case, it should be handled + in the processor's `__call__` method before calling this method. + """ + + def is_url(val) -> bool: + return isinstance(val, str) and val.startswith("http") + + def _is_valid_images_input_for_processor(imgs): + # If we have an list of images, make sure every image is valid + if isinstance(imgs, (list, tuple)): + for img in imgs: + if not _is_valid_images_input_for_processor(img): + return False + # If not a list or tuple, we have been given a single image or batched tensor of images + elif not (is_valid_image(imgs) or is_url(imgs)): + return False + return True + + def _is_valid_text_input_for_processor(t): + if isinstance(t, str): + # Strings are fine + return True + elif isinstance(t, (list, tuple)): + # List are fine as long as they are... + if len(t) == 0: + # ... not empty + return False + for t_s in t: + return _is_valid_text_input_for_processor(t_s) + return False + + def _is_valid(input, validator): + return validator(input) or input is None + + images_is_valid = _is_valid(images, _is_valid_images_input_for_processor) + images_is_text = _is_valid_text_input_for_processor(images) + + text_is_valid = _is_valid(text, _is_valid_text_input_for_processor) + text_is_images = _is_valid_images_input_for_processor(text) + # Handle cases where both inputs are valid + if images_is_valid and text_is_valid: + return images, text + + # Handle cases where inputs need to and can be swapped + if (images is None and text_is_images) or (text is None and images_is_text) or (images_is_text and text_is_images): + logger.warning_once( + "You may have used the wrong order for inputs. `images` should be passed before `text`. " + "The `images` and `text` inputs will be swapped. This behavior will be deprecated in transformers v4.47." + ) + return text, images + + raise ValueError("Invalid input type. Check that `images` and/or `text` are valid inputs.") + + +ProcessorMixin.push_to_hub = copy_func(ProcessorMixin.push_to_hub) +if ProcessorMixin.push_to_hub.__doc__ is not None: + ProcessorMixin.push_to_hub.__doc__ = ProcessorMixin.push_to_hub.__doc__.format( + object="processor", object_class="AutoProcessor", object_files="processor files" + ) diff --git a/lib/python3.12/site-packages/transformers/py.typed b/lib/python3.12/site-packages/transformers/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/transformers/pytorch_utils.py b/lib/python3.12/site-packages/transformers/pytorch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ff3856c9d4573d9b44edeb9a359e8958bf31811b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/pytorch_utils.py @@ -0,0 +1,360 @@ +# Copyright 2022 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. +from __future__ import annotations + +import inspect +from functools import lru_cache, wraps +from typing import Callable + +import torch +from safetensors.torch import storage_ptr, storage_size +from torch import nn + +from .utils import is_torch_greater_or_equal, is_torch_xla_available, is_torchdynamo_compiling, logging + + +ALL_LAYERNORM_LAYERS = [nn.LayerNorm] + +logger = logging.get_logger(__name__) + +is_torch_greater_or_equal_than_2_6 = is_torch_greater_or_equal("2.6", accept_dev=True) +is_torch_greater_or_equal_than_2_4 = is_torch_greater_or_equal("2.4", accept_dev=True) +is_torch_greater_or_equal_than_2_3 = is_torch_greater_or_equal("2.3", accept_dev=True) +is_torch_greater_or_equal_than_2_2 = is_torch_greater_or_equal("2.2", accept_dev=True) +is_torch_greater_or_equal_than_2_1 = is_torch_greater_or_equal("2.1", accept_dev=True) + +# For backwards compatibility (e.g. some remote codes on Hub using those variables). +is_torch_greater_or_equal_than_2_0 = is_torch_greater_or_equal("2.0", accept_dev=True) +is_torch_greater_or_equal_than_1_13 = is_torch_greater_or_equal("1.13", accept_dev=True) +is_torch_greater_or_equal_than_1_12 = is_torch_greater_or_equal("1.12", accept_dev=True) + +# Cache this result has it's a C FFI call which can be pretty time-consuming +_torch_distributed_available = torch.distributed.is_available() + +if is_torch_greater_or_equal("2.5") and _torch_distributed_available: + pass + + +def softmax_backward_data(parent, grad_output, output, dim, self): + """ + A function that calls the internal `_softmax_backward_data` PyTorch method and that adjusts the arguments according + to the torch version detected. + """ + + from torch import _softmax_backward_data + + return _softmax_backward_data(grad_output, output, parent.dim, self.dtype) + + +def prune_linear_layer(layer: nn.Linear, index: torch.LongTensor, dim: int = 0) -> nn.Linear: + """ + Prune a linear layer to keep only entries in index. + + Used to remove heads. + + Args: + layer (`torch.nn.Linear`): The layer to prune. + index (`torch.LongTensor`): The indices to keep in the layer. + dim (`int`, *optional*, defaults to 0): The dimension on which to keep the indices. + + Returns: + `torch.nn.Linear`: The pruned layer as a new layer with `requires_grad=True`. + """ + index = index.to(layer.weight.device) + W = layer.weight.index_select(dim, index).detach().clone() + if layer.bias is not None: + if dim == 1: + b = layer.bias.detach().clone() + else: + b = layer.bias[index].detach().clone() + new_size = list(layer.weight.size()) + new_size[dim] = len(index) + new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device) + new_layer.weight.requires_grad = False + new_layer.weight.copy_(W.contiguous()) + new_layer.weight.requires_grad = True + if layer.bias is not None: + new_layer.bias.requires_grad = False + new_layer.bias.copy_(b.contiguous()) + new_layer.bias.requires_grad = True + return new_layer + + +class Conv1D(nn.Module): + """ + 1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2). + + Basically works like a linear layer but the weights are transposed. + + Args: + nf (`int`): The number of output features. + nx (`int`): The number of input features. + """ + + def __init__(self, nf, nx): + super().__init__() + self.nf = nf + self.nx = nx + self.weight = nn.Parameter(torch.empty(nx, nf)) + self.bias = nn.Parameter(torch.zeros(nf)) + nn.init.normal_(self.weight, std=0.02) + + def __repr__(self) -> str: + return "Conv1D(nf={nf}, nx={nx})".format(**self.__dict__) + + def forward(self, x): + size_out = x.size()[:-1] + (self.nf,) + x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight) + x = x.view(size_out) + return x + + +def prune_conv1d_layer(layer: Conv1D, index: torch.LongTensor, dim: int = 1) -> Conv1D: + """ + Prune a Conv1D layer to keep only entries in index. A Conv1D work as a Linear layer (see e.g. BERT) but the weights + are transposed. + + Used to remove heads. + + Args: + layer ([`~pytorch_utils.Conv1D`]): The layer to prune. + index (`torch.LongTensor`): The indices to keep in the layer. + dim (`int`, *optional*, defaults to 1): The dimension on which to keep the indices. + + Returns: + [`~pytorch_utils.Conv1D`]: The pruned layer as a new layer with `requires_grad=True`. + """ + index = index.to(layer.weight.device) + W = layer.weight.index_select(dim, index).detach().clone() + if dim == 0: + b = layer.bias.detach().clone() + else: + b = layer.bias[index].detach().clone() + new_size = list(layer.weight.size()) + new_size[dim] = len(index) + new_layer = Conv1D(new_size[1], new_size[0]).to(layer.weight.device) + new_layer.weight.requires_grad = False + new_layer.weight.copy_(W.contiguous()) + new_layer.weight.requires_grad = True + new_layer.bias.requires_grad = False + new_layer.bias.copy_(b.contiguous()) + new_layer.bias.requires_grad = True + return new_layer + + +def prune_layer(layer: nn.Linear | Conv1D, index: torch.LongTensor, dim: int | None = None) -> nn.Linear | Conv1D: + """ + Prune a Conv1D or linear layer to keep only entries in index. + + Used to remove heads. + + Args: + layer (`Union[torch.nn.Linear, Conv1D]`): The layer to prune. + index (`torch.LongTensor`): The indices to keep in the layer. + dim (`int`, *optional*): The dimension on which to keep the indices. + + Returns: + `torch.nn.Linear` or [`~pytorch_utils.Conv1D`]: The pruned layer as a new layer with `requires_grad=True`. + """ + if isinstance(layer, nn.Linear): + return prune_linear_layer(layer, index, dim=0 if dim is None else dim) + elif isinstance(layer, Conv1D): + return prune_conv1d_layer(layer, index, dim=1 if dim is None else dim) + else: + raise ValueError(f"Can't prune layer of class {layer.__class__}") + + +def apply_chunking_to_forward( + forward_fn: Callable[..., torch.Tensor], + chunk_size: int, + chunk_dim: int, + *input_tensors, +) -> torch.Tensor: + """ + This function chunks the `input_tensors` into smaller input tensor parts of size `chunk_size` over the dimension + `chunk_dim`. It then applies a layer `forward_fn` to each chunk independently to save memory. + + If the `forward_fn` is independent across the `chunk_dim` this function will yield the same result as directly + applying `forward_fn` to `input_tensors`. + + Args: + forward_fn (`Callable[..., torch.Tensor]`): + The forward function of the model. + chunk_size (`int`): + The chunk size of a chunked tensor: `num_chunks = len(input_tensors[0]) / chunk_size`. + chunk_dim (`int`): + The dimension over which the `input_tensors` should be chunked. + input_tensors (`Tuple[torch.Tensor]`): + The input tensors of `forward_fn` which will be chunked + + Returns: + `torch.Tensor`: A tensor with the same shape as the `forward_fn` would have given if applied`. + + + Examples: + + ```python + # rename the usual forward() fn to forward_chunk() + def forward_chunk(self, hidden_states): + hidden_states = self.decoder(hidden_states) + return hidden_states + + + # implement a chunked forward function + def forward(self, hidden_states): + return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states) + ```""" + + assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors" + + # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility + num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters) + if num_args_in_forward_chunk_fn != len(input_tensors): + raise ValueError( + f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input " + "tensors are given" + ) + + if chunk_size > 0: + tensor_shape = input_tensors[0].shape[chunk_dim] + for input_tensor in input_tensors: + if input_tensor.shape[chunk_dim] != tensor_shape: + raise ValueError( + f"All input tenors have to be of the same shape: {tensor_shape}, " + f"found shape {input_tensor.shape[chunk_dim]}" + ) + + if input_tensors[0].shape[chunk_dim] % chunk_size != 0: + raise ValueError( + f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk " + f"size {chunk_size}" + ) + + num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size + + # chunk input tensor into tuples + input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors) + # apply forward fn to every tuple + output_chunks = tuple(forward_fn(*input_tensors_chunk) for input_tensors_chunk in zip(*input_tensors_chunks)) + # concatenate output at same dimension + return torch.cat(output_chunks, dim=chunk_dim) + + return forward_fn(*input_tensors) + + +def find_pruneable_heads_and_indices( + heads: list[int], n_heads: int, head_size: int, already_pruned_heads: set[int] +) -> tuple[set[int], torch.LongTensor]: + """ + Finds the heads and their indices taking `already_pruned_heads` into account. + + Args: + heads (`List[int]`): List of the indices of heads to prune. + n_heads (`int`): The number of heads in the model. + head_size (`int`): The size of each head. + already_pruned_heads (`Set[int]`): A set of already pruned heads. + + Returns: + `Tuple[Set[int], torch.LongTensor]`: A tuple with the indices of heads to prune taking `already_pruned_heads` + into account and the indices of rows/columns to keep in the layer weight. + """ + mask = torch.ones(n_heads, head_size) + heads = set(heads) - already_pruned_heads # Convert to set and remove already pruned heads + for head in heads: + # Compute how many pruned heads are before the head and move the index accordingly + head = head - sum(1 if h < head else 0 for h in already_pruned_heads) + mask[head] = 0 + mask = mask.view(-1).contiguous().eq(1) + index: torch.LongTensor = torch.arange(len(mask))[mask].long() + return heads, index + + +def meshgrid(*tensors: torch.Tensor | list[torch.Tensor], indexing: str | None = None) -> tuple[torch.Tensor, ...]: + """ + Wrapper around torch.meshgrid to avoid warning messages about the introduced `indexing` argument. + + Reference: https://pytorch.org/docs/1.13/generated/torch.meshgrid.html + """ + return torch.meshgrid(*tensors, indexing=indexing) + + +def id_tensor_storage(tensor: torch.Tensor) -> tuple[torch.device, int, int]: + """ + 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. + """ + if tensor.device.type == "xla" and is_torch_xla_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, storage_size(tensor) + + +def isin_mps_friendly(elements: torch.Tensor, test_elements: torch.Tensor | int) -> torch.Tensor: + """ + Same as `torch.isin` without flags, but MPS-friendly. We can remove this function when we stop supporting + torch <= 2.3. See https://github.com/pytorch/pytorch/issues/77764#issuecomment-2067838075 + + Args: + elements (`torch.Tensor`): Input elements + test_elements (`torch.Tensor` or `int`): The elements to check against. + + Returns: + `torch.Tensor`: A boolean tensor of the same shape as `elements` that is True for `elements` in `test_elements` + and False otherwise + """ + + if elements.device.type == "mps" and not is_torch_greater_or_equal_than_2_4: + test_elements = torch.tensor(test_elements) + if test_elements.ndim == 0: + test_elements = test_elements.unsqueeze(0) + return elements.tile(test_elements.shape[0], 1).eq(test_elements.unsqueeze(1)).sum(dim=0).bool().squeeze() + else: + # Note: don't use named arguments in `torch.isin`, see https://github.com/pytorch/pytorch/issues/126045 + return torch.isin(elements, test_elements) + + +def compile_compatible_method_lru_cache(*lru_args, **lru_kwargs): + """ + LRU cache decorator from standard functools library, but with a workaround to disable + caching when torchdynamo is compiling. Expected to work with class methods. + """ + + def decorator(func): + @wraps(func) + def wrapper(self, *args, **kwargs): + if not is_torchdynamo_compiling(): + # Cache the function only if the model is not being compiled + # check if the function is already cached, otherwise create it + if not hasattr(self, f"_cached_{func.__name__}"): + self.__setattr__( + f"_cached_{func.__name__}", lru_cache(*lru_args, **lru_kwargs)(func.__get__(self)) + ) + return self.__getattribute__(f"_cached_{func.__name__}")(*args, **kwargs) + else: + # Otherwise, just call the original function + return func(self, *args, **kwargs) + + return wrapper + + return decorator diff --git a/lib/python3.12/site-packages/transformers/quantizers/__init__.py b/lib/python3.12/site-packages/transformers/quantizers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7117bc2b5d8045380f582ac09338af3f3f3e3f61 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2024 The HuggingFace Inc. 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. +from .auto import AutoHfQuantizer, AutoQuantizationConfig, register_quantization_config, register_quantizer +from .base import HfQuantizer +from .quantizers_utils import get_module_from_name diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e6f94daf328ac14f9761b9d557d2afa82f00a20c Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/auto.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/auto.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..95b34df615733ed2a8495351fc3116904025a260 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/auto.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d7dc47aa733a0187df723815ded4009d55051b8d Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_aqlm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_aqlm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c21172859fe3ad619ff5bcd56a74aa15f9790416 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_aqlm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_awq.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_awq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f613519a784ae41d1bd6c729bff55391b858154 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_awq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_bitnet.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_bitnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b692378163ed3bf01788b6715e1a485412f53cb Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_bitnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_bnb_4bit.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_bnb_4bit.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..55d9d28a160fecac1a4f4965062505413b901e39 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_bnb_4bit.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_bnb_8bit.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_bnb_8bit.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30a0641dca6b9847a2d86ea186c5f9d695f7955a Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_bnb_8bit.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_compressed_tensors.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_compressed_tensors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce548ed39899d2712bcd9b8379e9855120ec271d Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_compressed_tensors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_eetq.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_eetq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1d929d8b2a46a185e10c73bb6f2effb7ec022d5 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_eetq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_fbgemm_fp8.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_fbgemm_fp8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8793d2cb4d8de3a49c0aab3e4e5c85872d6917de Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_fbgemm_fp8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_finegrained_fp8.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_finegrained_fp8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d105801e45b6d1ef40c0a60db68ecd03c8da9eee Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_finegrained_fp8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_gptq.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_gptq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a498fc9e22c49203f60f2d4346638c45ceab51db Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_gptq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_higgs.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_higgs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..039974991a0f8abbffb8ed43018bc19e058aa8ba Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_higgs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_hqq.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_hqq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9f8bb119fe006e31c674dcd15ac60a3753174ee Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_hqq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_quanto.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_quanto.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d750a80d301df1d5ab7f85f581e936701b9845e1 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_quanto.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_quark.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_quark.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c33f18ae73f9b7f39bb62174b22c9c273cd5444 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_quark.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_spqr.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_spqr.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..38e952433c20b8066a797d9c6d94cce0dfe78955 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_spqr.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_torchao.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_torchao.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..937f2655beca5ab0fcb2a75027c35a28f4c12c46 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_torchao.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_vptq.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_vptq.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eab362e3214448c37234127d4318c09e12b12c3d Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizer_vptq.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizers_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizers_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f56a0bf682af8f8835586c24b7d3e43b71240ef7 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/quantizers/__pycache__/quantizers_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/quantizers/auto.py b/lib/python3.12/site-packages/transformers/quantizers/auto.py new file mode 100644 index 0000000000000000000000000000000000000000..9d24b3539530b5d1ae03972de776f1d610c703d9 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/auto.py @@ -0,0 +1,266 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# Modifications Copyright (C) 2025, Advanced Micro Devices, Inc. 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. +import warnings +from typing import Dict, Optional, Union + +from ..models.auto.configuration_auto import AutoConfig +from ..utils import logging +from ..utils.quantization_config import ( + AqlmConfig, + AwqConfig, + BitNetConfig, + BitsAndBytesConfig, + CompressedTensorsConfig, + EetqConfig, + FbgemmFp8Config, + FineGrainedFP8Config, + GPTQConfig, + HiggsConfig, + HqqConfig, + QuantizationConfigMixin, + QuantizationMethod, + QuantoConfig, + QuarkConfig, + SpQRConfig, + TorchAoConfig, + VptqConfig, +) +from .base import HfQuantizer +from .quantizer_aqlm import AqlmHfQuantizer +from .quantizer_awq import AwqQuantizer +from .quantizer_bitnet import BitNetHfQuantizer +from .quantizer_bnb_4bit import Bnb4BitHfQuantizer +from .quantizer_bnb_8bit import Bnb8BitHfQuantizer +from .quantizer_compressed_tensors import CompressedTensorsHfQuantizer +from .quantizer_eetq import EetqHfQuantizer +from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer +from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer +from .quantizer_gptq import GptqHfQuantizer +from .quantizer_higgs import HiggsHfQuantizer +from .quantizer_hqq import HqqHfQuantizer +from .quantizer_quanto import QuantoHfQuantizer +from .quantizer_quark import QuarkHfQuantizer +from .quantizer_spqr import SpQRHfQuantizer +from .quantizer_torchao import TorchAoHfQuantizer +from .quantizer_vptq import VptqHfQuantizer + + +AUTO_QUANTIZER_MAPPING = { + "awq": AwqQuantizer, + "bitsandbytes_4bit": Bnb4BitHfQuantizer, + "bitsandbytes_8bit": Bnb8BitHfQuantizer, + "gptq": GptqHfQuantizer, + "aqlm": AqlmHfQuantizer, + "quanto": QuantoHfQuantizer, + "quark": QuarkHfQuantizer, + "eetq": EetqHfQuantizer, + "higgs": HiggsHfQuantizer, + "hqq": HqqHfQuantizer, + "compressed-tensors": CompressedTensorsHfQuantizer, + "fbgemm_fp8": FbgemmFp8HfQuantizer, + "torchao": TorchAoHfQuantizer, + "bitnet": BitNetHfQuantizer, + "vptq": VptqHfQuantizer, + "spqr": SpQRHfQuantizer, + "fp8": FineGrainedFP8HfQuantizer, +} + +AUTO_QUANTIZATION_CONFIG_MAPPING = { + "awq": AwqConfig, + "bitsandbytes_4bit": BitsAndBytesConfig, + "bitsandbytes_8bit": BitsAndBytesConfig, + "eetq": EetqConfig, + "gptq": GPTQConfig, + "aqlm": AqlmConfig, + "quanto": QuantoConfig, + "quark": QuarkConfig, + "hqq": HqqConfig, + "compressed-tensors": CompressedTensorsConfig, + "fbgemm_fp8": FbgemmFp8Config, + "higgs": HiggsConfig, + "torchao": TorchAoConfig, + "bitnet": BitNetConfig, + "vptq": VptqConfig, + "spqr": SpQRConfig, + "fp8": FineGrainedFP8Config, +} + +logger = logging.get_logger(__name__) + + +class AutoQuantizationConfig: + """ + The Auto-HF quantization config class that takes care of automatically dispatching to the correct + quantization config given a quantization config stored in a dictionary. + """ + + @classmethod + def from_dict(cls, quantization_config_dict: Dict): + quant_method = quantization_config_dict.get("quant_method", None) + # We need a special care for bnb models to make sure everything is BC .. + if quantization_config_dict.get("load_in_8bit", False) or quantization_config_dict.get("load_in_4bit", False): + suffix = "_4bit" if quantization_config_dict.get("load_in_4bit", False) else "_8bit" + quant_method = QuantizationMethod.BITS_AND_BYTES + suffix + elif quant_method is None: + raise ValueError( + "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized" + ) + + if quant_method not in AUTO_QUANTIZATION_CONFIG_MAPPING.keys(): + raise ValueError( + f"Unknown quantization type, got {quant_method} - supported types are:" + f" {list(AUTO_QUANTIZER_MAPPING.keys())}" + ) + + target_cls = AUTO_QUANTIZATION_CONFIG_MAPPING[quant_method] + return target_cls.from_dict(quantization_config_dict) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + model_config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) + if getattr(model_config, "quantization_config", None) is None: + raise ValueError( + f"Did not found a `quantization_config` in {pretrained_model_name_or_path}. Make sure that the model is correctly quantized." + ) + quantization_config_dict = model_config.quantization_config + quantization_config = cls.from_dict(quantization_config_dict) + # Update with potential kwargs that are passed through from_pretrained. + quantization_config.update(**kwargs) + return quantization_config + + +class AutoHfQuantizer: + """ + The Auto-HF quantizer class that takes care of automatically instantiating to the correct + `HfQuantizer` given the `QuantizationConfig`. + """ + + @classmethod + def from_config(cls, quantization_config: Union[QuantizationConfigMixin, Dict], **kwargs): + # Convert it to a QuantizationConfig if the q_config is a dict + if isinstance(quantization_config, dict): + quantization_config = AutoQuantizationConfig.from_dict(quantization_config) + + quant_method = quantization_config.quant_method + + # Again, we need a special care for bnb as we have a single quantization config + # class for both 4-bit and 8-bit quantization + if quant_method == QuantizationMethod.BITS_AND_BYTES: + if quantization_config.load_in_8bit: + quant_method += "_8bit" + else: + quant_method += "_4bit" + + if quant_method not in AUTO_QUANTIZER_MAPPING.keys(): + raise ValueError( + f"Unknown quantization type, got {quant_method} - supported types are:" + f" {list(AUTO_QUANTIZER_MAPPING.keys())}" + ) + + target_cls = AUTO_QUANTIZER_MAPPING[quant_method] + return target_cls(quantization_config, **kwargs) + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): + quantization_config = AutoQuantizationConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) + return cls.from_config(quantization_config) + + @classmethod + def merge_quantization_configs( + cls, + quantization_config: Union[dict, QuantizationConfigMixin], + quantization_config_from_args: Optional[QuantizationConfigMixin], + ): + """ + handles situations where both quantization_config from args and quantization_config from model config are present. + """ + if quantization_config_from_args is not None: + warning_msg = ( + "You passed `quantization_config` or equivalent parameters to `from_pretrained` but the model you're loading" + " already has a `quantization_config` attribute. The `quantization_config` from the model will be used." + ) + else: + warning_msg = "" + + if isinstance(quantization_config, dict): + quantization_config = AutoQuantizationConfig.from_dict(quantization_config) + + if ( + isinstance(quantization_config, (GPTQConfig, AwqConfig, FbgemmFp8Config, CompressedTensorsConfig)) + and quantization_config_from_args is not None + ): + # special case for GPTQ / AWQ / FbgemmFp8 config collision + loading_attr_dict = quantization_config_from_args.get_loading_attributes() + for attr, val in loading_attr_dict.items(): + setattr(quantization_config, attr, val) + + warning_msg += f"However, loading attributes (e.g. {list(loading_attr_dict.keys())}) will be overwritten with the one you passed to `from_pretrained`. The rest will be ignored." + + if warning_msg != "": + warnings.warn(warning_msg) + + return quantization_config + + @staticmethod + def supports_quant_method(quantization_config_dict): + quant_method = quantization_config_dict.get("quant_method", None) + if quantization_config_dict.get("load_in_8bit", False) or quantization_config_dict.get("load_in_4bit", False): + suffix = "_4bit" if quantization_config_dict.get("load_in_4bit", False) else "_8bit" + quant_method = QuantizationMethod.BITS_AND_BYTES + suffix + elif quant_method is None: + raise ValueError( + "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized" + ) + + if quant_method not in AUTO_QUANTIZATION_CONFIG_MAPPING.keys(): + logger.warning( + f"Unknown quantization type, got {quant_method} - supported types are:" + f" {list(AUTO_QUANTIZER_MAPPING.keys())}. Hence, we will skip the quantization. " + "To remove the warning, you can delete the quantization_config attribute in config.json" + ) + return False + return True + + +def register_quantization_config(method: str): + """Register a custom quantization configuration.""" + + def register_config_fn(cls): + if method in AUTO_QUANTIZATION_CONFIG_MAPPING: + raise ValueError(f"Config '{method}' already registered") + + if not issubclass(cls, QuantizationConfigMixin): + raise ValueError("Config must extend QuantizationConfigMixin") + + AUTO_QUANTIZATION_CONFIG_MAPPING[method] = cls + return cls + + return register_config_fn + + +def register_quantizer(name: str): + """Register a custom quantizer.""" + + def register_quantizer_fn(cls): + if name in AUTO_QUANTIZER_MAPPING: + raise ValueError(f"Quantizer '{name}' already registered") + + if not issubclass(cls, HfQuantizer): + raise ValueError("Quantizer must extend HfQuantizer") + + AUTO_QUANTIZER_MAPPING[name] = cls + return cls + + return register_quantizer_fn diff --git a/lib/python3.12/site-packages/transformers/quantizers/base.py b/lib/python3.12/site-packages/transformers/quantizers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..a780dca7548f49ae3dac3fee57cec79748287187 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/base.py @@ -0,0 +1,340 @@ +# Copyright 2024 The HuggingFace Inc. 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. +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from ..utils import is_torch_available +from ..utils.quantization_config import QuantizationConfigMixin, QuantizationMethod +from .quantizers_utils import get_module_from_name + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +if is_torch_available(): + import torch + from torch.nn import ModuleList +else: + ModuleList = str + + +class HfQuantizer(ABC): + """ + Abstract class of the HuggingFace quantizer. Supports for now quantizing HF transformers models for inference and/or quantization. + This class is used only for transformers.PreTrainedModel.from_pretrained and cannot be easily used outside the scope of that method + yet. + + Attributes + quantization_config (`transformers.utils.quantization_config.QuantizationConfigMixin`): + The quantization config that defines the quantization parameters of your model that you want to quantize. + modules_to_not_convert (`List[str]`, *optional*): + The list of module names to not convert when quantizing the model. + required_packages (`List[str]`, *optional*): + The list of required pip packages to install prior to using the quantizer + requires_calibration (`bool`): + Whether the quantization method requires to calibrate the model before using it. + requires_parameters_quantization (`bool`): + Whether the quantization method requires to create a new Parameter. For example, for bitsandbytes, it is + required to create a new xxxParameter in order to properly quantize the model. + """ + + requires_calibration = False + required_packages = None + requires_parameters_quantization = False + + def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): + self.quantization_config = quantization_config + + # -- Handle extra kwargs below -- + self.modules_to_not_convert = kwargs.pop("modules_to_not_convert", []) + self.pre_quantized = kwargs.pop("pre_quantized", True) + + if not self.pre_quantized and self.requires_calibration: + raise ValueError( + f"The quantization method {quantization_config.quant_method} does require the model to be pre-quantized." + f" You explicitly passed `pre_quantized=False` meaning your model weights are not quantized. Make sure to " + f"pass `pre_quantized=True` while knowing what you are doing." + ) + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + """ + Some quantization methods require to explicitly set the dtype of the model to a + target dtype. You need to override this method in case you want to make sure that behavior is + preserved + + Args: + torch_dtype (`torch.dtype`): + The input dtype that is passed in `from_pretrained` + """ + return torch_dtype + + def update_device_map(self, device_map: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: + """ + Override this method if you want to pass a override the existing device map with a new + one. E.g. for bitsandbytes, since `accelerate` is a hard requirement, if no device_map is + passed, the device_map is set to `"auto"`` + + Args: + device_map (`Union[dict, str]`, *optional*): + The device_map that is passed through the `from_pretrained` method. + """ + return device_map + + def adjust_target_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + """ + Override this method if you want to adjust the `target_dtype` variable used in `from_pretrained` + to compute the device_map in case the device_map is a `str`. E.g. for bitsandbytes we force-set `target_dtype` + to `torch.int8` and for 4-bit we pass a custom enum `accelerate.CustomDtype.int4`. + + Args: + torch_dtype (`torch.dtype`, *optional*): + The torch_dtype that is used to compute the device_map. + """ + return torch_dtype + + def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]: + """ + Override this method if you want to adjust the `missing_keys`. + + Args: + missing_keys (`List[str]`, *optional*): + The list of missing keys in the checkpoint compared to the state dict of the model + """ + return missing_keys + + def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str) -> List[str]: + """ + Override this method if you want to adjust the `unexpected_keys`. + + Args: + unexpected_keys (`List[str]`, *optional*): + The list of unexpected keys in the checkpoint compared to the state dict of the model + """ + return unexpected_keys + + def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: + """ + Override this method if you want to adjust the `missing_keys` after loading the model params, + but before the model is post-processed. + + Args: + missing_keys (`List[str]`, *optional*): + The list of missing keys in the checkpoint compared to the state dict of the model + """ + return missing_keys + + def update_expected_keys(self, model, expected_keys: List[str], loaded_keys: List[str]) -> List[str]: + """ + Override this method if you want to adjust the `update_expected_keys`. + + Args: + expected_keys (`List[str]`, *optional*): + The list of the expected keys in the initialized model. + loaded_keys (`List[str]`, *optional*): + The list of the loaded keys in the checkpoint. + """ + return expected_keys + + def get_special_dtypes_update(self, model, torch_dtype: "torch.dtype") -> Dict[str, "torch.dtype"]: + """ + returns dtypes for modules that are not quantized - used for the computation of the device_map in case + one passes a str as a device_map. The method will use the `modules_to_not_convert` that is modified + in `_process_model_before_weight_loading`. + + Args: + model (`~transformers.PreTrainedModel`): + The model to quantize + torch_dtype (`torch.dtype`): + The dtype passed in `from_pretrained` method. + """ + + return { + name: torch_dtype + for name, _ in model.named_parameters() + if any(m in name for m in self.modules_to_not_convert) + } + + def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + """adjust max_memory argument for infer_auto_device_map() if extra memory is needed for quantization""" + return max_memory + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ) -> bool: + """ + checks if a loaded state_dict component is part of quantized param + some validation; only defined if + requires_parameters_quantization == True for quantization methods that require to create a new parameters + for quantization. + """ + return False + + def create_quantized_param(self, *args, **kwargs) -> "torch.nn.Parameter": + """ + takes needed components from state_dict and creates quantized param; only applicable if + requires_parameters_quantization == True + """ + if not self.requires_parameters_quantization: + raise AttributeError( + f"`.create_quantized_param()` method is not supported by quantizer class {self.__class__.__name__}." + ) + + def validate_environment(self, *args, **kwargs): + """ + This method is used to potentially check for potential conflicts with arguments that are + passed in `from_pretrained`. You need to define it for all future quantizers that are integrated with transformers. + If no explicit check are needed, simply return nothing. + """ + return + + def update_tp_plan(self, config): + "updates the tp plan for the scales" + return config + + def preprocess_model(self, model: "PreTrainedModel", **kwargs): + """ + Setting model attributes and/or converting model before weights loading. At this point + the model should be initialized on the meta device so you can freely manipulate the skeleton + of the model in order to replace modules in-place. Make sure to override the abstract method `_process_model_before_weight_loading`. + + Args: + model (`~transformers.PreTrainedModel`): + The model to quantize + kwargs (`dict`, *optional*): + The keyword arguments that are passed along `_process_model_before_weight_loading`. + """ + model.is_quantized = True + model.quantization_method = self.quantization_config.quant_method + self._convert_model_for_quantization(model) + return self._process_model_before_weight_loading(model, **kwargs) + + def postprocess_model(self, model: "PreTrainedModel", **kwargs): + """ + Post-process the model post weights loading. + Make sure to override the abstract method `_process_model_after_weight_loading`. + + Args: + model (`~transformers.PreTrainedModel`): + The model to quantize + kwargs (`dict`, *optional*): + The keyword arguments that are passed along `_process_model_after_weight_loading`. + """ + return self._process_model_after_weight_loading(model, **kwargs) + + def dequantize(self, model): + """ + Potentially dequantize the model to retrive the original model, with some loss in accuracy / performance. + Note not all quantization schemes support this. + """ + model = self._dequantize(model) + + # Delete quantizer and quantization config + del model.hf_quantizer + del model.config.quantization_config + del model.config._pre_quantization_dtype + model.is_quantized = False + + return model + + def _dequantize(self, model): + raise NotImplementedError( + f"{self.quantization_config.quant_method} has no implementation of `dequantize`, please raise an issue on GitHub." + ) + + @staticmethod + def get_modules_to_not_convert( + model: "PreTrainedModel", + skip_modules: Optional[List[str]] = None, + keep_in_fp32_modules: Optional[List[str]] = None, + ): + from ..integrations import get_keys_to_not_convert + + modules_to_not_convert = [] + if skip_modules is None: + modules_to_not_convert = get_keys_to_not_convert(model) + else: + modules_to_not_convert = skip_modules + + if keep_in_fp32_modules is not None: + modules_to_not_convert.extend(keep_in_fp32_modules) + + return modules_to_not_convert + + @property + def is_qat_trainable(self) -> bool: + """Flag indicating whether the quantized model can carry out quantization aware training""" + return False + + @property + def is_compileable(self) -> bool: + """Flag indicating whether the quantized model can be compiled""" + return False + + @abstractmethod + def _process_model_before_weight_loading(self, model, **kwargs): ... + + @abstractmethod + def _process_model_after_weight_loading(self, model, **kwargs): ... + + @abstractmethod + def is_serializable(self, safe_serialization=None): ... + + @property + @abstractmethod + def is_trainable(self): ... + + def _convert_model_for_quantization(self, model): + from accelerate import init_empty_weights + + for name, module in model.named_modules(): + module_class_name = module.__class__.__name__ + if ( + module_class_name in MODULES_TO_PATCH_FOR_QUANTIZATION.keys() + and self.quantization_config.quant_method == QuantizationMethod.COMPRESSED_TENSORS + ): + with init_empty_weights(): + parent_module, name = get_module_from_name(model, name) + parent_module._modules[name] = MODULES_TO_PATCH_FOR_QUANTIZATION[module_class_name]( + model.config.get_text_config() + ) + + +class SequentialLlama4TextExperts(ModuleList): + """ + A module that implements a compressed version of a list of expert modules. + This is specifically designed to work with Llama4TextExperts in MoE layers. + """ + + def __init__(self, config): + from transformers.models.llama4.modeling_llama4 import Llama4TextMLP + + super().__init__([Llama4TextMLP(config) for _ in range(config.num_local_experts)]) + self.num_experts = config.num_local_experts + + def forward( + self, + hidden_states: "torch.Tensor", + ) -> "torch.Tensor": + hidden_states = hidden_states.reshape(self.num_experts, -1, hidden_states.shape[-1]) + routed_out = torch.zeros_like(hidden_states) + for expert_idx in range(self.num_experts): + routed_out[expert_idx] = self[expert_idx](hidden_states[expert_idx]) + return routed_out + + +MODULES_TO_PATCH_FOR_QUANTIZATION = {"Llama4TextExperts": SequentialLlama4TextExperts} diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_aqlm.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_aqlm.py new file mode 100644 index 0000000000000000000000000000000000000000..9d1d6f7e89f1e9aa2454a1fd56c81524660f9079 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_aqlm.py @@ -0,0 +1,97 @@ +# Copyright 2024 The HuggingFace Inc. 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. +import importlib +from typing import TYPE_CHECKING, Optional + +from packaging import version + +from .base import HfQuantizer + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..integrations import replace_with_aqlm_linear +from ..utils import is_accelerate_available, is_aqlm_available, is_torch_available, logging +from ..utils.quantization_config import QuantizationConfigMixin + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class AqlmHfQuantizer(HfQuantizer): + """ + Quantizer of the AQLM method. Enables the loading of prequantized models. + """ + + requires_calibration = True + required_packages = ["aqlm"] + optimum_quantizer = None + + def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): + super().__init__(quantization_config, **kwargs) + self.quantization_config = quantization_config + + def validate_environment(self, *args, **kwargs): + if not is_accelerate_available(): + raise ImportError("Using `aqlm` quantization requires Accelerate: `pip install accelerate`") + + if not is_aqlm_available(): + raise ImportError("Using `aqlm` quantization requires AQLM: `pip install aqlm[gpu,cpu]`") + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + if torch.cuda.is_available(): + torch_dtype = torch.float16 + logger.info( + "CUDA available. Assuming AQLM inference on GPU and loading the model in `torch.float16`. To overwrite it, set `torch_dtype` manually." + ) + else: + torch_dtype = torch.float32 + logger.info( + "CUDA is unavailable. Assuming AQLM inference on CPU and loading the model in `torch.float32`. To overwrite it, set `torch_dtype` manually." + ) + return torch_dtype + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + **kwargs, + ): + replace_with_aqlm_linear( + model, + quantization_config=self.quantization_config, + linear_weights_not_to_quantize=self.quantization_config.linear_weights_not_to_quantize, + ) + model.config.quantization_config = self.quantization_config + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + return model + + @property + def is_trainable(self, model: Optional["PreTrainedModel"] = None): + aqlm_supports_training = version.parse(importlib.metadata.version("aqlm")) >= version.parse("1.0.2") + if aqlm_supports_training: + return True + else: + logger.warning( + f"Currently installed `aqlm` version ({importlib.metadata.version('aqlm')}) doesn't support training. If you wish to train a quantized model, please update `aqlm` with `pip install aqlm>=1.0.2`" + ) + return False + + def is_serializable(self, safe_serialization=None): + return True diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_awq.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_awq.py new file mode 100644 index 0000000000000000000000000000000000000000..28460ac38edb5e899689df7259d95b5d733bce6d --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_awq.py @@ -0,0 +1,153 @@ +# Copyright 2024 The HuggingFace Inc. 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. +import importlib.metadata +from typing import TYPE_CHECKING, List, Optional + +from packaging import version + +from .base import HfQuantizer + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..utils import is_accelerate_available, is_auto_awq_available, is_torch_available, logging +from ..utils.quantization_config import AWQLinearVersion + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class AwqQuantizer(HfQuantizer): + """ + 4-bit quantization for Activation-aware Weight Quantization(AWQ) (https://arxiv.org/abs/2306.00978) + """ + + # AWQ requires data callibration - we support only inference + requires_calibration = True + + required_packages = ["awq", "accelerate"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + + def validate_environment(self, device_map, **kwargs): + if not is_auto_awq_available(): + raise ImportError("Loading an AWQ quantized model requires auto-awq library (`pip install autoawq`)") + + if not is_accelerate_available(): + raise ImportError("Loading an AWQ quantized model requires accelerate (`pip install accelerate`)") + + if self.quantization_config.version == AWQLinearVersion.GEMM and not torch.cuda.is_available(): + logger.warning_once("No CUDA found, replace GEMM with IPEX version to support non-cuda AWQ model.") + self.quantization_config.version = AWQLinearVersion.IPEX + + if self.quantization_config.version == AWQLinearVersion.IPEX: + if version.parse(importlib.metadata.version("autoawq")) < version.parse("0.2.6"): + raise RuntimeError( + "To use IPEX backend, you need autoawq>0.2.6. Please install the latest version or from source." + ) + if device_map is None: + logger.warning_once( + "You have loaded an AWQ model without setting device_map, please set 'cpu' or 'xpu' or 'auto'" + ) + elif isinstance(device_map, dict) and "disk" in device_map.values(): + raise ValueError( + "You are attempting to load an IPEX version AWQ model with a device_map that contains disk device." + " This is not supported. Please make sure only cpu and xpu in the device_map." + ) + else: + if not torch.cuda.is_available(): + raise RuntimeError( + "GPU is required to run AWQ quantized model. You can use IPEX version AWQ if you have an Intel CPU" + ) + + if device_map is None: + logger.warning_once( + "You have loaded an AWQ model on CPU and have a CUDA device available, make sure to set " + "your model on a GPU device in order to run your model." + ) + elif device_map is not None: + if isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()): + raise ValueError( + "You are attempting to load an AWQ model with a device_map that contains a CPU or disk device." + " This is not supported. Please remove the CPU or disk device from the device_map." + ) + + def update_torch_dtype(self, torch_dtype): + if torch_dtype is None: + torch_dtype = torch.float16 + logger.info("Loading the model in `torch.float16`. To overwrite it, set `torch_dtype` manually.") + elif torch_dtype != torch.float16: + logger.warning("We suggest you to set `torch_dtype=torch.float16` for better efficiency with AWQ.") + return torch_dtype + + def _process_model_before_weight_loading( + self, model: "PreTrainedModel", keep_in_fp32_modules: Optional[List[str]] = None, **kwargs + ): + from ..integrations import replace_quantization_scales, replace_with_awq_linear + + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules + ) + + model, has_been_replaced = replace_with_awq_linear( + model, quantization_config=self.quantization_config, modules_to_not_convert=self.modules_to_not_convert + ) + + model = replace_quantization_scales(model, model.config.model_type) + + if not has_been_replaced: + logger.warning( + "You are loading an AWQ model but no linear modules were found in your model." + " Please double check your model architecture, or submit an issue on github if you think this is a bug." + ) + + def _process_model_after_weight_loading(self, model, **kwargs): + if self.quantization_config.do_fuse: + from ..integrations import fuse_awq_modules + + model = fuse_awq_modules(model, self.quantization_config) + model._awq_is_fused = True # TODO: consider storing this flag in model.config instead + + if self.quantization_config.version == AWQLinearVersion.EXLLAMA: + from ..integrations import post_init_awq_exllama_modules + + model = post_init_awq_exllama_modules(model, self.quantization_config.exllama_config) + + if self.quantization_config.version == AWQLinearVersion.IPEX: + from ..integrations import post_init_awq_ipex_modules + + model = post_init_awq_ipex_modules(model) + + def is_serializable(self, safe_serialization=None): + # AWQ through auto-awq has been always serializable, except if the model is fused. + if self.quantization_config.do_fuse: + logger.warning("You cannot save an AWQ model that uses fused modules!") + return False + + if self.quantization_config.version == AWQLinearVersion.EXLLAMA: + logger.warning("You cannot save an AWQ model that uses Exllama backend!") + return False + + return True + + @property + def is_trainable(self): + # AWQ supports PEFT fine-tuning from version 0.2.0 + MIN_AWQ_VERSION_FOR_PEFT = "0.2.0" + return version.parse(importlib.metadata.version("autoawq")) >= version.parse(MIN_AWQ_VERSION_FOR_PEFT) diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_bitnet.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_bitnet.py new file mode 100644 index 0000000000000000000000000000000000000000..e56bb161ac449853f0fb12bc86aed1a9c75b88ee --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_bitnet.py @@ -0,0 +1,113 @@ +# Copyright 2024 The HuggingFace Inc. 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. +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +from .base import HfQuantizer + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..utils import is_accelerate_available, is_torch_available, logging + + +if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) + + +class BitNetHfQuantizer(HfQuantizer): + """ + 1.58-bit quantization from BitNet quantization method: + Before loading: it converts the linear layers into BitLinear layers during loading. + + Checkout the paper introducing this method : https://arxiv.org/pdf/2402.17764 + """ + + requires_parameters_quantization = False + requires_calibration = True + + required_packages = ["accelerate"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + self.quantization_config = quantization_config + + def validate_environment(self, *args, **kwargs): + if not is_accelerate_available(): + raise ImportError("Loading a BitNet quantized model requires accelerate (`pip install accelerate`)") + + if kwargs.get("from_tf", False) or kwargs.get("from_flax", False): + raise ValueError( + "Loading ternary weights from tf/flax is currently not supported, please make" + " sure the weights are in PyTorch format." + ) + + if not torch.cuda.is_available(): + logger.warning_once( + "You don't have a GPU available to load the model, the inference will be slow because of weight unpacking" + ) + return + + device_map = kwargs.get("device_map", None) + if device_map is None: + logger.warning_once( + "You have loaded a BitNet model on CPU and have a CUDA device available, make sure to set " + "your model on a GPU device in order to run your model." + ) + elif device_map is not None: + if isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()): + raise ValueError( + "You are attempting to load a BitNet model with a device_map that contains a CPU or disk device." + "This is not supported. Please remove the CPU or disk device from the device_map." + ) + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + return model + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + keep_in_fp32_modules: Optional[List[str]] = None, + **kwargs, + ): + from ..integrations import replace_with_bitnet_linear + + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules + ) + + model = replace_with_bitnet_linear( + model, + modules_to_not_convert=self.modules_to_not_convert, + quantization_config=self.quantization_config, + pre_quantized=self.pre_quantized, + ) + + def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + max_memory = {key: val * 0.90 for key, val in max_memory.items()} + return max_memory + + def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": + target_dtype = torch.int8 + return target_dtype + + def is_serializable(self, safe_serialization=None): + return True + + @property + def is_trainable(self) -> bool: + return False diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_bnb_4bit.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_bnb_4bit.py new file mode 100644 index 0000000000000000000000000000000000000000..7fb9176c4673a009914600ae47068c22826739d2 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_bnb_4bit.py @@ -0,0 +1,358 @@ +# Copyright 2024 The HuggingFace Inc. 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. +import importlib +from functools import cached_property +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from packaging import version + +from .base import HfQuantizer +from .quantizers_utils import get_module_from_name + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..utils import ( + ACCELERATE_MIN_VERSION, + is_accelerate_available, + is_bitsandbytes_available, + is_torch_available, + is_torch_hpu_available, + is_torch_npu_available, + is_torch_xpu_available, + logging, +) + + +if is_torch_available(): + import torch + + from ..pytorch_utils import Conv1D + +logger = logging.get_logger(__name__) + + +class Bnb4BitHfQuantizer(HfQuantizer): + """ + 4-bit quantization from bitsandbytes.py quantization method: + before loading: converts transformer layers into Linear4bit during loading: load 16bit weight and pass to the + layer object after: quantizes individual weights in Linear4bit into 4bit at the first .cuda() call + saving: + from state dict, as usual; saves weights and `quant_state` components + loading: + need to locate `quant_state` components and pass to Param4bit constructor + """ + + use_keep_in_fp32_modules = True + requires_parameters_quantization = True + requires_calibration = False + + required_packages = ["bitsandbytes", "accelerate"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + + if self.quantization_config.llm_int8_skip_modules is not None: + self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules + + def validate_environment(self, *args, **kwargs): + if not is_accelerate_available(): + raise ImportError( + f"Using `bitsandbytes` 4-bit quantization requires Accelerate: `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`" + ) + if not is_bitsandbytes_available(): + raise ImportError( + "Using `bitsandbytes` 4-bit quantization requires the latest version of bitsandbytes: `pip install -U bitsandbytes`" + ) + + from ..integrations import validate_bnb_backend_availability + from ..utils import is_bitsandbytes_multi_backend_available + + bnb_multibackend_is_enabled = is_bitsandbytes_multi_backend_available() + validate_bnb_backend_availability(raise_exception=True) + + if kwargs.get("from_tf", False) or kwargs.get("from_flax", False): + raise ValueError( + "Converting into 4-bit or 8-bit weights from tf/flax weights is currently not supported, please make" + " sure the weights are in PyTorch format." + ) + + device_map = kwargs.get("device_map", None) + if ( + device_map is not None + and isinstance(device_map, dict) + and not self.quantization_config.llm_int8_enable_fp32_cpu_offload + ): + device_map_without_lm_head = { + key: device_map[key] for key in device_map.keys() if key not in self.modules_to_not_convert + } + if set(device_map.values()) == {"cpu"} and bnb_multibackend_is_enabled: + pass + elif "cpu" in device_map_without_lm_head.values() or "disk" in device_map_without_lm_head.values(): + raise ValueError( + "Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the " + "quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules " + "in 32-bit, you need to set `llm_int8_enable_fp32_cpu_offload=True` and pass a custom `device_map` to " + "`from_pretrained`. Check " + "https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu " + "for more details. " + ) + + if version.parse(importlib.metadata.version("bitsandbytes")) < version.parse("0.39.0"): + raise ValueError( + "You have a version of `bitsandbytes` that is not compatible with 4bit inference and training" + " make sure you have the latest version of `bitsandbytes` installed" + ) + + def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": + if version.parse(importlib.metadata.version("accelerate")) > version.parse("0.19.0"): + from accelerate.utils import CustomDtype + + if target_dtype != torch.int8: + logger.info("target_dtype {target_dtype} is replaced by `CustomDtype.INT4` for 4-bit BnB quantization") + return CustomDtype.INT4 + else: + raise ValueError( + "You are using `device_map='auto'` on a 4bit loaded version of the model. To automatically compute" + " the appropriate device map, you should upgrade your `accelerate` library," + "`pip install --upgrade accelerate` or install it from source to support fp4 auto device map" + "calculation. You may encounter unexpected behavior, or pass your own device map" + ) + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ) -> bool: + import bitsandbytes as bnb + + module, tensor_name = get_module_from_name(model, param_name) + if isinstance(module._parameters.get(tensor_name, None), bnb.nn.Params4bit): + # Add here check for loaded components' dtypes once serialization is implemented + return True + elif isinstance(module, bnb.nn.Linear4bit) and tensor_name == "bias": + # bias could be loaded by regular set_module_tensor_to_device() from accelerate, + # but it would wrongly use uninitialized weight there. + return True + else: + return False + + def create_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + state_dict: Dict[str, Any], + unexpected_keys: Optional[List[str]] = None, + ): + """ + combines logic from _load_state_dict_into_meta_model and .integrations.bitsandbytes.py::set_module_quantized_tensor_to_device() + """ + import bitsandbytes as bnb + + module, tensor_name = get_module_from_name(model, param_name) + + if tensor_name not in module._parameters: + raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.") + + old_value = getattr(module, tensor_name) + + # `torch.Tensor.to()` is not supported by `torch_npu` (see this [issue](https://github.com/Ascend/pytorch/issues/16)). + if isinstance(target_device, int) and is_torch_npu_available(): + target_device = f"npu:{target_device}" + if tensor_name == "bias": + if param_value is None: + new_value = old_value.to(target_device) + else: + new_value = param_value.to(target_device) + + new_value = torch.nn.Parameter(new_value, requires_grad=old_value.requires_grad) + module._parameters[tensor_name] = new_value + return + + if not isinstance(module._parameters[tensor_name], bnb.nn.Params4bit): + raise ValueError("this function only loads `Linear4bit components`") + if ( + old_value.device == torch.device("meta") + and target_device not in ["meta", torch.device("meta")] + and param_value is None + ): + raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {target_device}.") + + # construct `new_value` for the module._parameters[tensor_name]: + if self.pre_quantized: + # 4bit loading. Collecting components for restoring quantized weight + # This can be expanded to make a universal call for any quantized weight loading + + if not self.is_serializable: + raise ValueError( + "Detected int4 weights but the version of bitsandbytes is not compatible with int4 serialization. " + "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`." + ) + + if (param_name + ".quant_state.bitsandbytes__fp4" not in state_dict) and ( + param_name + ".quant_state.bitsandbytes__nf4" not in state_dict + ): + raise ValueError( + f"Supplied state dict for {param_name} does not contain `bitsandbytes__*` and possibly other `quantized_stats` components." + ) + + quantized_stats = {} + for k, v in state_dict.items(): + if param_name + "." in k: + quantized_stats[k] = v + if unexpected_keys is not None and k in unexpected_keys: + unexpected_keys.remove(k) + + param_kwargs = {} + if self.is_bnb_supports_quant_storage_module: + param_kwargs["module"] = module + + new_value = bnb.nn.Params4bit.from_prequantized( + data=param_value, + quantized_stats=quantized_stats, + requires_grad=False, + device=target_device, + **param_kwargs, + ) + else: + new_value = param_value.to("cpu") + + # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization. + # Since weights are saved in the correct "orientation", we skip transposing when loading. + if issubclass(module.source_cls, Conv1D): + new_value = new_value.T + + kwargs = old_value.__dict__ + new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(target_device) + + module._parameters[tensor_name] = new_value + + # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer.adjust_max_memory + def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + # need more space for buffers that are created during quantization + max_memory = {key: val * 0.90 for key, val in max_memory.items()} + return max_memory + + # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer.update_torch_dtype + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + # We force the `dtype` to be float16, this is a requirement from `bitsandbytes` + logger.info( + "Overriding torch_dtype=%s with `torch_dtype=torch.float16` due to " + "requirements of `bitsandbytes` to enable model loading in 8-bit or 4-bit. " + "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass" + " torch_dtype=torch.float16 to remove this warning.", + torch_dtype, + ) + torch_dtype = torch.float16 + return torch_dtype + + def update_device_map(self, device_map): + if device_map is None: + if torch.cuda.is_available(): + device_map = {"": torch.cuda.current_device()} + elif is_torch_npu_available(): + device_map = {"": f"npu:{torch.npu.current_device()}"} + elif is_torch_hpu_available(): + device_map = {"": f"hpu:{torch.hpu.current_device()}"} + elif is_torch_xpu_available(): + device_map = {"": f"xpu:{torch.xpu.current_device()}"} + else: + device_map = {"": "cpu"} + logger.info( + "The device_map was not initialized. " + f"Setting device_map to {device_map}. " + "If you want to use the model for inference, please set device_map ='auto' " + ) + return device_map + + # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer._process_model_before_weight_loading + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + device_map, + keep_in_fp32_modules: Optional[List[str]] = None, + **kwargs, + ): + from ..integrations import replace_with_bnb_linear + + llm_int8_enable_fp32_cpu_offload = self.quantization_config.llm_int8_enable_fp32_cpu_offload + + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules + ) + + # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk` + if isinstance(device_map, dict) and len(device_map.keys()) > 1: + keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]] + + if len(keys_on_cpu) > 0 and not llm_int8_enable_fp32_cpu_offload: + raise ValueError( + "If you want to offload some keys to `cpu` or `disk`, you need to set " + "`llm_int8_enable_fp32_cpu_offload=True`. Note that these modules will not be " + " converted to 8-bit but kept in 32-bit." + ) + self.modules_to_not_convert.extend(keys_on_cpu) + + model = replace_with_bnb_linear( + model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config + ) + # TODO: consider bringing replace_with_bnb_linear() code from ..integrations/bitsandbyter.py to here + + model.config.quantization_config = self.quantization_config + + # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer._process_model_after_weight_loading with 8bit->4bit + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + model.is_loaded_in_4bit = True + model.is_4bit_serializable = self.is_serializable() + return model + + def is_serializable(self, safe_serialization=None): + _is_4bit_serializable = version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse("0.41.3") + + if not _is_4bit_serializable: + logger.warning( + "You are calling `save_pretrained` to a 4-bit converted model, but your `bitsandbytes` version doesn't support it. " + "If you want to save 4-bit models, make sure to have `bitsandbytes>=0.41.3` installed." + ) + return False + + return True + + @cached_property + def is_bnb_supports_quant_storage_module(self) -> bool: + """ + determines if the current version of bitsandbytes supports + the `module` parameter in `Params4bit.from_prequantized` + :return: + """ + return version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse("0.43.3") + + @property + def is_trainable(self) -> bool: + return True + + def _dequantize(self, model): + from ..integrations import dequantize_and_replace + + model = dequantize_and_replace( + model, self.modules_to_not_convert, quantization_config=self.quantization_config + ) + return model diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_bnb_8bit.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_bnb_8bit.py new file mode 100644 index 0000000000000000000000000000000000000000..cac339b16b9a099f52cc8715eb0b0ecddcff54fd --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_bnb_8bit.py @@ -0,0 +1,303 @@ +# Copyright 2024 The HuggingFace Inc. 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. +import importlib +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from packaging import version + +from .base import HfQuantizer + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..utils import ( + ACCELERATE_MIN_VERSION, + is_accelerate_available, + is_bitsandbytes_available, + is_torch_available, + is_torch_xpu_available, + logging, +) +from .quantizers_utils import get_module_from_name + + +if is_torch_available(): + import torch + + from ..pytorch_utils import Conv1D + +logger = logging.get_logger(__name__) + + +class Bnb8BitHfQuantizer(HfQuantizer): + """ + 8-bit quantization from bitsandbytes quantization method: + before loading: converts transformer layers into Linear8bitLt during loading: load 16bit weight and pass to the + layer object after: quantizes individual weights in Linear8bitLt into 8bit at fitst .cuda() call + saving: + from state dict, as usual; saves weights and 'SCB' component + loading: + need to locate SCB component and pass to the Linear8bitLt object + """ + + use_keep_in_fp32_modules = True + requires_parameters_quantization = True + requires_calibration = False + + required_packages = ["bitsandbytes", "accelerate"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + + if self.quantization_config.llm_int8_skip_modules is not None: + self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules + + def validate_environment(self, *args, **kwargs): + if not is_accelerate_available(): + raise ImportError( + f"Using `bitsandbytes` 8-bit quantization requires Accelerate: `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`" + ) + if not is_bitsandbytes_available(): + raise ImportError( + "Using `bitsandbytes` 8-bit quantization requires the latest version of bitsandbytes: `pip install -U bitsandbytes`" + ) + + from ..integrations import validate_bnb_backend_availability + from ..utils import is_bitsandbytes_multi_backend_available + + bnb_multibackend_is_enabled = is_bitsandbytes_multi_backend_available() + validate_bnb_backend_availability(raise_exception=True) + + if kwargs.get("from_tf", False) or kwargs.get("from_flax", False): + raise ValueError( + "Converting into 4-bit or 8-bit weights from tf/flax weights is currently not supported, please make" + " sure the weights are in PyTorch format." + ) + + device_map = kwargs.get("device_map", None) + if ( + device_map is not None + and isinstance(device_map, dict) + and not self.quantization_config.llm_int8_enable_fp32_cpu_offload + ): + device_map_without_lm_head = { + key: device_map[key] for key in device_map.keys() if key not in self.modules_to_not_convert + } + if set(device_map.values()) == {"cpu"} and bnb_multibackend_is_enabled: + pass + elif "cpu" in device_map_without_lm_head.values() or "disk" in device_map_without_lm_head.values(): + raise ValueError( + "Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the " + "quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules " + "in 32-bit, you need to set `llm_int8_enable_fp32_cpu_offload=True` and pass a custom `device_map` to " + "`from_pretrained`. Check " + "https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu " + "for more details. " + ) + + if version.parse(importlib.metadata.version("bitsandbytes")) < version.parse("0.37.2"): + raise ValueError( + "You have a version of `bitsandbytes` that is not compatible with 8bit inference and training" + " make sure you have the latest version of `bitsandbytes` installed" + ) + + def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + # need more space for buffers that are created during quantization + max_memory = {key: val * 0.90 for key, val in max_memory.items()} + return max_memory + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + # We force the `dtype` to be float16, this is a requirement from `bitsandbytes` + logger.info( + "Overriding torch_dtype=%s with `torch_dtype=torch.float16` due to " + "requirements of `bitsandbytes` to enable model loading in 8-bit or 4-bit. " + "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass" + " torch_dtype=torch.float16 to remove this warning.", + torch_dtype, + ) + torch_dtype = torch.float16 + return torch_dtype + + def update_device_map(self, device_map): + if device_map is None: + if torch.cuda.is_available(): + device_map = {"": torch.cuda.current_device()} + elif is_torch_xpu_available(): + device_map = {"": f"xpu:{torch.xpu.current_device()}"} + else: + device_map = {"": "cpu"} + logger.info( + "The device_map was not initialized. " + f"Setting device_map to {device_map}. " + "If you want to use the model for inference, please set device_map ='auto' " + ) + return device_map + + def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": + if target_dtype != torch.int8: + logger.info("target_dtype {target_dtype} is replaced by `torch.int8` for 8-bit BnB quantization") + return torch.int8 + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ): + import bitsandbytes as bnb + + module, tensor_name = get_module_from_name(model, param_name) + if isinstance(module._parameters.get(tensor_name, None), bnb.nn.Int8Params): + if self.pre_quantized: + if param_name.replace("weight", "SCB") not in state_dict.keys(): + raise ValueError("Missing quantization component `SCB`") + if param_value.dtype != torch.int8: + raise ValueError( + f"Incompatible dtype `{param_value.dtype}` when loading 8-bit prequantized weight. Expected `torch.int8`." + ) + return True + return False + + def create_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + state_dict: Dict[str, Any], + unexpected_keys: Optional[List[str]] = None, + ): + """ + combines logic from _load_state_dict_into_meta_model and .integrations.bitsandbytes.py::set_module_quantized_tensor_to_device() + needs aux items from state dicts, if found - removes them from unexpected_keys + """ + import bitsandbytes as bnb + + fp16_statistics_key = param_name.replace("weight", "SCB") + fp16_weights_format_key = param_name.replace("weight", "weight_format") + + fp16_statistics = state_dict.get(fp16_statistics_key, None) + fp16_weights_format = state_dict.get(fp16_weights_format_key, None) + + module, tensor_name = get_module_from_name(model, param_name) + if tensor_name not in module._parameters: + raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.") + + old_value = getattr(module, tensor_name) + + if not isinstance(module._parameters[tensor_name], bnb.nn.Int8Params): + raise ValueError(f"Parameter `{tensor_name}` should only be a `bnb.nn.Int8Params` instance.") + if ( + old_value.device == torch.device("meta") + and target_device not in ["meta", torch.device("meta")] + and param_value is None + ): + raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {target_device}.") + + new_value = param_value.to("cpu") + if self.pre_quantized and not self.is_serializable(): + raise ValueError( + "Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. " + "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`." + ) + + # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization. + # Since weights are saved in the correct "orientation", we skip transposing when loading. + if issubclass(module.source_cls, Conv1D): + if fp16_statistics is None: + new_value = new_value.T + + kwargs = old_value.__dict__ + new_value = bnb.nn.Int8Params(new_value, requires_grad=False, **kwargs).to(target_device) + + module._parameters[tensor_name] = new_value + if fp16_statistics is not None: + setattr(module.weight, "SCB", fp16_statistics.to(target_device)) + if unexpected_keys is not None: + unexpected_keys.remove(fp16_statistics_key) + + # We just need to pop the `weight_format` keys from the state dict to remove unneeded + # messages. The correct format is correctly retrieved during the first forward pass. + if fp16_weights_format is not None and unexpected_keys is not None: + unexpected_keys.remove(fp16_weights_format_key) + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + model.is_loaded_in_8bit = True + model.is_8bit_serializable = self.is_serializable() + return model + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + device_map, + keep_in_fp32_modules: Optional[List[str]] = None, + **kwargs, + ): + from ..integrations import replace_with_bnb_linear + + llm_int8_enable_fp32_cpu_offload = self.quantization_config.llm_int8_enable_fp32_cpu_offload + + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules + ) + + # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk` + if isinstance(device_map, dict) and len(device_map.keys()) > 1: + keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]] + + if len(keys_on_cpu) > 0 and not llm_int8_enable_fp32_cpu_offload: + raise ValueError( + "If you want to offload some keys to `cpu` or `disk`, you need to set " + "`llm_int8_enable_fp32_cpu_offload=True`. Note that these modules will not be " + " converted to 8-bit but kept in 32-bit." + ) + self.modules_to_not_convert.extend(keys_on_cpu) + + model = replace_with_bnb_linear( + model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config + ) + # TODO: consider bringing replace_with_bnb_linear() code from ..integrations/bitsandbyter.py to here + + model.config.quantization_config = self.quantization_config + + def is_serializable(self, safe_serialization=None): + _bnb_supports_8bit_serialization = version.parse(importlib.metadata.version("bitsandbytes")) > version.parse( + "0.37.2" + ) + + if not _bnb_supports_8bit_serialization: + logger.warning( + "You are calling `save_pretrained` to a 8-bit converted model, but your `bitsandbytes` version doesn't support it. " + "If you want to save 8-bit models, make sure to have `bitsandbytes>0.37.2` installed. You will most likely face errors or" + " unexpected behaviours." + ) + return False + + return True + + @property + def is_trainable(self) -> bool: + return version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse("0.37.0") + + def _dequantize(self, model): + from ..integrations import dequantize_and_replace + + model = dequantize_and_replace( + model, self.modules_to_not_convert, quantization_config=self.quantization_config + ) + return model diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_compressed_tensors.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_compressed_tensors.py new file mode 100644 index 0000000000000000000000000000000000000000..4e45abf95340b51138983a3a7132ab2f98b295de --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_compressed_tensors.py @@ -0,0 +1,173 @@ +# Copyright 2024 The HuggingFace Inc. 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. + + +import os +import re +from typing import List + +from ..utils import is_compressed_tensors_available, is_torch_available, logging +from ..utils.quantization_config import CompressedTensorsConfig +from .base import HfQuantizer + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class CompressedTensorsHfQuantizer(HfQuantizer): + """ + Quantizer for the compressed_tensors package. Loads and restores models to + quantized state with compressed_tensors + """ + + requires_calibration = True + required_packages = ["compressed_tensors"] + + def __init__(self, quantization_config: CompressedTensorsConfig, **kwargs): + super().__init__(quantization_config, **kwargs) + + if not is_compressed_tensors_available(): + raise ImportError( + "Using `compressed_tensors` quantized models requires the compressed-tensors library: " + "`pip install compressed-tensors`" + ) + + # Call post_init here to ensure proper config setup when `run_compressed` + # is provided directly via CompressedTensorsConfig, and to avoid duplicate logging. + + quantization_config.post_init() + from compressed_tensors.compressors import ModelCompressor + + self.compressor = ModelCompressor.from_compression_config(quantization_config) + self.run_compressed = quantization_config.run_compressed + self.quantization_config = quantization_config + + def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: + """ + Update missing keys after loading the model. This is necessary for compressed tensors + to load the model correctly. We expect weights to be present in missing keys. + The weight's are re-constructed by ModelCompressor in _process_model_after_weight_loading + + This function cleans up expected missing keys and returns the remaining missing keys + """ + + if self.run_compressed: + return missing_keys + + # We expect some keys to be missing for + # compresed models + # This is fine as the weights are reconstructed by ModelCompressor + # in _process_model_after_weight_loading + + expected_missing_keys = self.compressor.get_missing_module_keys(model) + return [ + key for key in missing_keys if not any(re.match(f".*{pattern}", key) for pattern in expected_missing_keys) + ] + + def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str) -> List[str]: + """ + Override this method if you want to adjust the `unexpected_keys`. + + Args: + unexpected_keys (`List[str]`, *optional*): + The list of unexpected keys in the checkpoint compared to the state dict of the model + """ + + if self.run_compressed: + return unexpected_keys + + # We expect some unexpected keys in model + # safetensors file for compressed models + keys_to_ignore = self.compressor.get_unexpected_file_keys(model) + return [key for key in unexpected_keys if not any(re.match(f".*{pattern}", key) for pattern in keys_to_ignore)] + + def validate_environment(self, *args, **kwargs): + if not is_compressed_tensors_available(): + raise ImportError( + "Using `compressed_tensors` quantized models requires the compressed-tensors library: " + "`pip install compressed-tensors`" + ) + if not is_torch_available(): + # torch already should be installed as part of compressed tensors + raise ImportError("torch is required for using compressed-tensors quantization") + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + logger.info("Loading model using torch.float16 for compressed-tensors quantization") + torch_dtype = torch.float16 + elif torch_dtype != torch.float16: + logger.info( + "We suggest you to set `torch_dtype=torch.float16` for better efficiency with compressed_tensors." + ) + return torch_dtype + + def _process_model_before_weight_loading(self, model, **kwargs): + from compressed_tensors.quantization import apply_quantization_config + + ct_quantization_config = self.compressor.quantization_config + + if self.run_compressed: + apply_quantization_config(model, ct_quantization_config, run_compressed=True) + elif not self.quantization_config.is_quantization_compressed: + apply_quantization_config(model, ct_quantization_config) + + def _process_model_after_weight_loading(self, model, **kwargs): + """Decompress loaded model if necessary - need for qat""" + + if ( + self.quantization_config.is_quantization_compressed and not self.run_compressed + ) or self.quantization_config.is_sparsification_compressed: + config = kwargs.get("config", None) + cache_path = config._name_or_path + + if not os.path.exists(cache_path): + from transformers.utils import cached_file + + config_file_path = cached_file(cache_path, "config.json") + cache_path = os.path.sep.join(config_file_path.split(os.path.sep)[:-1]) + + if self.quantization_config.is_quantization_compressed and not self.run_compressed: + from compressed_tensors.quantization import QuantizationStatus + + self.compressor.quantization_config.quantization_status = QuantizationStatus.FROZEN + self.compressor.decompress(model_path=cache_path, model=model) + + def update_tp_plan(self, config): + additional_plan = { + "layers.*.feed_forward.experts.*.gate_proj.weight": "local_colwise", + "layers.*.feed_forward.experts.*.gate_proj.weight_scale": "local_colwise", + "layers.*.feed_forward.experts.*.up_proj.weight": "local_colwise", + "layers.*.feed_forward.experts.*.up_proj.weight_scale": "local_colwise", + "layers.*.feed_forward.experts.*.down_proj.weight": "local_rowwise", + } + if config.get_text_config() is not None and config.get_text_config().base_model_tp_plan is not None: + config.get_text_config().base_model_tp_plan.update(additional_plan) + + return config + + @property + def is_trainable(self): + return True + + def is_qat_trainable(self) -> bool: + """Loaded Models can carry out quantization aware training""" + # models need to be decompressed carry out qat + return not self.run_compressed or not self.quantization_config.is_quantization_compressed + + def is_serializable(self, safe_serialization=None) -> bool: + """Models quantized using compressed tensors can be saved to disk""" + return True diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_eetq.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_eetq.py new file mode 100644 index 0000000000000000000000000000000000000000..988f90789ac3751ca27a9ce4d1ed71762fd7c434 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_eetq.py @@ -0,0 +1,181 @@ +# Copyright 2024 The HuggingFace Inc. 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. +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +from .base import HfQuantizer + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..utils import is_accelerate_available, is_eetq_available, is_torch_available, logging +from .quantizers_utils import get_module_from_name + + +if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) + + +class EetqHfQuantizer(HfQuantizer): + """ + 8-bit quantization from EETQ quantization method: + before loading: converts transformer layers into W8A16Linear during loading: load 16bit weight and pass to the + layer object after: quantizes individual weights in Linear8bitLt into 8bit at first .cuda() call + """ + + requires_parameters_quantization = True + requires_calibration = False + + required_packages = ["eetq", "accelerate"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + self.quantization_config = quantization_config + + def validate_environment(self, *args, **kwargs): + if not is_eetq_available(): + raise ImportError( + "Using `eetq` 8-bit quantization requires eetq." + "Please install the latest version of eetq from : https://github.com/NetEase-FuXi/EETQ" + ) + + try: + import eetq # noqa: F401 + except ImportError as exc: + if "shard_checkpoint" in str(exc): + # EETQ 1.0.0 is currently broken with the latest transformers because it tries to import the removed + # shard_checkpoint function, see https://github.com/NetEase-FuXi/EETQ/issues/34. + # TODO: Update message once eetq releases a fix + raise ImportError( + "You are using a version of EETQ that is incompatible with the current transformers version. " + "Either downgrade transformers to <= v4.46.3 or, if available, upgrade EETQ to > v1.0.0." + ) from exc + else: + raise + + if not is_accelerate_available(): + raise ImportError("Loading an EETQ quantized model requires accelerate (`pip install accelerate`)") + + if kwargs.get("from_tf", False) or kwargs.get("from_flax", False): + raise ValueError( + "Converting into 8-bit weights from tf/flax weights is currently not supported, please make" + " sure the weights are in PyTorch format." + ) + + if not torch.cuda.is_available(): + raise RuntimeError("No GPU found. A GPU is needed for quantization.") + + device_map = kwargs.get("device_map", None) + if device_map is None: + logger.warning_once( + "You have loaded an EETQ model on CPU and have a CUDA device available, make sure to set " + "your model on a GPU device in order to run your model." + ) + elif device_map is not None: + if isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()): + raise ValueError( + "You are attempting to load an EETQ model with a device_map that contains a CPU or disk device." + " This is not supported. Please remove the CPU or disk device from the device_map." + ) + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + torch_dtype = torch.float16 + logger.info( + "Overriding torch_dtype=%s with `torch_dtype=torch.float16` due to " + "requirements of `eetq` to enable model loading in 8-bit. " + "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass" + " torch_dtype=torch.float16 to remove this warning.", + torch_dtype, + ) + elif torch_dtype != torch.float16: + logger.info("We suggest you to set `torch_dtype=torch.float16` for better efficiency with EETQ.") + return torch_dtype + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ): + from eetq import EetqLinear + + module, tensor_name = get_module_from_name(model, param_name) + + if isinstance(module, EetqLinear): + if self.pre_quantized or tensor_name == "bias": + if tensor_name == "weight" and param_value.dtype != torch.int8: + raise ValueError("Expect quantized weights but got an unquantized weight") + return False + else: + if tensor_name == "weight_scale": + raise ValueError("Expect unquantized weights but got a quantized weight_scale") + return True + return False + + def create_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + state_dict: Dict[str, Any], + unexpected_keys: Optional[List[str]] = None, + ): + """ + quantizes weights into qweight and weight_scales + """ + from eetq import quantize_and_preprocess_weights + + module, tensor_name = get_module_from_name(model, param_name) + new_value, weight_scale = quantize_and_preprocess_weights(param_value) + + module._buffers[tensor_name] = new_value.to(target_device) + module.register("weight_scales", weight_scale.to(target_device)) + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + return model + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + keep_in_fp32_modules: Optional[List[str]] = None, + **kwargs, + ): + from ..integrations import replace_with_eetq_linear + + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules + ) + + model = replace_with_eetq_linear( + model, + modules_to_not_convert=self.modules_to_not_convert, + quantization_config=self.quantization_config, + pre_quantized=self.pre_quantized, + ) + + model.config.quantization_config = self.quantization_config + + def is_serializable(self, safe_serialization=None): + return True + + @property + def is_trainable(self) -> bool: + return True diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_fbgemm_fp8.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_fbgemm_fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..f0fa8f063bb452c2c49fc042a27b507e89ba5d9a --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_fbgemm_fp8.py @@ -0,0 +1,252 @@ +# Copyright 2024 The HuggingFace Inc. 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. +import importlib +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +from packaging import version + +from .base import HfQuantizer + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..utils import is_accelerate_available, is_fbgemm_gpu_available, is_torch_available, logging +from .quantizers_utils import get_module_from_name + + +if is_torch_available(): + import torch + + +logger = logging.get_logger(__name__) + + +class FbgemmFp8HfQuantizer(HfQuantizer): + """ + FP8 quantization using fbgemm kernels + """ + + requires_parameters_quantization = True + requires_calibration = False + + required_packages = ["fbgemm-gpu", "accelerate"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + self.quantization_config = quantization_config + + def validate_environment(self, *args, **kwargs): + if not is_torch_available() or version.parse(importlib.metadata.version("torch")) < version.parse("2.1.0"): + raise ImportError( + "Using fbgemm fp8 quantization requires torch > 2.1.0" + "Please install the latest version of torch ( pip install --upgrade torch )" + ) + if not is_fbgemm_gpu_available(): + raise ImportError( + "Using fbgemm fp8 quantization requires fbgemm-gpu library" + "Please install the latest version of fbgemm-gpu library by following : https://pytorch.org/FBGEMM/fbgemm_gpu-development/InstallationInstructions.html#fbgemm-gpu-install-libraries" + ) + + if not is_accelerate_available("0.32.2"): + raise ImportError( + "Loading an FP8 quantized model requires accelerate > 0.32.1 (`pip install --upgrade accelerate`)" + ) + + if not torch.cuda.is_available(): + raise RuntimeError("Using FP8 quantized models with fbgemm kernels requires a GPU") + + compute_capability = torch.cuda.get_device_capability() + major, minor = compute_capability + if major < 9: + raise ValueError( + "FP8 quantized models is only supported on GPUs with compute capability >= 9.0 (e.g H100)" + ) + + device_map = kwargs.get("device_map", None) + if device_map is None: + logger.warning_once( + "You have loaded an FP8 model on CPU and have a CUDA device available, make sure to set " + "your model on a GPU device in order to run your model. To remove this warning, pass device_map = 'cuda'. " + ) + elif device_map is not None: + if ( + not self.pre_quantized + and isinstance(device_map, dict) + and ("cpu" in device_map.values() or "disk" in device_map.values()) + ): + raise ValueError( + "You are attempting to load an FP8 model with a device_map that contains a CPU or disk device." + "This is not supported when the model is quantized on the fly. " + "Please use a quantized checkpoint or remove the CPU or disk device from the device_map." + ) + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + torch_dtype = torch.bfloat16 + logger.info( + "Overriding torch_dtype=%s with `torch_dtype=torch.bloat16` due to " + "requirements of `fbgemm-gpu` to enable model loading in fp8. " + "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass" + " torch_dtype=torch.bfloat16 to remove this warning.", + torch_dtype, + ) + elif torch_dtype == torch.float16: + raise ValueError( + "You cannot use FP8 with torch_dtype=torch.float16.We recommend you passing torch_dtype=torch.bfloat16" + ) + return torch_dtype + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ): + from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts + + module, tensor_name = get_module_from_name(model, param_name) + + if isinstance(module, FbgemmFp8Linear): + if self.pre_quantized or tensor_name == "bias": + if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn: + raise ValueError("Expect quantized weights but got an unquantized weight") + return False + else: + if tensor_name == "weight_scale": + raise ValueError("Expect unquantized weights but got a quantized weight_scale") + return True + if isinstance(module, FbgemmFp8Llama4TextExperts): + if self.pre_quantized or tensor_name == "bias": + return False + else: + if tensor_name == "gate_up_proj_scale" or tensor_name == "down_proj_scale": + raise ValueError("Expect unquantized weights but got a quantized weight_scale") + return True + return False + + def create_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + state_dict: Dict[str, Any], + unexpected_keys: Optional[List[str]] = None, + ): + """ + Quantizes weights into weight and weight_scale + """ + + from ..integrations import FbgemmFp8Llama4TextExperts + + module, tensor_name = get_module_from_name(model, param_name) + if isinstance(module, FbgemmFp8Llama4TextExperts): + if tensor_name == "gate_up_proj": + # Process each expert separately + # Transpose the second and third dimension + transposed_param = param_value.transpose(1, 2) + + # Reshape to 2D for quantization + original_shape = transposed_param.shape + flattened_param = transposed_param.reshape(-1, original_shape[-1]) + + # Quantize using per row instead of per column + new_value_flat, weight_scale_flat = torch.ops.fbgemm.quantize_fp8_per_row(flattened_param) + + # Reshape back to original dimensions + new_value = new_value_flat.reshape(original_shape) + new_value = new_value.transpose(1, 2) + weight_scale = weight_scale_flat.reshape(original_shape[0], 1, original_shape[1]) + elif tensor_name == "down_proj": + # Process each expert separately + # Transpose the weights for proper quantization + transposed_param = param_value.transpose(1, 2) + + # Reshape to 2D for quantization + original_shape = transposed_param.shape + flattened_param = transposed_param.reshape(-1, original_shape[-1]) + + # Quantize using per column + new_value_flat, weight_scale_flat = torch.ops.fbgemm.quantize_fp8_per_row(flattened_param) + + # Reshape back to original dimensions + new_value = new_value_flat.reshape(original_shape) + new_value = new_value.transpose(1, 2) + weight_scale = weight_scale_flat.reshape(original_shape[0], original_shape[1], 1) + + module._parameters[f"{tensor_name}_scale"] = torch.nn.Parameter(weight_scale.to(target_device)) + else: + new_value, weight_scale = torch.ops.fbgemm.quantize_fp8_per_row(param_value) + module._parameters[f"{tensor_name}_scale"] = torch.nn.Parameter( + weight_scale.view(weight_scale.shape[0], 1).to(target_device) + ) + + module._parameters[tensor_name] = torch.nn.Parameter(new_value.to(target_device)) + + if unexpected_keys is not None and param_name in unexpected_keys: + unexpected_keys.remove(param_name) + del param_name + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + return model + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + keep_in_fp32_modules: Optional[List[str]] = None, + **kwargs, + ): + from ..integrations import replace_with_fbgemm_fp8_linear + + tp_plan = model._tp_plan + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules + ) + + config = model.config + model = replace_with_fbgemm_fp8_linear( + model, + modules_to_not_convert=self.modules_to_not_convert, + quantization_config=self.quantization_config, + pre_quantized=self.pre_quantized, + config=config, + tp_plan=tp_plan, + ) + + model.config.quantization_config = self.quantization_config + + def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]: + from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts + + not_missing_keys = [] + for name, module in model.named_modules(): + if isinstance(module, FbgemmFp8Linear) or isinstance(module, FbgemmFp8Llama4TextExperts): + for missing in missing_keys: + if ( + (name in missing or name in f"{prefix}.{missing}") + and not missing.endswith(".weight") + and not missing.endswith(".bias") + ): + not_missing_keys.append(missing) + return [k for k in missing_keys if k not in not_missing_keys] + + def is_serializable(self, safe_serialization=None): + return True + + @property + def is_trainable(self) -> bool: + return False diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_finegrained_fp8.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_finegrained_fp8.py new file mode 100644 index 0000000000000000000000000000000000000000..16ce7f6a9eb633959afbf2bcac0a4c18378304f4 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_finegrained_fp8.py @@ -0,0 +1,206 @@ +import importlib +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +from packaging import version + +from ..utils import is_accelerate_available, is_torch_available, logging +from .base import HfQuantizer +from .quantizers_utils import get_module_from_name + + +if is_torch_available(): + import torch + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +logger = logging.get_logger(__name__) + + +class FineGrainedFP8HfQuantizer(HfQuantizer): + """ + FP8 quantization implementation supporting both standard and MoE models. + Supports both e4m3fn formats based on platform. + """ + + requires_parameters_quantization = True + requires_calibration = False + required_packages = ["accelerate"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + self.quantization_config = quantization_config + + def validate_environment(self, *args, **kwargs): + if not is_torch_available() or version.parse(importlib.metadata.version("torch")) < version.parse("2.1.0"): + raise ImportError( + "Using fp8 quantization requires torch >= 2.1.0" + "Please install the latest version of torch ( pip install --upgrade torch )" + ) + + if not is_accelerate_available(): + raise ImportError("Loading an FP8 quantized model requires accelerate (`pip install accelerate`)") + + if kwargs.get("from_tf", False) or kwargs.get("from_flax", False): + raise ValueError( + "Converting into FP8 weights from tf/flax weights is currently not supported, " + "please make sure the weights are in PyTorch format." + ) + + if not torch.cuda.is_available(): + raise RuntimeError("No GPU found. A GPU is needed for FP8 quantization.") + + compute_capability = torch.cuda.get_device_capability() + major, minor = compute_capability + if (major < 8) or (major == 8 and minor < 9): + raise ValueError( + "FP8 quantized models is only supported on GPUs with compute capability >= 8.9 (e.g 4090/H100)" + f", actual = `{major}.{minor}`" + ) + + device_map = kwargs.get("device_map", None) + if device_map is None: + logger.warning_once( + "You have loaded an FP8 model on CPU and have a CUDA device available, make sure to set " + "your model on a GPU device in order to run your model. To remove this warning, pass device_map = 'cuda'. " + ) + elif device_map is not None: + if ( + not self.pre_quantized + and isinstance(device_map, dict) + and ("cpu" in device_map.values() or "disk" in device_map.values()) + ): + raise ValueError( + "You are attempting to load an FP8 model with a device_map that contains a cpu/disk device." + "This is not supported when the model is quantized on the fly. " + "Please use a quantized checkpoint or remove the cpu/disk device from the device_map." + ) + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + logger.info("Setting torch_dtype to torch.float32 as no torch_dtype was specified in from_pretrained") + torch_dtype = torch.float32 + return torch_dtype + + def create_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + state_dict: Dict[str, Any], + unexpected_keys: Optional[List[str]] = None, + ): + """ + Quantizes weights to FP8 format using Block-wise quantization + """ + from accelerate.utils import set_module_tensor_to_device + + set_module_tensor_to_device(model, param_name, target_device, param_value) + + module, tensor_name = get_module_from_name(model, param_name) + + # Get FP8 min/max values + fp8_min = torch.finfo(torch.float8_e4m3fn).min + fp8_max = torch.finfo(torch.float8_e4m3fn).max + + block_size_m, block_size_n = self.quantization_config.weight_block_size + + rows, cols = param_value.shape[-2:] + + if rows % block_size_m != 0 or cols % block_size_n != 0: + raise ValueError( + f"Matrix dimensions ({rows}, {cols}) must be divisible by block sizes ({block_size_m}, {block_size_n})" + ) + param_value_orig_shape = param_value.shape + + param_value = param_value.reshape( + -1, rows // block_size_m, block_size_m, cols // block_size_n, block_size_n + ).permute(0, 1, 3, 2, 4) + + # Calculate scaling factor for each block + max_abs = torch.amax(torch.abs(param_value), dim=(-1, -2)) + scale = fp8_max / max_abs + scale_orig_shape = scale.shape + scale = scale.unsqueeze(-1).unsqueeze(-1) + + # Quantize the weights + quantized_param = torch.clamp(param_value * scale, min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) + + quantized_param = quantized_param.permute(0, 1, 3, 2, 4) + # Reshape back to matrix shape + quantized_param = quantized_param.reshape(param_value_orig_shape) + + # Reshape scale to match the number of blocks + scale = scale.reshape(scale_orig_shape).squeeze().reciprocal() + + module._buffers[tensor_name] = quantized_param.to(target_device) + module._buffers["weight_scale_inv"] = scale.to(target_device) + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ): + from ..integrations.finegrained_fp8 import FP8Linear + + module, tensor_name = get_module_from_name(model, param_name) + + if isinstance(module, FP8Linear): + if self.pre_quantized or tensor_name == "bias": + if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn: + raise ValueError("Expect quantized weights but got an unquantized weight") + return False + else: + if tensor_name == "weight_scale_inv": + raise ValueError("Expect unquantized weights but got a quantized weight_scale") + return True + return False + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + keep_in_fp32_modules: Optional[List[str]] = None, + **kwargs, + ): + from ..integrations.finegrained_fp8 import replace_with_fp8_linear + + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules + ) + + model = replace_with_fp8_linear( + model, + modules_to_not_convert=self.modules_to_not_convert, + quantization_config=self.quantization_config, + ) + + model.config.quantization_config = self.quantization_config + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + return model + + def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]: + from ..integrations import FP8Linear + + not_missing_keys = [] + for name, module in model.named_modules(): + if isinstance(module, FP8Linear): + for missing in missing_keys: + if ( + (name in missing or name in f"{prefix}.{missing}") + and not missing.endswith(".weight") + and not missing.endswith(".bias") + ): + not_missing_keys.append(missing) + return [k for k in missing_keys if k not in not_missing_keys] + + def is_serializable(self, safe_serialization=None): + return True + + @property + def is_trainable(self) -> bool: + return False diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_gptq.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_gptq.py new file mode 100644 index 0000000000000000000000000000000000000000..0fc35bea1635edd957566f53005ff65a58d57282 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_gptq.py @@ -0,0 +1,124 @@ +# Copyright 2024 The HuggingFace Inc. 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. +import importlib +from typing import TYPE_CHECKING, Optional + +from packaging import version + +from .base import HfQuantizer + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..utils import is_auto_gptq_available, is_gptqmodel_available, is_optimum_available, is_torch_available, logging +from ..utils.quantization_config import GPTQConfig, QuantizationConfigMixin + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class GptqHfQuantizer(HfQuantizer): + """ + Quantizer of the GPTQ method - for GPTQ the quantizer support calibration of the model through + `auto_gptq` or `gptqmodel` package. Quantization is done under the hood for users if they load a non-prequantized model. + """ + + requires_calibration = False + required_packages = ["optimum", "auto_gptq", "gptqmodel"] + optimum_quantizer = None + + def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): + super().__init__(quantization_config, **kwargs) + + if not is_optimum_available(): + raise ImportError("Loading a GPTQ quantized model requires optimum (`pip install optimum`)") + from optimum.gptq import GPTQQuantizer + + self.optimum_quantizer = GPTQQuantizer.from_dict(self.quantization_config.to_dict_optimum()) + + def validate_environment(self, *args, **kwargs): + if not is_optimum_available(): + raise ImportError("Loading a GPTQ quantized model requires optimum (`pip install optimum`)") + if is_auto_gptq_available() and is_gptqmodel_available(): + logger.warning("Detected gptqmodel and auto-gptq, will use gptqmodel") + + gptq_supports_cpu = ( + is_auto_gptq_available() + and version.parse(importlib.metadata.version("auto-gptq")) > version.parse("0.4.2") + ) or is_gptqmodel_available() + if not gptq_supports_cpu and not torch.cuda.is_available(): + raise RuntimeError("GPU is required to quantize or run quantize model.") + elif not (is_auto_gptq_available() or is_gptqmodel_available()): + raise ImportError( + "Loading a GPTQ quantized model requires gptqmodel (`pip install gptqmodel`) or auto-gptq (`pip install auto-gptq`) library. " + ) + elif is_auto_gptq_available() and version.parse(importlib.metadata.version("auto_gptq")) < version.parse( + "0.4.2" + ): + raise ImportError( + "You need a version of auto_gptq >= 0.4.2 to use GPTQ: `pip install --upgrade auto-gptq` or use gptqmodel by `pip install gptqmodel>=1.4.3`." + ) + elif is_gptqmodel_available() and ( + version.parse(importlib.metadata.version("gptqmodel")) < version.parse("1.4.3") + or version.parse(importlib.metadata.version("optimum")) < version.parse("1.23.99") + ): + raise ImportError("The gptqmodel version should be >= 1.4.3, optimum version should >= 1.24.0") + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + torch_dtype = torch.float16 + logger.info("Loading the model in `torch.float16`. To overwrite it, set `torch_dtype` manually.") + elif torch_dtype != torch.float16: + logger.info("We suggest you to set `torch_dtype=torch.float16` for better efficiency with GPTQ.") + return torch_dtype + + def update_device_map(self, device_map): + if device_map is None: + device_map = {"": torch.device("cpu")} + # Only with auto-gptq do not support CPU, we should move the model to cuda if available. + if not is_gptqmodel_available() and device_map in ("cpu", {"": torch.device("cpu")}): + device_map == {"": 0} + return device_map + + def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs): + if model.__class__.main_input_name != "input_ids": + raise RuntimeError("We can only quantize pure text model.") + + if self.pre_quantized: + # compat: latest optimum has gptqmodel refactor + if version.parse(importlib.metadata.version("optimum")) <= version.parse("1.23.99"): + model = self.optimum_quantizer.convert_model(model) + else: + model = self.optimum_quantizer.convert_model(model, **kwargs) + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + if self.pre_quantized: + model = self.optimum_quantizer.post_init_model(model) + else: + if self.quantization_config.tokenizer is None: + self.quantization_config.tokenizer = model.name_or_path + + self.optimum_quantizer.quantize_model(model, self.quantization_config.tokenizer) + model.config.quantization_config = GPTQConfig.from_dict(self.optimum_quantizer.to_dict()) + + @property + def is_trainable(self, model: Optional["PreTrainedModel"] = None): + return True + + def is_serializable(self, safe_serialization=None): + return True diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_higgs.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_higgs.py new file mode 100644 index 0000000000000000000000000000000000000000..83c102f16c3e303514d856508e4dc37d87163ffb --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_higgs.py @@ -0,0 +1,202 @@ +# Copyright 2024 The HuggingFace Inc. 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. +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +from ..utils.logging import tqdm +from .base import HfQuantizer +from .quantizers_utils import get_module_from_name + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..utils import is_accelerate_available, is_flute_available, is_hadamard_available, is_torch_available, logging +from ..utils.quantization_config import QuantizationConfigMixin + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class HiggsHfQuantizer(HfQuantizer): + """ + Quantizer of the HIGGS method. Enables the loading of prequantized models and in-flight quantization of full-precision models. + """ + + requires_calibration = False + requires_parameters_quantization = True + required_packages = ["flute-kernel", "fast_hadamard_transform"] + + def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): + super().__init__(quantization_config, **kwargs) + self.quantization_config = quantization_config + + def validate_environment(self, device_map, **kwargs): + if not torch.cuda.is_available(): + raise NotImplementedError("HIGGS quantization is only supported on GPU. Please use a different quantizer.") + + if not is_accelerate_available(): + raise ImportError("Using `higgs` quantization requires Accelerate: `pip install accelerate`") + + if not is_flute_available(): + raise ImportError("Using `higgs` quantization requires FLUTE: `pip install flute-kernel>=0.3.0`") + + if not is_hadamard_available(): + raise ImportError( + "Using `higgs` quantization requires fast_hadamard_transform: `pip install fast_hadamard_transform`" + ) + + if device_map is None: + raise ValueError( + "You are attempting to load a HIGGS model without setting device_map." + " Please set device_map comprised of 'cuda' devices." + ) + elif isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()): + raise ValueError( + "You are attempting to load a HIGGS model with a device_map that contains a CPU or disk device." + " This is not supported. Please remove the CPU or disk device from the device_map." + ) + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + logger.info("`torch_dtype` is None. Setting `torch_dtype=torch.float16` for FLUTE compatibility.") + torch_dtype = torch.float16 + elif torch_dtype != torch.float16 and torch_dtype != torch.bfloat16: + raise ValueError( + f"Invalid `torch_dtype` {torch_dtype}. HIGGS quantization only supports `torch_dtype=torch.float16` or `torch_dtype=torch.bfloat16`." + ) + + return torch_dtype + + def create_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + state_dict: Dict[str, Any], + unexpected_keys: Optional[List[str]] = None, + ): + from ..integrations import quantize_with_higgs + + """ + Quantizes weights into weight and weight_scale + """ + flute_dict = quantize_with_higgs( + param_value.to(target_device), + self.quantization_config.bits, + self.quantization_config.p, + self.quantization_config.group_size, + self.quantization_config.hadamard_size, + ) + del param_value + + module, _ = get_module_from_name(model, param_name) + module_name = ".".join(param_name.split(".")[:-1]) + for key, value in flute_dict.items(): + if key in module._parameters: + module._parameters[key] = torch.nn.Parameter(value, requires_grad=False) + elif key in module._buffers: + module._buffers[key] = torch.nn.Buffer(value) + elif key == "tune_metadata": + module.tune_metadata = value + self.quantization_config.tune_metadata[module_name] = value.to_dict() + else: + raise ValueError(f"Unexpected key {key} in module {module}") + + if unexpected_keys is not None and param_name in unexpected_keys: + unexpected_keys.remove(param_name) + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + **kwargs, + ): + from ..integrations import replace_with_higgs_linear + + replace_with_higgs_linear( + model, + quantization_config=self.quantization_config, + ) + model.config.quantization_config = self.quantization_config + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + from flute.tune import TuneMetaData, maybe_tune_and_repack + from flute.utils import make_workspace_streamk + + from ..integrations import HiggsLinear + + flute_workspaces = {} + flute_modules = {name: module for name, module in model.named_modules() if isinstance(module, HiggsLinear)} + for name, module in tqdm(flute_modules.items(), desc="Repacking HIGGS modules", leave=False): + # Every HiggsLinear needs a "workspace": a buffer for the unpacking operation. + # This buffer needs to be on the same device as the weights, but can be reused across modules otherwise. + if module.weight.device not in flute_workspaces: + flute_workspaces[module.weight.device] = make_workspace_streamk(device=module.weight.device) + module.workspace = flute_workspaces[module.weight.device] + + # FLUTE weights are packed in a way that is optimized for a specific number of SMs (GPU streaming multiprocessors). + # If the model is loaded on a different device than the one it was saved on, we need to repack the weights. + module.tune_metadata = TuneMetaData.from_dict(self.quantization_config.tune_metadata[name]) + module.weight.data, module.tune_metadata = maybe_tune_and_repack( + weight=module.weight.data, + scales=module.scales.data, + metadata=module.tune_metadata, + ) + self.quantization_config.tune_metadata[name] = module.tune_metadata.to_dict() + + def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]: + from ..integrations import HiggsLinear + + higgs_names = {name for name, module in model.named_modules() if isinstance(module, HiggsLinear)} + + def should_update(key: str) -> bool: + if key.endswith(".weight") or key.endswith(".bias"): + return False + full_key = f"{prefix}.{key}" + return any(name in key or name in full_key for name in higgs_names) + + return [key for key in missing_keys if not should_update(key)] + + @property + def is_trainable(self, model: Optional["PreTrainedModel"] = None): + return False + + def is_serializable(self, safe_serialization=None): + return True + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ) -> bool: + from ..integrations import HiggsLinear + + module, tensor_name = get_module_from_name(model, param_name) + if isinstance(module, HiggsLinear) and tensor_name == "weight" and param_value.dtype != torch.int16: + # Only quantize weights of HiggsLinear modules that are not already quantized + return True + else: + return False + + def _dequantize(self, model): + from ..integrations import dequantize_higgs + + model = dequantize_higgs(model) + return model diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_hqq.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_hqq.py new file mode 100644 index 0000000000000000000000000000000000000000..4adc323f958ebfcd7de8fd5d62d8fdcb99606582 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_hqq.py @@ -0,0 +1,322 @@ +# 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. + +from typing import TYPE_CHECKING, Any, Dict, List + +from ..integrations import prepare_for_hqq_linear +from ..utils import is_accelerate_available, is_hqq_available, is_torch_available, logging +from .base import HfQuantizer +from .quantizers_utils import get_module_from_name + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + + +if is_accelerate_available(): + from accelerate.hooks import remove_hook_from_module + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +# Finds the parent of a node module named "name" +def find_parent(model, name): + module_tree = name.split(".")[:-1] + parent = model + for m in module_tree: + parent = parent._modules[m] + return parent + + +class HqqHfQuantizer(HfQuantizer): + """ + HQQ quantizer base HF class. + nn.Linear modules are first tagged with quant_config in _process_model_before_weight_loading(). + The actual quantization and offloading to the GPU is done in check_quantized_param(). + """ + + use_keep_in_fp32_modules = False + requires_parameters_quantization = True + requires_calibration = False + required_packages = ["hqq"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + self.torch_dtype = None + self.using_multi_gpu = False + + def validate_environment(self, *args, **kwargs): + if not (is_hqq_available()): + raise ImportError( + "A valid HQQ version (>=0.2.1) is not available. Please follow the instructions to install it: `https://github.com/mobiusml/hqq/`." + ) + + if kwargs.get("from_tf", False) or kwargs.get("from_flax", False): + raise ValueError( + "Converting weights from tf/flax weights is currently not supported, please make" + " sure the weights are in PyTorch format." + ) + + if not torch.cuda.is_available(): + raise RuntimeError("No GPU found. A GPU is needed for quantization.") + + if self.torch_dtype is None: + if "torch_dtype" in kwargs: + self.torch_dtype = kwargs["torch_dtype"] + else: + self.torch_dtype = torch.float32 + logger.info("Setting torch_dtype to torch.float32 as the default value since it was not specified.") + + device_map = kwargs.get("device_map", None) + if isinstance(device_map, dict): + if "cpu" in device_map.values() or "disk" in device_map.values(): + raise ValueError( + "You are attempting to use an HQQ model with a device_map that contains a CPU or disk device." + " This is not supported. Please remove the CPU or disk device from the device_map." + ) + else: + self.using_multi_gpu = len(set(device_map.values())) > 1 + + def update_missing_keys( + self, model: "PreTrainedModel", missing_keys: List[str], prefix: str, **kwargs + ) -> List[str]: + if self.pre_quantized: + return [key for key in missing_keys if ("weight" not in key)] + else: + return missing_keys + + # Adds missing keys for HQQLinear modules that are loaded but the model with initialized with torch.nn.Linear + def update_expected_keys( + self, model: "PreTrainedModel", expected_keys: List[str], loaded_keys: List[str] + ) -> List[str]: + if not self.pre_quantized: + return expected_keys + + # Collects all quantizable (linear) layers + def _find_hqq_quantizable_layers(model, layers): + for name, module in model.named_children(): + if isinstance(module, (torch.nn.Linear)): + layers.add(module.name) + _find_hqq_quantizable_layers(module, layers) + + new_keys = set(expected_keys) + if is_hqq_available(): + from hqq.core.quantize import HQQLinear + + # Name modules + for name, module in model.named_modules(): + module.name = name + + # valid modules are Linear layers that have HQQLinear state_dict. We ignore skip_modules and any layers with Linear state_dict() params + _valid_modules = set() + _find_hqq_quantizable_layers(model, _valid_modules) + + # Remove skipped modules + _skipped_modules = set() + for _module in _valid_modules: + for _skip_module in model.config.quantization_config["skip_modules"]: + if _skip_module in _module: + _skipped_modules.add(_module) + _valid_modules -= _skipped_modules + + # Append new expected layers based on _ref_keys + _ref_keys = HQQLinear( + linear_layer=None, quant_config=None, compute_dtype=torch.float16, device="cpu" + ).state_dict_keys() - {"bias"} + + # Clean-up + _rm_keys = set() + for key in new_keys: + if any(_module in key for _module in _valid_modules): + _rm_keys.add(key) + new_keys -= _rm_keys + # At this point, new_keys contains all the keys of the layers that are NOT HQQLinear or torch.nn.Linear + + # Re-populate Linear/HQQLinear + for _module in _valid_modules: + if _module + ".weight" in loaded_keys: + new_keys.add(_module + ".weight") + else: + new_keys.update({_module + "." + _ref_key for _ref_key in _ref_keys}) + if _module + ".bias" in loaded_keys: + new_keys.add(_module + ".bias") + + return list(new_keys) + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ) -> bool: + if is_hqq_available(): + from hqq.core.quantize import HQQLinear + module, tensor_name = get_module_from_name(model, param_name) + + if self.pre_quantized: + return ( + (isinstance(module, torch.nn.Linear) or isinstance(module, HQQLinear)) + and tensor_name != "weight" + and tensor_name != "bias" + ) + else: + # we need a special path for bias since hqq overwrote load_state_dict for this layer + return ( + isinstance(module, torch.nn.Linear) + and tensor_name == "weight" + or (isinstance(module, HQQLinear) and tensor_name == "bias") + ) + + def create_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + state_dict: Dict[str, Any], + unexpected_keys: List[str], + ): + """ + Each nn.Linear layer is processed here. + We first check if the corresponding module state_dict contains already HQQ quantized parameters. + If not, we create a temp linear layer with the module state_dict params and use it for quantization + """ + + if is_hqq_available(): + from hqq.core.quantize import HQQLinear + + module, tensor_name = get_module_from_name(model, param_name) + layer_name = ".".join(param_name.split(".")[:-1]) + parent_module = find_parent(model, layer_name) + node = layer_name.split(".")[-1] + + if tensor_name == "bias": + # this should already be set + return + + # set module state_dict + module_state_dict = {} + for k, v in state_dict.items(): + if layer_name + "." in k: + module_state_dict[k.split(".")[-1]] = v + if unexpected_keys is not None and k in unexpected_keys: + unexpected_keys.remove(k) + + if self.pre_quantized: + if isinstance(module, HQQLinear): + return + else: + hqq_layer = HQQLinear( + linear_layer=None, + quant_config=None, + compute_dtype=self.torch_dtype, + device=target_device, + ) + + hqq_layer.load_state_dict(module_state_dict) + + if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor): + hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias) + + if self.using_multi_gpu: + hqq_layer = self._patch_layer_for_multigpu(hqq_layer) + + setattr(parent_module, node, hqq_layer) + + # cleanup + del module.__dict__, module + torch.cuda.empty_cache() + return + + # Step 1: populate module with weight/bias from module state dict + for key in module_state_dict: + setattr(module, key, torch.nn.Parameter(module_state_dict[key])) + + # Step 2: Replace module with either HQQLinear or move it to device. We do this via setattr on the parent as doing on it on the module + # directly doesn't work. + quant_config = model.config.quantization_config["quant_config"] + skip_modules = model.config.quantization_config["skip_modules"] + module_tag = ".".join(module.name.split(".")[-2:]) + module_quant_config = None + if "weight_quant_params" in quant_config: + module_quant_config = quant_config + elif module_tag in quant_config: + module_quant_config = quant_config[module_tag] + + for skip_module in skip_modules: + if skip_module in module.name: + module_quant_config = None + break + + if module_quant_config is not None: + hqq_layer = HQQLinear( + module, + quant_config=module_quant_config, + compute_dtype=self.torch_dtype, + device=target_device, + del_orig=True, + ) + + if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor): + hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias) + + if self.using_multi_gpu: + hqq_layer = self._patch_layer_for_multigpu(hqq_layer) + + setattr(parent_module, node, hqq_layer) + + else: + module = module.to(dtype=self.torch_dtype, device=target_device) + setattr(parent_module, node, module) + + torch.cuda.empty_cache() + + # Remove accelerate hook and uses a simpler forward pass. Otherwise, this breaks with multi-gpu + def _patch_layer_for_multigpu(self, hqq_layer): + hqq_layer = remove_hook_from_module(hqq_layer) + + def forward_with_device(self, x): + out = torch.matmul(x.to(self.device), self.dequantize().t()) + if self.bias is not None: + out += self.bias + return out + + hqq_layer.forward = lambda x: forward_with_device(hqq_layer, x) + return hqq_layer + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + **kwargs, + ): + # Add the corresponding quant_config to each valid module. This allows us to do the actual nn.Linear -> HQQLinear conversion in create_quantized_param(). + # prepare_for_hqq_linear() also sets the right quantization config inside the model (model.config.quantization_config) and the layers (hqq_layer.quant_config) + model = prepare_for_hqq_linear(model, quantization_config=self.quantization_config) + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + model.is_hqq_quantized = True + model.is_hqq_serializable = self.is_serializable() + return model + + def is_serializable(self, safe_serialization=None): + return True + + @property + def is_trainable(self) -> bool: + return True diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_quanto.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_quanto.py new file mode 100644 index 0000000000000000000000000000000000000000..be760f0d430c16c5a81036972689ebca553bb20f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_quanto.py @@ -0,0 +1,201 @@ +# Copyright 2024 The HuggingFace Inc. 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. +import importlib +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from packaging import version + +from .base import HfQuantizer +from .quantizers_utils import get_module_from_name + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..utils import ( + is_accelerate_available, + is_optimum_quanto_available, + is_torch_available, + logging, +) +from ..utils.quantization_config import QuantoConfig + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class QuantoHfQuantizer(HfQuantizer): + """ + Quantizer for the quanto library + """ + + required_packages = ["quanto", "accelerate"] + requires_parameters_quantization = True + requires_calibration = False + + def __init__(self, quantization_config: QuantoConfig, **kwargs): + super().__init__(quantization_config, **kwargs) + self.post_init() + + def post_init(self): + r""" + Safety checker + """ + if self.quantization_config.activations is not None and not self.pre_quantized: + raise ValueError( + "We don't support quantizing the activations with transformers library." + "Use quanto library for more complex use cases such as activations quantization, calibration and quantization aware training." + ) + + def validate_environment(self, *args, **kwargs): + if not is_optimum_quanto_available(): + raise ImportError( + "Loading an optimum-quanto quantized model requires optimum-quanto library (`pip install optimum-quanto`)" + ) + if not is_accelerate_available(): + raise ImportError( + "Loading an optimum-quanto quantized model requires accelerate library (`pip install accelerate`)" + ) + + def update_device_map(self, device_map): + if device_map is None: + device_map = {"": "cpu"} + logger.info( + "The device_map was not initialized. " + "Setting device_map to {'':'cpu'}. " + "If you want to use the model for inference, please set device_map ='auto'" + ) + return device_map + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + logger.info("You did not specify `torch_dtype` in `from_pretrained`. Setting it to `torch.float32`.") + torch_dtype = torch.float32 + return torch_dtype + + def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]: + if is_optimum_quanto_available(): + from optimum.quanto import QModuleMixin + + not_missing_keys = [] + for name, module in model.named_modules(): + if isinstance(module, QModuleMixin): + for missing in missing_keys: + if ( + (name in missing or name in f"{prefix}.{missing}") + and not missing.endswith(".weight") + and not missing.endswith(".bias") + ): + not_missing_keys.append(missing) + return [k for k in missing_keys if k not in not_missing_keys] + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ) -> bool: + """ + Check if a parameter needs to be quantized. + """ + if is_optimum_quanto_available(): + from optimum.quanto import QModuleMixin + + device_map = kwargs.get("device_map", None) + param_device = kwargs.get("param_device", None) + # we don't quantize the model if the module is going to be offloaded to the cpu + if device_map is not None and param_device is not None: + device_map_values = set(device_map.values()) + if param_device == "cpu" and len(device_map_values) > 1: + if not (device_map_values == {"cpu"} or device_map_values == {"cpu", "disk"}): + return False + + module, tensor_name = get_module_from_name(model, param_name) + # We only quantize the weights and the bias is not quantized. + if isinstance(module, QModuleMixin) and "weight" in tensor_name: + # if the weights are quantized, don't need to recreate it again with `create_quantized_param` + return not module.frozen + else: + return False + + def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + max_memory = {key: val * 0.90 for key, val in max_memory.items()} + return max_memory + + def create_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + *args, + **kwargs, + ): + """ + Create the quantized parameter by calling .freeze() after setting it to the module. + """ + from accelerate.utils import set_module_tensor_to_device + + set_module_tensor_to_device(model, param_name, target_device, param_value) + module, _ = get_module_from_name(model, param_name) + module.freeze() + module.weight.requires_grad = False + + def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": + if version.parse(importlib.metadata.version("accelerate")) > version.parse("0.27.0"): + from accelerate.utils import CustomDtype + + mapping = { + "int8": torch.int8, + "float8": CustomDtype.FP8, + "int4": CustomDtype.INT4, + "int2": CustomDtype.INT2, + } + target_dtype = mapping[self.quantization_config.weights] + return target_dtype + else: + raise ValueError( + "You are using `device_map='auto'` on an optimum-quanto quantized model. To automatically compute" + " the appropriate device map, you should upgrade your `accelerate` library," + "`pip install --upgrade accelerate` or install it from source." + ) + + def _process_model_before_weight_loading( + self, model: "PreTrainedModel", keep_in_fp32_modules: Optional[List[str]] = None, **kwargs + ): + from ..integrations import replace_with_quanto_layers + + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules + ) + + model, _ = replace_with_quanto_layers( + model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config + ) + model.config.quantization_config = self.quantization_config + + def _process_model_after_weight_loading(self, model, **kwargs): + return model + + @property + def is_trainable(self, model: Optional["PreTrainedModel"] = None): + return True + + def is_serializable(self, safe_serialization=None): + return False diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_quark.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_quark.py new file mode 100644 index 0000000000000000000000000000000000000000..374360b1cb8f20724d74fe3d59cb78b0b47b2912 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_quark.py @@ -0,0 +1,113 @@ +# coding=utf-8 +# Copyright 2025 Advanced Micro Devices, Inc. and The HuggingFace Inc. 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. + +from typing import TYPE_CHECKING, Any, Dict + +from ..file_utils import is_torch_available +from .base import HfQuantizer + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + + if is_torch_available(): + import torch + +from ..utils import is_accelerate_available, is_quark_available, logging + + +if is_accelerate_available(): + from accelerate.utils import set_module_tensor_to_device + +logger = logging.get_logger(__name__) + + +CHECKPOINT_KEYS = { + "weight_scale": "weight_quantizer.scale", + "bias_scale": "bias_quantizer.scale", + "input_scale": "input_quantizer.scale", + "output_scale": "output_quantizer.scale", + "weight_zero_point": "weight_quantizer.zero_point", + "bias_zero_point": "bias_quantizer.zero_point", + "input_zero_point": "input_quantizer.zero_point", + "output_zero_point": "output_quantizer.zero_point", +} + + +class QuarkHfQuantizer(HfQuantizer): + """ + Quark quantizer (https://quark.docs.amd.com/latest/). + """ + + requires_calibration = True # On-the-fly quantization with quark is not supported for now. + required_packages = ["quark"] + + # Checkpoints are expected to be already quantized when loading a quark model. However, as some keys from + # the checkpoint might mismatch the model parameters keys, we use the `create_quantized_param` method + # to load the checkpoints, remapping the keys. + requires_parameters_quantization = True + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + + self.json_export_config = quantization_config.json_export_config + + def validate_environment(self, *args, **kwargs): + if not is_quark_available(): + raise ImportError( + "Loading a Quark quantized model requires the `quark` library but it was not found in the environment. Please refer to https://quark.docs.amd.com/latest/install.html." + ) + + def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs): + from quark.torch.export.api import _map_to_quark + + _map_to_quark( + model, + self.quantization_config.quant_config, + pack_method=self.json_export_config.pack_method, + custom_mode=self.quantization_config.custom_mode, + ) + + return model + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ) -> bool: + return True + + def create_quantized_param( + self, model, param, param_name, param_device, state_dict, unexpected_keys + ) -> "torch.nn.Parameter": + postfix = param_name.split(".")[-1] + + if postfix in CHECKPOINT_KEYS: + param_name = param_name.replace(postfix, CHECKPOINT_KEYS[postfix]) + + set_module_tensor_to_device(model, param_name, param_device, value=param) + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + return model + + def is_serializable(self, safe_serialization=None): + return False + + @property + def is_trainable(self): + return False diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_spqr.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_spqr.py new file mode 100644 index 0000000000000000000000000000000000000000..4cf1193edbf2cc676704db06fda7505c824548c0 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_spqr.py @@ -0,0 +1,90 @@ +# Copyright 2024 The HuggingFace Inc. 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/lic enses/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 typing import TYPE_CHECKING, List, Optional + +from .base import HfQuantizer + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..integrations import replace_with_spqr_linear +from ..utils import is_accelerate_available, is_spqr_available, is_torch_available, logging +from ..utils.quantization_config import QuantizationConfigMixin + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class SpQRHfQuantizer(HfQuantizer): + """ + Quantizer of the SpQR method. Enables the loading of prequantized models. + """ + + requires_calibration = True + + def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): + super().__init__(quantization_config, **kwargs) + self.quantization_config = quantization_config + + def validate_environment(self, *args, **kwargs): + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run SpQR quantized model.") + + if not is_accelerate_available(): + raise ImportError("Using `spqr` quantization requires Accelerate: `pip install accelerate`") + + if not is_spqr_available(): + raise ImportError("Using `spqr` quantization requires SpQR: `pip install spqr_quant[gpu]`") + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + torch_dtype = torch.float16 + logger.info("Assuming SpQR inference on GPU and loading the model in `torch.float16`.") + elif torch_dtype != torch.float16: + raise ValueError( + "You cannot use any type other than torch.float16 for SpQR. Please either leave it None or set it to" + "torch.float16 explicitly." + ) + return torch_dtype + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + keep_in_fp32_modules: Optional[List[str]] = None, + **kwargs, + ): + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules + ) + + replace_with_spqr_linear( + model, + quantization_config=self.quantization_config, + modules_to_not_convert=self.modules_to_not_convert, + ) + model.config.quantization_config = self.quantization_config + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + return model + + @property + def is_trainable(self, model: Optional["PreTrainedModel"] = None): + return False + + def is_serializable(self, safe_serialization=None): + return True diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_torchao.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_torchao.py new file mode 100644 index 0000000000000000000000000000000000000000..ac8155f22767dc77d06343fe661b675cc02eb2da --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_torchao.py @@ -0,0 +1,286 @@ +# Copyright 2024 The HuggingFace Inc. 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. +import importlib +import re +import types +from typing import TYPE_CHECKING, Optional, Union + +from packaging import version + +from .base import HfQuantizer +from .quantizers_utils import get_module_from_name + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from typing import Any, Dict, List + +from ..utils import is_torch_available, is_torchao_available, logging +from ..utils.quantization_config import TorchAoConfig + + +if is_torch_available(): + import torch + import torch.nn as nn + +logger = logging.get_logger(__name__) + + +def fuzzy_match_size(config_name: str) -> Optional[str]: + """ + Extract the size digit from strings like "4weight", "8weight". + Returns the digit as an integer if found, otherwise None. + """ + config_name = config_name.lower() + + str_match = re.search(r"(\d)weight", config_name) + + if str_match: + return str_match.group(1) + + return None + + +# Finds the parent of a node module named "name" +def find_parent(model, name): + module_tree = name.split(".")[:-1] + parent = model + for m in module_tree: + parent = parent._modules[m] + return parent + + +def _quantization_type(weight): + from torchao.dtypes import AffineQuantizedTensor + from torchao.quantization.linear_activation_quantized_tensor import LinearActivationQuantizedTensor + + if isinstance(weight, AffineQuantizedTensor): + return f"{weight.__class__.__name__}({weight._quantization_type()})" + + if isinstance(weight, LinearActivationQuantizedTensor): + return f"{weight.__class__.__name__}(activation={weight.input_quant_func}, weight={_quantization_type(weight.original_weight_tensor)})" + + +def _linear_extra_repr(self): + weight = _quantization_type(self.weight) + if weight is None: + return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight=None" + else: + return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight={weight}" + + +class TorchAoHfQuantizer(HfQuantizer): + """ + Quantizer for torchao: https://github.com/pytorch/ao/ + """ + + requires_parameters_quantization = True + requires_calibration = False + required_packages = ["torchao"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + + def validate_environment(self, *args, **kwargs): + if not is_torchao_available(): + raise ImportError("Loading an torchao quantized model requires torchao library (`pip install torchao`)") + + self.offload = False + device_map = kwargs.get("device_map", None) + if isinstance(device_map, dict): + if "cpu" in device_map.values() or "disk" in device_map.values(): + if self.pre_quantized: + raise ValueError( + "You are attempting to perform cpu/disk offload with a pre-quantized torchao model " + "This is not supported yet . Please remove the CPU or disk device from the device_map." + ) + else: + self.offload = True + if self.pre_quantized: + weights_only = kwargs.get("weights_only", None) + if weights_only: + torch_version = version.parse(importlib.metadata.version("torch")) + if torch_version < version.parse("2.5.0"): + raise RuntimeError( + f"In order to use torchao pre-quantized model, you need to have torch>=2.5.0. However, the current version is {torch_version}." + f" You can also set with `weights_only=False` in `from_pretrained` if you don't want to update torch" + ) + + def update_torch_dtype(self, torch_dtype): + if self.quantization_config.quant_type == "int4_weight_only": + if torch_dtype is not None and torch_dtype != torch.bfloat16: + logger.warning_once( + f"Setting torch_dtype to {torch_dtype} for int4_weight_only quantization, but only bfloat16 is supported right now. Please set the torch_dtype to bfloat16." + ) + if torch_dtype is None: + logger.warning_once( + "Setting torch_dtype to torch.bfloat16 for int4_weight_only quantization since only bfloat16 is supported right now. Please set torch_dtype=torch.bfloat16 to remove this warning." + ) + torch_dtype = torch.bfloat16 + if self.quantization_config.quant_type == "int8_dynamic_activation_int8_weight": + if torch_dtype is None: + logger.info( + "Setting torch_dtype to torch.float32 for int8_dynamic_activation_int8_weight quantization as no torch_dtype was specified in from_pretrained" + ) + # we need to set the torch_dtype, otherwise we have dtype mismatch when performing the quantized linear op + torch_dtype = torch.float32 + return torch_dtype + + def adjust_target_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if version.parse(importlib.metadata.version("accelerate")) > version.parse("0.19.0"): + from accelerate.utils import CustomDtype + + # Import AOBaseConfig directly since we know we have the right version + if self.quantization_config._get_ao_version() > version.Version("0.9.0"): + from torchao.core.config import AOBaseConfig + + quant_type = self.quantization_config.quant_type + if isinstance(quant_type, AOBaseConfig): + # Extract size digit using fuzzy match on the class name + config_name = quant_type.__class__.__name__ + size_digit = fuzzy_match_size(config_name) + + # Map the extracted digit to appropriate dtype + if size_digit == "4": + return CustomDtype.INT4 + else: + # Default to int8 + return torch.int8 + + # Original mapping for non-AOBaseConfig types + map_to_target_dtype = { + "int4_weight_only": CustomDtype.INT4, + "int8_weight_only": torch.int8, + "int8_dynamic_activation_int8_weight": torch.int8, + "autoquant": None, + } + return map_to_target_dtype[self.quantization_config.quant_type] + else: + raise ValueError( + "You are using `device_map='auto'` on a torchao quantized model. To automatically compute" + " the appropriate device map, you should upgrade your `accelerate` library with " + "`pip install --upgrade accelerate`" + ) + + def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + # need more space for the quantization parameters (e.g. scale). Tested with int4 wo and group size = 128 + max_memory = {key: val * 0.9 for key, val in max_memory.items()} + return max_memory + + def _process_model_before_weight_loading( + self, model: "PreTrainedModel", keep_in_fp32_modules: Optional[List[str]] = None, **kwargs + ): + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules + ) + return + + def check_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], + **kwargs, + ) -> bool: + if self.quantization_config.quant_type == "autoquant": + return False + + param_device = kwargs.pop("param_device", None) + # check if the param_name is not in self.modules_to_not_convert + if any((key + "." in param_name) or (key == param_name) for key in self.modules_to_not_convert): + return False + elif param_device == "cpu" and self.offload: + # We don't quantize weights that we offload + return False + else: + # we only quantize the weight of nn.Linear + module, tensor_name = get_module_from_name(model, param_name) + return isinstance(module, torch.nn.Linear) and (tensor_name == "weight") + + def create_quantized_param( + self, + model: "PreTrainedModel", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + state_dict: Dict[str, Any], + unexpected_keys: List[str], + ): + """ + Each nn.Linear layer that needs to be quantized is processsed here. + First, we set the value the weight tensor, then we move it to the target device. Finally, we quantize the module. + """ + if self.quantization_config.quant_type == "autoquant": + return + + from torchao.quantization import quantize_ + + module, tensor_name = get_module_from_name(model, param_name) + if self.pre_quantized: + module._parameters[tensor_name] = torch.nn.Parameter(param_value.to(device=target_device)) + if isinstance(module, nn.Linear): + module.extra_repr = types.MethodType(_linear_extra_repr, module) + else: + assert isinstance(self.quantization_config, TorchAoConfig) + module._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device) + quantize_(module, self.quantization_config.get_apply_tensor_subclass()) + + def _process_model_after_weight_loading(self, model, **kwargs): + """No process required for torchao quantized model""" + if self.quantization_config.quant_type == "autoquant": + from torchao import autoquant + from torchao.quantization import ALL_AUTOQUANT_CLASS_LIST + + model = torch.compile(model, mode="max-autotune") + model = autoquant( + model, + qtensor_class_list=ALL_AUTOQUANT_CLASS_LIST, + set_inductor_config=False, + **self.quantization_config.quant_type_kwargs, + ) + return model + return + + def is_serializable(self, safe_serialization=None) -> bool: + if safe_serialization: + logger.warning( + "torchao quantized model does not support safe serialization, please set `safe_serialization` to False" + ) + return False + _is_torchao_serializable = version.parse(importlib.metadata.version("huggingface_hub")) >= version.parse( + "0.25.0" + ) + if not _is_torchao_serializable: + logger.warning("torchao quantized model is only serializable after huggingface_hub >= 0.25.0 ") + if self.offload and self.quantization_config.modules_to_not_convert is None: + logger.warning( + "The model contains offloaded modules and these modules are not quantized. We don't recommend saving the model as we won't be able to reload them." + "If you want to specify modules to not quantize, please specify modules_to_not_convert in the quantization_config." + ) + return False + return _is_torchao_serializable + + @property + def is_trainable(self) -> bool: + supported_quant_types_for_training = [ + "int8_weight_only", + "int8_dynamic_activation_int8_weight", + ] + return self.quantization_config.quant_type in supported_quant_types_for_training + + @property + def is_compileable(self) -> bool: + return True diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizer_vptq.py b/lib/python3.12/site-packages/transformers/quantizers/quantizer_vptq.py new file mode 100644 index 0000000000000000000000000000000000000000..85483357448f15df66dc77b4e0e9158970ab1764 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizer_vptq.py @@ -0,0 +1,99 @@ +# Copyright 2024 The HuggingFace Inc. 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. +from typing import TYPE_CHECKING, List, Optional + +from .base import HfQuantizer + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + +from ..utils import is_accelerate_available, is_torch_available, is_vptq_available, logging +from ..utils.quantization_config import QuantizationConfigMixin + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class VptqHfQuantizer(HfQuantizer): + """ + Quantizer of the VPTQ method. Enables the loading of prequantized models. + """ + + requires_calibration = True + required_packages = ["vptq"] + + def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): + super().__init__(quantization_config, **kwargs) + self.quantization_config = quantization_config + + def validate_environment(self, *args, **kwargs): + if not is_accelerate_available(): + raise ImportError("Using `vptq` quantization requires Accelerate: `pip install accelerate`") + + if not is_vptq_available(): + raise ImportError("Using `vptq` quantization requires VPTQ>=0.0.4: `pip install -U vptq`") + + def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype": + if torch_dtype is None: + if torch.cuda.is_available(): + torch_dtype = torch.float16 + logger.info( + "CUDA available. Assuming VPTQ inference on GPU and loading the model in `torch.float16`. To overwrite it, set `torch_dtype` manually." + ) + else: + import vptq + + device_availability = getattr(vptq, "device_availability", lambda device: False) + if device_availability("cpu") is True: + raise RuntimeError("No GPU found. Please wait for the next release of VPTQ to use CPU inference") + torch_dtype = torch.float32 + logger.info("No GPU found. Assuming VPTQ inference on CPU and loading the model in `torch.float32`.") + return torch_dtype + + def _process_model_before_weight_loading( + self, + model: "PreTrainedModel", + keep_in_fp32_modules: Optional[List[str]] = None, + **kwargs, + ): + """ + we don't have param like modules_to_not_convert to indicate which layers should not be quantized + because `quantization_config` include the layers that should be quantized + """ + from ..integrations import replace_with_vptq_linear + + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules + ) + + replace_with_vptq_linear( + model, + quantization_config=self.quantization_config, + modules_to_not_convert=self.modules_to_not_convert, + ) + model.config.quantization_config = self.quantization_config + + def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): + return model + + @property + def is_trainable(self, model: Optional["PreTrainedModel"] = None): + return False + + def is_serializable(self, safe_serialization=None): + return True diff --git a/lib/python3.12/site-packages/transformers/quantizers/quantizers_utils.py b/lib/python3.12/site-packages/transformers/quantizers/quantizers_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bbed8317049f9e2530c03dd40730e9c79cb879b2 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/quantizers/quantizers_utils.py @@ -0,0 +1,21 @@ +# Copyright 2024 The HuggingFace Inc. 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. +from typing import Any, Tuple + + +def get_module_from_name(module, tensor_name: str) -> Tuple[Any, str]: + if "." in tensor_name: + module_name, tensor_name = tensor_name.rsplit(".", 1) + module = module.get_submodule(module_name) + return module, tensor_name diff --git a/lib/python3.12/site-packages/transformers/safetensors_conversion.py b/lib/python3.12/site-packages/transformers/safetensors_conversion.py new file mode 100644 index 0000000000000000000000000000000000000000..f1612d3ea57c98fd1d383887cfbeb4e2882d3963 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/safetensors_conversion.py @@ -0,0 +1,105 @@ +from typing import Optional + +import requests +from huggingface_hub import Discussion, HfApi, get_repo_discussions + +from .utils import cached_file, http_user_agent, logging + + +logger = logging.get_logger(__name__) + + +def previous_pr(api: HfApi, model_id: str, pr_title: str, token: str) -> Optional["Discussion"]: + main_commit = api.list_repo_commits(model_id, token=token)[0].commit_id + for discussion in get_repo_discussions(repo_id=model_id, token=token): + if discussion.title == pr_title and discussion.status == "open" and discussion.is_pull_request: + commits = api.list_repo_commits(model_id, revision=discussion.git_reference, token=token) + + if main_commit == commits[1].commit_id: + return discussion + return None + + +def spawn_conversion(token: str, private: bool, model_id: str): + logger.info("Attempting to convert .bin model on the fly to safetensors.") + + safetensors_convert_space_url = "https://safetensors-convert.hf.space" + sse_url = f"{safetensors_convert_space_url}/call/run" + + def start(_sse_connection): + for line in _sse_connection.iter_lines(): + line = line.decode() + if line.startswith("event:"): + status = line[7:] + logger.debug(f"Safetensors conversion status: {status}") + + if status == "complete": + return + elif status == "heartbeat": + logger.debug("Heartbeat") + else: + logger.debug(f"Unknown status {status}") + else: + logger.debug(line) + + data = {"data": [model_id, private, token]} + + result = requests.post(sse_url, stream=True, json=data).json() + event_id = result["event_id"] + + with requests.get(f"{sse_url}/{event_id}", stream=True) as sse_connection: + try: + logger.debug("Spawning safetensors automatic conversion.") + start(sse_connection) + except Exception as e: + logger.warning(f"Error during conversion: {repr(e)}") + + +def get_conversion_pr_reference(api: HfApi, model_id: str, **kwargs): + private = api.model_info(model_id).private + + logger.info("Attempting to create safetensors variant") + pr_title = "Adding `safetensors` variant of this model" + token = kwargs.get("token") + + # This looks into the current repo's open PRs to see if a PR for safetensors was already open. If so, it + # returns it. It checks that the PR was opened by the bot and not by another user so as to prevent + # security breaches. + pr = previous_pr(api, model_id, pr_title, token=token) + + if pr is None or (not private and pr.author != "SFconvertbot"): + spawn_conversion(token, private, model_id) + pr = previous_pr(api, model_id, pr_title, token=token) + else: + logger.info("Safetensors PR exists") + + sha = f"refs/pr/{pr.num}" + + return sha + + +def auto_conversion(pretrained_model_name_or_path: str, ignore_errors_during_conversion=False, **cached_file_kwargs): + try: + api = HfApi(token=cached_file_kwargs.get("token"), headers={"user-agent": http_user_agent()}) + sha = get_conversion_pr_reference(api, pretrained_model_name_or_path, **cached_file_kwargs) + + if sha is None: + return None, None + cached_file_kwargs["revision"] = sha + del cached_file_kwargs["_commit_hash"] + + # This is an additional HEAD call that could be removed if we could infer sharded/non-sharded from the PR + # description. + sharded = api.file_exists( + pretrained_model_name_or_path, + "model.safetensors.index.json", + revision=sha, + token=cached_file_kwargs.get("token"), + ) + filename = "model.safetensors.index.json" if sharded else "model.safetensors" + + resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs) + return resolved_archive_file, sha, sharded + except Exception as e: + if not ignore_errors_during_conversion: + raise e diff --git a/lib/python3.12/site-packages/transformers/sagemaker/__init__.py b/lib/python3.12/site-packages/transformers/sagemaker/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..98fe38de89cd025911d03669f9e22b03ab0768bd --- /dev/null +++ b/lib/python3.12/site-packages/transformers/sagemaker/__init__.py @@ -0,0 +1,16 @@ +# Copyright 2021 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. + +from .trainer_sm import SageMakerTrainer +from .training_args_sm import SageMakerTrainingArguments, is_sagemaker_dp_enabled diff --git a/lib/python3.12/site-packages/transformers/sagemaker/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/transformers/sagemaker/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..869992acbf9409caf08459957d57a33ac600bb1b Binary files /dev/null and b/lib/python3.12/site-packages/transformers/sagemaker/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/sagemaker/__pycache__/trainer_sm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/sagemaker/__pycache__/trainer_sm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa4567080f65ca8a0aec0ccad0222f7fe536b9c1 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/sagemaker/__pycache__/trainer_sm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/sagemaker/__pycache__/training_args_sm.cpython-312.pyc b/lib/python3.12/site-packages/transformers/sagemaker/__pycache__/training_args_sm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe9a9862eda562687f86c3f687d7aaf1a1e7afee Binary files /dev/null and b/lib/python3.12/site-packages/transformers/sagemaker/__pycache__/training_args_sm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/sagemaker/trainer_sm.py b/lib/python3.12/site-packages/transformers/sagemaker/trainer_sm.py new file mode 100644 index 0000000000000000000000000000000000000000..6ab4e01acdbcd3ade1afc2339a75850bc538bd7a --- /dev/null +++ b/lib/python3.12/site-packages/transformers/sagemaker/trainer_sm.py @@ -0,0 +1,30 @@ +# Copyright 2021 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. +import warnings + +from ..trainer import Trainer +from ..utils import logging + + +logger = logging.get_logger(__name__) + + +class SageMakerTrainer(Trainer): + def __init__(self, args=None, **kwargs): + warnings.warn( + "`SageMakerTrainer` is deprecated and will be removed in v5 of Transformers. You can use `Trainer` " + "instead.", + FutureWarning, + ) + super().__init__(args=args, **kwargs) diff --git a/lib/python3.12/site-packages/transformers/sagemaker/training_args_sm.py b/lib/python3.12/site-packages/transformers/sagemaker/training_args_sm.py new file mode 100644 index 0000000000000000000000000000000000000000..3daac7859b550de31f211a5e7c9938d8d557fc4c --- /dev/null +++ b/lib/python3.12/site-packages/transformers/sagemaker/training_args_sm.py @@ -0,0 +1,136 @@ +# Copyright 2021 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. + +import importlib.util +import json +import os +import warnings +from dataclasses import dataclass, field + +import torch + +from ..training_args import TrainingArguments +from ..utils import cached_property, is_sagemaker_dp_enabled, logging + + +logger = logging.get_logger(__name__) + +# TODO: should be moved to `utils` after refactoring of SageMakerTrainer + + +def is_sagemaker_model_parallel_available(): + # Get the sagemaker specific mp parameters from smp_options variable. + smp_options = os.getenv("SM_HP_MP_PARAMETERS", "{}") + try: + # Parse it and check the field "partitions" is included, it is required for model parallel. + smp_options = json.loads(smp_options) + if "partitions" not in smp_options: + return False + except json.JSONDecodeError: + return False + + # Get the sagemaker specific framework parameters from mpi_options variable. + mpi_options = os.getenv("SM_FRAMEWORK_PARAMS", "{}") + try: + # Parse it and check the field "sagemaker_distributed_dataparallel_enabled". + mpi_options = json.loads(mpi_options) + if not mpi_options.get("sagemaker_mpi_enabled", False): + return False + except json.JSONDecodeError: + return False + # Lastly, check if the `smdistributed` module is present. + return importlib.util.find_spec("smdistributed") is not None + + +if is_sagemaker_model_parallel_available(): + import smdistributed.modelparallel.torch as smp + + smp.init() + + +@dataclass +class SageMakerTrainingArguments(TrainingArguments): + mp_parameters: str = field( + default="", + metadata={"help": "Used by the SageMaker launcher to send mp-specific args. Ignored in SageMakerTrainer"}, + ) + + def __post_init__(self): + super().__post_init__() + warnings.warn( + "`SageMakerTrainingArguments` is deprecated and will be removed in v5 of Transformers. You can use " + "`TrainingArguments` instead.", + FutureWarning, + ) + + @cached_property + def _setup_devices(self) -> "torch.device": + logger.info("PyTorch: setting up devices") + if torch.distributed.is_available() and torch.distributed.is_initialized() and self.local_rank == -1: + logger.warning( + "torch.distributed process group is initialized, but local_rank == -1. " + "In order to use Torch DDP, launch your script with `python -m torch.distributed.launch" + ) + if self.no_cuda: + device = torch.device("cpu") + self._n_gpu = 0 + elif is_sagemaker_model_parallel_available(): + local_rank = smp.local_rank() + device = torch.device("cuda", local_rank) + self._n_gpu = 1 + elif is_sagemaker_dp_enabled(): + import smdistributed.dataparallel.torch.torch_smddp # noqa: F401 + + torch.distributed.init_process_group(backend="smddp", timeout=self.ddp_timeout_delta) + self.local_rank = int(os.getenv("SMDATAPARALLEL_LOCAL_RANK")) + device = torch.device("cuda", self.local_rank) + self._n_gpu = 1 + elif self.local_rank == -1: + # if n_gpu is > 1 we'll use nn.DataParallel. + # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0` + # Explicitly set CUDA to the first (index 0) CUDA device, otherwise `set_device` will + # trigger an error that a device index is missing. Index 0 takes into account the + # GPUs available in the environment, so `CUDA_VISIBLE_DEVICES=1,2` with `cuda:0` + # will use the first GPU in that env, i.e. GPU#1 + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + # Sometimes the line in the postinit has not been run before we end up here, so just checking we're not at + # the default value. + self._n_gpu = torch.cuda.device_count() + else: + # Here, we'll use torch.distributed. + # Initializes the distributed backend which will take care of synchronizing nodes/GPUs + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group(backend="nccl", timeout=self.ddp_timeout_delta) + device = torch.device("cuda", self.local_rank) + self._n_gpu = 1 + + if device.type == "cuda": + torch.cuda.set_device(device) + + return device + + @property + def world_size(self): + if is_sagemaker_model_parallel_available(): + return smp.dp_size() + + return super().world_size + + @property + def place_model_on_device(self): + return not is_sagemaker_model_parallel_available() + + @property + def _no_sync_in_gradient_accumulation(self): + return False diff --git a/lib/python3.12/site-packages/transformers/testing_utils.py b/lib/python3.12/site-packages/transformers/testing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3ee55e893a332314c574b2c7684ff563f8a51b50 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/testing_utils.py @@ -0,0 +1,3155 @@ +# Copyright 2020 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. + +import collections +import contextlib +import copy +import doctest +import functools +import gc +import importlib +import inspect +import logging +import multiprocessing +import os +import re +import shlex +import shutil +import subprocess +import sys +import tempfile +import threading +import time +import unittest +from collections import UserDict, defaultdict +from collections.abc import Generator, Iterable, Iterator, Mapping +from dataclasses import MISSING, fields +from functools import cache, wraps +from io import StringIO +from pathlib import Path +from typing import Any, Callable, Optional, Union +from unittest import mock +from unittest.mock import patch + +import huggingface_hub.utils +import requests +import urllib3 +from huggingface_hub import delete_repo +from packaging import version + +from transformers import Trainer +from transformers import logging as transformers_logging + +from .integrations import ( + is_clearml_available, + is_optuna_available, + is_ray_available, + is_sigopt_available, + is_swanlab_available, + is_tensorboard_available, + is_wandb_available, +) +from .integrations.deepspeed import is_deepspeed_available +from .utils import ( + ACCELERATE_MIN_VERSION, + GGUF_MIN_VERSION, + is_accelerate_available, + is_apex_available, + is_apollo_torch_available, + is_aqlm_available, + is_auto_awq_available, + is_auto_gptq_available, + is_av_available, + is_bitsandbytes_available, + is_bitsandbytes_multi_backend_available, + is_bs4_available, + is_compressed_tensors_available, + is_cv2_available, + is_cython_available, + is_detectron2_available, + is_eetq_available, + is_essentia_available, + is_faiss_available, + is_fbgemm_gpu_available, + is_flash_attn_2_available, + is_flax_available, + is_flute_available, + is_fsdp_available, + is_ftfy_available, + is_g2p_en_available, + is_galore_torch_available, + is_gguf_available, + is_gptqmodel_available, + is_grokadamw_available, + is_hadamard_available, + is_hqq_available, + is_ipex_available, + is_jieba_available, + is_jinja_available, + is_jumanpp_available, + is_keras_nlp_available, + is_levenshtein_available, + is_librosa_available, + is_liger_kernel_available, + is_lomo_available, + is_natten_available, + is_nltk_available, + is_onnx_available, + is_optimum_available, + is_optimum_quanto_available, + is_pandas_available, + is_peft_available, + is_phonemizer_available, + is_pretty_midi_available, + is_pyctcdecode_available, + is_pytesseract_available, + is_pytest_available, + is_pytorch_quantization_available, + is_quark_available, + is_rjieba_available, + is_sacremoses_available, + is_safetensors_available, + is_schedulefree_available, + is_scipy_available, + is_sentencepiece_available, + is_seqio_available, + is_soundfile_available, + is_spacy_available, + is_spqr_available, + is_sudachi_available, + is_sudachi_projection_available, + is_tensorflow_probability_available, + is_tensorflow_text_available, + is_tf2onnx_available, + is_tf_available, + is_tiktoken_available, + is_timm_available, + is_tokenizers_available, + is_torch_available, + is_torch_bf16_available_on_device, + is_torch_bf16_cpu_available, + is_torch_bf16_gpu_available, + is_torch_deterministic, + is_torch_fp16_available_on_device, + is_torch_greater_or_equal, + is_torch_hpu_available, + is_torch_mlu_available, + is_torch_neuroncore_available, + is_torch_npu_available, + is_torch_sdpa_available, + is_torch_tensorrt_fx_available, + is_torch_tf32_available, + is_torch_xla_available, + is_torch_xpu_available, + is_torchao_available, + is_torchaudio_available, + is_torchdynamo_available, + is_torchvision_available, + is_vision_available, + is_vptq_available, + strtobool, +) + + +if is_accelerate_available(): + from accelerate.state import AcceleratorState, PartialState + from accelerate.utils.imports import is_fp8_available + + +if is_pytest_available(): + from _pytest.doctest import ( + Module, + _get_checker, + _get_continue_on_failure, + _get_runner, + _is_mocked, + _patch_unwrap_mock_aware, + get_optionflags, + ) + from _pytest.outcomes import skip + from _pytest.pathlib import import_path + from pytest import DoctestItem +else: + Module = object + DoctestItem = object + + +SMALL_MODEL_IDENTIFIER = "julien-c/bert-xsmall-dummy" +DUMMY_UNKNOWN_IDENTIFIER = "julien-c/dummy-unknown" +DUMMY_DIFF_TOKENIZER_IDENTIFIER = "julien-c/dummy-diff-tokenizer" +# Used to test Auto{Config, Model, Tokenizer} model_type detection. + +# Used to test the hub +USER = "__DUMMY_TRANSFORMERS_USER__" +ENDPOINT_STAGING = "https://hub-ci.huggingface.co" + +# Not critical, only usable on the sandboxed CI instance. +TOKEN = "hf_94wBhPGp6KrrTH3KDchhKpRxZwd6dmHWLL" + +if is_torch_available(): + import torch + + IS_ROCM_SYSTEM = torch.version.hip is not None + IS_CUDA_SYSTEM = torch.version.cuda is not None +else: + IS_ROCM_SYSTEM = False + IS_CUDA_SYSTEM = False + +logger = transformers_logging.get_logger(__name__) + + +def parse_flag_from_env(key, default=False): + try: + value = os.environ[key] + except KeyError: + # KEY isn't set, default to `default`. + _value = default + else: + # KEY is set, convert it to True or False. + try: + _value = strtobool(value) + except ValueError: + # More values are supported, but let's keep the message simple. + raise ValueError(f"If set, {key} must be yes or no.") + return _value + + +def parse_int_from_env(key, default=None): + try: + value = os.environ[key] + except KeyError: + _value = default + else: + try: + _value = int(value) + except ValueError: + raise ValueError(f"If set, {key} must be a int.") + return _value + + +_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False) +_run_custom_tokenizers = parse_flag_from_env("RUN_CUSTOM_TOKENIZERS", default=False) +_run_staging = parse_flag_from_env("HUGGINGFACE_CO_STAGING", default=False) +_run_pipeline_tests = parse_flag_from_env("RUN_PIPELINE_TESTS", default=True) +_run_agent_tests = parse_flag_from_env("RUN_AGENT_TESTS", default=False) +_run_third_party_device_tests = parse_flag_from_env("RUN_THIRD_PARTY_DEVICE_TESTS", default=False) + + +def is_staging_test(test_case): + """ + Decorator marking a test as a staging test. + + Those tests will run using the staging environment of huggingface.co instead of the real model hub. + """ + if not _run_staging: + return unittest.skip(reason="test is staging test")(test_case) + else: + try: + import pytest # We don't need a hard dependency on pytest in the main library + except ImportError: + return test_case + else: + return pytest.mark.is_staging_test()(test_case) + + +def is_pipeline_test(test_case): + """ + Decorator marking a test as a pipeline test. If RUN_PIPELINE_TESTS is set to a falsy value, those tests will be + skipped. + """ + if not _run_pipeline_tests: + return unittest.skip(reason="test is pipeline test")(test_case) + else: + try: + import pytest # We don't need a hard dependency on pytest in the main library + except ImportError: + return test_case + else: + return pytest.mark.is_pipeline_test()(test_case) + + +def is_agent_test(test_case): + """ + Decorator marking a test as an agent test. If RUN_TOOL_TESTS is set to a falsy value, those tests will be skipped. + """ + if not _run_agent_tests: + return unittest.skip(reason="test is an agent test")(test_case) + else: + try: + import pytest # We don't need a hard dependency on pytest in the main library + except ImportError: + return test_case + else: + return pytest.mark.is_agent_test()(test_case) + + +def slow(test_case): + """ + Decorator marking a test as slow. + + Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them. + + """ + return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case) + + +def tooslow(test_case): + """ + Decorator marking a test as too slow. + + Slow tests are skipped while they're in the process of being fixed. No test should stay tagged as "tooslow" as + these will not be tested by the CI. + + """ + return unittest.skip(reason="test is too slow")(test_case) + + +def skip_if_not_implemented(test_func): + @functools.wraps(test_func) + def wrapper(*args, **kwargs): + try: + return test_func(*args, **kwargs) + except NotImplementedError as e: + raise unittest.SkipTest(f"Test skipped due to NotImplementedError: {e}") + + return wrapper + + +def apply_skip_if_not_implemented(cls): + """ + Class decorator to apply @skip_if_not_implemented to all test methods. + """ + for attr_name in dir(cls): + if attr_name.startswith("test_"): + attr = getattr(cls, attr_name) + if callable(attr): + setattr(cls, attr_name, skip_if_not_implemented(attr)) + return cls + + +def custom_tokenizers(test_case): + """ + Decorator marking a test for a custom tokenizer. + + Custom tokenizers require additional dependencies, and are skipped by default. Set the RUN_CUSTOM_TOKENIZERS + environment variable to a truthy value to run them. + """ + return unittest.skipUnless(_run_custom_tokenizers, "test of custom tokenizers")(test_case) + + +def require_bs4(test_case): + """ + Decorator marking a test that requires BeautifulSoup4. These tests are skipped when BeautifulSoup4 isn't installed. + """ + return unittest.skipUnless(is_bs4_available(), "test requires BeautifulSoup4")(test_case) + + +def require_galore_torch(test_case): + """ + Decorator marking a test that requires GaLore. These tests are skipped when GaLore isn't installed. + https://github.com/jiaweizzhao/GaLore + """ + return unittest.skipUnless(is_galore_torch_available(), "test requires GaLore")(test_case) + + +def require_apollo_torch(test_case): + """ + Decorator marking a test that requires GaLore. These tests are skipped when APOLLO isn't installed. + https://github.com/zhuhanqing/APOLLO + """ + return unittest.skipUnless(is_apollo_torch_available(), "test requires APOLLO")(test_case) + + +def require_lomo(test_case): + """ + Decorator marking a test that requires LOMO. These tests are skipped when LOMO-optim isn't installed. + https://github.com/OpenLMLab/LOMO + """ + return unittest.skipUnless(is_lomo_available(), "test requires LOMO")(test_case) + + +def require_grokadamw(test_case): + """ + Decorator marking a test that requires GrokAdamW. These tests are skipped when GrokAdamW isn't installed. + """ + return unittest.skipUnless(is_grokadamw_available(), "test requires GrokAdamW")(test_case) + + +def require_schedulefree(test_case): + """ + Decorator marking a test that requires schedulefree. These tests are skipped when schedulefree isn't installed. + https://github.com/facebookresearch/schedule_free + """ + return unittest.skipUnless(is_schedulefree_available(), "test requires schedulefree")(test_case) + + +def require_cv2(test_case): + """ + Decorator marking a test that requires OpenCV. + + These tests are skipped when OpenCV isn't installed. + + """ + return unittest.skipUnless(is_cv2_available(), "test requires OpenCV")(test_case) + + +def require_levenshtein(test_case): + """ + Decorator marking a test that requires Levenshtein. + + These tests are skipped when Levenshtein isn't installed. + + """ + return unittest.skipUnless(is_levenshtein_available(), "test requires Levenshtein")(test_case) + + +def require_nltk(test_case): + """ + Decorator marking a test that requires NLTK. + + These tests are skipped when NLTK isn't installed. + + """ + return unittest.skipUnless(is_nltk_available(), "test requires NLTK")(test_case) + + +def require_accelerate(test_case, min_version: str = ACCELERATE_MIN_VERSION): + """ + Decorator marking a test that requires accelerate. These tests are skipped when accelerate isn't installed. + """ + return unittest.skipUnless( + is_accelerate_available(min_version), f"test requires accelerate version >= {min_version}" + )(test_case) + + +def require_gguf(test_case, min_version: str = GGUF_MIN_VERSION): + """ + Decorator marking a test that requires ggguf. These tests are skipped when gguf isn't installed. + """ + return unittest.skipUnless(is_gguf_available(min_version), f"test requires gguf version >= {min_version}")( + test_case + ) + + +def require_fsdp(test_case, min_version: str = "1.12.0"): + """ + Decorator marking a test that requires fsdp. These tests are skipped when fsdp isn't installed. + """ + return unittest.skipUnless(is_fsdp_available(min_version), f"test requires torch version >= {min_version}")( + test_case + ) + + +def require_g2p_en(test_case): + """ + Decorator marking a test that requires g2p_en. These tests are skipped when SentencePiece isn't installed. + """ + return unittest.skipUnless(is_g2p_en_available(), "test requires g2p_en")(test_case) + + +def require_safetensors(test_case): + """ + Decorator marking a test that requires safetensors. These tests are skipped when safetensors isn't installed. + """ + return unittest.skipUnless(is_safetensors_available(), "test requires safetensors")(test_case) + + +def require_rjieba(test_case): + """ + Decorator marking a test that requires rjieba. These tests are skipped when rjieba isn't installed. + """ + return unittest.skipUnless(is_rjieba_available(), "test requires rjieba")(test_case) + + +def require_jieba(test_case): + """ + Decorator marking a test that requires jieba. These tests are skipped when jieba isn't installed. + """ + return unittest.skipUnless(is_jieba_available(), "test requires jieba")(test_case) + + +def require_jinja(test_case): + """ + Decorator marking a test that requires jinja. These tests are skipped when jinja isn't installed. + """ + return unittest.skipUnless(is_jinja_available(), "test requires jinja")(test_case) + + +def require_tf2onnx(test_case): + return unittest.skipUnless(is_tf2onnx_available(), "test requires tf2onnx")(test_case) + + +def require_onnx(test_case): + return unittest.skipUnless(is_onnx_available(), "test requires ONNX")(test_case) + + +def require_timm(test_case): + """ + Decorator marking a test that requires Timm. + + These tests are skipped when Timm isn't installed. + + """ + return unittest.skipUnless(is_timm_available(), "test requires Timm")(test_case) + + +def require_natten(test_case): + """ + Decorator marking a test that requires NATTEN. + + These tests are skipped when NATTEN isn't installed. + + """ + return unittest.skipUnless(is_natten_available(), "test requires natten")(test_case) + + +def require_torch(test_case): + """ + Decorator marking a test that requires PyTorch. + + These tests are skipped when PyTorch isn't installed. + + """ + return unittest.skipUnless(is_torch_available(), "test requires PyTorch")(test_case) + + +def require_torch_greater_or_equal(version: str): + """ + Decorator marking a test that requires PyTorch version >= `version`. + + These tests are skipped when PyTorch version is less than `version`. + """ + + def decorator(test_case): + return unittest.skipUnless(is_torch_greater_or_equal(version), f"test requires PyTorch version >= {version}")( + test_case + ) + + return decorator + + +def require_flash_attn(test_case): + """ + Decorator marking a test that requires Flash Attention. + + These tests are skipped when Flash Attention isn't installed. + + """ + return unittest.skipUnless(is_flash_attn_2_available(), "test requires Flash Attention")(test_case) + + +def require_torch_sdpa(test_case): + """ + Decorator marking a test that requires PyTorch's SDPA. + + These tests are skipped when requirements are not met (torch version). + """ + return unittest.skipUnless(is_torch_sdpa_available(), "test requires PyTorch SDPA")(test_case) + + +def require_read_token(fn): + """ + A decorator that loads the HF token for tests that require to load gated models. + """ + token = os.getenv("HF_HUB_READ_TOKEN") + + @wraps(fn) + def _inner(*args, **kwargs): + if token is not None: + with patch("huggingface_hub.utils._headers.get_token", return_value=token): + return fn(*args, **kwargs) + else: # Allow running locally with the default token env variable + return fn(*args, **kwargs) + + return _inner + + +def require_peft(test_case): + """ + Decorator marking a test that requires PEFT. + + These tests are skipped when PEFT isn't installed. + + """ + return unittest.skipUnless(is_peft_available(), "test requires PEFT")(test_case) + + +def require_torchvision(test_case): + """ + Decorator marking a test that requires Torchvision. + + These tests are skipped when Torchvision isn't installed. + + """ + return unittest.skipUnless(is_torchvision_available(), "test requires Torchvision")(test_case) + + +def require_torch_or_tf(test_case): + """ + Decorator marking a test that requires PyTorch or TensorFlow. + + These tests are skipped when neither PyTorch not TensorFlow is installed. + + """ + return unittest.skipUnless(is_torch_available() or is_tf_available(), "test requires PyTorch or TensorFlow")( + test_case + ) + + +def require_intel_extension_for_pytorch(test_case): + """ + Decorator marking a test that requires Intel Extension for PyTorch. + + These tests are skipped when Intel Extension for PyTorch isn't installed or it does not match current PyTorch + version. + + """ + return unittest.skipUnless( + is_ipex_available(), + "test requires Intel Extension for PyTorch to be installed and match current PyTorch version, see" + " https://github.com/intel/intel-extension-for-pytorch", + )(test_case) + + +def require_tensorflow_probability(test_case): + """ + Decorator marking a test that requires TensorFlow probability. + + These tests are skipped when TensorFlow probability isn't installed. + + """ + return unittest.skipUnless(is_tensorflow_probability_available(), "test requires TensorFlow probability")( + test_case + ) + + +def require_torchaudio(test_case): + """ + Decorator marking a test that requires torchaudio. These tests are skipped when torchaudio isn't installed. + """ + return unittest.skipUnless(is_torchaudio_available(), "test requires torchaudio")(test_case) + + +def require_tf(test_case): + """ + Decorator marking a test that requires TensorFlow. These tests are skipped when TensorFlow isn't installed. + """ + return unittest.skipUnless(is_tf_available(), "test requires TensorFlow")(test_case) + + +def require_flax(test_case): + """ + Decorator marking a test that requires JAX & Flax. These tests are skipped when one / both are not installed + """ + return unittest.skipUnless(is_flax_available(), "test requires JAX & Flax")(test_case) + + +def require_sentencepiece(test_case): + """ + Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed. + """ + return unittest.skipUnless(is_sentencepiece_available(), "test requires SentencePiece")(test_case) + + +def require_sacremoses(test_case): + """ + Decorator marking a test that requires Sacremoses. These tests are skipped when Sacremoses isn't installed. + """ + return unittest.skipUnless(is_sacremoses_available(), "test requires Sacremoses")(test_case) + + +def require_seqio(test_case): + """ + Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed. + """ + return unittest.skipUnless(is_seqio_available(), "test requires Seqio")(test_case) + + +def require_scipy(test_case): + """ + Decorator marking a test that requires Scipy. These tests are skipped when SentencePiece isn't installed. + """ + return unittest.skipUnless(is_scipy_available(), "test requires Scipy")(test_case) + + +def require_tokenizers(test_case): + """ + Decorator marking a test that requires 🤗 Tokenizers. These tests are skipped when 🤗 Tokenizers isn't installed. + """ + return unittest.skipUnless(is_tokenizers_available(), "test requires tokenizers")(test_case) + + +def require_tensorflow_text(test_case): + """ + Decorator marking a test that requires tensorflow_text. These tests are skipped when tensroflow_text isn't + installed. + """ + return unittest.skipUnless(is_tensorflow_text_available(), "test requires tensorflow_text")(test_case) + + +def require_keras_nlp(test_case): + """ + Decorator marking a test that requires keras_nlp. These tests are skipped when keras_nlp isn't installed. + """ + return unittest.skipUnless(is_keras_nlp_available(), "test requires keras_nlp")(test_case) + + +def require_pandas(test_case): + """ + Decorator marking a test that requires pandas. These tests are skipped when pandas isn't installed. + """ + return unittest.skipUnless(is_pandas_available(), "test requires pandas")(test_case) + + +def require_pytesseract(test_case): + """ + Decorator marking a test that requires PyTesseract. These tests are skipped when PyTesseract isn't installed. + """ + return unittest.skipUnless(is_pytesseract_available(), "test requires PyTesseract")(test_case) + + +def require_pytorch_quantization(test_case): + """ + Decorator marking a test that requires PyTorch Quantization Toolkit. These tests are skipped when PyTorch + Quantization Toolkit isn't installed. + """ + return unittest.skipUnless(is_pytorch_quantization_available(), "test requires PyTorch Quantization Toolkit")( + test_case + ) + + +def require_vision(test_case): + """ + Decorator marking a test that requires the vision dependencies. These tests are skipped when torchaudio isn't + installed. + """ + return unittest.skipUnless(is_vision_available(), "test requires vision")(test_case) + + +def require_ftfy(test_case): + """ + Decorator marking a test that requires ftfy. These tests are skipped when ftfy isn't installed. + """ + return unittest.skipUnless(is_ftfy_available(), "test requires ftfy")(test_case) + + +def require_spacy(test_case): + """ + Decorator marking a test that requires SpaCy. These tests are skipped when SpaCy isn't installed. + """ + return unittest.skipUnless(is_spacy_available(), "test requires spacy")(test_case) + + +def require_torch_multi_gpu(test_case): + """ + Decorator marking a test that requires a multi-GPU CUDA setup (in PyTorch). These tests are skipped on a machine without + multiple CUDA GPUs. + + To run *only* the multi_gpu tests, assuming all test names contain multi_gpu: $ pytest -sv ./tests -k "multi_gpu" + """ + if not is_torch_available(): + return unittest.skip(reason="test requires PyTorch")(test_case) + + import torch + + return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple CUDA GPUs")(test_case) + + +def require_torch_multi_accelerator(test_case): + """ + Decorator marking a test that requires a multi-accelerator (in PyTorch). These tests are skipped on a machine + without multiple accelerators. To run *only* the multi_accelerator tests, assuming all test names contain + multi_accelerator: $ pytest -sv ./tests -k "multi_accelerator" + """ + if not is_torch_available(): + return unittest.skip(reason="test requires PyTorch")(test_case) + + return unittest.skipUnless(backend_device_count(torch_device) > 1, "test requires multiple accelerators")( + test_case + ) + + +def require_torch_non_multi_gpu(test_case): + """ + Decorator marking a test that requires 0 or 1 GPU setup (in PyTorch). + """ + if not is_torch_available(): + return unittest.skip(reason="test requires PyTorch")(test_case) + + import torch + + return unittest.skipUnless(torch.cuda.device_count() < 2, "test requires 0 or 1 GPU")(test_case) + + +def require_torch_non_multi_accelerator(test_case): + """ + Decorator marking a test that requires 0 or 1 accelerator setup (in PyTorch). + """ + if not is_torch_available(): + return unittest.skip(reason="test requires PyTorch")(test_case) + + return unittest.skipUnless(backend_device_count(torch_device) < 2, "test requires 0 or 1 accelerator")(test_case) + + +def require_torch_up_to_2_gpus(test_case): + """ + Decorator marking a test that requires 0 or 1 or 2 GPU setup (in PyTorch). + """ + if not is_torch_available(): + return unittest.skip(reason="test requires PyTorch")(test_case) + + import torch + + return unittest.skipUnless(torch.cuda.device_count() < 3, "test requires 0 or 1 or 2 GPUs")(test_case) + + +def require_torch_up_to_2_accelerators(test_case): + """ + Decorator marking a test that requires 0 or 1 or 2 accelerator setup (in PyTorch). + """ + if not is_torch_available(): + return unittest.skip(reason="test requires PyTorch")(test_case) + + return unittest.skipUnless(backend_device_count(torch_device) < 3, "test requires 0 or 1 or 2 accelerators")( + test_case + ) + + +def require_torch_xla(test_case): + """ + Decorator marking a test that requires TorchXLA (in PyTorch). + """ + return unittest.skipUnless(is_torch_xla_available(), "test requires TorchXLA")(test_case) + + +def require_torch_neuroncore(test_case): + """ + Decorator marking a test that requires NeuronCore (in PyTorch). + """ + return unittest.skipUnless(is_torch_neuroncore_available(check_device=False), "test requires PyTorch NeuronCore")( + test_case + ) + + +def require_torch_npu(test_case): + """ + Decorator marking a test that requires NPU (in PyTorch). + """ + return unittest.skipUnless(is_torch_npu_available(), "test requires PyTorch NPU")(test_case) + + +def require_torch_multi_npu(test_case): + """ + Decorator marking a test that requires a multi-NPU setup (in PyTorch). These tests are skipped on a machine without + multiple NPUs. + + To run *only* the multi_npu tests, assuming all test names contain multi_npu: $ pytest -sv ./tests -k "multi_npu" + """ + if not is_torch_npu_available(): + return unittest.skip(reason="test requires PyTorch NPU")(test_case) + + return unittest.skipUnless(torch.npu.device_count() > 1, "test requires multiple NPUs")(test_case) + + +def require_non_hpu(test_case): + """ + Decorator marking a test that should be skipped for HPU. + """ + return unittest.skipUnless(torch_device != "hpu", "test requires a non-HPU")(test_case) + + +def require_torch_xpu(test_case): + """ + Decorator marking a test that requires XPU (in PyTorch). + + These tests are skipped when XPU backend is not available. XPU backend might be available either via stock + PyTorch (>=2.4) or via Intel Extension for PyTorch. In the latter case, if IPEX is installed, its version + must match match current PyTorch version. + """ + return unittest.skipUnless(is_torch_xpu_available(), "test requires XPU device")(test_case) + + +def require_non_xpu(test_case): + """ + Decorator marking a test that should be skipped for XPU. + """ + return unittest.skipUnless(torch_device != "xpu", "test requires a non-XPU")(test_case) + + +def require_torch_multi_xpu(test_case): + """ + Decorator marking a test that requires a multi-XPU setup (in PyTorch). These tests are skipped on a machine without + multiple XPUs. + + To run *only* the multi_xpu tests, assuming all test names contain multi_xpu: $ pytest -sv ./tests -k "multi_xpu" + """ + if not is_torch_xpu_available(): + return unittest.skip(reason="test requires PyTorch XPU")(test_case) + + return unittest.skipUnless(torch.xpu.device_count() > 1, "test requires multiple XPUs")(test_case) + + +def require_torch_multi_hpu(test_case): + """ + Decorator marking a test that requires a multi-HPU setup (in PyTorch). These tests are skipped on a machine without + multiple HPUs. + + To run *only* the multi_hpu tests, assuming all test names contain multi_hpu: $ pytest -sv ./tests -k "multi_hpu" + """ + if not is_torch_hpu_available(): + return unittest.skip(reason="test requires PyTorch HPU")(test_case) + + return unittest.skipUnless(torch.hpu.device_count() > 1, "test requires multiple HPUs")(test_case) + + +if is_torch_available(): + # Set env var CUDA_VISIBLE_DEVICES="" to force cpu-mode + import torch + + if "TRANSFORMERS_TEST_BACKEND" in os.environ: + backend = os.environ["TRANSFORMERS_TEST_BACKEND"] + try: + _ = importlib.import_module(backend) + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + f"Failed to import `TRANSFORMERS_TEST_BACKEND` '{backend}'! This should be the name of an installed module. The original error (look up to see its" + f" traceback):\n{e}" + ) from e + + if "TRANSFORMERS_TEST_DEVICE" in os.environ: + torch_device = os.environ["TRANSFORMERS_TEST_DEVICE"] + if torch_device == "cuda" and not torch.cuda.is_available(): + raise ValueError( + f"TRANSFORMERS_TEST_DEVICE={torch_device}, but CUDA is unavailable. Please double-check your testing environment." + ) + if torch_device == "xpu" and not is_torch_xpu_available(): + raise ValueError( + f"TRANSFORMERS_TEST_DEVICE={torch_device}, but XPU is unavailable. Please double-check your testing environment." + ) + if torch_device == "npu" and not is_torch_npu_available(): + raise ValueError( + f"TRANSFORMERS_TEST_DEVICE={torch_device}, but NPU is unavailable. Please double-check your testing environment." + ) + if torch_device == "mlu" and not is_torch_mlu_available(): + raise ValueError( + f"TRANSFORMERS_TEST_DEVICE={torch_device}, but MLU is unavailable. Please double-check your testing environment." + ) + if torch_device == "hpu" and not is_torch_hpu_available(): + raise ValueError( + f"TRANSFORMERS_TEST_DEVICE={torch_device}, but HPU is unavailable. Please double-check your testing environment." + ) + + try: + # try creating device to see if provided device is valid + _ = torch.device(torch_device) + except RuntimeError as e: + raise RuntimeError( + f"Unknown testing device specified by environment variable `TRANSFORMERS_TEST_DEVICE`: {torch_device}" + ) from e + elif torch.cuda.is_available(): + torch_device = "cuda" + elif _run_third_party_device_tests and is_torch_npu_available(): + torch_device = "npu" + elif _run_third_party_device_tests and is_torch_mlu_available(): + torch_device = "mlu" + elif _run_third_party_device_tests and is_torch_hpu_available(): + torch_device = "hpu" + elif _run_third_party_device_tests and is_torch_xpu_available(): + torch_device = "xpu" + else: + torch_device = "cpu" +else: + torch_device = None + +if is_tf_available(): + import tensorflow as tf + +if is_flax_available(): + import jax + + jax_device = jax.default_backend() +else: + jax_device = None + + +def require_torchdynamo(test_case): + """Decorator marking a test that requires TorchDynamo""" + return unittest.skipUnless(is_torchdynamo_available(), "test requires TorchDynamo")(test_case) + + +def require_torchao(test_case): + """Decorator marking a test that requires torchao""" + return unittest.skipUnless(is_torchao_available(), "test requires torchao")(test_case) + + +def require_torchao_version_greater_or_equal(torchao_version): + def decorator(test_case): + correct_torchao_version = is_torchao_available() and version.parse( + version.parse(importlib.metadata.version("torchao")).base_version + ) >= version.parse(torchao_version) + return unittest.skipUnless( + correct_torchao_version, f"Test requires torchao with the version greater than {torchao_version}." + )(test_case) + + return decorator + + +def require_torch_tensorrt_fx(test_case): + """Decorator marking a test that requires Torch-TensorRT FX""" + return unittest.skipUnless(is_torch_tensorrt_fx_available(), "test requires Torch-TensorRT FX")(test_case) + + +def require_torch_gpu(test_case): + """Decorator marking a test that requires CUDA and PyTorch.""" + return unittest.skipUnless(torch_device == "cuda", "test requires CUDA")(test_case) + + +def require_torch_large_gpu(test_case, memory: float = 20): + """Decorator marking a test that requires a CUDA GPU with more than `memory` GiB of memory.""" + if torch_device != "cuda": + return unittest.skip(reason=f"test requires a CUDA GPU with more than {memory} GiB of memory")(test_case) + + return unittest.skipUnless( + torch.cuda.get_device_properties(0).total_memory / 1024**3 > memory, + f"test requires a GPU with more than {memory} GiB of memory", + )(test_case) + + +def require_torch_gpu_if_bnb_not_multi_backend_enabled(test_case): + """ + Decorator marking a test that requires a GPU if bitsandbytes multi-backend feature is not enabled. + """ + if is_bitsandbytes_available() and is_bitsandbytes_multi_backend_available(): + return test_case + return require_torch_gpu(test_case) + + +def require_torch_accelerator(test_case): + """Decorator marking a test that requires an accessible accelerator and PyTorch.""" + return unittest.skipUnless(torch_device is not None and torch_device != "cpu", "test requires accelerator")( + test_case + ) + + +def require_torch_fp16(test_case): + """Decorator marking a test that requires a device that supports fp16""" + return unittest.skipUnless( + is_torch_fp16_available_on_device(torch_device), "test requires device with fp16 support" + )(test_case) + + +def require_fp8(test_case): + """Decorator marking a test that requires supports for fp8""" + return unittest.skipUnless(is_accelerate_available() and is_fp8_available(), "test requires fp8 support")( + test_case + ) + + +def require_torch_bf16(test_case): + """Decorator marking a test that requires a device that supports bf16""" + return unittest.skipUnless( + is_torch_bf16_available_on_device(torch_device), "test requires device with bf16 support" + )(test_case) + + +def require_torch_bf16_gpu(test_case): + """Decorator marking a test that requires torch>=1.10, using Ampere GPU or newer arch with cuda>=11.0""" + return unittest.skipUnless( + is_torch_bf16_gpu_available(), + "test requires torch>=1.10, using Ampere GPU or newer arch with cuda>=11.0", + )(test_case) + + +def require_torch_bf16_cpu(test_case): + """Decorator marking a test that requires torch>=1.10, using CPU.""" + return unittest.skipUnless( + is_torch_bf16_cpu_available(), + "test requires torch>=1.10, using CPU", + )(test_case) + + +def require_deterministic_for_xpu(test_case): + if is_torch_xpu_available(): + return unittest.skipUnless(is_torch_deterministic(), "test requires torch to use deterministic algorithms")( + test_case + ) + else: + return test_case + + +def require_torch_tf32(test_case): + """Decorator marking a test that requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7.""" + return unittest.skipUnless( + is_torch_tf32_available(), "test requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7" + )(test_case) + + +def require_detectron2(test_case): + """Decorator marking a test that requires detectron2.""" + return unittest.skipUnless(is_detectron2_available(), "test requires `detectron2`")(test_case) + + +def require_faiss(test_case): + """Decorator marking a test that requires faiss.""" + return unittest.skipUnless(is_faiss_available(), "test requires `faiss`")(test_case) + + +def require_optuna(test_case): + """ + Decorator marking a test that requires optuna. + + These tests are skipped when optuna isn't installed. + + """ + return unittest.skipUnless(is_optuna_available(), "test requires optuna")(test_case) + + +def require_ray(test_case): + """ + Decorator marking a test that requires Ray/tune. + + These tests are skipped when Ray/tune isn't installed. + + """ + return unittest.skipUnless(is_ray_available(), "test requires Ray/tune")(test_case) + + +def require_sigopt(test_case): + """ + Decorator marking a test that requires SigOpt. + + These tests are skipped when SigOpt isn't installed. + + """ + return unittest.skipUnless(is_sigopt_available(), "test requires SigOpt")(test_case) + + +def require_swanlab(test_case): + """ + Decorator marking a test that requires swanlab. + + These tests are skipped when swanlab isn't installed. + + """ + return unittest.skipUnless(is_swanlab_available(), "test requires swanlab")(test_case) + + +def require_wandb(test_case): + """ + Decorator marking a test that requires wandb. + + These tests are skipped when wandb isn't installed. + + """ + return unittest.skipUnless(is_wandb_available(), "test requires wandb")(test_case) + + +def require_clearml(test_case): + """ + Decorator marking a test requires clearml. + + These tests are skipped when clearml isn't installed. + + """ + return unittest.skipUnless(is_clearml_available(), "test requires clearml")(test_case) + + +def require_soundfile(test_case): + """ + Decorator marking a test that requires soundfile + + These tests are skipped when soundfile isn't installed. + + """ + return unittest.skipUnless(is_soundfile_available(), "test requires soundfile")(test_case) + + +def require_deepspeed(test_case): + """ + Decorator marking a test that requires deepspeed + """ + return unittest.skipUnless(is_deepspeed_available(), "test requires deepspeed")(test_case) + + +def require_apex(test_case): + """ + Decorator marking a test that requires apex + """ + return unittest.skipUnless(is_apex_available(), "test requires apex")(test_case) + + +def require_aqlm(test_case): + """ + Decorator marking a test that requires aqlm + """ + return unittest.skipUnless(is_aqlm_available(), "test requires aqlm")(test_case) + + +def require_vptq(test_case): + """ + Decorator marking a test that requires vptq + """ + return unittest.skipUnless(is_vptq_available(), "test requires vptq")(test_case) + + +def require_spqr(test_case): + """ + Decorator marking a test that requires spqr + """ + return unittest.skipUnless(is_spqr_available(), "test requires spqr")(test_case) + + +def require_eetq(test_case): + """ + Decorator marking a test that requires eetq + """ + eetq_available = is_eetq_available() + if eetq_available: + try: + import eetq # noqa: F401 + except ImportError as exc: + if "shard_checkpoint" in str(exc): + # EETQ 1.0.0 is currently broken with the latest transformers because it tries to import the removed + # shard_checkpoint function, see https://github.com/NetEase-FuXi/EETQ/issues/34. + # TODO: Remove once eetq releases a fix and this release is used in CI + eetq_available = False + return unittest.skipUnless(eetq_available, "test requires eetq")(test_case) + + +def require_av(test_case): + """ + Decorator marking a test that requires av + """ + return unittest.skipUnless(is_av_available(), "test requires av")(test_case) + + +def require_bitsandbytes(test_case): + """ + Decorator marking a test that requires the bitsandbytes library. Will be skipped when the library or its hard dependency torch is not installed. + """ + if is_bitsandbytes_available() and is_torch_available(): + try: + import pytest + + return pytest.mark.bitsandbytes(test_case) + except ImportError: + return test_case + else: + return unittest.skip(reason="test requires bitsandbytes and torch")(test_case) + + +def require_optimum(test_case): + """ + Decorator for optimum dependency + """ + return unittest.skipUnless(is_optimum_available(), "test requires optimum")(test_case) + + +def require_tensorboard(test_case): + """ + Decorator for `tensorboard` dependency + """ + return unittest.skipUnless(is_tensorboard_available(), "test requires tensorboard") + + +def require_gptq(test_case): + """ + Decorator for auto_gptq dependency + """ + return unittest.skipUnless( + is_gptqmodel_available() or is_auto_gptq_available(), "test requires gptqmodel or auto-gptq" + )(test_case) + + +def require_hqq(test_case): + """ + Decorator for hqq dependency + """ + return unittest.skipUnless(is_hqq_available(), "test requires hqq")(test_case) + + +def require_auto_awq(test_case): + """ + Decorator for auto_awq dependency + """ + return unittest.skipUnless(is_auto_awq_available(), "test requires autoawq")(test_case) + + +def require_optimum_quanto(test_case): + """ + Decorator for quanto dependency + """ + return unittest.skipUnless(is_optimum_quanto_available(), "test requires optimum-quanto")(test_case) + + +def require_compressed_tensors(test_case): + """ + Decorator for compressed_tensors dependency + """ + return unittest.skipUnless(is_compressed_tensors_available(), "test requires compressed_tensors")(test_case) + + +def require_fbgemm_gpu(test_case): + """ + Decorator for fbgemm_gpu dependency + """ + return unittest.skipUnless(is_fbgemm_gpu_available(), "test requires fbgemm-gpu")(test_case) + + +def require_quark(test_case): + """ + Decorator for quark dependency + """ + return unittest.skipUnless(is_quark_available(), "test requires quark")(test_case) + + +def require_flute_hadamard(test_case): + """ + Decorator marking a test that requires higgs and hadamard + """ + return unittest.skipUnless( + is_flute_available() and is_hadamard_available(), "test requires flute and fast_hadamard_transform" + )(test_case) + + +def require_phonemizer(test_case): + """ + Decorator marking a test that requires phonemizer + """ + return unittest.skipUnless(is_phonemizer_available(), "test requires phonemizer")(test_case) + + +def require_pyctcdecode(test_case): + """ + Decorator marking a test that requires pyctcdecode + """ + return unittest.skipUnless(is_pyctcdecode_available(), "test requires pyctcdecode")(test_case) + + +def require_librosa(test_case): + """ + Decorator marking a test that requires librosa + """ + return unittest.skipUnless(is_librosa_available(), "test requires librosa")(test_case) + + +def require_liger_kernel(test_case): + """ + Decorator marking a test that requires liger_kernel + """ + return unittest.skipUnless(is_liger_kernel_available(), "test requires liger_kernel")(test_case) + + +def require_essentia(test_case): + """ + Decorator marking a test that requires essentia + """ + return unittest.skipUnless(is_essentia_available(), "test requires essentia")(test_case) + + +def require_pretty_midi(test_case): + """ + Decorator marking a test that requires pretty_midi + """ + return unittest.skipUnless(is_pretty_midi_available(), "test requires pretty_midi")(test_case) + + +def cmd_exists(cmd): + return shutil.which(cmd) is not None + + +def require_usr_bin_time(test_case): + """ + Decorator marking a test that requires `/usr/bin/time` + """ + return unittest.skipUnless(cmd_exists("/usr/bin/time"), "test requires /usr/bin/time")(test_case) + + +def require_sudachi(test_case): + """ + Decorator marking a test that requires sudachi + """ + return unittest.skipUnless(is_sudachi_available(), "test requires sudachi")(test_case) + + +def require_sudachi_projection(test_case): + """ + Decorator marking a test that requires sudachi_projection + """ + return unittest.skipUnless(is_sudachi_projection_available(), "test requires sudachi which supports projection")( + test_case + ) + + +def require_jumanpp(test_case): + """ + Decorator marking a test that requires jumanpp + """ + return unittest.skipUnless(is_jumanpp_available(), "test requires jumanpp")(test_case) + + +def require_cython(test_case): + """ + Decorator marking a test that requires jumanpp + """ + return unittest.skipUnless(is_cython_available(), "test requires cython")(test_case) + + +def require_tiktoken(test_case): + """ + Decorator marking a test that requires TikToken. These tests are skipped when TikToken isn't installed. + """ + return unittest.skipUnless(is_tiktoken_available(), "test requires TikToken")(test_case) + + +def get_gpu_count(): + """ + Return the number of available gpus (regardless of whether torch, tf or jax is used) + """ + if is_torch_available(): + import torch + + return torch.cuda.device_count() + elif is_tf_available(): + import tensorflow as tf + + return len(tf.config.list_physical_devices("GPU")) + elif is_flax_available(): + import jax + + return jax.device_count() + else: + return 0 + + +def get_tests_dir(append_path=None): + """ + Args: + append_path: optional path to append to the tests dir path + + Return: + The full path to the `tests` dir, so that the tests can be invoked from anywhere. Optionally `append_path` is + joined after the `tests` dir the former is provided. + + """ + # this function caller's __file__ + caller__file__ = inspect.stack()[1][1] + tests_dir = os.path.abspath(os.path.dirname(caller__file__)) + + while not tests_dir.endswith("tests"): + tests_dir = os.path.dirname(tests_dir) + + if append_path: + return os.path.join(tests_dir, append_path) + else: + return tests_dir + + +def get_steps_per_epoch(trainer: Trainer) -> int: + training_args = trainer.args + train_dataloader = trainer.get_train_dataloader() + + initial_training_values = trainer.set_initial_training_values( + args=training_args, + dataloader=train_dataloader, + total_train_batch_size=training_args.per_device_train_batch_size, + ) + steps_per_epoch = initial_training_values[1] + + return steps_per_epoch + + +def evaluate_side_effect_factory( + side_effect_values: list[dict[str, float]], +) -> Generator[dict[str, float], None, None]: + """ + Function that returns side effects for the _evaluate method. + Used when we're unsure of exactly how many times _evaluate will be called. + """ + yield from side_effect_values + + while True: + yield side_effect_values[-1] + + +# +# Helper functions for dealing with testing text outputs +# The original code came from: +# https://github.com/fastai/fastai/blob/master/tests/utils/text.py + + +# When any function contains print() calls that get overwritten, like progress bars, +# a special care needs to be applied, since under pytest -s captured output (capsys +# or contextlib.redirect_stdout) contains any temporary printed strings, followed by +# \r's. This helper function ensures that the buffer will contain the same output +# with and without -s in pytest, by turning: +# foo bar\r tar mar\r final message +# into: +# final message +# it can handle a single string or a multiline buffer +def apply_print_resets(buf): + return re.sub(r"^.*\r", "", buf, 0, re.M) + + +def assert_screenout(out, what): + out_pr = apply_print_resets(out).lower() + match_str = out_pr.find(what.lower()) + assert match_str != -1, f"expecting to find {what} in output: f{out_pr}" + + +def set_model_tester_for_less_flaky_test(test_case): + target_num_hidden_layers = 1 + # TODO (if possible): Avoid exceptional cases + exceptional_classes = [ + "ZambaModelTester", + "Zamba2ModelTester", + "RwkvModelTester", + "AriaVisionText2TextModelTester", + "GPTNeoModelTester", + "DPTModelTester", + ] + if test_case.model_tester.__class__.__name__ in exceptional_classes: + target_num_hidden_layers = None + if hasattr(test_case.model_tester, "out_features") or hasattr(test_case.model_tester, "out_indices"): + target_num_hidden_layers = None + + if hasattr(test_case.model_tester, "num_hidden_layers") and target_num_hidden_layers is not None: + test_case.model_tester.num_hidden_layers = target_num_hidden_layers + if ( + hasattr(test_case.model_tester, "vision_config") + and "num_hidden_layers" in test_case.model_tester.vision_config + and target_num_hidden_layers is not None + ): + test_case.model_tester.vision_config = copy.deepcopy(test_case.model_tester.vision_config) + if isinstance(test_case.model_tester.vision_config, dict): + test_case.model_tester.vision_config["num_hidden_layers"] = 1 + else: + test_case.model_tester.vision_config.num_hidden_layers = 1 + if ( + hasattr(test_case.model_tester, "text_config") + and "num_hidden_layers" in test_case.model_tester.text_config + and target_num_hidden_layers is not None + ): + test_case.model_tester.text_config = copy.deepcopy(test_case.model_tester.text_config) + if isinstance(test_case.model_tester.text_config, dict): + test_case.model_tester.text_config["num_hidden_layers"] = 1 + else: + test_case.model_tester.text_config.num_hidden_layers = 1 + + # A few model class specific handling + + # For Albert + if hasattr(test_case.model_tester, "num_hidden_groups"): + test_case.model_tester.num_hidden_groups = test_case.model_tester.num_hidden_layers + + +def set_config_for_less_flaky_test(config): + target_attrs = [ + "rms_norm_eps", + "layer_norm_eps", + "norm_eps", + "norm_epsilon", + "layer_norm_epsilon", + "batch_norm_eps", + ] + for target_attr in target_attrs: + setattr(config, target_attr, 1.0) + + # norm layers (layer/group norm, etc.) could cause flaky tests when the tensors have very small variance. + # (We don't need the original epsilon values to check eager/sdpa matches) + attrs = ["text_config", "vision_config", "text_encoder", "audio_encoder", "decoder"] + for attr in attrs: + if hasattr(config, attr): + for target_attr in target_attrs: + setattr(getattr(config, attr), target_attr, 1.0) + + +def set_model_for_less_flaky_test(model): + # Another way to make sure norm layers have desired epsilon. (Some models don't set it from its config.) + target_names = ( + "LayerNorm", + "GroupNorm", + "BatchNorm", + "RMSNorm", + "BatchNorm2d", + "BatchNorm1d", + "BitGroupNormActivation", + "WeightStandardizedConv2d", + ) + target_attrs = ["eps", "epsilon", "variance_epsilon"] + if is_torch_available() and isinstance(model, torch.nn.Module): + for module in model.modules(): + if type(module).__name__.endswith(target_names): + for attr in target_attrs: + if hasattr(module, attr): + setattr(module, attr, 1.0) + + +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 + + +# in tests it's the best to capture only the stream that's wanted, otherwise +# it's easy to miss things, so unless you need to capture both streams, use the +# subclasses below (less typing). Or alternatively, configure `CaptureStd` to +# disable the stream you don't need to test. + + +class CaptureStdout(CaptureStd): + """Same as CaptureStd but captures only stdout""" + + def __init__(self, replay=True): + super().__init__(err=False, replay=replay) + + +class CaptureStderr(CaptureStd): + """Same as CaptureStd but captures only stderr""" + + def __init__(self, replay=True): + super().__init__(out=False, replay=replay) + + +class CaptureLogger: + """ + Context manager to capture `logging` streams + + Args: + logger: 'logging` logger object + + Returns: + The captured output is available via `self.out` + + Example: + + ```python + >>> from transformers import logging + >>> from transformers.testing_utils import CaptureLogger + + >>> msg = "Testing 1, 2, 3" + >>> logging.set_verbosity_info() + >>> logger = logging.get_logger("transformers.models.bart.tokenization_bart") + >>> with CaptureLogger(logger) as cl: + ... logger.info(msg) + >>> assert cl.out, msg + "\n" + ``` + """ + + def __init__(self, logger): + self.logger = logger + self.io = StringIO() + self.sh = logging.StreamHandler(self.io) + self.out = "" + + def __enter__(self): + self.logger.addHandler(self.sh) + return self + + def __exit__(self, *exc): + self.logger.removeHandler(self.sh) + self.out = self.io.getvalue() + + def __repr__(self): + return f"captured: {self.out}\n" + + +@contextlib.contextmanager +def LoggingLevel(level): + """ + This is a context manager to temporarily change transformers modules logging level to the desired value and have it + restored to the original setting at the end of the scope. + + Example: + + ```python + with LoggingLevel(logging.INFO): + AutoModel.from_pretrained("openai-community/gpt2") # calls logger.info() several times + ``` + """ + orig_level = transformers_logging.get_verbosity() + try: + transformers_logging.set_verbosity(level) + yield + finally: + transformers_logging.set_verbosity(orig_level) + + +class TemporaryHubRepo: + """Create a temporary Hub repository and return its `RepoUrl` object. This is similar to + `tempfile.TemporaryDirectory` and can be used as a context manager. For example: + + with TemporaryHubRepo(token=self._token) as temp_repo: + ... + + Upon exiting the context, the repository and everything contained in it are removed. + + Example: + + ```python + with TemporaryHubRepo(token=self._token) as temp_repo: + model.push_to_hub(tmp_repo.repo_id, token=self._token) + ``` + """ + + def __init__(self, namespace: Optional[str] = None, token: Optional[str] = None) -> None: + self.token = token + with tempfile.TemporaryDirectory() as tmp_dir: + repo_id = Path(tmp_dir).name + if namespace is not None: + repo_id = f"{namespace}/{repo_id}" + self.repo_url = huggingface_hub.create_repo(repo_id, token=self.token) + + def __enter__(self): + return self.repo_url + + def __exit__(self, exc, value, tb): + delete_repo(repo_id=self.repo_url.repo_id, token=self.token, missing_ok=True) + + +@contextlib.contextmanager +# adapted from https://stackoverflow.com/a/64789046/9201239 +def ExtendSysPath(path: Union[str, os.PathLike]) -> Iterator[None]: + """ + Temporary add given path to `sys.path`. + + Usage : + + ```python + with ExtendSysPath("/path/to/dir"): + mymodule = importlib.import_module("mymodule") + ``` + """ + + path = os.fspath(path) + try: + sys.path.insert(0, path) + yield + finally: + sys.path.remove(path) + + +class TestCasePlus(unittest.TestCase): + """ + This class extends *unittest.TestCase* with additional features. + + Feature 1: A set of fully resolved important file and dir path accessors. + + In tests often we need to know where things are relative to the current test file, and it's not trivial since the + test could be invoked from more than one directory or could reside in sub-directories with different depths. This + class solves this problem by sorting out all the basic paths and provides easy accessors to them: + + - `pathlib` objects (all fully resolved): + + - `test_file_path` - the current test file path (=`__file__`) + - `test_file_dir` - the directory containing the current test file + - `tests_dir` - the directory of the `tests` test suite + - `examples_dir` - the directory of the `examples` test suite + - `repo_root_dir` - the directory of the repository + - `src_dir` - the directory of `src` (i.e. where the `transformers` sub-dir resides) + + - stringified paths---same as above but these return paths as strings, rather than `pathlib` objects: + + - `test_file_path_str` + - `test_file_dir_str` + - `tests_dir_str` + - `examples_dir_str` + - `repo_root_dir_str` + - `src_dir_str` + + Feature 2: Flexible auto-removable temporary dirs which are guaranteed to get removed at the end of test. + + 1. Create a unique temporary dir: + + ```python + def test_whatever(self): + tmp_dir = self.get_auto_remove_tmp_dir() + ``` + + `tmp_dir` will contain the path to the created temporary dir. It will be automatically removed at the end of the + test. + + + 2. Create a temporary dir of my choice, ensure it's empty before the test starts and don't + empty it after the test. + + ```python + def test_whatever(self): + tmp_dir = self.get_auto_remove_tmp_dir("./xxx") + ``` + + This is useful for debug when you want to monitor a specific directory and want to make sure the previous tests + didn't leave any data in there. + + 3. You can override the first two options by directly overriding the `before` and `after` args, leading to the + following behavior: + + `before=True`: the temporary dir will always be cleared at the beginning of the test. + + `before=False`: if the temporary dir already existed, any existing files will remain there. + + `after=True`: the temporary dir will always be deleted at the end of the test. + + `after=False`: the temporary dir will always be left intact at the end of the test. + + Note 1: In order to run the equivalent of `rm -r` safely, only subdirs of the project repository checkout are + allowed if an explicit `tmp_dir` is used, so that by mistake no `/tmp` or similar important part of the filesystem + will get nuked. i.e. please always pass paths that start with `./` + + Note 2: Each test can register multiple temporary dirs and they all will get auto-removed, unless requested + otherwise. + + Feature 3: Get a copy of the `os.environ` object that sets up `PYTHONPATH` specific to the current test suite. This + is useful for invoking external programs from the test suite - e.g. distributed training. + + + ```python + def test_whatever(self): + env = self.get_env() + ```""" + + def setUp(self): + # get_auto_remove_tmp_dir feature: + self.teardown_tmp_dirs = [] + + # figure out the resolved paths for repo_root, tests, examples, etc. + self._test_file_path = inspect.getfile(self.__class__) + path = Path(self._test_file_path).resolve() + self._test_file_dir = path.parents[0] + for up in [1, 2, 3]: + tmp_dir = path.parents[up] + if (tmp_dir / "src").is_dir() and (tmp_dir / "tests").is_dir(): + break + if tmp_dir: + self._repo_root_dir = tmp_dir + else: + raise ValueError(f"can't figure out the root of the repo from {self._test_file_path}") + self._tests_dir = self._repo_root_dir / "tests" + self._examples_dir = self._repo_root_dir / "examples" + self._src_dir = self._repo_root_dir / "src" + + @property + def test_file_path(self): + return self._test_file_path + + @property + def test_file_path_str(self): + return str(self._test_file_path) + + @property + def test_file_dir(self): + return self._test_file_dir + + @property + def test_file_dir_str(self): + return str(self._test_file_dir) + + @property + def tests_dir(self): + return self._tests_dir + + @property + def tests_dir_str(self): + return str(self._tests_dir) + + @property + def examples_dir(self): + return self._examples_dir + + @property + def examples_dir_str(self): + return str(self._examples_dir) + + @property + def repo_root_dir(self): + return self._repo_root_dir + + @property + def repo_root_dir_str(self): + return str(self._repo_root_dir) + + @property + def src_dir(self): + return self._src_dir + + @property + def src_dir_str(self): + return str(self._src_dir) + + def get_env(self): + """ + Return a copy of the `os.environ` object that sets up `PYTHONPATH` correctly, depending on the test suite it's + invoked from. This is useful for invoking external programs from the test suite - e.g. distributed training. + + It always inserts `./src` first, then `./tests` or `./examples` depending on the test suite type and finally + the preset `PYTHONPATH` if any (all full resolved paths). + + """ + env = os.environ.copy() + paths = [self.src_dir_str] + if "/examples" in self.test_file_dir_str: + paths.append(self.examples_dir_str) + else: + paths.append(self.tests_dir_str) + paths.append(env.get("PYTHONPATH", "")) + + env["PYTHONPATH"] = ":".join(paths) + return env + + def get_auto_remove_tmp_dir(self, tmp_dir=None, before=None, after=None): + """ + Args: + tmp_dir (`string`, *optional*): + if `None`: + + - a unique temporary path will be created + - sets `before=True` if `before` is `None` + - sets `after=True` if `after` is `None` + else: + + - `tmp_dir` will be created + - sets `before=True` if `before` is `None` + - sets `after=False` if `after` is `None` + before (`bool`, *optional*): + If `True` and the `tmp_dir` already exists, make sure to empty it right away if `False` and the + `tmp_dir` already exists, any existing files will remain there. + after (`bool`, *optional*): + If `True`, delete the `tmp_dir` at the end of the test if `False`, leave the `tmp_dir` and its contents + intact at the end of the test. + + Returns: + tmp_dir(`string`): either the same value as passed via *tmp_dir* or the path to the auto-selected tmp dir + """ + if tmp_dir is not None: + # defining the most likely desired behavior for when a custom path is provided. + # this most likely indicates the debug mode where we want an easily locatable dir that: + # 1. gets cleared out before the test (if it already exists) + # 2. is left intact after the test + if before is None: + before = True + if after is None: + after = False + + # using provided path + path = Path(tmp_dir).resolve() + + # to avoid nuking parts of the filesystem, only relative paths are allowed + if not tmp_dir.startswith("./"): + raise ValueError( + f"`tmp_dir` can only be a relative path, i.e. `./some/path`, but received `{tmp_dir}`" + ) + + # ensure the dir is empty to start with + if before is True and path.exists(): + shutil.rmtree(tmp_dir, ignore_errors=True) + + path.mkdir(parents=True, exist_ok=True) + + else: + # defining the most likely desired behavior for when a unique tmp path is auto generated + # (not a debug mode), here we require a unique tmp dir that: + # 1. is empty before the test (it will be empty in this situation anyway) + # 2. gets fully removed after the test + if before is None: + before = True + if after is None: + after = True + + # using unique tmp dir (always empty, regardless of `before`) + tmp_dir = tempfile.mkdtemp() + + if after is True: + # register for deletion + self.teardown_tmp_dirs.append(tmp_dir) + + return tmp_dir + + def python_one_liner_max_rss(self, one_liner_str): + """ + Runs the passed python one liner (just the code) and returns how much max cpu memory was used to run the + program. + + Args: + one_liner_str (`string`): + a python one liner code that gets passed to `python -c` + + Returns: + max cpu memory bytes used to run the program. This value is likely to vary slightly from run to run. + + Requirements: + this helper needs `/usr/bin/time` to be installed (`apt install time`) + + Example: + + ``` + one_liner_str = 'from transformers import AutoModel; AutoModel.from_pretrained("google-t5/t5-large")' + max_rss = self.python_one_liner_max_rss(one_liner_str) + ``` + """ + + if not cmd_exists("/usr/bin/time"): + raise ValueError("/usr/bin/time is required, install with `apt install time`") + + cmd = shlex.split(f"/usr/bin/time -f %M python -c '{one_liner_str}'") + with CaptureStd() as cs: + execute_subprocess_async(cmd, env=self.get_env()) + # returned data is in KB so convert to bytes + max_rss = int(cs.err.split("\n")[-2].replace("stderr: ", "")) * 1024 + return max_rss + + def tearDown(self): + # get_auto_remove_tmp_dir feature: remove registered temp dirs + for path in self.teardown_tmp_dirs: + shutil.rmtree(path, ignore_errors=True) + self.teardown_tmp_dirs = [] + if is_accelerate_available(): + AcceleratorState._reset_state() + PartialState._reset_state() + + # delete all the env variables having `ACCELERATE` in them + for k in list(os.environ.keys()): + if "ACCELERATE" in k: + del os.environ[k] + + +def mockenv(**kwargs): + """ + this is a convenience wrapper, that allows this :: + + @mockenv(RUN_SLOW=True, USE_TF=False) def test_something(): + run_slow = os.getenv("RUN_SLOW", False) use_tf = os.getenv("USE_TF", False) + + """ + return mock.patch.dict(os.environ, kwargs) + + +# from https://stackoverflow.com/a/34333710/9201239 +@contextlib.contextmanager +def mockenv_context(*remove, **update): + """ + Temporarily updates the `os.environ` dictionary in-place. Similar to mockenv + + The `os.environ` dictionary is updated in-place so that the modification is sure to work in all situations. + + Args: + remove: Environment variables to remove. + update: Dictionary of environment variables and values to add/update. + """ + env = os.environ + update = update or {} + remove = remove or [] + + # List of environment variables being updated or removed. + stomped = (set(update.keys()) | set(remove)) & set(env.keys()) + # Environment variables and values to restore on exit. + update_after = {k: env[k] for k in stomped} + # Environment variables and values to remove on exit. + remove_after = frozenset(k for k in update if k not in env) + + try: + env.update(update) + [env.pop(k, None) for k in remove] + yield + finally: + env.update(update_after) + [env.pop(k) for k in remove_after] + + +# --- pytest conf functions --- # + +# to avoid multiple invocation from tests/conftest.py and examples/conftest.py - make sure it's called only once +pytest_opt_registered = {} + + +def pytest_addoption_shared(parser): + """ + This function is to be called from `conftest.py` via `pytest_addoption` wrapper that has to be defined there. + + It allows loading both `conftest.py` files at once without causing a failure due to adding the same `pytest` + option. + + """ + option = "--make-reports" + if option not in pytest_opt_registered: + parser.addoption( + option, + action="store", + default=False, + help="generate report files. The value of this option is used as a prefix to report names", + ) + pytest_opt_registered[option] = 1 + + +def pytest_terminal_summary_main(tr, id): + """ + Generate multiple reports at the end of test suite run - each report goes into a dedicated file in the current + directory. The report files are prefixed with the test suite name. + + This function emulates --duration and -rA pytest arguments. + + This function is to be called from `conftest.py` via `pytest_terminal_summary` wrapper that has to be defined + there. + + Args: + - tr: `terminalreporter` passed from `conftest.py` + - id: unique id like `tests` or `examples` that will be incorporated into the final reports filenames - this is + needed as some jobs have multiple runs of pytest, so we can't have them overwrite each other. + + NB: this functions taps into a private _pytest API and while unlikely, it could break should pytest do internal + changes - also it calls default internal methods of terminalreporter which can be hijacked by various `pytest-` + plugins and interfere. + + """ + from _pytest.config import create_terminal_writer + + if not len(id): + id = "tests" + + config = tr.config + orig_writer = config.get_terminal_writer() + orig_tbstyle = config.option.tbstyle + orig_reportchars = tr.reportchars + + dir = f"reports/{id}" + Path(dir).mkdir(parents=True, exist_ok=True) + report_files = { + k: f"{dir}/{k}.txt" + for k in [ + "durations", + "errors", + "failures_long", + "failures_short", + "failures_line", + "passes", + "stats", + "summary_short", + "warnings", + ] + } + + # custom durations report + # note: there is no need to call pytest --durations=XX to get this separate report + # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/runner.py#L66 + dlist = [] + for replist in tr.stats.values(): + for rep in replist: + if hasattr(rep, "duration"): + dlist.append(rep) + if dlist: + dlist.sort(key=lambda x: x.duration, reverse=True) + with open(report_files["durations"], "w") as f: + durations_min = 0.05 # sec + f.write("slowest durations\n") + for i, rep in enumerate(dlist): + if rep.duration < durations_min: + f.write(f"{len(dlist) - i} durations < {durations_min} secs were omitted") + break + f.write(f"{rep.duration:02.2f}s {rep.when:<8} {rep.nodeid}\n") + + def summary_failures_short(tr): + # expecting that the reports were --tb=long (default) so we chop them off here to the last frame + reports = tr.getreports("failed") + if not reports: + return + tr.write_sep("=", "FAILURES SHORT STACK") + for rep in reports: + msg = tr._getfailureheadline(rep) + tr.write_sep("_", msg, red=True, bold=True) + # chop off the optional leading extra frames, leaving only the last one + longrepr = re.sub(r".*_ _ _ (_ ){10,}_ _ ", "", rep.longreprtext, 0, re.M | re.S) + tr._tw.line(longrepr) + # note: not printing out any rep.sections to keep the report short + + # use ready-made report funcs, we are just hijacking the filehandle to log to a dedicated file each + # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/terminal.py#L814 + # note: some pytest plugins may interfere by hijacking the default `terminalreporter` (e.g. + # pytest-instafail does that) + + # report failures with line/short/long styles + config.option.tbstyle = "auto" # full tb + with open(report_files["failures_long"], "w") as f: + tr._tw = create_terminal_writer(config, f) + tr.summary_failures() + + # config.option.tbstyle = "short" # short tb + with open(report_files["failures_short"], "w") as f: + tr._tw = create_terminal_writer(config, f) + summary_failures_short(tr) + + config.option.tbstyle = "line" # one line per error + with open(report_files["failures_line"], "w") as f: + tr._tw = create_terminal_writer(config, f) + tr.summary_failures() + + with open(report_files["errors"], "w") as f: + tr._tw = create_terminal_writer(config, f) + tr.summary_errors() + + with open(report_files["warnings"], "w") as f: + tr._tw = create_terminal_writer(config, f) + tr.summary_warnings() # normal warnings + tr.summary_warnings() # final warnings + + tr.reportchars = "wPpsxXEf" # emulate -rA (used in summary_passes() and short_test_summary()) + + # Skip the `passes` report, as it starts to take more than 5 minutes, and sometimes it timeouts on CircleCI if it + # takes > 10 minutes (as this part doesn't generate any output on the terminal). + # (also, it seems there is no useful information in this report, and we rarely need to read it) + # with open(report_files["passes"], "w") as f: + # tr._tw = create_terminal_writer(config, f) + # tr.summary_passes() + + with open(report_files["summary_short"], "w") as f: + tr._tw = create_terminal_writer(config, f) + tr.short_test_summary() + + with open(report_files["stats"], "w") as f: + tr._tw = create_terminal_writer(config, f) + tr.summary_stats() + + # restore: + tr._tw = orig_writer + tr.reportchars = orig_reportchars + config.option.tbstyle = orig_tbstyle + + +# --- distributed testing functions --- # + +# adapted from https://stackoverflow.com/a/59041913/9201239 +import asyncio # noqa + + +class _RunOutput: + def __init__(self, returncode, stdout, stderr): + self.returncode = returncode + self.stdout = stdout + self.stderr = stderr + + +async def _read_stream(stream, callback): + while True: + line = await stream.readline() + if line: + callback(line) + else: + break + + +async def _stream_subprocess(cmd, env=None, stdin=None, timeout=None, quiet=False, echo=False) -> _RunOutput: + if echo: + print("\nRunning: ", " ".join(cmd)) + + p = await asyncio.create_subprocess_exec( + cmd[0], + *cmd[1:], + stdin=stdin, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + env=env, + ) + + # note: there is a warning for a possible deadlock when using `wait` with huge amounts of data in the pipe + # https://docs.python.org/3/library/asyncio-subprocess.html#asyncio.asyncio.subprocess.Process.wait + # + # If it starts hanging, will need to switch to the following code. The problem is that no data + # will be seen until it's done and if it hangs for example there will be no debug info. + # out, err = await p.communicate() + # return _RunOutput(p.returncode, out, err) + + out = [] + err = [] + + def tee(line, sink, pipe, label=""): + line = line.decode("utf-8").rstrip() + sink.append(line) + if not quiet: + print(label, line, file=pipe) + + # XXX: the timeout doesn't seem to make any difference here + await asyncio.wait( + [ + asyncio.create_task(_read_stream(p.stdout, lambda l: tee(l, out, sys.stdout, label="stdout:"))), + asyncio.create_task(_read_stream(p.stderr, lambda l: tee(l, err, sys.stderr, label="stderr:"))), + ], + timeout=timeout, + ) + return _RunOutput(await p.wait(), out, err) + + +def execute_subprocess_async(cmd, env=None, stdin=None, timeout=180, quiet=False, echo=True) -> _RunOutput: + loop = asyncio.get_event_loop() + result = loop.run_until_complete( + _stream_subprocess(cmd, env=env, stdin=stdin, timeout=timeout, quiet=quiet, echo=echo) + ) + + cmd_str = " ".join(cmd) + if result.returncode > 0: + stderr = "\n".join(result.stderr) + raise RuntimeError( + f"'{cmd_str}' failed with returncode {result.returncode}\n\n" + f"The combined stderr from workers follows:\n{stderr}" + ) + + # check that the subprocess actually did run and produced some output, should the test rely on + # the remote side to do the testing + if not result.stdout and not result.stderr: + raise RuntimeError(f"'{cmd_str}' produced no output.") + + return result + + +def pytest_xdist_worker_id(): + """ + Returns an int value of worker's numerical id under `pytest-xdist`'s concurrent workers `pytest -n N` regime, or 0 + if `-n 1` or `pytest-xdist` isn't being used. + """ + worker = os.environ.get("PYTEST_XDIST_WORKER", "gw0") + worker = re.sub(r"^gw", "", worker, 0, re.M) + return int(worker) + + +def get_torch_dist_unique_port(): + """ + Returns a port number that can be fed to `torch.distributed.launch`'s `--master_port` argument. + + Under `pytest-xdist` it adds a delta number based on a worker id so that concurrent tests don't try to use the same + port at once. + """ + port = 29500 + uniq_delta = pytest_xdist_worker_id() + return port + uniq_delta + + +def nested_simplify(obj, decimals=3): + """ + Simplifies an object by rounding float numbers, and downcasting tensors/numpy arrays to get simple equality test + within tests. + """ + import numpy as np + + if isinstance(obj, list): + return [nested_simplify(item, decimals) for item in obj] + if isinstance(obj, tuple): + return tuple([nested_simplify(item, decimals) for item in obj]) + elif isinstance(obj, np.ndarray): + return nested_simplify(obj.tolist()) + elif isinstance(obj, Mapping): + return {nested_simplify(k, decimals): nested_simplify(v, decimals) for k, v in obj.items()} + elif isinstance(obj, (str, int, np.int64)): + return obj + elif obj is None: + return obj + elif is_torch_available() and isinstance(obj, torch.Tensor): + return nested_simplify(obj.tolist(), decimals) + elif is_tf_available() and tf.is_tensor(obj): + return nested_simplify(obj.numpy().tolist()) + elif isinstance(obj, float): + return round(obj, decimals) + elif isinstance(obj, (np.int32, np.float32, np.float16)): + return nested_simplify(obj.item(), decimals) + else: + raise Exception(f"Not supported: {type(obj)}") + + +def check_json_file_has_correct_format(file_path): + with open(file_path) as f: + lines = f.readlines() + if len(lines) == 1: + # length can only be 1 if dict is empty + assert lines[0] == "{}" + else: + # otherwise make sure json has correct format (at least 3 lines) + assert len(lines) >= 3 + # each key one line, ident should be 2, min length is 3 + assert lines[0].strip() == "{" + for line in lines[1:-1]: + left_indent = len(lines[1]) - len(lines[1].lstrip()) + assert left_indent == 2 + assert lines[-1].strip() == "}" + + +def to_2tuple(x): + if isinstance(x, collections.abc.Iterable): + return x + return (x, x) + + +# These utils relate to ensuring the right error message is received when running scripts +class SubprocessCallException(Exception): + pass + + +def run_command(command: list[str], return_stdout=False): + """ + Runs `command` with `subprocess.check_output` and will potentially return the `stdout`. Will also properly capture + if an error occurred while running `command` + """ + try: + output = subprocess.check_output(command, stderr=subprocess.STDOUT) + if return_stdout: + if hasattr(output, "decode"): + output = output.decode("utf-8") + return output + except subprocess.CalledProcessError as e: + raise SubprocessCallException( + f"Command `{' '.join(command)}` failed with the following error:\n\n{e.output.decode()}" + ) from e + + +class RequestCounter: + """ + Helper class that will count all requests made online. + + Might not be robust if urllib3 changes its logging format but should be good enough for us. + + Usage: + ```py + with RequestCounter() as counter: + _ = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bert") + assert counter["GET"] == 0 + assert counter["HEAD"] == 1 + assert counter.total_calls == 1 + ``` + """ + + def __enter__(self): + self._counter = defaultdict(int) + self._thread_id = threading.get_ident() + self._extra_info = [] + + def patched_with_thread_info(func): + def wrap(*args, **kwargs): + self._extra_info.append(threading.get_ident()) + return func(*args, **kwargs) + + return wrap + + self.patcher = patch.object( + urllib3.connectionpool.log, "debug", side_effect=patched_with_thread_info(urllib3.connectionpool.log.debug) + ) + self.mock = self.patcher.start() + return self + + def __exit__(self, *args, **kwargs) -> None: + assert len(self.mock.call_args_list) == len(self._extra_info) + for thread_id, call in zip(self._extra_info, self.mock.call_args_list): + if thread_id != self._thread_id: + continue + # code 307: the URL being requested by the user has moved to a temporary location + if call.args[-2] == 307: + continue + log = call.args[0] % call.args[1:] + for method in ("HEAD", "GET", "POST", "PUT", "DELETE", "CONNECT", "OPTIONS", "TRACE", "PATCH"): + if method in log: + self._counter[method] += 1 + break + self.patcher.stop() + + def __getitem__(self, key: str) -> int: + return self._counter[key] + + @property + def total_calls(self) -> int: + return sum(self._counter.values()) + + +def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None, description: Optional[str] = None): + """ + To decorate flaky tests. They will be retried on failures. + + Please note that our push tests use `pytest-rerunfailures`, which prompts the CI to rerun certain types of + failed tests. More specifically, if the test exception contains any substring in `FLAKY_TEST_FAILURE_PATTERNS` + (in `.circleci/create_circleci_config.py`), it will be rerun. If you find a recurrent pattern of failures, + expand `FLAKY_TEST_FAILURE_PATTERNS` in our CI configuration instead of using `is_flaky`. + + Args: + max_attempts (`int`, *optional*, defaults to 5): + The maximum number of attempts to retry the flaky test. + wait_before_retry (`float`, *optional*): + If provided, will wait that number of seconds before retrying the test. + description (`str`, *optional*): + A string to describe the situation (what / where / why is flaky, link to GH issue/PR comments, errors, + etc.) + """ + + def decorator(test_func_ref): + @functools.wraps(test_func_ref) + def wrapper(*args, **kwargs): + retry_count = 1 + + while retry_count < max_attempts: + try: + return test_func_ref(*args, **kwargs) + + except Exception as err: + logger.error(f"Test failed with {err} at try {retry_count}/{max_attempts}.") + if wait_before_retry is not None: + time.sleep(wait_before_retry) + retry_count += 1 + + return test_func_ref(*args, **kwargs) + + return wrapper + + return decorator + + +def hub_retry(max_attempts: int = 5, wait_before_retry: Optional[float] = 2): + """ + To decorate tests that download from the Hub. They can fail due to a + variety of network issues such as timeouts, connection resets, etc. + + Args: + max_attempts (`int`, *optional*, defaults to 5): + The maximum number of attempts to retry the flaky test. + wait_before_retry (`float`, *optional*, defaults to 2): + If provided, will wait that number of seconds before retrying the test. + """ + + def decorator(test_func_ref): + @functools.wraps(test_func_ref) + def wrapper(*args, **kwargs): + retry_count = 1 + + while retry_count < max_attempts: + try: + return test_func_ref(*args, **kwargs) + # We catch all exceptions related to network issues from requests + except ( + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + requests.exceptions.ReadTimeout, + requests.exceptions.HTTPError, + requests.exceptions.RequestException, + ) as err: + logger.error( + f"Test failed with {err} at try {retry_count}/{max_attempts} as it couldn't connect to the specified Hub repository." + ) + if wait_before_retry is not None: + time.sleep(wait_before_retry) + retry_count += 1 + + return test_func_ref(*args, **kwargs) + + return wrapper + + return decorator + + +def run_first(test_case): + """ + Decorator marking a test with order(1). When pytest-order plugin is installed, tests marked with this decorator + are garanteed to run first. + + This is especially useful in some test settings like on a Gaudi instance where a Gaudi device can only be used by a + single process at a time. So we make sure all tests that run in a subprocess are launched first, to avoid device + allocation conflicts. + """ + import pytest + + return pytest.mark.order(1)(test_case) + + +def run_test_in_subprocess(test_case, target_func, inputs=None, timeout=None): + """ + To run a test in a subprocess. In particular, this can avoid (GPU) memory issue. + + Args: + test_case (`unittest.TestCase`): + The test that will run `target_func`. + target_func (`Callable`): + The function implementing the actual testing logic. + inputs (`dict`, *optional*, defaults to `None`): + The inputs that will be passed to `target_func` through an (input) queue. + timeout (`int`, *optional*, defaults to `None`): + The timeout (in seconds) that will be passed to the input and output queues. If not specified, the env. + variable `PYTEST_TIMEOUT` will be checked. If still `None`, its value will be set to `600`. + """ + if timeout is None: + timeout = int(os.environ.get("PYTEST_TIMEOUT", 600)) + + start_methohd = "spawn" + ctx = multiprocessing.get_context(start_methohd) + + input_queue = ctx.Queue(1) + output_queue = ctx.JoinableQueue(1) + + # We can't send `unittest.TestCase` to the child, otherwise we get issues regarding pickle. + input_queue.put(inputs, timeout=timeout) + + process = ctx.Process(target=target_func, args=(input_queue, output_queue, timeout)) + process.start() + # Kill the child process if we can't get outputs from it in time: otherwise, the hanging subprocess prevents + # the test to exit properly. + try: + results = output_queue.get(timeout=timeout) + output_queue.task_done() + except Exception as e: + process.terminate() + test_case.fail(e) + process.join(timeout=timeout) + + if results["error"] is not None: + test_case.fail(f"{results['error']}") + + +def run_test_using_subprocess(func): + """ + To decorate a test to run in a subprocess using the `subprocess` module. This could avoid potential GPU memory + issues (GPU OOM or a test that causes many subsequential failing with `CUDA error: device-side assert triggered`). + """ + import pytest + + @functools.wraps(func) + def wrapper(*args, **kwargs): + if os.getenv("_INSIDE_SUB_PROCESS", None) == "1": + func(*args, **kwargs) + else: + test = " ".join(os.environ.get("PYTEST_CURRENT_TEST").split(" ")[:-1]) + try: + import copy + + env = copy.deepcopy(os.environ) + env["_INSIDE_SUB_PROCESS"] = "1" + # This prevents the entries in `short test summary info` given by the subprocess being truncated. so the + # full information can be passed to the parent pytest process. + # See: https://docs.pytest.org/en/stable/explanation/ci.html + env["CI"] = "true" + + # If not subclass of `unitTest.TestCase` and `pytestconfig` is used: try to grab and use the arguments + if "pytestconfig" in kwargs: + command = list(kwargs["pytestconfig"].invocation_params.args) + for idx, x in enumerate(command): + if x in kwargs["pytestconfig"].args: + test = test.split("::")[1:] + command[idx] = "::".join([f"{func.__globals__['__file__']}"] + test) + command = [f"{sys.executable}", "-m", "pytest"] + command + command = [x for x in command if x not in ["--no-summary"]] + # Otherwise, simply run the test with no option at all + else: + command = [f"{sys.executable}", "-m", "pytest", f"{test}"] + + subprocess.run(command, env=env, check=True, capture_output=True) + except subprocess.CalledProcessError as e: + exception_message = e.stdout.decode() + lines = exception_message.split("\n") + # Add a first line with more informative information instead of just `= test session starts =`. + # This makes the `short test summary info` section more useful. + if "= test session starts =" in lines[0]: + text = "" + for line in lines[1:]: + if line.startswith("FAILED "): + text = line[len("FAILED ") :] + text = "".join(text.split(" - ")[1:]) + elif line.startswith("=") and line.endswith("=") and " failed in " in line: + break + elif len(text) > 0: + text += f"\n{line}" + text = "(subprocess) " + text + lines = [text] + lines + exception_message = "\n".join(lines) + raise pytest.fail(exception_message, pytrace=False) + + return wrapper + + +""" +The following contains utils to run the documentation tests without having to overwrite any files. + +The `preprocess_string` function adds `# doctest: +IGNORE_RESULT` markers on the fly anywhere a `load_dataset` call is +made as a print would otherwise fail the corresponding line. + +To skip cuda tests, make sure to call `SKIP_CUDA_DOCTEST=1 pytest --doctest-modules +""" + + +def preprocess_string(string, skip_cuda_tests): + """Prepare a docstring or a `.md` file to be run by doctest. + + The argument `string` would be the whole file content if it is a `.md` file. For a python file, it would be one of + its docstring. In each case, it may contain multiple python code examples. If `skip_cuda_tests` is `True` and a + cuda stuff is detective (with a heuristic), this method will return an empty string so no doctest will be run for + `string`. + """ + codeblock_pattern = r"(```(?:python|py)\s*\n\s*>>> )(.*?```)" + codeblocks = re.split(codeblock_pattern, string, flags=re.DOTALL) + is_cuda_found = False + for i, codeblock in enumerate(codeblocks): + if "load_dataset(" in codeblock and "# doctest: +IGNORE_RESULT" not in codeblock: + codeblocks[i] = re.sub(r"(>>> .*load_dataset\(.*)", r"\1 # doctest: +IGNORE_RESULT", codeblock) + if ( + (">>>" in codeblock or "..." in codeblock) + and re.search(r"cuda|to\(0\)|device=0", codeblock) + and skip_cuda_tests + ): + is_cuda_found = True + break + + modified_string = "" + if not is_cuda_found: + modified_string = "".join(codeblocks) + + return modified_string + + +class HfDocTestParser(doctest.DocTestParser): + """ + Overwrites the DocTestParser from doctest to properly parse the codeblocks that are formatted with black. This + means that there are no extra lines at the end of our snippets. The `# doctest: +IGNORE_RESULT` marker is also + added anywhere a `load_dataset` call is made as a print would otherwise fail the corresponding line. + + Tests involving cuda are skipped base on a naive pattern that should be updated if it is not enough. + """ + + # This regular expression is used to find doctest examples in a + # string. It defines three groups: `source` is the source code + # (including leading indentation and prompts); `indent` is the + # indentation of the first (PS1) line of the source code; and + # `want` is the expected output (including leading indentation). + # fmt: off + _EXAMPLE_RE = re.compile(r''' + # Source consists of a PS1 line followed by zero or more PS2 lines. + (?P + (?:^(?P [ ]*) >>> .*) # PS1 line + (?:\n [ ]* \.\.\. .*)*) # PS2 lines + \n? + # Want consists of any non-blank lines that do not start with PS1. + (?P (?:(?![ ]*$) # Not a blank line + (?![ ]*>>>) # Not a line starting with PS1 + # !!!!!!!!!!! HF Specific !!!!!!!!!!! + (?:(?!```).)* # Match any character except '`' until a '```' is found (this is specific to HF because black removes the last line) + # !!!!!!!!!!! HF Specific !!!!!!!!!!! + (?:\n|$) # Match a new line or end of string + )*) + ''', re.MULTILINE | re.VERBOSE + ) + # fmt: on + + # !!!!!!!!!!! HF Specific !!!!!!!!!!! + skip_cuda_tests: bool = bool(os.environ.get("SKIP_CUDA_DOCTEST", False)) + # !!!!!!!!!!! HF Specific !!!!!!!!!!! + + def parse(self, string, name=""): + """ + Overwrites the `parse` method to incorporate a skip for CUDA tests, and remove logs and dataset prints before + calling `super().parse` + """ + string = preprocess_string(string, self.skip_cuda_tests) + return super().parse(string, name) + + +class HfDoctestModule(Module): + """ + Overwrites the `DoctestModule` of the pytest package to make sure the HFDocTestParser is used when discovering + tests. + """ + + def collect(self) -> Iterable[DoctestItem]: + class MockAwareDocTestFinder(doctest.DocTestFinder): + """A hackish doctest finder that overrides stdlib internals to fix a stdlib bug. + + https://github.com/pytest-dev/pytest/issues/3456 https://bugs.python.org/issue25532 + """ + + def _find_lineno(self, obj, source_lines): + """Doctest code does not take into account `@property`, this + is a hackish way to fix it. https://bugs.python.org/issue17446 + + Wrapped Doctests will need to be unwrapped so the correct line number is returned. This will be + reported upstream. #8796 + """ + if isinstance(obj, property): + obj = getattr(obj, "fget", obj) + + if hasattr(obj, "__wrapped__"): + # Get the main obj in case of it being wrapped + obj = inspect.unwrap(obj) + + # Type ignored because this is a private function. + return super()._find_lineno( # type:ignore[misc] + obj, + source_lines, + ) + + def _find(self, tests, obj, name, module, source_lines, globs, seen) -> None: + if _is_mocked(obj): + return + with _patch_unwrap_mock_aware(): + # Type ignored because this is a private function. + super()._find( # type:ignore[misc] + tests, obj, name, module, source_lines, globs, seen + ) + + if self.path.name == "conftest.py": + module = self.config.pluginmanager._importconftest( + self.path, + self.config.getoption("importmode"), + rootpath=self.config.rootpath, + ) + else: + try: + module = import_path( + self.path, + root=self.config.rootpath, + mode=self.config.getoption("importmode"), + ) + except ImportError: + if self.config.getvalue("doctest_ignore_import_errors"): + skip("unable to import module %r" % self.path) + else: + raise + + # !!!!!!!!!!! HF Specific !!!!!!!!!!! + finder = MockAwareDocTestFinder(parser=HfDocTestParser()) + # !!!!!!!!!!! HF Specific !!!!!!!!!!! + optionflags = get_optionflags(self) + runner = _get_runner( + verbose=False, + optionflags=optionflags, + checker=_get_checker(), + continue_on_failure=_get_continue_on_failure(self.config), + ) + for test in finder.find(module, module.__name__): + if test.examples: # skip empty doctests and cuda + yield DoctestItem.from_parent(self, name=test.name, runner=runner, dtest=test) + + +def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, Callable], *args, **kwargs): + if device not in dispatch_table: + return dispatch_table["default"](*args, **kwargs) + + fn = dispatch_table[device] + + # Some device agnostic functions return values. Need to guard against `None` + # instead at user level. + if fn is None: + return None + return fn(*args, **kwargs) + + +if is_torch_available(): + # Mappings from device names to callable functions to support device agnostic + # testing. + BACKEND_MANUAL_SEED = { + "cuda": torch.cuda.manual_seed, + "cpu": torch.manual_seed, + "default": torch.manual_seed, + } + BACKEND_EMPTY_CACHE = { + "cuda": torch.cuda.empty_cache, + "cpu": None, + "default": None, + } + BACKEND_DEVICE_COUNT = { + "cuda": torch.cuda.device_count, + "cpu": lambda: 0, + "default": lambda: 1, + } +else: + BACKEND_MANUAL_SEED = {"default": None} + BACKEND_EMPTY_CACHE = {"default": None} + BACKEND_DEVICE_COUNT = {"default": lambda: 0} + +if is_torch_hpu_available(): + BACKEND_MANUAL_SEED["hpu"] = torch.hpu.manual_seed + BACKEND_DEVICE_COUNT["hpu"] = torch.hpu.device_count + +if is_torch_mlu_available(): + BACKEND_EMPTY_CACHE["mlu"] = torch.mlu.empty_cache + BACKEND_MANUAL_SEED["mlu"] = torch.mlu.manual_seed + BACKEND_DEVICE_COUNT["mlu"] = torch.mlu.device_count + +if is_torch_npu_available(): + BACKEND_EMPTY_CACHE["npu"] = torch.npu.empty_cache + BACKEND_MANUAL_SEED["npu"] = torch.npu.manual_seed + BACKEND_DEVICE_COUNT["npu"] = torch.npu.device_count + +if is_torch_xpu_available(): + BACKEND_EMPTY_CACHE["xpu"] = torch.xpu.empty_cache + BACKEND_MANUAL_SEED["xpu"] = torch.xpu.manual_seed + BACKEND_DEVICE_COUNT["xpu"] = torch.xpu.device_count + +if is_torch_xla_available(): + BACKEND_EMPTY_CACHE["xla"] = torch.cuda.empty_cache + BACKEND_MANUAL_SEED["xla"] = torch.cuda.manual_seed + BACKEND_DEVICE_COUNT["xla"] = torch.cuda.device_count + + +def backend_manual_seed(device: str, seed: int): + return _device_agnostic_dispatch(device, BACKEND_MANUAL_SEED, seed) + + +def backend_empty_cache(device: str): + return _device_agnostic_dispatch(device, BACKEND_EMPTY_CACHE) + + +def backend_device_count(device: str): + return _device_agnostic_dispatch(device, BACKEND_DEVICE_COUNT) + + +if is_torch_available(): + # If `TRANSFORMERS_TEST_DEVICE_SPEC` is enabled we need to import extra entries + # into device to function mappings. + if "TRANSFORMERS_TEST_DEVICE_SPEC" in os.environ: + device_spec_path = os.environ["TRANSFORMERS_TEST_DEVICE_SPEC"] + if not Path(device_spec_path).is_file(): + raise ValueError( + f"Specified path to device spec file is not a file or not found. Received '{device_spec_path}" + ) + + # Try to strip extension for later import – also verifies we are importing a + # python file. + device_spec_dir, _ = os.path.split(os.path.realpath(device_spec_path)) + sys.path.append(device_spec_dir) + try: + import_name = device_spec_path[: device_spec_path.index(".py")] + except ValueError as e: + raise ValueError(f"Provided device spec file was not a Python file! Received '{device_spec_path}") from e + + device_spec_module = importlib.import_module(import_name) + + # Imported file must contain `DEVICE_NAME`. If it doesn't, terminate early. + try: + device_name = device_spec_module.DEVICE_NAME + except AttributeError as e: + raise AttributeError("Device spec file did not contain `DEVICE_NAME`") from e + + if "TRANSFORMERS_TEST_DEVICE" in os.environ and torch_device != device_name: + msg = f"Mismatch between environment variable `TRANSFORMERS_TEST_DEVICE` '{torch_device}' and device found in spec '{device_name}'\n" + msg += "Either unset `TRANSFORMERS_TEST_DEVICE` or ensure it matches device spec name." + raise ValueError(msg) + + torch_device = device_name + + def update_mapping_from_spec(device_fn_dict: dict[str, Callable], attribute_name: str): + try: + # Try to import the function directly + spec_fn = getattr(device_spec_module, attribute_name) + device_fn_dict[torch_device] = spec_fn + except AttributeError as e: + # If the function doesn't exist, and there is no default, throw an error + if "default" not in device_fn_dict: + raise AttributeError( + f"`{attribute_name}` not found in '{device_spec_path}' and no default fallback function found." + ) from e + + # Add one entry here for each `BACKEND_*` dictionary. + update_mapping_from_spec(BACKEND_MANUAL_SEED, "MANUAL_SEED_FN") + update_mapping_from_spec(BACKEND_EMPTY_CACHE, "EMPTY_CACHE_FN") + update_mapping_from_spec(BACKEND_DEVICE_COUNT, "DEVICE_COUNT_FN") + + +def compare_pipeline_output_to_hub_spec(output, hub_spec): + missing_keys = [] + unexpected_keys = [] + all_field_names = {field.name for field in fields(hub_spec)} + matching_keys = sorted([key for key in output.keys() if key in all_field_names]) + + # Fields with a MISSING default are required and must be in the output + for field in fields(hub_spec): + if field.default is MISSING and field.name not in output: + missing_keys.append(field.name) + + # All output keys must match either a required or optional field in the Hub spec + for output_key in output: + if output_key not in all_field_names: + unexpected_keys.append(output_key) + + if missing_keys or unexpected_keys: + error = ["Pipeline output does not match Hub spec!"] + if matching_keys: + error.append(f"Matching keys: {matching_keys}") + if missing_keys: + error.append(f"Missing required keys in pipeline output: {missing_keys}") + if unexpected_keys: + error.append(f"Keys in pipeline output that are not in Hub spec: {unexpected_keys}") + raise KeyError("\n".join(error)) + + +@require_torch +def cleanup(device: str, gc_collect=False): + if gc_collect: + gc.collect() + backend_empty_cache(device) + torch._dynamo.reset() + + +# Type definition of key used in `Expectations` class. +DeviceProperties = tuple[Union[str, None], Union[int, None]] + + +@cache +def get_device_properties() -> DeviceProperties: + """ + Get environment device properties. + """ + if IS_CUDA_SYSTEM or IS_ROCM_SYSTEM: + import torch + + major, _ = torch.cuda.get_device_capability() + if IS_ROCM_SYSTEM: + return ("rocm", major) + else: + return ("cuda", major) + else: + return (torch_device, None) + + +class Expectations(UserDict[DeviceProperties, Any]): + def get_expectation(self) -> Any: + """ + Find best matching expectation based on environment device properties. + """ + return self.find_expectation(get_device_properties()) + + @staticmethod + def is_default(key: DeviceProperties) -> bool: + return all(p is None for p in key) + + @staticmethod + def score(key: DeviceProperties, other: DeviceProperties) -> int: + """ + Returns score indicating how similar two instances of the `Properties` tuple are. + Points are calculated using bits, but documented as int. + Rules are as follows: + * Matching `type` gives 8 points. + * Semi-matching `type`, for example cuda and rocm, gives 4 points. + * Matching `major` (compute capability major version) gives 2 points. + * Default expectation (if present) gives 1 points. + """ + (device_type, major) = key + (other_device_type, other_major) = other + + score = 0b0 + if device_type == other_device_type: + score |= 0b1000 + elif device_type in ["cuda", "rocm"] and other_device_type in ["cuda", "rocm"]: + score |= 0b100 + + if major == other_major and other_major is not None: + score |= 0b10 + + if Expectations.is_default(other): + score |= 0b1 + + return int(score) + + def find_expectation(self, key: DeviceProperties = (None, None)) -> Any: + """ + Find best matching expectation based on provided device properties. + """ + (result_key, result) = max(self.data.items(), key=lambda x: Expectations.score(key, x[0])) + + if Expectations.score(key, result_key) == 0: + raise ValueError(f"No matching expectation found for {key}") + + return result + + def __repr__(self): + return f"{self.data}" diff --git a/lib/python3.12/site-packages/transformers/tf_utils.py b/lib/python3.12/site-packages/transformers/tf_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c3770cb1237f3c6ea3517bda3dd29b2b6f44f7de --- /dev/null +++ b/lib/python3.12/site-packages/transformers/tf_utils.py @@ -0,0 +1,294 @@ +# Copyright 2022 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. + +from typing import Optional, Union + +import numpy as np +import tensorflow as tf + +from .feature_extraction_utils import BatchFeature +from .tokenization_utils_base import BatchEncoding +from .utils import logging + + +logger = logging.get_logger(__name__) + + +def shape_list(tensor: Union[tf.Tensor, np.ndarray]) -> list[int]: + """ + Deal with dynamic shape in tensorflow cleanly. + + Args: + tensor (`tf.Tensor` or `np.ndarray`): The tensor we want the shape of. + + Returns: + `List[int]`: The shape of the tensor as a list. + """ + if isinstance(tensor, np.ndarray): + return list(tensor.shape) + + dynamic = tf.shape(tensor) + + if tensor.shape == tf.TensorShape(None): + return dynamic + + static = tensor.shape.as_list() + + return [dynamic[i] if s is None else s for i, s in enumerate(static)] + + +def stable_softmax(logits: tf.Tensor, axis: Optional[int] = None, name: Optional[str] = None) -> tf.Tensor: + """ + Stable wrapper that returns the same output as `tf.nn.softmax`, but that works reliably with XLA on CPU. It is + meant as a workaround for the [following issue](https://github.com/tensorflow/tensorflow/issues/55682), and will be + removed after it gets fixed. The arguments and outputs are the same as `tf.nn.softmax`, and relies on the fact that + `softmax(x) = softmax(x + c)` (see https://ogunlao.github.io/2020/04/26/you_dont_really_know_softmax.html). + + Args: + logits (`tf.Tensor`): + Must be one of the following types: half, float32, float64. + axis (`int`, *optional*): + The dimension softmax would be performed on. The default is -1 which indicates the last dimension. + name (`str`, *optional*): + A name for the operation. + + Returns: + `tf.Tensor`: + A Tensor. Has the same type and shape as logits. + """ + # TODO: When the issue linked above gets sorted, add a check on TF version here and use the original function if + # it has the fix. After we drop the support for unfixed versions, remove this function. + return tf.nn.softmax(logits=logits + 1e-9, axis=axis, name=name) + + +def functional_layernorm(inputs, weight, bias, epsilon=1e-5, axis=-1): + # This is a very simplified functional layernorm, designed to duplicate + # the functionality of PyTorch nn.functional.layer_norm when this is needed to port + # models in Transformers. + + if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(axis, int): + raise NotImplementedError("Only 1D weight and bias tensors are supported for now, with only a single axis.") + + # Get mean and variance on the axis to be normalized + mean, variance = tf.nn.moments(inputs, axes=[axis], keepdims=True) + + if axis != -1: + # Reshape scale and weight to have the same rank as inputs, but with 1 dimensions + # on every dimension except axis + shape = [1] * inputs.shape.rank + shape[axis] = shape_list(inputs)[axis] + weight = tf.reshape(weight, shape) + bias = tf.reshape(bias, shape) + + # Compute layer normalization using the batch_normalization + # function. + outputs = tf.nn.batch_normalization( + inputs, + mean, + variance, + offset=bias, + scale=weight, + variance_epsilon=epsilon, + ) + return outputs + + +def scaled_dot_product_attention( + query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale: Optional[float] = None +): + """TF equivalent for torch's nn.functional.scaled_dot_product_attention""" + if dropout_p != 0.0: + raise ValueError( + "Dropout is not supported in this implementation - file an issue " + "with Transformers and ping @Rocketknight1 if you need it for a port!" + ) + if is_causal and attn_mask is not None: + raise ValueError("You cannot specify an attn_mask and is_causal at the same time!") + if is_causal: + attn_mask = tf.ones((tf.shape(query)[-2], tf.shape(key)[-2]), dtype=tf.int32) + attn_mask = tf.experimental.numpy.tril(attn_mask, k=0) + if attn_mask is not None and (attn_mask.dtype.is_integer or attn_mask.dtype.is_bool): + # Convert boolean mask to a negative logit bias + attn_mask = tf.where(attn_mask > 0, tf.cast(0.0, query.dtype), tf.cast(-1000.0, query.dtype)) + logits = tf.einsum("...qd, ...kd -> ...qk", query, key) + if scale is None: + scale = tf.cast(tf.shape(key)[-1], logits.dtype) ** -0.5 + logits *= scale # scale by 1/sqrt(key_dim) + if attn_mask is not None: + logits += attn_mask + probs = tf.nn.softmax(logits) + return probs @ value + + +def flatten(input, start_dim=0, end_dim=-1): + # Replicates the behavior of torch.flatten in TF + + # If end_dim or start_dim is negative, count them from the end + if end_dim < 0: + end_dim += input.shape.rank + if start_dim < 0: + start_dim += input.shape.rank + + if start_dim == end_dim: + return input + + in_shape = tf.shape(input) + flattened_dim = tf.math.reduce_prod(in_shape[start_dim : end_dim + 1]) + out_shape = tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]], axis=0) + return tf.reshape(input, out_shape) + + +def invert_attention_mask(encoder_attention_mask: tf.Tensor) -> tf.Tensor: + """ + Invert an attention mask (e.g., switches 0. and 1.). + + Args: + encoder_attention_mask (`torch.Tensor`): An attention mask. + + Returns: + `tf.Tensor`: The inverted attention mask. + """ + if not isinstance(encoder_attention_mask, tf.Tensor): + encoder_attention_mask = tf.convert_to_tensor(encoder_attention_mask) # Catches stray NumPy inputs + if encoder_attention_mask.shape.rank == 3: + encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :] + if encoder_attention_mask.shape.rank == 2: + encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :] + # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition + # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow + # /transformer/transformer_layers.py#L270 + # encoder_extended_attention_mask = (encoder_extended_attention_mask == + # encoder_extended_attention_mask.transpose(-1, -2)) + encoder_extended_attention_mask = ( + tf.cast(1, encoder_attention_mask.dtype) - encoder_extended_attention_mask + ) * encoder_extended_attention_mask.dtype.min + + return encoder_extended_attention_mask + + +def check_embeddings_within_bounds(tensor: tf.Tensor, embed_dim: int, tensor_name: str = "input_ids") -> None: + """ + `tf.gather`, on which TF embedding layers are based, won't check positive out of bound indices on GPU, returning + zeros instead. This function adds a check against that dangerous silent behavior. + + Args: + tensor (`tf.Tensor`): The tensor of indices to check. + embed_dim (`int`): The embedding dimension. + tensor_name (`str`, *optional*): The name of the tensor to use in the error message. + """ + tf.debugging.assert_less( + tensor, + tf.cast(embed_dim, dtype=tensor.dtype), + message=( + f"The maximum value of {tensor_name} ({tf.math.reduce_max(tensor)}) must be smaller than the embedding " + f"layer's input dimension ({embed_dim}). The likely cause is some problem at tokenization time." + ), + ) + + +def save_attributes_to_hdf5_group(group, name, data): + """Saves attributes (data) of the specified name into the HDF5 group. + + This method deals with an inherent problem of HDF5 file which is not able to store data larger than + HDF5_OBJECT_HEADER_LIMIT bytes. + + Args: + group: A pointer to a HDF5 group. + name: A name of the attributes to save. + data: Attributes data to store. + + Raises: + RuntimeError: If any single attribute is too large to be saved. + + Copied from Keras to Transformers to avoid versioning issues. + """ + HDF5_OBJECT_HEADER_LIMIT = 64512 + # Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT` + # because in that case even chunking the array would not make the saving + # possible. + bad_attributes = [x for x in data if len(x) > HDF5_OBJECT_HEADER_LIMIT] + + # Expecting this to never be true. + if bad_attributes: + raise RuntimeError( + "The following attributes cannot be saved to HDF5 file because " + f"they are larger than {HDF5_OBJECT_HEADER_LIMIT} " + f"bytes: {bad_attributes}" + ) + + data_npy = np.asarray(data) + + num_chunks = 1 + chunked_data = np.array_split(data_npy, num_chunks) + + # This will never loop forever thanks to the test above. + while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data): + num_chunks += 1 + chunked_data = np.array_split(data_npy, num_chunks) + + if num_chunks > 1: + for chunk_id, chunk_data in enumerate(chunked_data): + group.attrs["%s%d" % (name, chunk_id)] = chunk_data + else: + group.attrs[name] = data + + +def load_attributes_from_hdf5_group(group, name): + """Loads attributes of the specified name from the HDF5 group. + + This method deals with an inherent problem of HDF5 file which is not able to store data larger than + HDF5_OBJECT_HEADER_LIMIT bytes. + + Args: + group: A pointer to a HDF5 group. + name: A name of the attributes to load. + + Returns: + data: Attributes data. + + Copied from Keras to Transformers to avoid versioning issues. + """ + if name in group.attrs: + data = [n.decode("utf8") if hasattr(n, "decode") else n for n in group.attrs[name]] + else: + data = [] + chunk_id = 0 + while "%s%d" % (name, chunk_id) in group.attrs: + data.extend( + [n.decode("utf8") if hasattr(n, "decode") else n for n in group.attrs["%s%d" % (name, chunk_id)]] + ) + chunk_id += 1 + return data + + +def expand_1d(data): + """Expands 1-dimensional `Tensor`s into 2-dimensional `Tensor`s. + Copied from Keras to here to avoid versioning issues.""" + + def _expand_single_1d_tensor(t): + if isinstance(t, tf.Tensor) and t.shape.rank == 1: + return tf.expand_dims(t, axis=-1) + return t + + return tf.nest.map_structure(_expand_single_1d_tensor, data) + + +def convert_batch_encoding(*args, **kwargs): + # Convert HF BatchEncoding/BatchFeature objects in the inputs to dicts that Keras understands + if args and isinstance(args[0], (BatchEncoding, BatchFeature)): + args = list(args) + args[0] = dict(args[0]) + elif "x" in kwargs and isinstance(kwargs["x"], (BatchEncoding, BatchFeature)): + kwargs["x"] = dict(kwargs["x"]) + return args, kwargs diff --git a/lib/python3.12/site-packages/transformers/time_series_utils.py b/lib/python3.12/site-packages/transformers/time_series_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3a5cf4f2f4d8f636e623c07ae3400ca5e17b5891 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/time_series_utils.py @@ -0,0 +1,225 @@ +# Copyright 2023 The HuggingFace Inc. team. +# Copyright 2018 Amazon.com, Inc. or its affiliates. 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. +""" +Time series distributional output classes and utilities. +""" + +from typing import Callable, Optional + +import torch +from torch import nn +from torch.distributions import ( + AffineTransform, + Distribution, + Independent, + NegativeBinomial, + Normal, + StudentT, + TransformedDistribution, +) + + +class AffineTransformed(TransformedDistribution): + def __init__(self, base_distribution: Distribution, loc=None, scale=None, event_dim=0): + self.scale = 1.0 if scale is None else scale + self.loc = 0.0 if loc is None else loc + + super().__init__(base_distribution, [AffineTransform(loc=self.loc, scale=self.scale, event_dim=event_dim)]) + + @property + def mean(self): + """ + Returns the mean of the distribution. + """ + return self.base_dist.mean * self.scale + self.loc + + @property + def variance(self): + """ + Returns the variance of the distribution. + """ + return self.base_dist.variance * self.scale**2 + + @property + def stddev(self): + """ + Returns the standard deviation of the distribution. + """ + return self.variance.sqrt() + + +class ParameterProjection(nn.Module): + def __init__( + self, in_features: int, args_dim: dict[str, int], domain_map: Callable[..., tuple[torch.Tensor]], **kwargs + ) -> None: + super().__init__(**kwargs) + self.args_dim = args_dim + self.proj = nn.ModuleList([nn.Linear(in_features, dim) for dim in args_dim.values()]) + self.domain_map = domain_map + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor]: + params_unbounded = [proj(x) for proj in self.proj] + + return self.domain_map(*params_unbounded) + + +class LambdaLayer(nn.Module): + def __init__(self, function): + super().__init__() + self.function = function + + def forward(self, x, *args): + return self.function(x, *args) + + +class DistributionOutput: + distribution_class: type + in_features: int + args_dim: dict[str, int] + + def __init__(self, dim: int = 1) -> None: + self.dim = dim + self.args_dim = {k: dim * self.args_dim[k] for k in self.args_dim} + + def _base_distribution(self, distr_args): + if self.dim == 1: + return self.distribution_class(*distr_args) + else: + return Independent(self.distribution_class(*distr_args), 1) + + def distribution( + self, + distr_args, + loc: Optional[torch.Tensor] = None, + scale: Optional[torch.Tensor] = None, + ) -> Distribution: + distr = self._base_distribution(distr_args) + if loc is None and scale is None: + return distr + else: + return AffineTransformed(distr, loc=loc, scale=scale, event_dim=self.event_dim) + + @property + def event_shape(self) -> tuple: + r""" + Shape of each individual event contemplated by the distributions that this object constructs. + """ + return () if self.dim == 1 else (self.dim,) + + @property + def event_dim(self) -> int: + r""" + Number of event dimensions, i.e., length of the `event_shape` tuple, of the distributions that this object + constructs. + """ + return len(self.event_shape) + + @property + def value_in_support(self) -> float: + r""" + A float that will have a valid numeric value when computing the log-loss of the corresponding distribution. By + default 0.0. This value will be used when padding data series. + """ + return 0.0 + + def get_parameter_projection(self, in_features: int) -> nn.Module: + r""" + Return the parameter projection layer that maps the input to the appropriate parameters of the distribution. + """ + return ParameterProjection( + in_features=in_features, + args_dim=self.args_dim, + domain_map=LambdaLayer(self.domain_map), + ) + + def domain_map(self, *args: torch.Tensor): + r""" + Converts arguments to the right shape and domain. The domain depends on the type of distribution, while the + correct shape is obtained by reshaping the trailing axis in such a way that the returned tensors define a + distribution of the right event_shape. + """ + raise NotImplementedError() + + @staticmethod + def squareplus(x: torch.Tensor) -> torch.Tensor: + r""" + Helper to map inputs to the positive orthant by applying the square-plus operation. Reference: + https://twitter.com/jon_barron/status/1387167648669048833 + """ + return (x + torch.sqrt(torch.square(x) + 4.0)) / 2.0 + + +class StudentTOutput(DistributionOutput): + """ + Student-T distribution output class. + """ + + args_dim: dict[str, int] = {"df": 1, "loc": 1, "scale": 1} + distribution_class: type = StudentT + + @classmethod + def domain_map(cls, df: torch.Tensor, loc: torch.Tensor, scale: torch.Tensor): + scale = cls.squareplus(scale).clamp_min(torch.finfo(scale.dtype).eps) + df = 2.0 + cls.squareplus(df) + return df.squeeze(-1), loc.squeeze(-1), scale.squeeze(-1) + + +class NormalOutput(DistributionOutput): + """ + Normal distribution output class. + """ + + args_dim: dict[str, int] = {"loc": 1, "scale": 1} + distribution_class: type = Normal + + @classmethod + def domain_map(cls, loc: torch.Tensor, scale: torch.Tensor): + scale = cls.squareplus(scale).clamp_min(torch.finfo(scale.dtype).eps) + return loc.squeeze(-1), scale.squeeze(-1) + + +class NegativeBinomialOutput(DistributionOutput): + """ + Negative Binomial distribution output class. + """ + + args_dim: dict[str, int] = {"total_count": 1, "logits": 1} + distribution_class: type = NegativeBinomial + + @classmethod + def domain_map(cls, total_count: torch.Tensor, logits: torch.Tensor): + total_count = cls.squareplus(total_count) + return total_count.squeeze(-1), logits.squeeze(-1) + + def _base_distribution(self, distr_args) -> Distribution: + total_count, logits = distr_args + if self.dim == 1: + return self.distribution_class(total_count=total_count, logits=logits) + else: + return Independent(self.distribution_class(total_count=total_count, logits=logits), 1) + + # Overwrites the parent class method. We cannot scale using the affine + # transformation since negative binomial should return integers. Instead + # we scale the parameters. + def distribution( + self, distr_args, loc: Optional[torch.Tensor] = None, scale: Optional[torch.Tensor] = None + ) -> Distribution: + total_count, logits = distr_args + + if scale is not None: + # See scaling property of Gamma. + logits += scale.log() + + return self._base_distribution((total_count, logits)) diff --git a/lib/python3.12/site-packages/transformers/tokenization_utils.py b/lib/python3.12/site-packages/transformers/tokenization_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..51382d776f0643b7526d9a634fb98d3b43380031 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/tokenization_utils.py @@ -0,0 +1,1133 @@ +# Copyright 2020 The HuggingFace Inc. team. +# +# 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. +""" +Tokenization classes for python tokenizers. For fast tokenizers (provided by HuggingFace's tokenizers library) see +tokenization_utils_fast.py +""" + +import bisect +import itertools +import re +import unicodedata +from collections import OrderedDict +from typing import Any, Optional, Union, overload + +from .tokenization_utils_base import ( + ENCODE_KWARGS_DOCSTRING, + ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING, + INIT_TOKENIZER_DOCSTRING, + AddedToken, + BatchEncoding, + EncodedInput, + EncodedInputPair, + PreTokenizedInput, + PreTokenizedInputPair, + PreTrainedTokenizerBase, + TextInput, + TextInputPair, + TruncationStrategy, +) +from .utils import PaddingStrategy, TensorType, add_end_docstrings, logging + + +logger = logging.get_logger(__name__) + +# Slow tokenizers are saved in a vocabulary plus three separated files +SPECIAL_TOKENS_MAP_FILE = "special_tokens_map.json" +ADDED_TOKENS_FILE = "added_tokens.json" +TOKENIZER_CONFIG_FILE = "tokenizer_config.json" + + +class Trie: + """ + Trie in Python. Creates a Trie out of a list of words. The trie is used to split on `added_tokens` in one pass + Loose reference https://en.wikipedia.org/wiki/Trie + """ + + def __init__(self, *args): + self.data = {} + self._tokens = set() + self._termination_char = "" + self.update(*args) + + def update(self, *args): + """ + Updates the Trie with new tokens provided as arguments. + + Args: + *args: Variable number of words to be added to the Trie. + """ + for token in tuple(*args): + self.add(token) + + def add(self, word: str): + """ + Passes over every char (utf-8 char) on word and recursively adds it to the internal `data` trie representation. + The special key `""` in `self._termination_char` is used to represent termination. + + This function is idempotent, adding twice the same word will leave the trie unchanged + + Example: + + ```python + >>> trie = Trie() + >>> trie.add("Hello 友達") + >>> trie.data + {"H": {"e": {"l": {"l": {"o": {" ": {"友": {"達": {"": 1}}}}}}}}} + + >>> trie.add("Hello") + >>> trie.data + {"H": {"e": {"l": {"l": {"o": {"": 1, " ": {"友": {"達": {"": 1}}}}}}}}} + ``` + """ + if not word: + # Prevent empty string + return + + self._tokens.add(word) + ref = self.data + for char in word: + ref[char] = ref.setdefault(char, {}) + ref = ref[char] + ref[self._termination_char] = 1 + + def split(self, text: str) -> list[str]: + """ + Will look for the words added to the trie within `text`. Output is the original string splitted along the + boundaries of the words found. + + This trie will match the longest possible word first ! + + Example: + + ```python + >>> trie = Trie() + >>> trie.split("[CLS] This is a extra_id_100") + ["[CLS] This is a extra_id_100"] + + >>> trie.add("[CLS]") + >>> trie.add("extra_id_1") + >>> trie.add("extra_id_100") + >>> trie.split("[CLS] This is a extra_id_100") + ["[CLS]", " This is a ", "extra_id_100"] + ``` + """ + # indexes are counted left of the chars index. + # "hello", index 0, is left of h, index 1 is between h and e. + # index 5 is right of the "o". + + # States are going to capture every possible start (indexes as above) + # as keys, and have as values, a pointer to the position in the trie + # where we're at. This is a partial match for now. + # This enables to keep track of multiple matches while we're iterating + # the string + # If the trie contains, "blowing", and "lower" and we encounter the + # string "blower", we need to split into ["b", "lower"]. + # This is where we need to keep track of multiple possible starts. + states = OrderedDict() + + # This will contain every indices where we need + # to cut. + # We force to cut at offset 0 and len(text) (added later) + offsets = [0] + + # This is used by the lookahead which needs to skip over + # some text where the full match exceeded the place in the initial + # for loop + skip = 0 + # Main loop, Giving this algorithm O(n) complexity + for current, current_char in enumerate(text): + if skip and current < skip: + # Prevents the lookahead for matching twice + # like extra_id_100 and id_100 + continue + + # This will track every state + # that stop matching, we need to stop tracking them. + # If we look at "lowball", we're going to match "l" (add it to states), "o", "w", then + # fail on "b", we need to remove 0 from the valid states. + to_remove = set() + # Whenever we found a match, we need to drop everything + # this is a greedy algorithm, it will match on the first found token + reset = False + + # In this case, we already have partial matches (But unfinished) + for start, trie_pointer in states.items(): + if "" in trie_pointer: + # This is a final match, we need to reset and + # store the results in `offsets`. + + # Lookahead to match longest first + # Important in case of extra_id_1 vs extra_id_100 + # Here we are also actively looking for other earlier partial + # matches + # "[CLS]", "L", we need to match CLS even if L is special + for lookstart, looktrie_pointer in states.items(): + if lookstart > start: + # This partial match is later, we can stop looking + break + elif lookstart < start: + # This partial match is earlier, the trie pointer + # was already updated, so index is + 1 + lookahead_index = current + 1 + end = current + 1 + else: + # Here lookstart == start and + # looktrie_pointer == trie_pointer + # It wasn't updated yet so indices are current ones + lookahead_index = current + end = current + next_char = text[lookahead_index] if lookahead_index < len(text) else None + if "" in looktrie_pointer: + start = lookstart + end = lookahead_index + skip = lookahead_index + + while next_char in looktrie_pointer: + looktrie_pointer = looktrie_pointer[next_char] + lookahead_index += 1 + if "" in looktrie_pointer: + start = lookstart + end = lookahead_index + skip = lookahead_index + + if lookahead_index == len(text): + # End of string + break + next_char = text[lookahead_index] + # End lookahead + + # Storing and resetting + offsets.append(start) + offsets.append(end) + reset = True + break + elif current_char in trie_pointer: + # The current character being looked at has a match within the trie + # update the pointer (it will be stored back into states later). + trie_pointer = trie_pointer[current_char] + + # Storing back the new pointer into the states. + # Partial matches got longer by one. + states[start] = trie_pointer + else: + # The new character has not match in the trie, we need + # to stop keeping track of this partial match. + # We can't do it directly within the loop because of how + # python iteration works + to_remove.add(start) + + # Either clearing the full start (we found a real match) + # Or clearing only the partial matches that didn't work. + if reset: + states = {} + else: + for start in to_remove: + del states[start] + + # If this character is a starting character within the trie + # start keeping track of this partial match. + if current >= skip and current_char in self.data: + states[current] = self.data[current_char] + + # We have a cut at the end with states. + for start, trie_pointer in states.items(): + if "" in trie_pointer: + # This is a final match, we need to reset and + # store the results in `offsets`. + end = len(text) + offsets.append(start) + offsets.append(end) + # Longest cut is always the one with lower start so the first + # item so we need to break. + break + + return self.cut_text(text, offsets) + + def cut_text(self, text, offsets): + # We have all the offsets now, we just need to do the actual splitting. + # We need to eventually add the first part of the string and the eventual + # last part. + offsets.append(len(text)) + tokens = [] + start = 0 + for end in offsets: + if start > end: + logger.error( + "There was a bug in Trie algorithm in tokenization. Attempting to recover. Please report it" + " anyway." + ) + continue + elif start == end: + # This might happen if there's a match at index 0 + # we're also preventing zero-width cuts in case of two + # consecutive matches + continue + tokens.append(text[start:end]) + start = end + + return tokens + + +class ExtensionsTrie(Trie): + def __init__(self, *args): + super().__init__(*args) + + def extensions(self, prefix: str): + """ + Generates all extensions of a given prefix token in the Trie. + + Example: + + ```python + >>> trie = Trie() + >>> trie.add("apple") + >>> trie.add("app") + >>> trie.add("application") + >>> trie.extensions("app") + ['app', 'apple', 'application'] + ``` + """ + prefix_node = self._get_node(prefix) + ret = self._collect_tokens(prefix_node) + return [prefix + token for token in ret] + + def _get_node(self, token: str) -> dict: + """ + Retrieves the node corresponding to the given token in the Trie. + + Args: + token (str): The token for which the corresponding node needs to be retrieved. + + Returns: + dict: The node in the Trie corresponding to the given token. + """ + node = self.data + for char in token: + if char not in node: + break + + node = node[char] + return node + + def _collect_tokens(self, node: dict) -> list: + """ + Generates all tokens in the Trie starting from a given node. + + Args: + node (dict): The node in the Trie from which tokens need to be generated. + + Returns: + list: List of tokens generated from the given node. + """ + tokens = [self._termination_char] if self._termination_char in node else [] + for token, subtrie_head in node.items(): + if token != self._termination_char: + subtokens = self._collect_tokens(subtrie_head) + tokens.extend([token + subtoken for subtoken in subtokens]) + return tokens + + +def _is_whitespace(char): + """Checks whether `char` is a whitespace character.""" + # \t, \n, and \r are technically control characters but we treat them + # as whitespace since they are generally considered as such. + if char == " " or char == "\t" or char == "\n" or char == "\r": + return True + cat = unicodedata.category(char) + if cat == "Zs": + return True + return False + + +def _is_control(char): + """Checks whether `char` is a control character.""" + # These are technically control characters but we count them as whitespace + # characters. + if char == "\t" or char == "\n" or char == "\r": + return False + cat = unicodedata.category(char) + if cat.startswith("C"): + return True + return False + + +def _is_punctuation(char): + """Checks whether `char` is a punctuation character.""" + cp = ord(char) + # We treat all non-letter/number ASCII as punctuation. + # Characters such as "^", "$", and "`" are not in the Unicode + # Punctuation class but we treat them as punctuation anyways, for + # consistency. + if (cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126): + return True + cat = unicodedata.category(char) + if cat.startswith("P"): + return True + return False + + +def _is_end_of_word(text): + """Checks whether the last character in text is one of a punctuation, control or whitespace character.""" + last_char = text[-1] + return bool(_is_control(last_char) | _is_punctuation(last_char) | _is_whitespace(last_char)) + + +def _is_start_of_word(text): + """Checks whether the first character in text is one of a punctuation, control or whitespace character.""" + first_char = text[0] + return bool(_is_control(first_char) | _is_punctuation(first_char) | _is_whitespace(first_char)) + + +def _insert_one_token_to_ordered_list(token_list: list[str], new_token: str): + """ + Inserts one token to an ordered list if it does not already exist. Note: token_list must be sorted. + """ + insertion_idx = bisect.bisect_left(token_list, new_token) + # Checks if new_token is already in the ordered token_list + if insertion_idx < len(token_list) and token_list[insertion_idx] == new_token: + # new_token is in token_list, don't add + return + else: + token_list.insert(insertion_idx, new_token) + + +@add_end_docstrings(INIT_TOKENIZER_DOCSTRING) +class PreTrainedTokenizer(PreTrainedTokenizerBase): + """ + Base class for all slow tokenizers. + + Inherits from [`~tokenization_utils_base.PreTrainedTokenizerBase`]. + + Handle all the shared methods for tokenization and special tokens as well as methods downloading/caching/loading + pretrained tokenizers as well as adding tokens to the vocabulary. + + This class also contain the added tokens in a unified way on top of all tokenizers so we don't have to handle the + specific vocabulary augmentation methods of the various underlying dictionary structures (BPE, sentencepiece...). + """ + + def __init__(self, **kwargs): + # 1. Init the parent class + + self.tokens_trie = Trie() + + # 2. init `_added_tokens_decoder` if child class did not + if not hasattr(self, "_added_tokens_decoder"): + self._added_tokens_decoder: dict[int, AddedToken] = {} + + # 3. if a `added_tokens_decoder` is passed, we are loading from a saved tokenizer, we overwrite + self._added_tokens_decoder.update(kwargs.pop("added_tokens_decoder", {})) + self._added_tokens_encoder: dict[str, int] = {k.content: v for v, k in self._added_tokens_decoder.items()} + + # 4 init the parent class + super().__init__(**kwargs) + + # 4. If some of the special tokens are not part of the vocab, we add them, at the end. + # the order of addition is the same as self.SPECIAL_TOKENS_ATTRIBUTES following `tokenizers` + self._add_tokens( + [token for token in self.all_special_tokens_extended if token not in self._added_tokens_encoder], + special_tokens=True, + ) + + self._decode_use_source_tokenizer = False + + @property + def is_fast(self) -> bool: + return False + + @property + def vocab_size(self) -> int: + """ + `int`: Size of the base vocabulary (without the added tokens). + """ + raise NotImplementedError + + @property + def added_tokens_encoder(self) -> dict[str, int]: + """ + Returns the sorted mapping from string to index. The added tokens encoder is cached for performance + optimisation in `self._added_tokens_encoder` for the slow tokenizers. + """ + return {k.content: v for v, k in sorted(self._added_tokens_decoder.items(), key=lambda item: item[0])} + + @property + def added_tokens_decoder(self) -> dict[int, AddedToken]: + """ + Returns the added tokens in the vocabulary as a dictionary of index to AddedToken. + + Returns: + `Dict[str, int]`: The added tokens. + """ + return dict(sorted(self._added_tokens_decoder.items(), key=lambda item: item[0])) + + @added_tokens_decoder.setter + def added_tokens_decoder(self, value: dict[int, Union[AddedToken, str]]) -> dict[int, AddedToken]: + # Always raise an error if string because users should define the behavior + for index, token in value.items(): + if not isinstance(token, (str, AddedToken)) or not isinstance(index, int): + raise TypeError( + f"The provided `added_tokens_decoder` has an element of type {index.__class__, token.__class__}, should be a dict of {int, Union[AddedToken, str]}" + ) + + self._added_tokens_decoder[index] = AddedToken(token) if isinstance(token, str) else token + self._added_tokens_encoder[str(token)] = index + self._update_total_vocab_size() + + def get_added_vocab(self) -> dict[str, int]: + """ + Returns the added tokens in the vocabulary as a dictionary of token to index. Results might be different from + the fast call because for now we always add the tokens even if they are already in the vocabulary. This is + something we should change. + + Returns: + `Dict[str, int]`: The added tokens. + """ + return self._added_tokens_encoder + + def __len__(self): + """ + Size of the full vocabulary with the added tokens. + """ + return self.total_vocab_size + + def _update_total_vocab_size(self): + """ + Update the size of the full vocabulary with the added tokens. Counts the `keys` and not the `values` because + otherwise if there is a hole in the vocab, we will add tokenizers at a wrong index. This operation is slow and + is only updated when adding tokens. + """ + self.total_vocab_size = len(self.get_vocab()) + + def _add_tokens(self, new_tokens: Union[list[str], list[AddedToken]], special_tokens: bool = False) -> int: + """ + Add a list of new tokens to the tokenizer class. If the new tokens are not in the vocabulary, they are added to + it with indices starting from length of the current vocabulary. Special tokens are sometimes already in the + vocab which is why they have to be handled specifically. + + Args: + new_tokens (`List[str]`or `List[tokenizers.AddedToken]`): + Token(s) to add in vocabulary. A token is counted as added if it's not already in the vocabulary + (tested by checking if the tokenizer assign the index of the `unk_token` to them). If a token is part + of the vocabulary then we simply mark this token as an `AddedToken` which allows to control the + stripping and normalization of this token. This is NOT possible in `tokenizers`. + special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the tokens should be added as special tokens. + + Returns: + `int`: The number of tokens actually added to the vocabulary. + + Examples: + + ```python + # Let's see how to increase the vocabulary of Bert model and tokenizer + tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased") + model = BertModel.from_pretrained("google-bert/bert-base-uncased") + + num_added_toks = tokenizer.add_tokens(["new_tok1", "my_new-tok2"]) + print("We have added", num_added_toks, "tokens") + # Note: resize_token_embeddings expects to receive the full size of the new vocabulary, i.e. the length of the tokenizer. + model.resize_token_embeddings(len(tokenizer)) + ```""" + added_tokens = 0 + if new_tokens is None: + return added_tokens + # TODO this is fairly slow to improve! + current_vocab = self.get_vocab().copy() + new_idx = len(current_vocab) # only call this once, len gives the last index + 1 + for token in new_tokens: + if not isinstance(token, (str, AddedToken)): + raise TypeError(f"Token {token} is not a string but a {type(token)}.") + if str(token) == "": + continue + if isinstance(token, str): + if token in self._added_tokens_encoder: + continue + else: + # very important for fast and slow equivalence! + is_special = token in self.all_special_tokens or special_tokens + token = AddedToken( + token, rstrip=False, lstrip=False, normalized=not is_special, special=is_special + ) + elif special_tokens: + # doing token.special=True changes the normalization! will fix in rust + # this is important and the only reason why the AddedTokens in each class are normalized by default + token.__setstate__({"special": True, "normalized": token.normalized}) + if token in self._added_tokens_decoder: + continue + if not token.special and token.normalized and getattr(self, "do_lower_case", False): + # Normalize if requested + token.content = token.content.lower() + if token.content not in current_vocab: + token_index = new_idx + added_tokens + current_vocab[token.content] = token_index + added_tokens += 1 + else: + token_index = current_vocab[token.content] + + if token.special and str(token) not in self.all_special_tokens: + self._special_tokens_map["additional_special_tokens"].append(token) + # the setter automatically updates the reverse map + self._added_tokens_decoder[token_index] = token + self._added_tokens_encoder[token.content] = token_index + if self.verbose: + logger.info(f"Adding {token} to the vocabulary") + + self._update_trie() + self._update_total_vocab_size() + return added_tokens + + def _update_trie(self, unique_no_split_tokens: Optional[str] = []): + for token in self._added_tokens_decoder.values(): + if token not in self.tokens_trie._tokens: + self.tokens_trie.add(token.content) + for token in unique_no_split_tokens: + if token not in self.tokens_trie._tokens: + self.tokens_trie.add(token) + + def num_special_tokens_to_add(self, pair: bool = False) -> int: + """ + Returns the number of added tokens when encoding a sequence with special tokens. + + + + This encodes a dummy input and checks the number of added tokens, and is therefore not efficient. Do not put + this inside your training loop. + + + + Args: + pair (`bool`, *optional*, defaults to `False`): + Whether the number of added tokens should be computed in the case of a sequence pair or a single + sequence. + + Returns: + `int`: Number of special tokens added to sequences. + """ + token_ids_0 = [] + token_ids_1 = [] + return len(self.build_inputs_with_special_tokens(token_ids_0, token_ids_1 if pair else None)) + + def tokenize(self, text: TextInput, **kwargs) -> list[str]: + """ + Converts a string into a sequence of tokens, using the tokenizer. + + Split in words for word-based vocabulary or sub-words for sub-word-based vocabularies + (BPE/SentencePieces/WordPieces). Takes care of added tokens. + + Args: + text (`str`): + The sequence to be encoded. + **kwargs (additional keyword arguments): + Passed along to the model-specific `prepare_for_tokenization` preprocessing method. + + Returns: + `List[str]`: The list of tokens. + """ + split_special_tokens = kwargs.pop("split_special_tokens", self.split_special_tokens) + + text, kwargs = self.prepare_for_tokenization(text, **kwargs) + + if kwargs: + logger.warning(f"Keyword arguments {kwargs} not recognized.") + + if hasattr(self, "do_lower_case") and self.do_lower_case: + # convert non-special tokens to lowercase. Might be super slow as well? + escaped_special_toks = [re.escape(s_tok) for s_tok in (self.all_special_tokens)] + escaped_special_toks += [ + re.escape(s_tok.content) + for s_tok in (self._added_tokens_decoder.values()) + if not s_tok.special and s_tok.normalized + ] + pattern = r"(" + r"|".join(escaped_special_toks) + r")|" + r"(.+?)" + text = re.sub(pattern, lambda m: m.groups()[0] or m.groups()[1].lower(), text) + + if split_special_tokens: + no_split_token = [] + tokens = [text] + else: + no_split_token = self._added_tokens_encoder.keys() # don't split on any of the added tokens + # "This is something else" + tokens = self.tokens_trie.split(text) + + # ["This is something", "", " else"] + for i, token in enumerate(tokens): + if token in no_split_token: + tok_extended = self._added_tokens_decoder.get(self._added_tokens_encoder[token], None) + left = tokens[i - 1] if i > 0 else None + right = tokens[i + 1] if i < len(tokens) - 1 else None + if isinstance(tok_extended, AddedToken): + if tok_extended.rstrip and right: + # A bit counter-intuitive but we strip the left of the string + # since tok_extended.rstrip means the special token is eating all white spaces on its right + tokens[i + 1] = right.lstrip() + # Strip white spaces on the left + if tok_extended.lstrip and left: + tokens[i - 1] = left.rstrip() # Opposite here + if tok_extended.single_word and left and left[-1] != " ": + tokens[i - 1] += token + tokens[i] = "" + elif tok_extended.single_word and right and right[0] != " ": + tokens[i + 1] = token + tokens[i + 1] + tokens[i] = "" + else: + raise ValueError( + f"{tok_extended} cannot be tokenized because it was not properly added" + f" to the tokenizer. This means that it is not an `AddedToken` but a {type(tok_extended)}" + ) + # ["This is something", "", "else"] + tokenized_text = [] + for token in tokens: + # Need to skip eventual empty (fully stripped) tokens + if not token: + continue + if token in no_split_token: + tokenized_text.append(token) + else: + tokenized_text.extend(self._tokenize(token)) + # ["This", " is", " something", "", "else"] + return tokenized_text + + def _tokenize(self, text, **kwargs): + """ + Converts a string into a sequence of tokens (string), using the tokenizer. Split in words for word-based + vocabulary or sub-words for sub-word-based vocabularies (BPE/SentencePieces/WordPieces). + + Do NOT take care of added tokens. + """ + raise NotImplementedError + + def convert_tokens_to_ids(self, tokens: Union[str, list[str]]) -> Union[int, list[int]]: + """ + Converts a token string (or a sequence of tokens) in a single integer id (or a sequence of ids), using the + vocabulary. + + Args: + tokens (`str` or `List[str]`): One or several token(s) to convert to token id(s). + + Returns: + `int` or `List[int]`: The token id or list of token ids. + """ + if tokens is None: + return None + + if isinstance(tokens, str): + return self._convert_token_to_id_with_added_voc(tokens) + + ids = [] + for token in tokens: + ids.append(self._convert_token_to_id_with_added_voc(token)) + return ids + + def _convert_token_to_id_with_added_voc(self, token): + if token is None: + return None + + if token in self._added_tokens_encoder: + return self._added_tokens_encoder[token] + return self._convert_token_to_id(token) + + def _convert_token_to_id(self, token): + raise NotImplementedError + + def _encode_plus( + self, + text: Union[TextInput, PreTokenizedInput, EncodedInput], + text_pair: Optional[Union[TextInput, PreTokenizedInput, EncodedInput]] = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + def get_input_ids(text): + if isinstance(text, str): + tokens = self.tokenize(text, **kwargs) + return self.convert_tokens_to_ids(tokens) + elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], str): + if is_split_into_words: + tokens = list( + itertools.chain(*(self.tokenize(t, is_split_into_words=True, **kwargs) for t in text)) + ) + return self.convert_tokens_to_ids(tokens) + else: + return self.convert_tokens_to_ids(text) + elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], int): + return text + else: + if is_split_into_words: + raise ValueError( + f"Input {text} is not valid. Should be a string or a list/tuple of strings when" + " `is_split_into_words=True`." + ) + else: + raise ValueError( + f"Input {text} is not valid. Should be a string, a list/tuple of strings or a list/tuple of" + " integers." + ) + + if return_offsets_mapping: + raise NotImplementedError( + "return_offset_mapping is not available when using Python tokenizers. " + "To use this feature, change your tokenizer to one deriving from " + "transformers.PreTrainedTokenizerFast. " + "More information on available tokenizers at " + "https://github.com/huggingface/transformers/pull/2674" + ) + + first_ids = get_input_ids(text) + second_ids = get_input_ids(text_pair) if text_pair is not None else None + + return self.prepare_for_model( + first_ids, + pair_ids=second_ids, + add_special_tokens=add_special_tokens, + padding=padding_strategy.value, + truncation=truncation_strategy.value, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + prepend_batch_axis=True, + return_attention_mask=return_attention_mask, + return_token_type_ids=return_token_type_ids, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + verbose=verbose, + ) + + def _batch_encode_plus( + self, + batch_text_or_text_pairs: Union[ + list[TextInput], + list[TextInputPair], + list[PreTokenizedInput], + list[PreTokenizedInputPair], + list[EncodedInput], + list[EncodedInputPair], + ], + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + def get_input_ids(text): + if isinstance(text, str): + tokens = self.tokenize(text, **kwargs) + return self.convert_tokens_to_ids(tokens) + elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], str): + if is_split_into_words: + tokens = list( + itertools.chain(*(self.tokenize(t, is_split_into_words=True, **kwargs) for t in text)) + ) + return self.convert_tokens_to_ids(tokens) + else: + return self.convert_tokens_to_ids(text) + elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], int): + return text + else: + raise ValueError( + "Input is not valid. Should be a string, a list/tuple of strings or a list/tuple of integers." + ) + + if return_offsets_mapping: + raise NotImplementedError( + "return_offset_mapping is not available when using Python tokenizers. " + "To use this feature, change your tokenizer to one deriving from " + "transformers.PreTrainedTokenizerFast." + ) + + input_ids = [] + for ids_or_pair_ids in batch_text_or_text_pairs: + if not isinstance(ids_or_pair_ids, (list, tuple)): + ids, pair_ids = ids_or_pair_ids, None + elif is_split_into_words and not isinstance(ids_or_pair_ids[0], (list, tuple)): + ids, pair_ids = ids_or_pair_ids, None + else: + ids, pair_ids = ids_or_pair_ids + + first_ids = get_input_ids(ids) + second_ids = get_input_ids(pair_ids) if pair_ids is not None else None + input_ids.append((first_ids, second_ids)) + + batch_outputs = self._batch_prepare_for_model( + input_ids, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + return_token_type_ids=return_token_type_ids, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + return_tensors=return_tensors, + verbose=verbose, + split_special_tokens=split_special_tokens, + ) + + return BatchEncoding(batch_outputs) + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def _batch_prepare_for_model( + self, + batch_ids_pairs: list[Union[PreTokenizedInputPair, tuple[list[int], None]]], + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[str] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + ) -> BatchEncoding: + """ + Prepares a sequence of input id, or a pair of sequences of inputs ids so that it can be used by the model. It + adds special tokens, truncates sequences if overflowing while taking into account the special tokens and + manages a moving window (with user defined stride) for overflowing tokens + + Args: + batch_ids_pairs: list of tokenized input ids or input ids pairs + """ + + batch_outputs = {} + for first_ids, second_ids in batch_ids_pairs: + outputs = self.prepare_for_model( + first_ids, + second_ids, + add_special_tokens=add_special_tokens, + padding=PaddingStrategy.DO_NOT_PAD.value, # we pad in batch afterward + truncation=truncation_strategy.value, + max_length=max_length, + stride=stride, + pad_to_multiple_of=None, # we pad in batch afterward + padding_side=None, # we pad in batch afterward + return_attention_mask=False, # we pad in batch afterward + return_token_type_ids=return_token_type_ids, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_length=return_length, + return_tensors=None, # We convert the whole batch to tensors at the end + prepend_batch_axis=False, + verbose=verbose, + split_special_tokens=split_special_tokens, + ) + + for key, value in outputs.items(): + if key not in batch_outputs: + batch_outputs[key] = [] + batch_outputs[key].append(value) + + batch_outputs = self.pad( + batch_outputs, + padding=padding_strategy.value, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + batch_outputs = BatchEncoding(batch_outputs, tensor_type=return_tensors) + + return batch_outputs + + def prepare_for_tokenization( + self, text: str, is_split_into_words: bool = False, **kwargs + ) -> tuple[str, dict[str, Any]]: + """ + Performs any necessary transformations before tokenization. + + This method should pop the arguments from kwargs and return the remaining `kwargs` as well. We test the + `kwargs` at the end of the encoding process to be sure all the arguments have been used. + + Args: + text (`str`): + The text to prepare. + is_split_into_words (`bool`, *optional*, defaults to `False`): + Whether or not the input is already pre-tokenized (e.g., split into words). If set to `True`, the + tokenizer assumes the input is already split into words (for instance, by splitting it on whitespace) + which it will tokenize. This is useful for NER or token classification. + kwargs (`Dict[str, Any]`, *optional*): + Keyword arguments to use for the tokenization. + + Returns: + `Tuple[str, Dict[str, Any]]`: The prepared text and the unused kwargs. + """ + return (text, kwargs) + + def get_special_tokens_mask( + self, token_ids_0: list, token_ids_1: Optional[list] = None, already_has_special_tokens: bool = False + ) -> list[int]: + """ + Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding + special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods. + + Args: + token_ids_0 (`List[int]`): + List of ids of the first sequence. + token_ids_1 (`List[int]`, *optional*): + List of ids of the second sequence. + already_has_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the token list is already formatted with special tokens for the model. + + Returns: + A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + if already_has_special_tokens: + if token_ids_1 is not None: + raise ValueError( + "You should not supply a second sequence if the provided sequence of " + "ids is already formatted with special tokens for the model." + ) + + return super().get_special_tokens_mask( + token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True + ) + return [0] * ((len(token_ids_1) if token_ids_1 else 0) + len(token_ids_0)) + + @overload + def convert_ids_to_tokens(self, ids: int, skip_special_tokens: bool = False) -> str: ... + + @overload + def convert_ids_to_tokens(self, ids: list[int], skip_special_tokens: bool = False) -> list[str]: ... + + def convert_ids_to_tokens( + self, ids: Union[int, list[int]], skip_special_tokens: bool = False + ) -> Union[str, list[str]]: + """ + Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and + added tokens. + + Args: + ids (`int` or `List[int]`): + The token id (or token ids) to convert to tokens. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + + Returns: + `str` or `List[str]`: The decoded token(s). + """ + if isinstance(ids, int): + if ids in self._added_tokens_decoder: + return self._added_tokens_decoder[ids].content + else: + return self._convert_id_to_token(ids) + tokens = [] + for index in ids: + index = int(index) + if skip_special_tokens and index in self.all_special_ids: + continue + if index in self._added_tokens_decoder: + tokens.append(self._added_tokens_decoder[index].content) + else: + tokens.append(self._convert_id_to_token(index)) + return tokens + + def _convert_id_to_token(self, index: int) -> str: + raise NotImplementedError + + def convert_tokens_to_string(self, tokens: list[str]) -> str: + return " ".join(tokens) + + def _decode( + self, + token_ids: Union[int, list[int]], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = None, + spaces_between_special_tokens: bool = True, + **kwargs, + ) -> str: + self._decode_use_source_tokenizer = kwargs.pop("use_source_tokenizer", False) + + filtered_tokens = self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens) + # If given is a single id, prevents splitting the string in upcoming loop + if isinstance(filtered_tokens, str): + filtered_tokens = [filtered_tokens] + + legacy_added_tokens = set(self._added_tokens_encoder.keys()) - set(self.all_special_tokens) | { + token for token in self.additional_special_tokens if self.convert_tokens_to_ids(token) >= self.vocab_size + } + # To avoid mixing byte-level and unicode for byte-level BPT + # we need to build string separately for added tokens and byte-level tokens + # cf. https://github.com/huggingface/transformers/issues/1133 + sub_texts = [] + current_sub_text = [] + # TODO @ArthurZ in version 5, special tokens should be handled in convert_tokens_to_string, while _convert_tokens_to_string + for token in filtered_tokens: + if skip_special_tokens and token in self.all_special_tokens: + continue + if token in legacy_added_tokens: + if current_sub_text: + string = self.convert_tokens_to_string(current_sub_text) + if len(string) > 0: + sub_texts.append(string) + current_sub_text = [] + sub_texts.append(token) + else: + current_sub_text.append(token) + if current_sub_text: + sub_texts.append(self.convert_tokens_to_string(current_sub_text)) + + if spaces_between_special_tokens: + text = " ".join(sub_texts) + else: + text = "".join(sub_texts) + + clean_up_tokenization_spaces = ( + clean_up_tokenization_spaces + if clean_up_tokenization_spaces is not None + else self.clean_up_tokenization_spaces + ) + if clean_up_tokenization_spaces: + clean_text = self.clean_up_tokenization(text) + return clean_text + else: + return text diff --git a/lib/python3.12/site-packages/transformers/tokenization_utils_base.py b/lib/python3.12/site-packages/transformers/tokenization_utils_base.py new file mode 100644 index 0000000000000000000000000000000000000000..8e047569e7ce56573949b4e03aa9f1057802bf90 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/tokenization_utils_base.py @@ -0,0 +1,4176 @@ +# coding=utf-8 +# Copyright 2020 The HuggingFace Inc. team. +# +# 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. +""" +Base classes common to both the slow and the fast tokenization classes: PreTrainedTokenizerBase (host all the user +fronting encoding methods) Special token mixing (host the special tokens logic) and BatchEncoding (wrap the dictionary +of output with special method for the Fast tokenizers) +""" + +import copy +import json +import os +import re +import warnings +from collections import UserDict +from collections.abc import Mapping, Sized +from contextlib import contextmanager +from dataclasses import dataclass +from inspect import isfunction +from typing import TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union + +import numpy as np +from packaging import version + +from . import __version__ +from .dynamic_module_utils import custom_object_save +from .utils import ( + ExplicitEnum, + PaddingStrategy, + PushToHubMixin, + TensorType, + add_end_docstrings, + add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, + cached_file, + copy_func, + download_url, + extract_commit_hash, + get_json_schema, + is_flax_available, + is_jax_tensor, + is_mlx_available, + is_numpy_array, + is_offline_mode, + is_protobuf_available, + is_remote_url, + is_tf_available, + is_tf_tensor, + is_tokenizers_available, + is_torch_available, + is_torch_device, + is_torch_tensor, + logging, + requires_backends, + to_py_obj, +) +from .utils.chat_template_utils import _compile_jinja_template, _render_with_assistant_indices +from .utils.import_utils import PROTOBUF_IMPORT_ERROR + + +if TYPE_CHECKING: + if is_torch_available(): + import torch + if is_tf_available(): + import tensorflow as tf + if is_flax_available(): + import jax.numpy as jnp # noqa: F401 + + +def import_protobuf_decode_error(error_message=""): + if is_protobuf_available(): + from google.protobuf.message import DecodeError + + return DecodeError + else: + raise ImportError(PROTOBUF_IMPORT_ERROR.format(error_message)) + + +if is_tokenizers_available(): + from tokenizers import AddedToken + from tokenizers import Encoding as EncodingFast +else: + + @dataclass(frozen=False, eq=True) + class AddedToken: + """ + AddedToken represents a token to be added to a Tokenizer An AddedToken can have special options defining the + way it should behave. + + The `normalized` will default to `not special` if it is not specified, similarly to the definition in + `tokenizers`. + """ + + def __init__( + self, content: str, single_word=False, lstrip=False, rstrip=False, special=False, normalized=None + ): + self.content = content + self.single_word = single_word + self.lstrip = lstrip + self.rstrip = rstrip + self.special = special + self.normalized = normalized if normalized is not None else not special + + def __getstate__(self): + return self.__dict__ + + def __str__(self): + return self.content + + @dataclass + class EncodingFast: + """This is dummy class because without the `tokenizers` library we don't have these objects anyway""" + + pass + + +logger = logging.get_logger(__name__) + +VERY_LARGE_INTEGER = int(1e30) # This is used to set the max input length for a model with infinite size input +LARGE_INTEGER = int(1e20) # This is used when we need something big but slightly smaller than VERY_LARGE_INTEGER + +# Define type aliases and NamedTuples +TextInput = str +PreTokenizedInput = List[str] +EncodedInput = List[int] +TextInputPair = Tuple[str, str] +PreTokenizedInputPair = Tuple[List[str], List[str]] +EncodedInputPair = Tuple[List[int], List[int]] + +# Define type aliases for text-related non-text modalities +AudioInput = Union["np.ndarray", "torch.Tensor", List["np.ndarray"], List["torch.Tensor"]] + +# Slow tokenizers used to be saved in three separated files +SPECIAL_TOKENS_MAP_FILE = "special_tokens_map.json" +ADDED_TOKENS_FILE = "added_tokens.json" +TOKENIZER_CONFIG_FILE = "tokenizer_config.json" +CHAT_TEMPLATE_FILE = "chat_template.jinja" + +# Fast tokenizers (provided by HuggingFace tokenizer's library) can be saved in a single file +FULL_TOKENIZER_FILE = "tokenizer.json" +_re_tokenizer_file = re.compile(r"tokenizer\.(.*)\.json") + + +class TruncationStrategy(ExplicitEnum): + """ + Possible values for the `truncation` argument in [`PreTrainedTokenizerBase.__call__`]. Useful for tab-completion in + an IDE. + """ + + ONLY_FIRST = "only_first" + ONLY_SECOND = "only_second" + LONGEST_FIRST = "longest_first" + DO_NOT_TRUNCATE = "do_not_truncate" + + +class CharSpan(NamedTuple): + """ + Character span in the original string. + + Args: + start (`int`): Index of the first character in the original string. + end (`int`): Index of the character following the last character in the original string. + """ + + start: int + end: int + + +class TokenSpan(NamedTuple): + """ + Token span in an encoded string (list of tokens). + + Args: + start (`int`): Index of the first token in the span. + end (`int`): Index of the token following the last token in the span. + """ + + start: int + end: int + + +class BatchEncoding(UserDict): + """ + Holds the output of the [`~tokenization_utils_base.PreTrainedTokenizerBase.__call__`], + [`~tokenization_utils_base.PreTrainedTokenizerBase.encode_plus`] and + [`~tokenization_utils_base.PreTrainedTokenizerBase.batch_encode_plus`] methods (tokens, attention_masks, etc). + + This class is derived from a python dictionary and can be used as a dictionary. In addition, this class exposes + utility methods to map from word/character space to token space. + + Args: + data (`dict`, *optional*): + Dictionary of lists/arrays/tensors returned by the `__call__`/`encode_plus`/`batch_encode_plus` methods + ('input_ids', 'attention_mask', etc.). + encoding (`tokenizers.Encoding` or `Sequence[tokenizers.Encoding]`, *optional*): + If the tokenizer is a fast tokenizer which outputs additional information like mapping from word/character + space to token space the `tokenizers.Encoding` instance or list of instance (for batches) hold this + information. + tensor_type (`Union[None, str, TensorType]`, *optional*): + You can give a tensor_type here to convert the lists of integers in PyTorch/TensorFlow/Numpy Tensors at + initialization. + prepend_batch_axis (`bool`, *optional*, defaults to `False`): + Whether or not to add a batch axis when converting to tensors (see `tensor_type` above). Note that this + parameter has an effect if the parameter `tensor_type` is set, *otherwise has no effect*. + n_sequences (`Optional[int]`, *optional*): + You can give a tensor_type here to convert the lists of integers in PyTorch/TensorFlow/Numpy Tensors at + initialization. + """ + + def __init__( + self, + data: Optional[Dict[str, Any]] = None, + encoding: Optional[Union[EncodingFast, Sequence[EncodingFast]]] = None, + tensor_type: Union[None, str, TensorType] = None, + prepend_batch_axis: bool = False, + n_sequences: Optional[int] = None, + ): + super().__init__(data) + + if isinstance(encoding, EncodingFast): + encoding = [encoding] + + self._encodings = encoding + + if n_sequences is None and encoding is not None and len(encoding): + n_sequences = encoding[0].n_sequences + + self._n_sequences = n_sequences + + self.convert_to_tensors(tensor_type=tensor_type, prepend_batch_axis=prepend_batch_axis) + + @property + def n_sequences(self) -> Optional[int]: + """ + `Optional[int]`: The number of sequences used to generate each sample from the batch encoded in this + [`BatchEncoding`]. Currently can be one of `None` (unknown), `1` (a single sentence) or `2` (a pair of + sentences) + """ + return self._n_sequences + + @property + def is_fast(self) -> bool: + """ + `bool`: Indicate whether this [`BatchEncoding`] was generated from the result of a [`PreTrainedTokenizerFast`] + or not. + """ + return self._encodings is not None + + def __getitem__(self, item: Union[int, str]) -> Union[Any, EncodingFast]: + """ + If the key is a string, returns the value of the dict associated to `key` ('input_ids', 'attention_mask', + etc.). + + If the key is an integer, get the `tokenizers.Encoding` for batch item with index `key`. + + If the key is a slice, returns the value of the dict associated to `key` ('input_ids', 'attention_mask', etc.) + with the constraint of slice. + """ + if isinstance(item, str): + return self.data[item] + elif self._encodings is not None: + return self._encodings[item] + elif isinstance(item, slice): + return {key: self.data[key][item] for key in self.data.keys()} + else: + raise KeyError( + "Invalid key. Only three types of key are available: " + "(1) string, (2) integers for backend Encoding, and (3) slices for data subsetting." + ) + + def __getattr__(self, item: str): + try: + return self.data[item] + except KeyError: + raise AttributeError + + def __getstate__(self): + return {"data": self.data, "encodings": self._encodings} + + def __setstate__(self, state): + if "data" in state: + self.data = state["data"] + + if "encodings" in state: + self._encodings = state["encodings"] + + def keys(self): + return self.data.keys() + + def values(self): + return self.data.values() + + def items(self): + return self.data.items() + + # After this point: + # Extended properties and methods only available for fast (Rust-based) tokenizers + # provided by HuggingFace tokenizers library. + + @property + def encodings(self) -> Optional[List[EncodingFast]]: + """ + `Optional[List[tokenizers.Encoding]]`: The list all encodings from the tokenization process. Returns `None` if + the input was tokenized through Python (i.e., not a fast) tokenizer. + """ + return self._encodings + + def tokens(self, batch_index: int = 0) -> List[str]: + """ + Return the list of tokens (sub-parts of the input strings after word/subword splitting and before conversion to + integer indices) at a given batch index (only works for the output of a fast tokenizer). + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `List[str]`: The list of tokens at that index. + """ + if not self._encodings: + raise ValueError( + "tokens() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`" + " class)." + ) + return self._encodings[batch_index].tokens + + def sequence_ids(self, batch_index: int = 0) -> List[Optional[int]]: + """ + Return a list mapping the tokens to the id of their original sentences: + + - `None` for special tokens added around or between sequences, + - `0` for tokens corresponding to words in the first sequence, + - `1` for tokens corresponding to words in the second sequence when a pair of sequences was jointly + encoded. + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `List[Optional[int]]`: A list indicating the sequence id corresponding to each token. Special tokens added + by the tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding + sequence. + """ + if not self._encodings: + raise ValueError( + "sequence_ids() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`" + " class)." + ) + return self._encodings[batch_index].sequence_ids + + def words(self, batch_index: int = 0) -> List[Optional[int]]: + """ + Return a list mapping the tokens to their actual word in the initial sentence for a fast tokenizer. + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `List[Optional[int]]`: A list indicating the word corresponding to each token. Special tokens added by the + tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding word + (several tokens will be mapped to the same word index if they are parts of that word). + """ + if not self._encodings: + raise ValueError( + "words() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`" + " class)." + ) + warnings.warn( + "`BatchEncoding.words()` property is deprecated and should be replaced with the identical, " + "but more self-explanatory `BatchEncoding.word_ids()` property.", + FutureWarning, + ) + return self.word_ids(batch_index) + + def word_ids(self, batch_index: int = 0) -> List[Optional[int]]: + """ + Return a list mapping the tokens to their actual word in the initial sentence for a fast tokenizer. + + Args: + batch_index (`int`, *optional*, defaults to 0): The index to access in the batch. + + Returns: + `List[Optional[int]]`: A list indicating the word corresponding to each token. Special tokens added by the + tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding word + (several tokens will be mapped to the same word index if they are parts of that word). + """ + if not self._encodings: + raise ValueError( + "word_ids() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`" + " class)." + ) + return self._encodings[batch_index].word_ids + + def token_to_sequence(self, batch_or_token_index: int, token_index: Optional[int] = None) -> int: + """ + Get the index of the sequence represented by the given token. In the general use case, this method returns `0` + for a single sequence or the first sequence of a pair, and `1` for the second sequence of a pair + + Can be called as: + + - `self.token_to_sequence(token_index)` if batch size is 1 + - `self.token_to_sequence(batch_index, token_index)` if batch size is greater than 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e., + words are defined by the user). In this case it allows to easily associate encoded tokens with provided + tokenized words. + + Args: + batch_or_token_index (`int`): + Index of the sequence in the batch. If the batch only comprises one sequence, this can be the index of + the token in the sequence. + token_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the token in the + sequence. + + Returns: + `int`: Index of the word in the input sequence. + """ + + if not self._encodings: + raise ValueError("token_to_sequence() is not available when using Python based tokenizers") + if token_index is not None: + batch_index = batch_or_token_index + else: + batch_index = 0 + token_index = batch_or_token_index + if batch_index < 0: + batch_index = self._batch_size + batch_index + if token_index < 0: + token_index = self._seq_len + token_index + return self._encodings[batch_index].token_to_sequence(token_index) + + def token_to_word(self, batch_or_token_index: int, token_index: Optional[int] = None) -> int: + """ + Get the index of the word corresponding (i.e. comprising) to an encoded token in a sequence of the batch. + + Can be called as: + + - `self.token_to_word(token_index)` if batch size is 1 + - `self.token_to_word(batch_index, token_index)` if batch size is greater than 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e., + words are defined by the user). In this case it allows to easily associate encoded tokens with provided + tokenized words. + + Args: + batch_or_token_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the token in the sequence. + token_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the token in the + sequence. + + Returns: + `int`: Index of the word in the input sequence. + """ + + if not self._encodings: + raise ValueError("token_to_word() is not available when using Python based tokenizers") + if token_index is not None: + batch_index = batch_or_token_index + else: + batch_index = 0 + token_index = batch_or_token_index + if batch_index < 0: + batch_index = self._batch_size + batch_index + if token_index < 0: + token_index = self._seq_len + token_index + return self._encodings[batch_index].token_to_word(token_index) + + def word_to_tokens( + self, batch_or_word_index: int, word_index: Optional[int] = None, sequence_index: int = 0 + ) -> Optional[TokenSpan]: + """ + Get the encoded token span corresponding to a word in a sequence of the batch. + + Token spans are returned as a [`~tokenization_utils_base.TokenSpan`] with: + + - **start** -- Index of the first token. + - **end** -- Index of the token following the last token. + + Can be called as: + + - `self.word_to_tokens(word_index, sequence_index: int = 0)` if batch size is 1 + - `self.word_to_tokens(batch_index, word_index, sequence_index: int = 0)` if batch size is greater or equal to + 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words + are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized + words. + + Args: + batch_or_word_index (`int`): + Index of the sequence in the batch. If the batch only comprises one sequence, this can be the index of + the word in the sequence. + word_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the + sequence. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided word index belongs to. + + Returns: + ([`~tokenization_utils_base.TokenSpan`], *optional*): Span of tokens in the encoded sequence. Returns + `None` if no tokens correspond to the word. This can happen especially when the token is a special token + that has been used to format the tokenization. For example when we add a class token at the very beginning + of the tokenization. + """ + + if not self._encodings: + raise ValueError("word_to_tokens() is not available when using Python based tokenizers") + if word_index is not None: + batch_index = batch_or_word_index + else: + batch_index = 0 + word_index = batch_or_word_index + if batch_index < 0: + batch_index = self._batch_size + batch_index + if word_index < 0: + word_index = self._seq_len + word_index + span = self._encodings[batch_index].word_to_tokens(word_index, sequence_index) + return TokenSpan(*span) if span is not None else None + + def token_to_chars(self, batch_or_token_index: int, token_index: Optional[int] = None) -> Optional[CharSpan]: + """ + Get the character span corresponding to an encoded token in a sequence of the batch. + + Character spans are returned as a [`~tokenization_utils_base.CharSpan`] with: + + - **start** -- Index of the first character in the original string associated to the token. + - **end** -- Index of the character following the last character in the original string associated to the + token. + + Can be called as: + + - `self.token_to_chars(token_index)` if batch size is 1 + - `self.token_to_chars(batch_index, token_index)` if batch size is greater or equal to 1 + + Args: + batch_or_token_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the token in the sequence. + token_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the token or tokens in + the sequence. + + Returns: + [`~tokenization_utils_base.CharSpan`]: Span of characters in the original string, or None, if the token + (e.g. , ) doesn't correspond to any chars in the origin string. + """ + + if not self._encodings: + raise ValueError("token_to_chars() is not available when using Python based tokenizers") + if token_index is not None: + batch_index = batch_or_token_index + else: + batch_index = 0 + token_index = batch_or_token_index + span_indices = self._encodings[batch_index].token_to_chars(token_index) + + return CharSpan(*span_indices) if span_indices is not None else None + + def char_to_token( + self, batch_or_char_index: int, char_index: Optional[int] = None, sequence_index: int = 0 + ) -> int: + """ + Get the index of the token in the encoded output comprising a character in the original string for a sequence + of the batch. + + Can be called as: + + - `self.char_to_token(char_index)` if batch size is 1 + - `self.char_to_token(batch_index, char_index)` if batch size is greater or equal to 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words + are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized + words. + + Args: + batch_or_char_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the word in the sequence + char_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the + sequence. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided character index belongs to. + + + Returns: + `int`: Index of the token, or None if the char index refers to a whitespace only token and whitespace is + trimmed with `trim_offsets=True`. + """ + + if not self._encodings: + raise ValueError("char_to_token() is not available when using Python based tokenizers") + if char_index is not None: + batch_index = batch_or_char_index + else: + batch_index = 0 + char_index = batch_or_char_index + return self._encodings[batch_index].char_to_token(char_index, sequence_index) + + def word_to_chars( + self, batch_or_word_index: int, word_index: Optional[int] = None, sequence_index: int = 0 + ) -> CharSpan: + """ + Get the character span in the original string corresponding to given word in a sequence of the batch. + + Character spans are returned as a CharSpan NamedTuple with: + + - start: index of the first character in the original string + - end: index of the character following the last character in the original string + + Can be called as: + + - `self.word_to_chars(word_index)` if batch size is 1 + - `self.word_to_chars(batch_index, word_index)` if batch size is greater or equal to 1 + + Args: + batch_or_word_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the word in the sequence + word_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the + sequence. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided word index belongs to. + + Returns: + `CharSpan` or `List[CharSpan]`: Span(s) of the associated character or characters in the string. CharSpan + are NamedTuple with: + + - start: index of the first character associated to the token in the original string + - end: index of the character following the last character associated to the token in the original + string + """ + + if not self._encodings: + raise ValueError("word_to_chars() is not available when using Python based tokenizers") + if word_index is not None: + batch_index = batch_or_word_index + else: + batch_index = 0 + word_index = batch_or_word_index + return CharSpan(*(self._encodings[batch_index].word_to_chars(word_index, sequence_index))) + + def char_to_word(self, batch_or_char_index: int, char_index: Optional[int] = None, sequence_index: int = 0) -> int: + """ + Get the word in the original string corresponding to a character in the original string of a sequence of the + batch. + + Can be called as: + + - `self.char_to_word(char_index)` if batch size is 1 + - `self.char_to_word(batch_index, char_index)` if batch size is greater than 1 + + This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words + are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized + words. + + Args: + batch_or_char_index (`int`): + Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of + the character in the original string. + char_index (`int`, *optional*): + If a batch index is provided in *batch_or_token_index*, this can be the index of the character in the + original string. + sequence_index (`int`, *optional*, defaults to 0): + If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0 + or 1) the provided character index belongs to. + + + Returns: + `int` or `List[int]`: Index or indices of the associated encoded token(s). + """ + + if not self._encodings: + raise ValueError("char_to_word() is not available when using Python based tokenizers") + if char_index is not None: + batch_index = batch_or_char_index + else: + batch_index = 0 + char_index = batch_or_char_index + return self._encodings[batch_index].char_to_word(char_index, sequence_index) + + def convert_to_tensors( + self, tensor_type: Optional[Union[str, TensorType]] = None, prepend_batch_axis: bool = False + ): + """ + Convert the inner content to tensors. + + Args: + tensor_type (`str` or [`~utils.TensorType`], *optional*): + The type of tensors to use. If `str`, should be one of the values of the enum [`~utils.TensorType`]. If + `None`, no modification is done. + prepend_batch_axis (`int`, *optional*, defaults to `False`): + Whether or not to add the batch dimension during the conversion. + """ + if tensor_type is None: + return self + + # Convert to TensorType + if not isinstance(tensor_type, TensorType): + tensor_type = TensorType(tensor_type) + + # Get a function reference for the correct framework + if tensor_type == TensorType.TENSORFLOW: + if not is_tf_available(): + raise ImportError( + "Unable to convert output to TensorFlow tensors format, TensorFlow is not installed." + ) + import tensorflow as tf + + as_tensor = tf.constant + is_tensor = tf.is_tensor + elif tensor_type == TensorType.PYTORCH: + if not is_torch_available(): + raise ImportError("Unable to convert output to PyTorch tensors format, PyTorch is not installed.") + import torch + + is_tensor = torch.is_tensor + + def as_tensor(value, dtype=None): + if isinstance(value, list) and isinstance(value[0], np.ndarray): + return torch.from_numpy(np.array(value)) + return torch.tensor(value) + + elif tensor_type == TensorType.JAX: + if not is_flax_available(): + raise ImportError("Unable to convert output to JAX tensors format, JAX is not installed.") + import jax.numpy as jnp # noqa: F811 + + as_tensor = jnp.array + is_tensor = is_jax_tensor + + elif tensor_type == TensorType.MLX: + if not is_mlx_available(): + raise ImportError("Unable to convert output to MLX tensors format, MLX is not installed.") + import mlx.core as mx + + as_tensor = mx.array + + def is_tensor(obj): + return isinstance(obj, mx.array) + else: + + def as_tensor(value, dtype=None): + if isinstance(value, (list, tuple)) and isinstance(value[0], (list, tuple, np.ndarray)): + value_lens = [len(val) for val in value] + if len(set(value_lens)) > 1 and dtype is None: + # we have a ragged list so handle explicitly + value = as_tensor([np.asarray(val) for val in value], dtype=object) + return np.asarray(value, dtype=dtype) + + is_tensor = is_numpy_array + + # Do the tensor conversion in batch + for key, value in self.items(): + try: + if prepend_batch_axis: + value = [value] + + if not is_tensor(value): + tensor = as_tensor(value) + + # Removing this for now in favor of controlling the shape with `prepend_batch_axis` + # # at-least2d + # if tensor.ndim > 2: + # tensor = tensor.squeeze(0) + # elif tensor.ndim < 2: + # tensor = tensor[None, :] + + self[key] = tensor + except Exception as e: + if key == "overflowing_tokens": + raise ValueError( + "Unable to create tensor returning overflowing tokens of different lengths. " + "Please see if a fast version of this tokenizer is available to have this feature available." + ) from e + raise ValueError( + "Unable to create tensor, you should probably activate truncation and/or padding with" + " 'padding=True' 'truncation=True' to have batched tensors with the same length. Perhaps your" + f" features (`{key}` in this case) have excessive nesting (inputs type `list` where type `int` is" + " expected)." + ) from e + + return self + + def to(self, device: Union[str, "torch.device"], *, non_blocking: bool = False) -> "BatchEncoding": + """ + Send all values to device by calling `v.to(device, non_blocking=non_blocking)` (PyTorch only). + + Args: + device (`str` or `torch.device`): The device to put the tensors on. + non_blocking (`bool`): Whether to perform the copy asynchronously. + + Returns: + [`BatchEncoding`]: The same instance after modification. + """ + requires_backends(self, ["torch"]) + import torch + + # This check catches things like APEX blindly calling "to" on all inputs to a module + # Otherwise it passes the casts down and casts the LongTensor containing the token idxs + # into a HalfTensor + if isinstance(device, str) or is_torch_device(device) or isinstance(device, int): + self.data = { + k: v.to(device=device, non_blocking=non_blocking) if isinstance(v, torch.Tensor) else v + for k, v in self.data.items() + } + else: + logger.warning(f"Attempting to cast a BatchEncoding to type {str(device)}. This is not supported.") + return self + + +class SpecialTokensMixin: + """ + A mixin derived by [`PreTrainedTokenizer`] and [`PreTrainedTokenizerFast`] to handle specific behaviors related to + special tokens. In particular, this class hold the attributes which can be used to directly access these special + tokens in a model-independent manner and allow to set and update the special tokens. + + Args: + bos_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the beginning of a sentence. + eos_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the end of a sentence. + unk_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing an out-of-vocabulary token. + sep_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token separating two different sentences in the same input (used by BERT for instance). + pad_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by + attention mechanisms or loss computation. + cls_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the class of the input (used by BERT for instance). + mask_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing a masked token (used by masked-language modeling pretraining objectives, like + BERT). + additional_special_tokens (tuple or list of `str` or `tokenizers.AddedToken`, *optional*): + A tuple or a list of additional tokens, which will be marked as `special`, meaning that they will be + skipped when decoding if `skip_special_tokens` is set to `True`. + """ + + SPECIAL_TOKENS_ATTRIBUTES = [ + "bos_token", + "eos_token", + "unk_token", + "sep_token", + "pad_token", + "cls_token", + "mask_token", + "additional_special_tokens", + ] + + def __init__(self, verbose=False, **kwargs): + self._pad_token_type_id = 0 + self.verbose = verbose + self._special_tokens_map = dict.fromkeys(self.SPECIAL_TOKENS_ATTRIBUTES) + self._special_tokens_map["additional_special_tokens"] = [] # for BC where it defaults to empty list + + # We directly set the hidden value to allow initialization with special tokens + # which are not yet in the vocabulary. Necessary for serialization/de-serialization + # TODO clean this up at some point (probably by switching to fast tokenizers) + + for key, value in kwargs.items(): + if value is None: + continue + if key in self.SPECIAL_TOKENS_ATTRIBUTES: + if key == "additional_special_tokens": + assert isinstance(value, (list, tuple)), f"Value {value} is not a list or tuple" + assert all(isinstance(t, (str, AddedToken)) for t in value), ( + "One of the tokens is not a string or an AddedToken" + ) + setattr(self, key, value) + elif isinstance(value, (str, AddedToken)): + setattr(self, key, value) + else: + raise TypeError(f"Special token {key} has to be either str or AddedToken but got: {type(value)}") + + def sanitize_special_tokens(self) -> int: + """ + The `sanitize_special_tokens` is now deprecated kept for backward compatibility and will be removed in + transformers v5. + """ + logger.warning_once("The `sanitize_special_tokens` will be removed in transformers v5.") + return self.add_tokens(self.all_special_tokens_extended, special_tokens=True) + + def add_special_tokens( + self, special_tokens_dict: Dict[str, Union[str, AddedToken]], replace_additional_special_tokens=True + ) -> int: + """ + Add a dictionary of special tokens (eos, pad, cls, etc.) to the encoder and link them to class attributes. If + special tokens are NOT in the vocabulary, they are added to it (indexed starting from the last index of the + current vocabulary). + + When adding new tokens to the vocabulary, you should make sure to also resize the token embedding matrix of the + model so that its embedding matrix matches the tokenizer. + + In order to do that, please use the [`~PreTrainedModel.resize_token_embeddings`] method. + + Using `add_special_tokens` will ensure your special tokens can be used in several ways: + + - Special tokens can be skipped when decoding using `skip_special_tokens = True`. + - Special tokens are carefully handled by the tokenizer (they are never split), similar to `AddedTokens`. + - You can easily refer to special tokens using tokenizer class attributes like `tokenizer.cls_token`. This + makes it easy to develop model-agnostic training and fine-tuning scripts. + + When possible, special tokens are already registered for provided pretrained models (for instance + [`BertTokenizer`] `cls_token` is already registered to be :obj*'[CLS]'* and XLM's one is also registered to be + `''`). + + Args: + special_tokens_dict (dictionary *str* to *str* or `tokenizers.AddedToken`): + Keys should be in the list of predefined special attributes: [`bos_token`, `eos_token`, `unk_token`, + `sep_token`, `pad_token`, `cls_token`, `mask_token`, `additional_special_tokens`]. + + Tokens are only added if they are not already in the vocabulary (tested by checking if the tokenizer + assign the index of the `unk_token` to them). + replace_additional_special_tokens (`bool`, *optional*,, defaults to `True`): + If `True`, the existing list of additional special tokens will be replaced by the list provided in + `special_tokens_dict`. Otherwise, `self._special_tokens_map["additional_special_tokens"]` is just extended. In the former + case, the tokens will NOT be removed from the tokenizer's full vocabulary - they are only being flagged + as non-special tokens. Remember, this only affects which tokens are skipped during decoding, not the + `added_tokens_encoder` and `added_tokens_decoder`. This means that the previous + `additional_special_tokens` are still added tokens, and will not be split by the model. + + Returns: + `int`: Number of tokens added to the vocabulary. + + Examples: + + ```python + # Let's see how to add a new classification token to GPT-2 + tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2") + model = GPT2Model.from_pretrained("openai-community/gpt2") + + special_tokens_dict = {"cls_token": ""} + + num_added_toks = tokenizer.add_special_tokens(special_tokens_dict) + print("We have added", num_added_toks, "tokens") + # Notice: resize_token_embeddings expect to receive the full size of the new vocabulary, i.e., the length of the tokenizer. + model.resize_token_embeddings(len(tokenizer)) + + assert tokenizer.cls_token == "" + ```""" + if not special_tokens_dict: + return 0 + + added_tokens = [] + for key, value in special_tokens_dict.items(): + assert key in self.SPECIAL_TOKENS_ATTRIBUTES, f"Key {key} is not a special token" + + if self.verbose: + logger.info(f"Assigning {value} to the {key} key of the tokenizer") + + if key == "additional_special_tokens": + assert isinstance(value, (list, tuple)) and all(isinstance(t, (str, AddedToken)) for t in value), ( + f"Tokens {value} for key {key} should all be str or AddedToken instances" + ) + + to_add = [] + for token in value: + if isinstance(token, str): + # for legacy purpose we default to stripping. `test_add_tokens_tokenizer` depends on this + token = AddedToken(token, rstrip=False, lstrip=False, normalized=False, special=True) + if not replace_additional_special_tokens and str(token) in self.additional_special_tokens: + continue + to_add.append(token) + if replace_additional_special_tokens and len(to_add) > 0: + setattr(self, key, list(to_add)) + else: + self._special_tokens_map["additional_special_tokens"].extend(to_add) + added_tokens += to_add + + else: + if not isinstance(value, (str, AddedToken)): + raise ValueError(f"Token {value} for key {key} should be a str or an AddedToken instance") + if isinstance(value, (str)): + # for legacy purpose we default to stripping. `False` depends on this + value = AddedToken(value, rstrip=False, lstrip=False, normalized=False, special=True) + if isinstance(value, AddedToken): + setattr(self, key, value) + if value not in added_tokens: + added_tokens.append(value) + + # if we are adding tokens that were not part of the vocab, we ought to add them + added_tokens = self.add_tokens(added_tokens, special_tokens=True) + return added_tokens + + def add_tokens( + self, new_tokens: Union[str, AddedToken, List[Union[str, AddedToken]]], special_tokens: bool = False + ) -> int: + """ + Add a list of new tokens to the tokenizer class. If the new tokens are not in the vocabulary, they are added to + it with indices starting from length of the current vocabulary and will be isolated before the tokenization + algorithm is applied. Added tokens and tokens from the vocabulary of the tokenization algorithm are therefore + not treated in the same way. + + Note, when adding new tokens to the vocabulary, you should make sure to also resize the token embedding matrix + of the model so that its embedding matrix matches the tokenizer. + + In order to do that, please use the [`~PreTrainedModel.resize_token_embeddings`] method. + + Args: + new_tokens (`str`, `tokenizers.AddedToken` or a list of *str* or `tokenizers.AddedToken`): + Tokens are only added if they are not already in the vocabulary. `tokenizers.AddedToken` wraps a string + token to let you personalize its behavior: whether this token should only match against a single word, + whether this token should strip all potential whitespaces on the left side, whether this token should + strip all potential whitespaces on the right side, etc. + special_tokens (`bool`, *optional*, defaults to `False`): + Can be used to specify if the token is a special token. This mostly change the normalization behavior + (special tokens like CLS or [MASK] are usually not lower-cased for instance). + + See details for `tokenizers.AddedToken` in HuggingFace tokenizers library. + + Returns: + `int`: Number of tokens added to the vocabulary. + + Examples: + + ```python + # Let's see how to increase the vocabulary of Bert model and tokenizer + tokenizer = BertTokenizerFast.from_pretrained("google-bert/bert-base-uncased") + model = BertModel.from_pretrained("google-bert/bert-base-uncased") + + num_added_toks = tokenizer.add_tokens(["new_tok1", "my_new-tok2"]) + print("We have added", num_added_toks, "tokens") + # Notice: resize_token_embeddings expect to receive the full size of the new vocabulary, i.e., the length of the tokenizer. + model.resize_token_embeddings(len(tokenizer)) + ```""" + if not new_tokens: + return 0 + + if not isinstance(new_tokens, (list, tuple)): + new_tokens = [new_tokens] + + return self._add_tokens(new_tokens, special_tokens=special_tokens) + + def _add_tokens(self, new_tokens: Union[List[str], List[AddedToken]], special_tokens: bool = False) -> int: + raise NotImplementedError + + @property + def pad_token_type_id(self) -> int: + """ + `int`: Id of the padding token type in the vocabulary. + """ + return self._pad_token_type_id + + def __setattr__(self, key, value): + key_without_id = key + key_is_special_id = key.endswith("_id") or key.endswith("_ids") + if key_is_special_id: + key_without_id = key[:-3] if not key.endswith("_ids") else key[:-4] + + if self.__dict__.get("_special_tokens_map", None) is not None and any( + name in self.__dict__["_special_tokens_map"] for name in [key, key_without_id] + ): + if key_is_special_id: + if value is not None: + value = ( + self.convert_ids_to_tokens(value) + if key != "additional_special_tokens" + else [self.convert_ids_to_tokens(val) for val in value] + ) + key = key_without_id + + if key != "additional_special_tokens" and not isinstance(value, (str, AddedToken)) and value is not None: + raise ValueError(f"Cannot set a non-string value as the {key}") + self._special_tokens_map[key] = value + else: + super().__setattr__(key, value) + + def __getattr__(self, key): + key_without_id = key + key_is_special_id = key.endswith("_id") or key.endswith("_ids") + if key_is_special_id: + key_without_id = key[:-3] if not key.endswith("_ids") else key[:-4] + + if self.__dict__.get("_special_tokens_map", None) is not None and any( + name in self.__dict__["_special_tokens_map"] for name in [key, key_without_id] + ): + _special_tokens_map = self.__dict__["_special_tokens_map"] + if not key_is_special_id: + if _special_tokens_map[key] is None: + if self.verbose: + logger.error(f"Using {key}, but it is not set yet.") + return None + value = _special_tokens_map[key] + return str(value) if key != "additional_special_tokens" else [str(tok) for tok in value] + else: + attr_as_tokens = getattr(self, key_without_id) + return self.convert_tokens_to_ids(attr_as_tokens) if attr_as_tokens is not None else None + + if key not in self.__dict__: + raise AttributeError(f"{self.__class__.__name__} has no attribute {key}") + else: + return super().__getattr__(key) + + @property + def special_tokens_map(self) -> Dict[str, Union[str, List[str]]]: + """ + `Dict[str, Union[str, List[str]]]`: A dictionary mapping special token class attributes (`cls_token`, + `unk_token`, etc.) to their values (`''`, `''`, etc.). + + Convert potential tokens of `tokenizers.AddedToken` type to string. + """ + set_attr = {} + for attr in self.SPECIAL_TOKENS_ATTRIBUTES: + attr_value = getattr(self, attr) + if attr_value: + set_attr[attr] = attr_value + return set_attr + + @property + def special_tokens_map_extended(self) -> Dict[str, Union[str, AddedToken, List[Union[str, AddedToken]]]]: + """ + `Dict[str, Union[str, tokenizers.AddedToken, List[Union[str, tokenizers.AddedToken]]]]`: A dictionary mapping + special token class attributes (`cls_token`, `unk_token`, etc.) to their values (`''`, `''`, etc.). + + Don't convert tokens of `tokenizers.AddedToken` type to string so they can be used to control more finely how + special tokens are tokenized. + """ + set_attr = {} + for attr in self.SPECIAL_TOKENS_ATTRIBUTES: + attr_value = self._special_tokens_map[attr] + if attr_value: + set_attr[attr] = attr_value + return set_attr + + @property + def all_special_tokens_extended(self) -> List[Union[str, AddedToken]]: + """ + `List[Union[str, tokenizers.AddedToken]]`: All the special tokens (`''`, `''`, etc.), the order has + nothing to do with the index of each tokens. If you want to know the correct indices, check + `self.added_tokens_encoder`. We can't create an order anymore as the keys are `AddedTokens` and not `Strings`. + + Don't convert tokens of `tokenizers.AddedToken` type to string so they can be used to control more finely how + special tokens are tokenized. + """ + all_tokens = [] + seen = set() + for value in self.special_tokens_map_extended.values(): + if isinstance(value, (list, tuple)): + tokens_to_add = [token for token in value if str(token) not in seen] + else: + tokens_to_add = [value] if str(value) not in seen else [] + seen.update(map(str, tokens_to_add)) + all_tokens.extend(tokens_to_add) + return all_tokens + + @property + def all_special_tokens(self) -> List[str]: + """ + `List[str]`: A list of the unique special tokens (`''`, `''`, ..., etc.). + + Convert tokens of `tokenizers.AddedToken` type to string. + """ + all_toks = [str(s) for s in self.all_special_tokens_extended] + return all_toks + + @property + def all_special_ids(self) -> List[int]: + """ + `List[int]`: List the ids of the special tokens(`''`, `''`, etc.) mapped to class attributes. + """ + all_toks = self.all_special_tokens + all_ids = self.convert_tokens_to_ids(all_toks) + return all_ids + + def _set_model_specific_special_tokens(self, special_tokens: List[str]): + """ + Adds new special tokens to the "SPECIAL_TOKENS_ATTRIBUTES" list which will be part + of "self.special_tokens" and saved as a special token in tokenizer's config. + This allows us to dynamically add new model-type specific tokens after initializing the tokenizer. + For example: if the model tokenizers is multimodal, we can support special image or audio tokens. + """ + self.SPECIAL_TOKENS_ATTRIBUTES = self.SPECIAL_TOKENS_ATTRIBUTES + list(special_tokens.keys()) + for key, value in special_tokens.items(): + if isinstance(value, (str, AddedToken)): + self._special_tokens_map[key] = value + else: + raise TypeError(f"Special token {key} has to be either str or AddedToken but got: {type(value)}") + + +ENCODE_KWARGS_DOCSTRING = r""" + add_special_tokens (`bool`, *optional*, defaults to `True`): + Whether or not to add special tokens when encoding the sequences. This will use the underlying + `PretrainedTokenizerBase.build_inputs_with_special_tokens` function, which defines which tokens are + automatically added to the input ids. This is useful if you want to add `bos` or `eos` tokens + automatically. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`): + Activates and controls truncation. Accepts the following values: + + - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or + to the maximum acceptable input length for the model if that argument is not provided. This will + truncate token by token, removing a token from the longest sequence in the pair if a pair of + sequences (or a batch of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths + greater than the model maximum admissible input size). + max_length (`int`, *optional*): + Controls the maximum length to use by one of the truncation/padding parameters. + + If left unset or set to `None`, this will use the predefined model maximum length if a maximum length + is required by one of the truncation/padding parameters. If the model has no specific maximum input + length (like XLNet) truncation/padding to a maximum length will be deactivated. + stride (`int`, *optional*, defaults to 0): + If set to a number along with `max_length`, the overflowing tokens returned when + `return_overflowing_tokens=True` will contain some tokens from the end of the truncated sequence + returned to provide some overlap between truncated and overflowing sequences. The value of this + argument defines the number of overlapping tokens. + is_split_into_words (`bool`, *optional*, defaults to `False`): + Whether or not the input is already pre-tokenized (e.g., split into words). If set to `True`, the + tokenizer assumes the input is already split into words (for instance, by splitting it on whitespace) + which it will tokenize. This is useful for NER or token classification. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. Requires `padding` to be activated. + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side (`str`, *optional*): + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. +""" + +ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING = r""" + return_token_type_ids (`bool`, *optional*): + Whether to return token type IDs. If left to the default, will return the token type IDs according to + the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are token type IDs?](../glossary#token-type-ids) + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) + return_overflowing_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to return overflowing token sequences. If a pair of sequences of input ids (or a batch + of pairs) is provided with `truncation_strategy = longest_first` or `True`, an error is raised instead + of returning overflowing tokens. + return_special_tokens_mask (`bool`, *optional*, defaults to `False`): + Whether or not to return special tokens mask information. + return_offsets_mapping (`bool`, *optional*, defaults to `False`): + Whether or not to return `(char_start, char_end)` for each token. + + This is only available on fast tokenizers inheriting from [`PreTrainedTokenizerFast`], if using + Python's tokenizer, this method will raise `NotImplementedError`. + return_length (`bool`, *optional*, defaults to `False`): + Whether or not to return the lengths of the encoded inputs. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + **kwargs: passed to the `self.tokenize()` method + + Return: + [`BatchEncoding`]: A [`BatchEncoding`] with the following fields: + + - **input_ids** -- List of token ids to be fed to a model. + + [What are input IDs?](../glossary#input-ids) + + - **token_type_ids** -- List of token type ids to be fed to a model (when `return_token_type_ids=True` or + if *"token_type_ids"* is in `self.model_input_names`). + + [What are token type IDs?](../glossary#token-type-ids) + + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when + `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names`). + + [What are attention masks?](../glossary#attention-mask) + + - **overflowing_tokens** -- List of overflowing tokens sequences (when a `max_length` is specified and + `return_overflowing_tokens=True`). + - **num_truncated_tokens** -- Number of tokens truncated (when a `max_length` is specified and + `return_overflowing_tokens=True`). + - **special_tokens_mask** -- List of 0s and 1s, with 1 specifying added special tokens and 0 specifying + regular sequence tokens (when `add_special_tokens=True` and `return_special_tokens_mask=True`). + - **length** -- The length of the inputs (when `return_length=True`) +""" + + +INIT_TOKENIZER_DOCSTRING = r""" + Class attributes (overridden by derived classes) + + - **vocab_files_names** (`Dict[str, str]`) -- A dictionary with, as keys, the `__init__` keyword name of each + vocabulary file required by the model, and as associated values, the filename for saving the associated file + (string). + - **pretrained_vocab_files_map** (`Dict[str, Dict[str, str]]`) -- A dictionary of dictionaries, with the + high-level keys being the `__init__` keyword name of each vocabulary file required by the model, the + low-level being the `short-cut-names` of the pretrained models with, as associated values, the `url` to the + associated pretrained vocabulary file. + - **model_input_names** (`List[str]`) -- A list of inputs expected in the forward pass of the model. + - **padding_side** (`str`) -- The default value for the side on which the model should have padding applied. + Should be `'right'` or `'left'`. + - **truncation_side** (`str`) -- The default value for the side on which the model should have truncation + applied. Should be `'right'` or `'left'`. + + Args: + model_max_length (`int`, *optional*): + The maximum length (in number of tokens) for the inputs to the transformer model. When the tokenizer is + loaded with [`~tokenization_utils_base.PreTrainedTokenizerBase.from_pretrained`], this will be set to the + value stored for the associated model in `max_model_input_sizes` (see above). If no value is provided, will + default to VERY_LARGE_INTEGER (`int(1e30)`). + padding_side (`str`, *optional*): + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + truncation_side (`str`, *optional*): + The side on which the model should have truncation applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + chat_template (`str`, *optional*): + A Jinja template string that will be used to format lists of chat messages. See + https://huggingface.co/docs/transformers/chat_templating for a full description. + model_input_names (`List[string]`, *optional*): + The list of inputs accepted by the forward pass of the model (like `"token_type_ids"` or + `"attention_mask"`). Default value is picked from the class attribute of the same name. + bos_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the beginning of a sentence. Will be associated to `self.bos_token` and + `self.bos_token_id`. + eos_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the end of a sentence. Will be associated to `self.eos_token` and + `self.eos_token_id`. + unk_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing an out-of-vocabulary token. Will be associated to `self.unk_token` and + `self.unk_token_id`. + sep_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token separating two different sentences in the same input (used by BERT for instance). Will be + associated to `self.sep_token` and `self.sep_token_id`. + pad_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by + attention mechanisms or loss computation. Will be associated to `self.pad_token` and `self.pad_token_id`. + cls_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing the class of the input (used by BERT for instance). Will be associated to + `self.cls_token` and `self.cls_token_id`. + mask_token (`str` or `tokenizers.AddedToken`, *optional*): + A special token representing a masked token (used by masked-language modeling pretraining objectives, like + BERT). Will be associated to `self.mask_token` and `self.mask_token_id`. + additional_special_tokens (tuple or list of `str` or `tokenizers.AddedToken`, *optional*): + A tuple or a list of additional special tokens. Add them here to ensure they are skipped when decoding with + `skip_special_tokens` is set to True. If they are not part of the vocabulary, they will be added at the end + of the vocabulary. + clean_up_tokenization_spaces (`bool`, *optional*, defaults to `True`): + Whether or not the model should cleanup the spaces that were added when splitting the input text during the + tokenization process. + split_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the special tokens should be split during the tokenization process. Passing will affect the + internal state of the tokenizer. The default behavior is to not split special tokens. This means that if + `` is the `bos_token`, then `tokenizer.tokenize("") = ['`]. Otherwise, if + `split_special_tokens=True`, then `tokenizer.tokenize("")` will be give `['<','s', '>']`. +""" + + +@add_end_docstrings(INIT_TOKENIZER_DOCSTRING) +class PreTrainedTokenizerBase(SpecialTokensMixin, PushToHubMixin): + """ + Base class for [`PreTrainedTokenizer`] and [`PreTrainedTokenizerFast`]. + + Handles shared (mostly boiler plate) methods for those two classes. + """ + + vocab_files_names: Dict[str, str] = {} + pretrained_vocab_files_map: Dict[str, Dict[str, str]] = {} + _auto_class: Optional[str] = None + + # first name has to correspond to main model input name + # to make sure `tokenizer.pad(...)` works correctly + model_input_names: List[str] = ["input_ids", "token_type_ids", "attention_mask"] + padding_side: str = "right" + truncation_side: str = "right" + slow_tokenizer_class = None + + def __init__(self, **kwargs): + # inputs and kwargs for saving and re-loading (see ``from_pretrained`` and ``save_pretrained``) + self.init_inputs = () + for key in kwargs: + if hasattr(self, key) and callable(getattr(self, key)): + raise AttributeError(f"{key} conflicts with the method {key} in {self.__class__.__name__}") + + self.init_kwargs = copy.deepcopy(kwargs) + self.name_or_path = kwargs.pop("name_or_path", "") + self._processor_class = kwargs.pop("processor_class", None) + + # For backward compatibility we fallback to set model_max_length from max_len if provided + model_max_length = kwargs.pop("model_max_length", kwargs.pop("max_len", None)) + self.model_max_length = model_max_length if model_max_length is not None else VERY_LARGE_INTEGER + + # Padding and truncation side are right by default and overridden in subclasses. If specified in the kwargs, it + # is changed. + self.padding_side = kwargs.pop("padding_side", self.padding_side) + if self.padding_side not in ["right", "left"]: + raise ValueError( + f"Padding side should be selected between 'right' and 'left', current value: {self.padding_side}" + ) + + self.truncation_side = kwargs.pop("truncation_side", self.truncation_side) + if self.truncation_side not in ["right", "left"]: + raise ValueError( + f"Truncation side should be selected between 'right' and 'left', current value: {self.truncation_side}" + ) + + self.model_input_names = kwargs.pop("model_input_names", self.model_input_names) + + # By default, cleaning tokenization spaces for both fast and slow tokenizers + self.clean_up_tokenization_spaces = kwargs.pop("clean_up_tokenization_spaces", False) + + # By default, do not split special tokens for both fast and slow tokenizers + self.split_special_tokens = kwargs.pop("split_special_tokens", False) + + self.deprecation_warnings = {} # Use to store when we have already noticed a deprecation warning (avoid overlogging). + self._in_target_context_manager = False + + # Stores a Jinja template that formats chat histories into tokenizable strings + self.chat_template = kwargs.pop("chat_template", None) + if isinstance(self.chat_template, (list, tuple)): + # Chat templates are stored as lists of dicts with fixed key names, + # we reconstruct that into a single dict while loading them. + self.chat_template = {template["name"]: template["template"] for template in self.chat_template} + + super().__init__(**kwargs) + + self.extra_special_tokens = kwargs.pop("extra_special_tokens", {}) + self._set_model_specific_special_tokens(special_tokens=self.extra_special_tokens) + + @property + def max_len_single_sentence(self) -> int: + """ + `int`: The maximum length of a sentence that can be fed to the model. + """ + return self.model_max_length - self.num_special_tokens_to_add(pair=False) + + @property + def max_len_sentences_pair(self) -> int: + """ + `int`: The maximum combined length of a pair of sentences that can be fed to the model. + """ + return self.model_max_length - self.num_special_tokens_to_add(pair=True) + + @max_len_single_sentence.setter + def max_len_single_sentence(self, value) -> int: + # For backward compatibility, allow to try to setup 'max_len_single_sentence'. + if value == self.model_max_length - self.num_special_tokens_to_add(pair=False) and self.verbose: + if not self.deprecation_warnings.get("max_len_single_sentence", False): + logger.warning( + "Setting 'max_len_single_sentence' is now deprecated. This value is automatically set up." + ) + self.deprecation_warnings["max_len_single_sentence"] = True + else: + raise ValueError( + "Setting 'max_len_single_sentence' is now deprecated. This value is automatically set up." + ) + + @max_len_sentences_pair.setter + def max_len_sentences_pair(self, value) -> int: + # For backward compatibility, allow to try to setup 'max_len_sentences_pair'. + if value == self.model_max_length - self.num_special_tokens_to_add(pair=True) and self.verbose: + if not self.deprecation_warnings.get("max_len_sentences_pair", False): + logger.warning( + "Setting 'max_len_sentences_pair' is now deprecated. This value is automatically set up." + ) + self.deprecation_warnings["max_len_sentences_pair"] = True + else: + raise ValueError("Setting 'max_len_sentences_pair' is now deprecated. This value is automatically set up.") + + def _set_processor_class(self, processor_class: str): + """Sets processor class as an attribute.""" + self._processor_class = processor_class + + @property + def added_tokens_decoder(self) -> Dict[int, AddedToken]: + raise NotImplementedError() + + def __repr__(self) -> str: + added_tokens_decoder_rep = "\n\t".join([f"{k}: {v.__repr__()}," for k, v in self.added_tokens_decoder.items()]) + return ( + f"{self.__class__.__name__}(name_or_path='{self.name_or_path}'," + f" vocab_size={self.vocab_size}, model_max_length={self.model_max_length}, is_fast={self.is_fast}," + f" padding_side='{self.padding_side}', truncation_side='{self.truncation_side}'," + f" special_tokens={self.special_tokens_map}, clean_up_tokenization_spaces={self.clean_up_tokenization_spaces}," + " added_tokens_decoder={\n\t" + added_tokens_decoder_rep + "\n}\n)" + ) + + def __len__(self) -> int: + raise NotImplementedError() + + def get_vocab(self) -> Dict[str, int]: + """ + Returns the vocabulary as a dictionary of token to index. + + `tokenizer.get_vocab()[token]` is equivalent to `tokenizer.convert_tokens_to_ids(token)` when `token` is in the + vocab. + + Returns: + `Dict[str, int]`: The vocabulary. + """ + raise NotImplementedError() + + def apply_chat_template( + self, + conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]]], + tools: Optional[List[Union[Dict, Callable]]] = None, + documents: Optional[List[Dict[str, str]]] = None, + chat_template: Optional[str] = None, + add_generation_prompt: bool = False, + continue_final_message: bool = False, + tokenize: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: bool = False, + max_length: Optional[int] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_dict: bool = False, + return_assistant_tokens_mask: bool = False, + tokenizer_kwargs: Optional[Dict[str, Any]] = None, + **kwargs, + ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]: + """ + Converts a list of dictionaries with `"role"` and `"content"` keys to a list of token + ids. This method is intended for use with chat models, and will read the tokenizer's chat_template attribute to + determine the format and control tokens to use when converting. + + Args: + conversation (Union[List[Dict[str, str]], List[List[Dict[str, str]]]]): A list of dicts + with "role" and "content" keys, representing the chat history so far. + tools (`List[Dict]`, *optional*): + A list of tools (callable functions) that will be accessible to the model. If the template does not + support function calling, this argument will have no effect. Each tool should be passed as a JSON Schema, + giving the name, description and argument types for the tool. See our + [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use) + for more information. + documents (`List[Dict[str, str]]`, *optional*): + A list of dicts representing documents that will be accessible to the model if it is performing RAG + (retrieval-augmented generation). If the template does not support RAG, this argument will have no + effect. We recommend that each document should be a dict containing "title" and "text" keys. Please + see the RAG section of the [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#arguments-for-RAG) + for examples of passing documents with chat templates. + chat_template (`str`, *optional*): + A Jinja template to use for this conversion. It is usually not necessary to pass anything to this + argument, as the model's template will be used by default. + add_generation_prompt (bool, *optional*): + If this is set, a prompt with the token(s) that indicate + the start of an assistant message will be appended to the formatted output. This is useful when you want to generate a response from the model. + Note that this argument will be passed to the chat template, and so it must be supported in the + template for this argument to have any effect. + continue_final_message (bool, *optional*): + If this is set, the chat will be formatted so that the final + message in the chat is open-ended, without any EOS tokens. The model will continue this message + rather than starting a new one. This allows you to "prefill" part of + the model's response for it. Cannot be used at the same time as `add_generation_prompt`. + tokenize (`bool`, defaults to `True`): + Whether to tokenize the output. If `False`, the output will be a string. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + truncation (`bool`, defaults to `False`): + Whether to truncate sequences at the maximum length. Has no effect if tokenize is `False`. + max_length (`int`, *optional*): + Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is `False`. If + not specified, the tokenizer's `max_length` attribute will be used as a default. + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors of a particular framework. Has no effect if tokenize is `False`. Acceptable + values are: + - `'tf'`: Return TensorFlow `tf.Tensor` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return NumPy `np.ndarray` objects. + - `'jax'`: Return JAX `jnp.ndarray` objects. + return_dict (`bool`, defaults to `False`): + Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`. + tokenizer_kwargs (`Dict[str: Any]`, *optional*): Additional kwargs to pass to the tokenizer. + return_assistant_tokens_mask (`bool`, defaults to `False`): + Whether to return a mask of the assistant generated tokens. For tokens generated by the assistant, + the mask will contain 1. For user and system tokens, the mask will contain 0. + This functionality is only available for chat templates that support it via the `{% generation %}` keyword. + **kwargs: Additional kwargs to pass to the template renderer. Will be accessible by the chat template. + + Returns: + `Union[List[int], Dict]`: A list of token ids representing the tokenized chat so far, including control tokens. This + output is ready to pass to the model, either directly or via methods like `generate()`. If `return_dict` is + set, will return a dict of tokenizer outputs instead. + """ + + if return_dict and not tokenize: + raise ValueError( + "`return_dict=True` is incompatible with `tokenize=False`, because there is no dict " + "of tokenizer outputs to return." + ) + + if return_assistant_tokens_mask and not return_dict: + raise ValueError("`return_assistant_tokens_mask=True` is incompatible with `return_dict=False`") + + if tokenizer_kwargs is None: + tokenizer_kwargs = {} + + chat_template = self.get_chat_template(chat_template, tools) + + if return_assistant_tokens_mask and not re.search(r"\{\%-?\s*generation\s*-?\%\}", chat_template): + logger.warning_once( + "return_assistant_tokens_mask==True but chat template does not contain `{% generation %}` keyword." + ) + + # Compilation function uses a cache to avoid recompiling the same template + compiled_template = _compile_jinja_template(chat_template) + + if isinstance(conversation, (list, tuple)) and ( + isinstance(conversation[0], (list, tuple)) or hasattr(conversation[0], "messages") + ): + conversations = conversation + is_batched = True + else: + conversations = [conversation] + is_batched = False + + if continue_final_message: + if add_generation_prompt: + raise ValueError( + "continue_final_message and add_generation_prompt are not compatible. Use continue_final_message when you want the model to continue the final message, and add_generation_prompt when you want to add a header that will prompt it to start a new assistant message instead." + ) + if return_assistant_tokens_mask: + raise ValueError("continue_final_message is not compatible with return_assistant_tokens_mask.") + + # We accept either JSON schemas or functions for tools. If we get functions, we convert them to schemas + if tools is not None: + tool_schemas = [] + for tool in tools: + if isinstance(tool, dict): + tool_schemas.append(tool) + elif isfunction(tool): + tool_schemas.append(get_json_schema(tool)) + else: + raise ValueError( + "Tools should either be a JSON schema, or a callable function with type hints " + "and a docstring suitable for auto-conversion to a schema." + ) + else: + tool_schemas = None + + if documents is not None: + for document in documents: + if not isinstance(document, dict): + raise TypeError("Documents should be a list of dicts with 'title' and 'text' keys!") + + rendered = [] + all_generation_indices = [] + template_kwargs = {**self.special_tokens_map, **kwargs} # kwargs overwrite special tokens if both are present + for chat in conversations: + if hasattr(chat, "messages"): + # Indicates it's a Conversation object + chat = chat.messages + if return_assistant_tokens_mask: + rendered_chat, generation_indices = _render_with_assistant_indices( + compiled_template=compiled_template, + messages=chat, + tools=tool_schemas, + documents=documents, + add_generation_prompt=add_generation_prompt, + **template_kwargs, + ) + all_generation_indices.append(generation_indices) + else: + rendered_chat = compiled_template.render( + messages=chat, + tools=tool_schemas, + documents=documents, + add_generation_prompt=add_generation_prompt, + **template_kwargs, + ) + if continue_final_message: + final_message = chat[-1]["content"] + if isinstance(final_message, (list, tuple)): + for content_block in reversed(final_message): + if "text" in content_block: + # Pick the last text block in the message (the first one we hit while iterating in reverse) + final_message = content_block["text"] + break + else: + raise ValueError( + "continue_final_message is set but we could not find any text to continue" + "in the final message!" + ) + if final_message.strip() not in rendered_chat: + raise ValueError( + "continue_final_message is set but the final message does not appear in the chat after " + "applying the chat template! This can happen if the chat template deletes portions of " + "the final message. Please verify the chat template and final message in your chat to " + "ensure they are compatible." + ) + final_msg_loc = rendered_chat.rindex(final_message.strip()) + if rendered_chat[final_msg_loc : final_msg_loc + len(final_message.lstrip())] == final_message: + # The template preserves spacing or the message doesn't have trailing spacing, so things are simple + rendered_chat = rendered_chat[: final_msg_loc + len(final_message.lstrip())] + else: + # The message has trailing spacing that was trimmed, so we must be more cautious + rendered_chat = rendered_chat[: final_msg_loc + len(final_message.strip())] + rendered.append(rendered_chat) + + if not is_batched: + rendered = rendered[0] + + if tokenize: + out = self( + rendered, + padding=padding, + truncation=truncation, + max_length=max_length, + add_special_tokens=False, + return_tensors=return_tensors, + **tokenizer_kwargs, + ) + if return_dict: + if return_assistant_tokens_mask: + assistant_masks = [] + if is_batched or return_tensors: + input_ids = out["input_ids"] + else: + input_ids = [out["input_ids"]] + for i in range(len(input_ids)): + current_mask = [0] * len(input_ids[i]) + for assistant_start_char, assistant_end_char in all_generation_indices[i]: + start_token = out.char_to_token(i, assistant_start_char) + end_token = out.char_to_token(i, assistant_end_char - 1) + if start_token is None: + # start_token is out of bounds maybe due to truncation. + break + for token_id in range(start_token, end_token + 1 if end_token else len(input_ids[i])): + current_mask[token_id] = 1 + assistant_masks.append(current_mask) + + if not is_batched and not return_tensors: + assistant_masks = assistant_masks[0] + + out["assistant_masks"] = assistant_masks + + if return_tensors: + out.convert_to_tensors(tensor_type=return_tensors) + + return out + else: + return out["input_ids"] + else: + return rendered + + def get_chat_template(self, chat_template: Optional[str] = None, tools: Optional[List[Dict]] = None) -> str: + """ + Retrieve the chat template string used for tokenizing chat messages. This template is used + internally by the `apply_chat_template` method and can also be used externally to retrieve the model's chat + template for better generation tracking. + + Args: + chat_template (`str`, *optional*): + A Jinja template or the name of a template to use for this conversion. + It is usually not necessary to pass anything to this argument, + as the model's template will be used by default. + tools (`List[Dict]`, *optional*): + A list of tools (callable functions) that will be accessible to the model. If the template does not + support function calling, this argument will have no effect. Each tool should be passed as a JSON Schema, + giving the name, description and argument types for the tool. See our + [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use) + for more information. + + Returns: + `str`: The chat template string. + """ + # First, handle the cases when the model has a dict of multiple templates + if isinstance(self.chat_template, dict): + template_dict = self.chat_template + if chat_template is not None and chat_template in template_dict: + # The user can pass the name of a template to the chat template argument instead of an entire template + chat_template = template_dict[chat_template] + elif chat_template is None: + if tools is not None and "tool_use" in template_dict: + chat_template = template_dict["tool_use"] + elif "default" in template_dict: + chat_template = template_dict["default"] + else: + raise ValueError( + "This model has multiple chat templates with no default specified! Please either pass a chat " + "template or the name of the template you wish to use to the `chat_template` argument. Available " + f"template names are {sorted(template_dict.keys())}." + ) + + elif chat_template is None: + # These are the cases when the model has a single template + # priority: `chat_template` argument > `tokenizer.chat_template` + if self.chat_template is not None: + chat_template = self.chat_template + else: + raise ValueError( + "Cannot use chat template functions because tokenizer.chat_template is not set and no template " + "argument was passed! For information about writing templates and setting the " + "tokenizer.chat_template attribute, please see the documentation at " + "https://huggingface.co/docs/transformers/main/en/chat_templating" + ) + + return chat_template + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Union[str, os.PathLike], + *init_inputs, + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + trust_remote_code=False, + **kwargs, + ): + r""" + Instantiate a [`~tokenization_utils_base.PreTrainedTokenizerBase`] (or a derived class) from a predefined + tokenizer. + + Args: + pretrained_model_name_or_path (`str` or `os.PathLike`): + Can be either: + + - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co. + - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved + using the [`~tokenization_utils_base.PreTrainedTokenizerBase.save_pretrained`] method, e.g., + `./my_model_directory/`. + - (**Deprecated**, not applicable to all derived classes) A path or url to a single saved vocabulary + file (if and only if the tokenizer only requires a single vocabulary file like Bert or XLNet), e.g., + `./my_model_directory/vocab.txt`. + cache_dir (`str` or `os.PathLike`, *optional*): + Path to a directory in which a downloaded predefined tokenizer vocabulary files should be cached if the + standard cache should not be used. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force the (re-)download the vocabulary files and override the cached versions if they + exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). + local_files_only (`bool`, *optional*, defaults to `False`): + Whether or not to only rely on local files and not to attempt to download any files. + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + subfolder (`str`, *optional*): + In case the relevant files are located inside a subfolder of the model repo on huggingface.co (e.g. for + facebook/rag-token-base), specify it here. + inputs (additional positional arguments, *optional*): + Will be passed along to the Tokenizer `__init__` method. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether or not to allow for custom models defined on the Hub in their own modeling files. This option + should only be set to `True` for repositories you trust and in which you have read the code, as it will + execute code present on the Hub on your local machine. + kwargs (additional keyword arguments, *optional*): + Will be passed to the Tokenizer `__init__` method. Can be used to set special tokens like `bos_token`, + `eos_token`, `unk_token`, `sep_token`, `pad_token`, `cls_token`, `mask_token`, + `additional_special_tokens`. See parameters in the `__init__` for more details. + + + + Passing `token=True` is required when you want to use a private model. + + + + Examples: + + ```python + # We can't instantiate directly the base class *PreTrainedTokenizerBase* so let's show our examples on a derived class: BertTokenizer + # Download vocabulary from huggingface.co and cache. + tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased") + + # Download vocabulary from huggingface.co (user-uploaded) and cache. + tokenizer = BertTokenizer.from_pretrained("dbmdz/bert-base-german-cased") + + # If vocabulary files are in a directory (e.g. tokenizer was saved using *save_pretrained('./test/saved_model/')*) + tokenizer = BertTokenizer.from_pretrained("./test/saved_model/") + + # If the tokenizer uses a single vocabulary file, you can point directly to this file + tokenizer = BertTokenizer.from_pretrained("./test/saved_model/my_vocab.txt") + + # You can link tokens to special vocabulary when instantiating + tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased", unk_token="") + # You should be sure '' is in the vocabulary when doing that. + # Otherwise use tokenizer.add_special_tokens({'unk_token': ''}) instead) + assert tokenizer.unk_token == "" + ```""" + resume_download = kwargs.pop("resume_download", None) + proxies = kwargs.pop("proxies", None) + use_auth_token = kwargs.pop("use_auth_token", None) + subfolder = kwargs.pop("subfolder", None) + from_pipeline = kwargs.pop("_from_pipeline", None) + from_auto_class = kwargs.pop("_from_auto", False) + commit_hash = kwargs.pop("_commit_hash", None) + gguf_file = kwargs.get("gguf_file", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + user_agent = {"file_type": "tokenizer", "from_auto_class": from_auto_class, "is_fast": "Fast" in cls.__name__} + if from_pipeline is not None: + user_agent["using_pipeline"] = from_pipeline + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + + pretrained_model_name_or_path = str(pretrained_model_name_or_path) + vocab_files = {} + init_configuration = {} + + is_local = os.path.isdir(pretrained_model_name_or_path) + single_file_id = None + if os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path): + if len(cls.vocab_files_names) > 1 and not gguf_file: + raise ValueError( + f"Calling {cls.__name__}.from_pretrained() with the path to a single file or url is not " + "supported for this tokenizer. Use a model identifier or the path to a directory instead." + ) + warnings.warn( + f"Calling {cls.__name__}.from_pretrained() with the path to a single file or url is deprecated and " + "won't be possible anymore in v5. Use a model identifier or the path to a directory instead.", + FutureWarning, + ) + file_id = list(cls.vocab_files_names.keys())[0] + + vocab_files[file_id] = pretrained_model_name_or_path + single_file_id = file_id + else: + if gguf_file: + vocab_files["vocab_file"] = gguf_file + else: + # At this point pretrained_model_name_or_path is either a directory or a model identifier name + additional_files_names = { + "added_tokens_file": ADDED_TOKENS_FILE, # kept only for legacy + "special_tokens_map_file": SPECIAL_TOKENS_MAP_FILE, # kept only for legacy + "tokenizer_config_file": TOKENIZER_CONFIG_FILE, + # tokenizer_file used to initialize a slow from a fast. Properly copy the `addedTokens` instead of adding in random orders + "tokenizer_file": FULL_TOKENIZER_FILE, + "chat_template_file": CHAT_TEMPLATE_FILE, + } + vocab_files = {**cls.vocab_files_names, **additional_files_names} + if "tokenizer_file" in vocab_files: + # Try to get the tokenizer config to see if there are versioned tokenizer files. + fast_tokenizer_file = FULL_TOKENIZER_FILE + resolved_config_file = cached_file( + pretrained_model_name_or_path, + TOKENIZER_CONFIG_FILE, + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + token=token, + revision=revision, + local_files_only=local_files_only, + subfolder=subfolder, + user_agent=user_agent, + _raise_exceptions_for_gated_repo=False, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + _commit_hash=commit_hash, + ) + commit_hash = extract_commit_hash(resolved_config_file, commit_hash) + if resolved_config_file is not None: + with open(resolved_config_file, encoding="utf-8") as reader: + tokenizer_config = json.load(reader) + if "fast_tokenizer_files" in tokenizer_config: + fast_tokenizer_file = get_fast_tokenizer_file(tokenizer_config["fast_tokenizer_files"]) + vocab_files["tokenizer_file"] = fast_tokenizer_file + + # Get files from url, cache, or disk depending on the case + resolved_vocab_files = {} + for file_id, file_path in vocab_files.items(): + if file_path is None: + resolved_vocab_files[file_id] = None + elif single_file_id == file_id: + if os.path.isfile(file_path): + resolved_vocab_files[file_id] = file_path + elif is_remote_url(file_path): + resolved_vocab_files[file_id] = download_url(file_path, proxies=proxies) + else: + resolved_vocab_files[file_id] = cached_file( + pretrained_model_name_or_path, + file_path, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _raise_exceptions_for_gated_repo=False, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + _commit_hash=commit_hash, + ) + commit_hash = extract_commit_hash(resolved_vocab_files[file_id], commit_hash) + + # If one passes a GGUF file path to `gguf_file` there is no need for this check as the tokenizer will be + # loaded directly from the GGUF file. + if all(full_file_name is None for full_file_name in resolved_vocab_files.values()) and not gguf_file: + raise EnvironmentError( + f"Can't load tokenizer for '{pretrained_model_name_or_path}'. If you were trying to load it from " + "'https://huggingface.co/models', make sure you don't have a local directory with the same name. " + f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory " + f"containing all relevant files for a {cls.__name__} tokenizer." + ) + + for file_id, file_path in vocab_files.items(): + if file_id not in resolved_vocab_files: + continue + + if is_local: + logger.info(f"loading file {file_path}") + else: + logger.info(f"loading file {file_path} from cache at {resolved_vocab_files[file_id]}") + + return cls._from_pretrained( + resolved_vocab_files, + pretrained_model_name_or_path, + init_configuration, + *init_inputs, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + _commit_hash=commit_hash, + _is_local=is_local, + trust_remote_code=trust_remote_code, + **kwargs, + ) + + @classmethod + def _from_pretrained( + cls, + resolved_vocab_files, + pretrained_model_name_or_path, + init_configuration, + *init_inputs, + token=None, + cache_dir=None, + local_files_only=False, + _commit_hash=None, + _is_local=False, + trust_remote_code=False, + **kwargs, + ): + # We instantiate fast tokenizers based on a slow tokenizer if we don't have access to the tokenizer.json + # file or if `from_slow` is set to True. + from_slow = kwargs.get("from_slow", False) + gguf_file = kwargs.get("gguf_file", None) + has_tokenizer_file = resolved_vocab_files.get("tokenizer_file", None) is not None + + # If one passes a GGUF file path to `gguf_file` there is no need for this check as the tokenizer will be + # loaded directly from the GGUF file. + if (from_slow or not has_tokenizer_file) and cls.slow_tokenizer_class is not None and not gguf_file: + slow_tokenizer = (cls.slow_tokenizer_class)._from_pretrained( + copy.deepcopy(resolved_vocab_files), + pretrained_model_name_or_path, + copy.deepcopy(init_configuration), + *init_inputs, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + _commit_hash=_commit_hash, + **(copy.deepcopy(kwargs)), + ) + else: + slow_tokenizer = None + + # Prepare tokenizer initialization kwargs + # Did we saved some inputs and kwargs to reload ? + tokenizer_config_file = resolved_vocab_files.pop("tokenizer_config_file", None) + if tokenizer_config_file is not None: + with open(tokenizer_config_file, encoding="utf-8") as tokenizer_config_handle: + init_kwargs = json.load(tokenizer_config_handle) + # First attempt. We get tokenizer_class from tokenizer_config to check mismatch between tokenizers. + config_tokenizer_class = init_kwargs.get("tokenizer_class") + init_kwargs.pop("tokenizer_class", None) + if not has_tokenizer_file: + init_kwargs.pop("tokenizer_file", None) + saved_init_inputs = init_kwargs.pop("init_inputs", ()) + if not init_inputs: + init_inputs = saved_init_inputs + else: + config_tokenizer_class = None + init_kwargs = init_configuration + + # If an independent chat template file exists, it takes priority over template entries in the tokenizer config + chat_template_file = resolved_vocab_files.pop("chat_template_file", None) + if chat_template_file is not None: + with open(chat_template_file) as chat_template_handle: + init_kwargs["chat_template"] = chat_template_handle.read() # Clobbers any template in the config + + if not _is_local: + if "auto_map" in init_kwargs: + # For backward compatibility with odl format. + if isinstance(init_kwargs["auto_map"], (tuple, list)): + init_kwargs["auto_map"] = {"AutoTokenizer": init_kwargs["auto_map"]} + init_kwargs["auto_map"] = add_model_info_to_auto_map( + init_kwargs["auto_map"], pretrained_model_name_or_path + ) + if "custom_pipelines" in init_kwargs: + init_kwargs["custom_pipelines"] = add_model_info_to_custom_pipelines( + init_kwargs["custom_pipelines"], pretrained_model_name_or_path + ) + + if config_tokenizer_class is None: + # Matt: This entire block is only used to decide if the tokenizer class matches the class in the repo. + # If not, it raises a warning, but otherwise continues. Since we mostly load tokenizers with + # AutoTokenizer these days, it seems like a lot of work (and a source of bugs) for little gain. + # Maybe we can just remove this entirely? + from .models.auto.configuration_auto import AutoConfig # tests_ignore + + # Second attempt. If we have not yet found tokenizer_class, let's try to use the config. + try: + config = AutoConfig.from_pretrained( + pretrained_model_name_or_path, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + trust_remote_code=trust_remote_code, + _commit_hash=_commit_hash, + ) + config_tokenizer_class = config.tokenizer_class + except (OSError, ValueError, KeyError): + # skip if an error occurred. + config = None + if config_tokenizer_class is None: + # Third attempt. If we have not yet found the original type of the tokenizer, + # we are loading we see if we can infer it from the type of the configuration file + from .models.auto.tokenization_auto import TOKENIZER_MAPPING_NAMES # tests_ignore + + if hasattr(config, "model_type"): + model_type = config.model_type + else: + # Fallback: use pattern matching on the string. + model_type = None + for pattern in TOKENIZER_MAPPING_NAMES.keys(): + if pattern in str(pretrained_model_name_or_path): + model_type = pattern + break + + if model_type is not None: + config_tokenizer_class, config_tokenizer_class_fast = TOKENIZER_MAPPING_NAMES.get( + model_type, (None, None) + ) + if config_tokenizer_class is None: + config_tokenizer_class = config_tokenizer_class_fast + + if config_tokenizer_class is not None: + if cls.__name__.replace("Fast", "") != config_tokenizer_class.replace("Fast", ""): + logger.warning( + "The tokenizer class you load from this checkpoint is not the same type as the class this" + " function is called from. It may result in unexpected tokenization. \nThe tokenizer class you" + f" load from this checkpoint is '{config_tokenizer_class}'. \nThe class this function is called" + f" from is '{cls.__name__}'." + ) + + # Update with newly provided kwargs + init_kwargs.update(kwargs) + + # Merge resolved_vocab_files arguments in init_kwargs. + added_tokens_file = resolved_vocab_files.pop("added_tokens_file", None) + special_tokens_map_file = resolved_vocab_files.pop("special_tokens_map_file", None) + for args_name, file_path in resolved_vocab_files.items(): + if args_name not in init_kwargs: + init_kwargs[args_name] = file_path + tokenizer_file = resolved_vocab_files.pop("tokenizer_file", None) + + if slow_tokenizer is not None: + init_kwargs["__slow_tokenizer"] = slow_tokenizer + init_kwargs["name_or_path"] = pretrained_model_name_or_path + + #### Handle tokenizer serialization of added and special tokens + added_tokens_decoder: Dict[int, AddedToken] = {} + added_tokens_map: Dict[str, AddedToken] = {} + # if we have info on the slow added tokens + if "added_tokens_decoder" in init_kwargs: + for idx, token in init_kwargs["added_tokens_decoder"].items(): + if isinstance(token, dict): + token = AddedToken(**token) + if isinstance(token, AddedToken): + added_tokens_decoder[int(idx)] = token + added_tokens_map[str(token)] = token + else: + raise ValueError( + f"Found a {token.__class__} in the saved `added_tokens_decoder`, should be a dictionary or an AddedToken instance" + ) + else: + # begin legacy: read the added_tokens_file and update kwargs with special_tokens_map if modified + if special_tokens_map_file is not None: + with open(special_tokens_map_file, encoding="utf-8") as special_tokens_map_handle: + special_tokens_map = json.load(special_tokens_map_handle) + for key, value in special_tokens_map.items(): + if key in kwargs and kwargs[key]: + # This value has already been redefined by the kwargs + # We keep this new value and ignore the one stored in the special_tokens_map_file + continue + if isinstance(value, dict): + value["special"] = True + value = AddedToken(**value) + elif key == "additional_special_tokens" and isinstance(value, list): + additional_special_tokens = init_kwargs.pop("additional_special_tokens", []) or [] + for token in value: + if isinstance(token, dict): + token["special"] = True + token = AddedToken(**token) + if token not in additional_special_tokens: + additional_special_tokens.append(token) + value = additional_special_tokens + init_kwargs[key] = value + + # slow -> slow|fast, legacy: convert the `"added_tokens.json"` file to `added_tokens_decoder`. + # this is for legacy purpose. We don't add the tokens after init for efficiency. + if added_tokens_file is not None: + special_tokens = [] + for key in cls.SPECIAL_TOKENS_ATTRIBUTES & init_kwargs.keys(): + if init_kwargs[key] is not None: + if key == "additional_special_tokens": + special_tokens += [str(token) for token in init_kwargs[key]] + else: + special_tokens.append(str(init_kwargs[key])) + + with open(added_tokens_file, encoding="utf-8") as added_tokens_handle: + added_tok_encoder = json.load(added_tokens_handle) + for str_token, index in added_tok_encoder.items(): + # if index not in added_tokens_decoder and str_token not in added_tokens_map: + special = str_token in special_tokens + added_tokens_decoder[index] = AddedToken( + str_token, rstrip=False, lstrip=False, normalized=not special, special=special + ) + added_tokens_map[str(token)] = added_tokens_decoder[index] + + # allows converting a fast -> slow: add the `tokenizer.json`'s `"added_tokens"` to the slow tokenizer + # if `tokenizer_config.json` is `None` + if tokenizer_file is not None: + # This is for slow so can be done before + with open(tokenizer_file, encoding="utf-8") as tokenizer_file_handle: + tokenizer_file_handle = json.load(tokenizer_file_handle) + added_tokens = tokenizer_file_handle.pop("added_tokens") + for serialized_tokens in added_tokens: + idx = serialized_tokens.pop("id") + added_tokens_decoder[idx] = AddedToken(**serialized_tokens) + added_tokens_map[str(added_tokens_decoder[idx])] = added_tokens_decoder[idx] + # end legacy + + # Passing AddedTokens and not strings to the class to prevent it from casting the string to a different AddedToken + # convert {'__type': 'AddedToken', 'content': '', 'lstrip': False, 'normalized': True, ...} to AddedTokens + init_kwargs["added_tokens_decoder"] = added_tokens_decoder + init_kwargs = cls.convert_added_tokens(init_kwargs, save=False) + for key in cls.SPECIAL_TOKENS_ATTRIBUTES & init_kwargs.keys(): + if added_tokens_map != {} and init_kwargs[key] is not None: + if key != "additional_special_tokens": + init_kwargs[key] = added_tokens_map.get(str(init_kwargs[key]), init_kwargs[key]) + + # Instantiate the tokenizer. + try: + tokenizer = cls(*init_inputs, **init_kwargs) + except import_protobuf_decode_error(): + logger.info( + "Unable to load tokenizer model from SPM, loading from TikToken will be attempted instead." + "(Google protobuf error: Tried to load SPM model with non-SPM vocab file).", + ) + return False + except RuntimeError as e: + if "sentencepiece_processor.cc" in str(e): + logger.info( + "Unable to load tokenizer model from SPM, loading from TikToken will be attempted instead." + "(SentencePiece RuntimeError: Tried to load SPM model with non-SPM vocab file).", + ) + return False + except OSError: + raise OSError( + "Unable to load vocabulary from file. " + "Please check that the provided vocabulary is accessible and not corrupted." + ) + + if added_tokens_decoder != {} and max(list(added_tokens_decoder.keys())[-1], 0) > tokenizer.vocab_size: + logger.info( + "Special tokens have been added in the vocabulary, make sure the associated word embeddings are" + " fine-tuned or trained." + ) + return tokenizer + + @staticmethod + def _eventually_correct_t5_max_length(pretrained_model_name_or_path, max_model_length, init_max_model_length): + # This method should be deleted in Transformers v5 + # Its only purpose is to potentially throw a warning + # that incorrectly defined max lengths of T5's tokenizer are used + # which we will correct in Transformers v5. + return max_model_length + + @classmethod + def convert_added_tokens(cls, obj: Union[AddedToken, Any], save=False, add_type_field=True): + if isinstance(obj, dict) and "__type" in obj and obj["__type"] == "AddedToken": + obj.pop("__type") + return AddedToken(**obj) + if isinstance(obj, AddedToken) and save: + obj = obj.__getstate__() + if add_type_field: + obj["__type"] = "AddedToken" + else: + # Don't save "special" for previous tokenizers + obj.pop("special") + return obj + elif isinstance(obj, (list, tuple)): + return [cls.convert_added_tokens(o, save=save, add_type_field=add_type_field) for o in obj] + elif isinstance(obj, dict): + return {k: cls.convert_added_tokens(v, save=save, add_type_field=add_type_field) for k, v in obj.items()} + return obj + + def save_pretrained( + self, + save_directory: Union[str, os.PathLike], + legacy_format: Optional[bool] = None, + filename_prefix: Optional[str] = None, + push_to_hub: bool = False, + **kwargs, + ) -> Tuple[str]: + """ + Save the full tokenizer state. + + + This method make sure the full tokenizer can then be re-loaded using the + [`~tokenization_utils_base.PreTrainedTokenizer.from_pretrained`] class method.. + + Warning,None This won't save modifications you may have applied to the tokenizer after the instantiation (for + instance, modifying `tokenizer.do_lower_case` after creation). + + Args: + save_directory (`str` or `os.PathLike`): The path to a directory where the tokenizer will be saved. + legacy_format (`bool`, *optional*): + Only applicable for a fast tokenizer. If unset (default), will save the tokenizer in the unified JSON + format as well as in legacy format if it exists, i.e. with tokenizer specific vocabulary and a separate + added_tokens files. + + If `False`, will only save the tokenizer in the unified JSON format. This format is incompatible with + "slow" tokenizers (not powered by the *tokenizers* library), so the tokenizer will not be able to be + loaded in the corresponding "slow" tokenizer. + + If `True`, will save the tokenizer in legacy format. If the "slow" tokenizer doesn't exits, a value + error is raised. + filename_prefix (`str`, *optional*): + A prefix to add to the names of the files saved by the tokenizer. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the + repository you want to push to with `repo_id` (will default to the name of `save_directory` in your + namespace). + kwargs (`Dict[str, Any]`, *optional*): + Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method. + + Returns: + A tuple of `str`: The files saved. + """ + use_auth_token = kwargs.pop("use_auth_token", None) + + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if kwargs.get("token", None) is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + kwargs["token"] = use_auth_token + + if os.path.isfile(save_directory): + logger.error(f"Provided path ({save_directory}) should be a directory, not a file") + return + + os.makedirs(save_directory, exist_ok=True) + + if push_to_hub: + commit_message = kwargs.pop("commit_message", None) + repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1]) + repo_id = self._create_repo(repo_id, **kwargs) + files_timestamps = self._get_files_timestamps(save_directory) + + special_tokens_map_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + SPECIAL_TOKENS_MAP_FILE + ) + tokenizer_config_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + TOKENIZER_CONFIG_FILE + ) + chat_template_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + CHAT_TEMPLATE_FILE + ) + + tokenizer_config = copy.deepcopy(self.init_kwargs) + + # Let's save the init kwargs + target_keys = set(self.init_kwargs.keys()) + # Let's save the special tokens map (only the strings) + target_keys.update(["model_max_length", "clean_up_tokenization_spaces"]) + + for k in target_keys: + if hasattr(self, k): + tokenizer_config[k] = getattr(self, k) + + # Let's make sure we properly save the special tokens + tokenizer_config.update(self.special_tokens_map) + if "extra_special_tokens" not in tokenizer_config: + tokenizer_config["extra_special_tokens"] = self.extra_special_tokens + tokenizer_config.update(self.extra_special_tokens) + + saved_raw_chat_template = False + if self.chat_template is not None: + if isinstance(self.chat_template, dict): + # Chat template dicts are saved to the config as lists of dicts with fixed key names. + # They will be reconstructed as a single dict during loading. + # We're trying to discourage chat template dicts, and they are always + # saved in the config, never as single files. + tokenizer_config["chat_template"] = [{"name": k, "template": v} for k, v in self.chat_template.items()] + elif kwargs.get("save_raw_chat_template", False): + with open(chat_template_file, "w", encoding="utf-8") as f: + f.write(self.chat_template) + saved_raw_chat_template = True + logger.info(f"chat template saved in {chat_template_file}") + if "chat_template" in tokenizer_config: + tokenizer_config.pop("chat_template") # To ensure it doesn't somehow end up in the config too + else: + tokenizer_config["chat_template"] = self.chat_template + + if len(self.init_inputs) > 0: + tokenizer_config["init_inputs"] = copy.deepcopy(self.init_inputs) + for file_id in self.vocab_files_names.keys(): + tokenizer_config.pop(file_id, None) + + # no typefields, this way old fast and slow can load it + tokenizer_config = self.convert_added_tokens(tokenizer_config, add_type_field=True, save=True) + + # Process added tokens separately: allows previous versions to ignore it! + added_tokens = {} + for key, value in self.added_tokens_decoder.items(): + added_tokens[key] = value.__getstate__() + tokenizer_config["added_tokens_decoder"] = added_tokens + + # Add tokenizer class to the tokenizer config to be able to reload it with from_pretrained + tokenizer_class = self.__class__.__name__ + # Remove the Fast at the end if we can save the slow tokenizer + if tokenizer_class.endswith("Fast") and getattr(self, "can_save_slow_tokenizer", False): + tokenizer_class = tokenizer_class[:-4] + tokenizer_config["tokenizer_class"] = tokenizer_class + if getattr(self, "_auto_map", None) is not None: + tokenizer_config["auto_map"] = self._auto_map + if getattr(self, "_processor_class", None) is not None: + tokenizer_config["processor_class"] = self._processor_class + + # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be + # loaded from the Hub. + if self._auto_class is not None: + custom_object_save(self, save_directory, config=tokenizer_config) + + # remove private information + if "name_or_path" in tokenizer_config: + tokenizer_config.pop("name_or_path") + tokenizer_config.pop("special_tokens_map_file", None) + tokenizer_config.pop("tokenizer_file", None) + if "device_map" in tokenizer_config: + tokenizer_config.pop("device_map") + + with open(tokenizer_config_file, "w", encoding="utf-8") as f: + out_str = json.dumps(tokenizer_config, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + f.write(out_str) + logger.info(f"tokenizer config file saved in {tokenizer_config_file}") + + # Sanitize AddedTokens in special_tokens_map + + # kept for forward compatibility, will be removed in transoformers 5. Typefields are not saved for FC, special should not be save either + write_dict = self.convert_added_tokens(self.special_tokens_map_extended, save=True, add_type_field=False) + with open(special_tokens_map_file, "w", encoding="utf-8") as f: + out_str = json.dumps(write_dict, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + f.write(out_str) + logger.info(f"Special tokens file saved in {special_tokens_map_file}") + + file_names = (tokenizer_config_file, special_tokens_map_file) + if saved_raw_chat_template: + file_names += (chat_template_file,) + + save_files = self._save_pretrained( + save_directory=save_directory, + file_names=file_names, + legacy_format=legacy_format, + filename_prefix=filename_prefix, + ) + + if push_to_hub: + self._upload_modified_files( + save_directory, + repo_id, + files_timestamps, + commit_message=commit_message, + token=kwargs.get("token"), + ) + + return save_files + + def _save_pretrained( + self, + save_directory: Union[str, os.PathLike], + file_names: Tuple[str], + legacy_format: Optional[bool] = None, + filename_prefix: Optional[str] = None, + ) -> Tuple[str]: + """ + Save a tokenizer using the slow-tokenizer/legacy format: vocabulary + added tokens. + + Fast tokenizers can also be saved in a unique JSON file containing {config + vocab + added-tokens} using the + specific [`~tokenization_utils_fast.PreTrainedTokenizerFast._save_pretrained`] + """ + if legacy_format is False: + raise ValueError( + "Only fast tokenizers (instances of PreTrainedTokenizerFast) can be saved in non legacy format." + ) + + save_directory = str(save_directory) + + added_tokens_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + ADDED_TOKENS_FILE + ) + # the new get_added_vocab() also returns special tokens and tokens that have an index < vocab_size + added_vocab = {tok: index for tok, index in self.added_tokens_encoder.items() if index >= self.vocab_size} + if added_vocab: + with open(added_tokens_file, "w", encoding="utf-8") as f: + out_str = json.dumps(added_vocab, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + f.write(out_str) + logger.info(f"added tokens file saved in {added_tokens_file}") + + vocab_files = self.save_vocabulary(save_directory, filename_prefix=filename_prefix) + + return file_names + vocab_files + (added_tokens_file,) + + def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]: + """ + Save only the vocabulary of the tokenizer (vocabulary + added tokens). + + This method won't save the configuration and special token mappings of the tokenizer. Use + [`~PreTrainedTokenizerFast._save_pretrained`] to save the whole state of the tokenizer. + + Args: + save_directory (`str`): + The directory in which to save the vocabulary. + filename_prefix (`str`, *optional*): + An optional prefix to add to the named of the saved files. + + Returns: + `Tuple(str)`: Paths to the files saved. + """ + raise NotImplementedError + + def tokenize(self, text: str, pair: Optional[str] = None, add_special_tokens: bool = False, **kwargs) -> List[str]: + """ + Converts a string into a sequence of tokens, replacing unknown tokens with the `unk_token`. + + Args: + text (`str`): + The sequence to be encoded. + pair (`str`, *optional*): + A second sequence to be encoded with the first. + add_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to add the special tokens associated with the corresponding model. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific encode method. See details in + [`~PreTrainedTokenizerBase.__call__`] + + Returns: + `List[str]`: The list of tokens. + """ + raise NotImplementedError + + @add_end_docstrings( + ENCODE_KWARGS_DOCSTRING, + """ + **kwargs: Passed along to the `.tokenize()` method. + """, + """ + Returns: + `List[int]`, `torch.Tensor`, `tf.Tensor` or `np.ndarray`: The tokenized ids of the text. + """, + ) + def encode( + self, + text: Union[TextInput, PreTokenizedInput, EncodedInput], + text_pair: Optional[Union[TextInput, PreTokenizedInput, EncodedInput]] = None, + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy, None] = None, + max_length: Optional[int] = None, + stride: int = 0, + padding_side: Optional[str] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + **kwargs, + ) -> List[int]: + """ + Converts a string to a sequence of ids (integer), using the tokenizer and vocabulary. + + Same as doing `self.convert_tokens_to_ids(self.tokenize(text))`. + + Args: + text (`str`, `List[str]` or `List[int]`): + The first sequence to be encoded. This can be a string, a list of strings (tokenized string using the + `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids` + method). + text_pair (`str`, `List[str]` or `List[int]`, *optional*): + Optional second sequence to be encoded. This can be a string, a list of strings (tokenized string using + the `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids` + method). + """ + encoded_inputs = self.encode_plus( + text, + text_pair=text_pair, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + padding_side=padding_side, + return_tensors=return_tensors, + **kwargs, + ) + + return encoded_inputs["input_ids"] + + def num_special_tokens_to_add(self, pair: bool = False) -> int: + raise NotImplementedError + + def _get_padding_truncation_strategies( + self, padding=False, truncation=None, max_length=None, pad_to_multiple_of=None, verbose=True, **kwargs + ): + """ + Find the correct padding/truncation strategy with backward compatibility for old arguments (truncation_strategy + and pad_to_max_length) and behaviors. + """ + old_truncation_strategy = kwargs.pop("truncation_strategy", "do_not_truncate") + old_pad_to_max_length = kwargs.pop("pad_to_max_length", False) + + # Backward compatibility for previous behavior, maybe we should deprecate it: + # If you only set max_length, it activates truncation for max_length + if max_length is not None and padding is False and truncation is None: + if verbose: + if not self.deprecation_warnings.get("Truncation-not-explicitly-activated", False): + logger.warning( + "Truncation was not explicitly activated but `max_length` is provided a specific value, please" + " use `truncation=True` to explicitly truncate examples to max length. Defaulting to" + " 'longest_first' truncation strategy. If you encode pairs of sequences (GLUE-style) with the" + " tokenizer you can select this strategy more precisely by providing a specific strategy to" + " `truncation`." + ) + self.deprecation_warnings["Truncation-not-explicitly-activated"] = True + truncation = "longest_first" + + # Get padding strategy + if padding is False and old_pad_to_max_length: + if verbose: + warnings.warn( + "The `pad_to_max_length` argument is deprecated and will be removed in a future version, " + "use `padding=True` or `padding='longest'` to pad to the longest sequence in the batch, or " + "use `padding='max_length'` to pad to a max length. In this case, you can give a specific " + "length with `max_length` (e.g. `max_length=45`) or leave max_length to None to pad to the " + "maximal input size of the model (e.g. 512 for Bert).", + FutureWarning, + ) + if max_length is None: + padding_strategy = PaddingStrategy.LONGEST + else: + padding_strategy = PaddingStrategy.MAX_LENGTH + elif padding is not False: + if padding is True: + if verbose: + if max_length is not None and ( + truncation is None or truncation is False or truncation == "do_not_truncate" + ): + warnings.warn( + "`max_length` is ignored when `padding`=`True` and there is no truncation strategy. " + "To pad to max length, use `padding='max_length'`." + ) + if old_pad_to_max_length is not False: + warnings.warn("Though `pad_to_max_length` = `True`, it is ignored because `padding`=`True`.") + padding_strategy = PaddingStrategy.LONGEST # Default to pad to the longest sequence in the batch + elif not isinstance(padding, PaddingStrategy): + padding_strategy = PaddingStrategy(padding) + elif isinstance(padding, PaddingStrategy): + padding_strategy = padding + else: + padding_strategy = PaddingStrategy.DO_NOT_PAD + + # Get truncation strategy + if truncation is None and old_truncation_strategy != "do_not_truncate": + if verbose: + warnings.warn( + "The `truncation_strategy` argument is deprecated and will be removed in a future version, use" + " `truncation=True` to truncate examples to a max length. You can give a specific length with" + " `max_length` (e.g. `max_length=45`) or leave max_length to None to truncate to the maximal input" + " size of the model (e.g. 512 for Bert). If you have pairs of inputs, you can give a specific" + " truncation strategy selected among `truncation='only_first'` (will only truncate the first" + " sentence in the pairs) `truncation='only_second'` (will only truncate the second sentence in the" + " pairs) or `truncation='longest_first'` (will iteratively remove tokens from the longest sentence" + " in the pairs).", + FutureWarning, + ) + truncation_strategy = TruncationStrategy(old_truncation_strategy) + elif truncation is not False and truncation is not None: + if truncation is True: + truncation_strategy = ( + TruncationStrategy.LONGEST_FIRST + ) # Default to truncate the longest sequences in pairs of inputs + elif not isinstance(truncation, TruncationStrategy): + truncation_strategy = TruncationStrategy(truncation) + elif isinstance(truncation, TruncationStrategy): + truncation_strategy = truncation + else: + truncation_strategy = TruncationStrategy.DO_NOT_TRUNCATE + + # Set max length if needed + if max_length is None: + if padding_strategy == PaddingStrategy.MAX_LENGTH: + if self.model_max_length > LARGE_INTEGER: + if verbose: + if not self.deprecation_warnings.get("Asking-to-pad-to-max_length", False): + logger.warning( + "Asking to pad to max_length but no maximum length is provided and the model has no" + " predefined maximum length. Default to no padding." + ) + self.deprecation_warnings["Asking-to-pad-to-max_length"] = True + padding_strategy = PaddingStrategy.DO_NOT_PAD + else: + max_length = self.model_max_length + + if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE: + if self.model_max_length > LARGE_INTEGER: + if verbose: + if not self.deprecation_warnings.get("Asking-to-truncate-to-max_length", False): + logger.warning( + "Asking to truncate to max_length but no maximum length is provided and the model has" + " no predefined maximum length. Default to no truncation." + ) + self.deprecation_warnings["Asking-to-truncate-to-max_length"] = True + truncation_strategy = TruncationStrategy.DO_NOT_TRUNCATE + else: + max_length = self.model_max_length + + # Test if we have a padding token + if padding_strategy != PaddingStrategy.DO_NOT_PAD and (self.pad_token is None or self.pad_token_id < 0): + raise ValueError( + "Asking to pad but the tokenizer does not have a padding token. " + "Please select a token to use as `pad_token` `(tokenizer.pad_token = tokenizer.eos_token e.g.)` " + "or add a new pad token via `tokenizer.add_special_tokens({'pad_token': '[PAD]'})`." + ) + + # Check that we will truncate to a multiple of pad_to_multiple_of if both are provided + if ( + truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE + and padding_strategy != PaddingStrategy.DO_NOT_PAD + and pad_to_multiple_of is not None + and max_length is not None + and (max_length % pad_to_multiple_of != 0) + ): + raise ValueError( + "Truncation and padding are both activated but " + f"truncation length ({max_length}) is not a multiple of pad_to_multiple_of ({pad_to_multiple_of})." + ) + + return padding_strategy, truncation_strategy, max_length, kwargs + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def __call__( + self, + text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput], None] = None, + text_pair: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None, + text_target: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput], None] = None, + text_pair_target: Optional[ + Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] + ] = None, + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy, None] = None, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """ + Main method to tokenize and prepare for the model one or several sequence(s) or one or several pair(s) of + sequences. + + Args: + text (`str`, `List[str]`, `List[List[str]]`, *optional*): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + text_pair (`str`, `List[str]`, `List[List[str]]`, *optional*): + The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings + (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set + `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + text_target (`str`, `List[str]`, `List[List[str]]`, *optional*): + The sequence or batch of sequences to be encoded as target texts. Each sequence can be a string or a + list of strings (pretokenized string). If the sequences are provided as list of strings (pretokenized), + you must set `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + text_pair_target (`str`, `List[str]`, `List[List[str]]`, *optional*): + The sequence or batch of sequences to be encoded as target texts. Each sequence can be a string or a + list of strings (pretokenized string). If the sequences are provided as list of strings (pretokenized), + you must set `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). + """ + # To avoid duplicating + all_kwargs = { + "add_special_tokens": add_special_tokens, + "padding": padding, + "truncation": truncation, + "max_length": max_length, + "stride": stride, + "is_split_into_words": is_split_into_words, + "pad_to_multiple_of": pad_to_multiple_of, + "padding_side": padding_side, + "return_tensors": return_tensors, + "return_token_type_ids": return_token_type_ids, + "return_attention_mask": return_attention_mask, + "return_overflowing_tokens": return_overflowing_tokens, + "return_special_tokens_mask": return_special_tokens_mask, + "return_offsets_mapping": return_offsets_mapping, + "return_length": return_length, + "split_special_tokens": kwargs.pop("split_special_tokens", self.split_special_tokens), + "verbose": verbose, + } + all_kwargs.update(kwargs) + if text is None and text_target is None: + raise ValueError("You need to specify either `text` or `text_target`.") + if text is not None: + # The context manager will send the inputs as normal texts and not text_target, but we shouldn't change the + # input mode in this case. + if not self._in_target_context_manager: + self._switch_to_input_mode() + encodings = self._call_one(text=text, text_pair=text_pair, **all_kwargs) + if text_target is not None: + self._switch_to_target_mode() + target_encodings = self._call_one(text=text_target, text_pair=text_pair_target, **all_kwargs) + # Leave back tokenizer in input mode + self._switch_to_input_mode() + + if text_target is None: + return encodings + elif text is None: + return target_encodings + else: + encodings["labels"] = target_encodings["input_ids"] + return encodings + + def _call_one( + self, + text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]], + text_pair: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None, + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy, None] = None, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + # Input type checking for clearer error + def _is_valid_text_input(t): + if isinstance(t, str): + # Strings are fine + return True + elif isinstance(t, (list, tuple)): + # List are fine as long as they are... + if len(t) == 0: + # ... empty + return True + elif isinstance(t[0], str): + # ... list of strings + return True + elif isinstance(t[0], (list, tuple)): + # ... list with an empty list or with a list of strings + return len(t[0]) == 0 or isinstance(t[0][0], str) + else: + return False + else: + return False + + if not _is_valid_text_input(text): + raise ValueError( + "text input must be of type `str` (single example), `List[str]` (batch or single pretokenized example) " + "or `List[List[str]]` (batch of pretokenized examples)." + ) + + if text_pair is not None and not _is_valid_text_input(text_pair): + raise ValueError( + "text input must be of type `str` (single example), `List[str]` (batch or single pretokenized example) " + "or `List[List[str]]` (batch of pretokenized examples)." + ) + + if is_split_into_words: + is_batched = isinstance(text, (list, tuple)) and text and isinstance(text[0], (list, tuple)) + else: + is_batched = isinstance(text, (list, tuple)) + + if is_batched: + if isinstance(text_pair, str): + raise TypeError( + "when tokenizing batches of text, `text_pair` must be a list or tuple with the same length as" + " `text`." + ) + if text_pair is not None and len(text) != len(text_pair): + raise ValueError( + f"batch length of `text`: {len(text)} does not match batch length of `text_pair`:" + f" {len(text_pair)}." + ) + batch_text_or_text_pairs = list(zip(text, text_pair)) if text_pair is not None else text + return self.batch_encode_plus( + batch_text_or_text_pairs=batch_text_or_text_pairs, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + split_special_tokens=split_special_tokens, + **kwargs, + ) + else: + return self.encode_plus( + text=text, + text_pair=text_pair, + add_special_tokens=add_special_tokens, + padding=padding, + truncation=truncation, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + split_special_tokens=split_special_tokens, + **kwargs, + ) + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def encode_plus( + self, + text: Union[TextInput, PreTokenizedInput, EncodedInput], + text_pair: Optional[Union[TextInput, PreTokenizedInput, EncodedInput]] = None, + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy, None] = None, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + **kwargs, + ) -> BatchEncoding: + """ + Tokenize and prepare for the model a sequence or a pair of sequences. + + + + This method is deprecated, `__call__` should be used instead. + + + + Args: + text (`str`, `List[str]` or (for non-fast tokenizers) `List[int]`): + The first sequence to be encoded. This can be a string, a list of strings (tokenized string using the + `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids` + method). + text_pair (`str`, `List[str]` or `List[int]`, *optional*): + Optional second sequence to be encoded. This can be a string, a list of strings (tokenized string using + the `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids` + method). + """ + + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + return self._encode_plus( + text=text, + text_pair=text_pair, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + split_special_tokens=kwargs.pop("split_special_tokens", self.split_special_tokens), + **kwargs, + ) + + def _encode_plus( + self, + text: Union[TextInput, PreTokenizedInput, EncodedInput], + text_pair: Optional[Union[TextInput, PreTokenizedInput, EncodedInput]] = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + raise NotImplementedError + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def batch_encode_plus( + self, + batch_text_or_text_pairs: Union[ + List[TextInput], + List[TextInputPair], + List[PreTokenizedInput], + List[PreTokenizedInputPair], + List[EncodedInput], + List[EncodedInputPair], + ], + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy, None] = None, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + """ + Tokenize and prepare for the model a list of sequences or a list of pairs of sequences. + + + + This method is deprecated, `__call__` should be used instead. + + + + Args: + batch_text_or_text_pairs (`List[str]`, `List[Tuple[str, str]]`, `List[List[str]]`, `List[Tuple[List[str], List[str]]]`, and for not-fast tokenizers, also `List[List[int]]`, `List[Tuple[List[int], List[int]]]`): + Batch of sequences or pair of sequences to be encoded. This can be a list of + string/string-sequences/int-sequences or a list of pair of string/string-sequences/int-sequence (see + details in `encode_plus`). + """ + + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + return self._batch_encode_plus( + batch_text_or_text_pairs=batch_text_or_text_pairs, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + is_split_into_words=is_split_into_words, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + split_special_tokens=split_special_tokens, + **kwargs, + ) + + def _batch_encode_plus( + self, + batch_text_or_text_pairs: Union[ + List[TextInput], + List[TextInputPair], + List[PreTokenizedInput], + List[PreTokenizedInputPair], + List[EncodedInput], + List[EncodedInputPair], + ], + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + raise NotImplementedError + + def pad( + self, + encoded_inputs: Union[ + BatchEncoding, + List[BatchEncoding], + Dict[str, EncodedInput], + Dict[str, List[EncodedInput]], + List[Dict[str, EncodedInput]], + ], + padding: Union[bool, str, PaddingStrategy] = True, + max_length: Optional[int] = None, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_attention_mask: Optional[bool] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + verbose: bool = True, + ) -> BatchEncoding: + """ + Pad a single encoded input or a batch of encoded inputs up to predefined length or to the max sequence length + in the batch. + + Padding side (left/right) padding token ids are defined at the tokenizer level (with `self.padding_side`, + `self.pad_token_id` and `self.pad_token_type_id`). + + Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the + text followed by a call to the `pad` method to get a padded encoding. + + + + If the `encoded_inputs` passed are dictionary of numpy arrays, PyTorch tensors or TensorFlow tensors, the + result will use the same type unless you provide a different tensor type with `return_tensors`. In the case of + PyTorch tensors, you will lose the specific device of your tensors however. + + + + Args: + encoded_inputs ([`BatchEncoding`], list of [`BatchEncoding`], `Dict[str, List[int]]`, `Dict[str, List[List[int]]` or `List[Dict[str, List[int]]]`): + Tokenized inputs. Can represent one input ([`BatchEncoding`] or `Dict[str, List[int]]`) or a batch of + tokenized inputs (list of [`BatchEncoding`], *Dict[str, List[List[int]]]* or *List[Dict[str, + List[int]]]*) so you can use this method during preprocessing as well as in a PyTorch Dataloader + collate function. + + Instead of `List[int]` you can have tensors (numpy arrays, PyTorch tensors or TensorFlow tensors), see + the note above for the return type. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`): + Select a strategy to pad the returned sequences (according to the model's padding side and padding + index) among: + + - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different + lengths). + max_length (`int`, *optional*): + Maximum length of the returned list and optionally padding length (see above). + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. + + This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side (`str`, *optional*): + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_attention_mask (`bool`, *optional*): + Whether to return the attention mask. If left to the default, will return the attention mask according + to the specific tokenizer's default, defined by the `return_outputs` attribute. + + [What are attention masks?](../glossary#attention-mask) + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + verbose (`bool`, *optional*, defaults to `True`): + Whether or not to print more information and warnings. + """ + if self.__class__.__name__.endswith("Fast"): + if not self.deprecation_warnings.get("Asking-to-pad-a-fast-tokenizer", False): + logger.warning_advice( + f"You're using a {self.__class__.__name__} tokenizer. Please note that with a fast tokenizer," + " using the `__call__` method is faster than using a method to encode the text followed by a call" + " to the `pad` method to get a padded encoding." + ) + self.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True + + # If we have a list of dicts, let's convert it in a dict of lists + # We do this to allow using this method as a collate_fn function in PyTorch Dataloader + if isinstance(encoded_inputs, (list, tuple)) and isinstance(encoded_inputs[0], Mapping): + encoded_inputs = {key: [example[key] for example in encoded_inputs] for key in encoded_inputs[0].keys()} + + # The model's main input name, usually `input_ids`, has been passed for padding + if self.model_input_names[0] not in encoded_inputs: + raise ValueError( + "You should supply an encoding or a list of encodings to this method " + f"that includes {self.model_input_names[0]}, but you provided {list(encoded_inputs.keys())}" + ) + + required_input = encoded_inputs[self.model_input_names[0]] + + if required_input is None or (isinstance(required_input, Sized) and len(required_input) == 0): + if return_attention_mask: + encoded_inputs["attention_mask"] = [] + return encoded_inputs + + # If we have PyTorch/TF/NumPy tensors/arrays as inputs, we cast them as python objects + # and rebuild them afterwards if no return_tensors is specified + # Note that we lose the specific device the tensor may be on for PyTorch + + first_element = required_input[0] + if isinstance(first_element, (list, tuple)): + # first_element might be an empty list/tuple in some edge cases so we grab the first non empty element. + for item in required_input: + if len(item) != 0: + first_element = item[0] + break + # At this state, if `first_element` is still a list/tuple, it's an empty one so there is nothing to do. + if not isinstance(first_element, (int, list, tuple)): + if is_tf_tensor(first_element): + return_tensors = "tf" if return_tensors is None else return_tensors + elif is_torch_tensor(first_element): + return_tensors = "pt" if return_tensors is None else return_tensors + elif isinstance(first_element, np.ndarray): + return_tensors = "np" if return_tensors is None else return_tensors + else: + raise ValueError( + f"type of {first_element} unknown: {type(first_element)}. " + "Should be one of a python, numpy, pytorch or tensorflow object." + ) + + for key, value in encoded_inputs.items(): + encoded_inputs[key] = to_py_obj(value) + + # Convert padding_strategy in PaddingStrategy + padding_strategy, _, max_length, _ = self._get_padding_truncation_strategies( + padding=padding, max_length=max_length, verbose=verbose + ) + + required_input = encoded_inputs[self.model_input_names[0]] + if required_input and not isinstance(required_input[0], (list, tuple)): + encoded_inputs = self._pad( + encoded_inputs, + max_length=max_length, + padding_strategy=padding_strategy, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + return BatchEncoding(encoded_inputs, tensor_type=return_tensors) + + batch_size = len(required_input) + assert all(len(v) == batch_size for v in encoded_inputs.values()), ( + "Some items in the output dictionary have a different batch size than others." + ) + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = max(len(inputs) for inputs in required_input) + padding_strategy = PaddingStrategy.MAX_LENGTH + + batch_outputs = {} + for i in range(batch_size): + inputs = {k: v[i] for k, v in encoded_inputs.items()} + outputs = self._pad( + inputs, + max_length=max_length, + padding_strategy=padding_strategy, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + for key, value in outputs.items(): + if key not in batch_outputs: + batch_outputs[key] = [] + batch_outputs[key].append(value) + + return BatchEncoding(batch_outputs, tensor_type=return_tensors) + + def create_token_type_ids_from_sequences( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """ + Create the token type IDs corresponding to the sequences passed. [What are token type + IDs?](../glossary#token-type-ids) + + Should be overridden in a subclass if the model has a special way of building those. + + Args: + token_ids_0 (`List[int]`): The first tokenized sequence. + token_ids_1 (`List[int]`, *optional*): The second tokenized sequence. + + Returns: + `List[int]`: The token type ids. + """ + if token_ids_1 is None: + return len(token_ids_0) * [0] + return [0] * len(token_ids_0) + [1] * len(token_ids_1) + + def build_inputs_with_special_tokens( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + """ + Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and + adding special tokens. + + This implementation does not add special tokens and this method should be overridden in a subclass. + + Args: + token_ids_0 (`List[int]`): The first tokenized sequence. + token_ids_1 (`List[int]`, *optional*): The second tokenized sequence. + + Returns: + `List[int]`: The model input with special tokens. + """ + if token_ids_1 is None: + return token_ids_0 + return token_ids_0 + token_ids_1 + + @add_end_docstrings(ENCODE_KWARGS_DOCSTRING, ENCODE_PLUS_ADDITIONAL_KWARGS_DOCSTRING) + def prepare_for_model( + self, + ids: List[int], + pair_ids: Optional[List[int]] = None, + add_special_tokens: bool = True, + padding: Union[bool, str, PaddingStrategy] = False, + truncation: Union[bool, str, TruncationStrategy, None] = None, + max_length: Optional[int] = None, + stride: int = 0, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + prepend_batch_axis: bool = False, + **kwargs, + ) -> BatchEncoding: + """ + Prepares a sequence of input id, or a pair of sequences of inputs ids so that it can be used by the model. It + adds special tokens, truncates sequences if overflowing while taking into account the special tokens and + manages a moving window (with user defined stride) for overflowing tokens. Please Note, for *pair_ids* + different than `None` and *truncation_strategy = longest_first* or `True`, it is not possible to return + overflowing tokens. Such a combination of arguments will raise an error. + + Args: + ids (`List[int]`): + Tokenized input ids of the first sequence. Can be obtained from a string by chaining the `tokenize` and + `convert_tokens_to_ids` methods. + pair_ids (`List[int]`, *optional*): + Tokenized input ids of the second sequence. Can be obtained from a string by chaining the `tokenize` + and `convert_tokens_to_ids` methods. + """ + + # Backward compatibility for 'truncation_strategy', 'pad_to_max_length' + padding_strategy, truncation_strategy, max_length, kwargs = self._get_padding_truncation_strategies( + padding=padding, + truncation=truncation, + max_length=max_length, + pad_to_multiple_of=pad_to_multiple_of, + verbose=verbose, + **kwargs, + ) + + pair = bool(pair_ids is not None) + len_ids = len(ids) + len_pair_ids = len(pair_ids) if pair else 0 + + if return_token_type_ids and not add_special_tokens: + raise ValueError( + "Asking to return token_type_ids while setting add_special_tokens to False " + "results in an undefined behavior. Please set add_special_tokens to True or " + "set return_token_type_ids to None." + ) + + if ( + return_overflowing_tokens + and truncation_strategy == TruncationStrategy.LONGEST_FIRST + and pair_ids is not None + ): + raise ValueError( + "Not possible to return overflowing tokens for pair of sequences with the " + "`longest_first`. Please select another truncation strategy than `longest_first`, " + "for instance `only_second` or `only_first`." + ) + + # Load from model defaults + if return_token_type_ids is None: + return_token_type_ids = "token_type_ids" in self.model_input_names + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + encoded_inputs = {} + + # Compute the total size of the returned encodings + total_len = len_ids + len_pair_ids + (self.num_special_tokens_to_add(pair=pair) if add_special_tokens else 0) + + # Truncation: Handle max sequence length + overflowing_tokens = [] + if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE and max_length and total_len > max_length: + ids, pair_ids, overflowing_tokens = self.truncate_sequences( + ids, + pair_ids=pair_ids, + num_tokens_to_remove=total_len - max_length, + truncation_strategy=truncation_strategy, + stride=stride, + ) + + if return_overflowing_tokens: + encoded_inputs["overflowing_tokens"] = overflowing_tokens + encoded_inputs["num_truncated_tokens"] = total_len - max_length + + # Add special tokens + if add_special_tokens: + sequence = self.build_inputs_with_special_tokens(ids, pair_ids) + token_type_ids = self.create_token_type_ids_from_sequences(ids, pair_ids) + else: + sequence = ids + pair_ids if pair else ids + token_type_ids = [0] * len(ids) + ([0] * len(pair_ids) if pair else []) + + # Build output dictionary + encoded_inputs["input_ids"] = sequence + if return_token_type_ids: + encoded_inputs["token_type_ids"] = token_type_ids + if return_special_tokens_mask: + if add_special_tokens: + encoded_inputs["special_tokens_mask"] = self.get_special_tokens_mask(ids, pair_ids) + else: + encoded_inputs["special_tokens_mask"] = [0] * len(sequence) + + # Check lengths + self._eventual_warn_about_too_long_sequence(encoded_inputs["input_ids"], max_length, verbose) + + # Padding + if padding_strategy != PaddingStrategy.DO_NOT_PAD or return_attention_mask: + encoded_inputs = self.pad( + encoded_inputs, + max_length=max_length, + padding=padding_strategy.value, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_attention_mask=return_attention_mask, + ) + + if return_length: + encoded_inputs["length"] = len(encoded_inputs["input_ids"]) + + batch_outputs = BatchEncoding( + encoded_inputs, tensor_type=return_tensors, prepend_batch_axis=prepend_batch_axis + ) + + return batch_outputs + + def truncate_sequences( + self, + ids: List[int], + pair_ids: Optional[List[int]] = None, + num_tokens_to_remove: int = 0, + truncation_strategy: Union[str, TruncationStrategy] = "longest_first", + stride: int = 0, + ) -> Tuple[List[int], List[int], List[int]]: + """ + Truncates a sequence pair in-place following the strategy. + + Args: + ids (`List[int]`): + Tokenized input ids of the first sequence. Can be obtained from a string by chaining the `tokenize` and + `convert_tokens_to_ids` methods. + pair_ids (`List[int]`, *optional*): + Tokenized input ids of the second sequence. Can be obtained from a string by chaining the `tokenize` + and `convert_tokens_to_ids` methods. + num_tokens_to_remove (`int`, *optional*, defaults to 0): + Number of tokens to remove using the truncation strategy. + truncation_strategy (`str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `'longest_first'`): + The strategy to follow for truncation. Can be: + + - `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will truncate + token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a + batch of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths greater + than the model maximum admissible input size). + stride (`int`, *optional*, defaults to 0): + If set to a positive number, the overflowing tokens returned will contain some tokens from the main + sequence returned. The value of this argument defines the number of additional tokens. + + Returns: + `Tuple[List[int], List[int], List[int]]`: The truncated `ids`, the truncated `pair_ids` and the list of + overflowing tokens. Note: The *longest_first* strategy returns empty list of overflowing tokens if a pair + of sequences (or a batch of pairs) is provided. + """ + if num_tokens_to_remove <= 0: + return ids, pair_ids, [] + + if not isinstance(truncation_strategy, TruncationStrategy): + truncation_strategy = TruncationStrategy(truncation_strategy) + + overflowing_tokens = [] + if truncation_strategy == TruncationStrategy.ONLY_FIRST or ( + truncation_strategy == TruncationStrategy.LONGEST_FIRST and pair_ids is None + ): + if len(ids) > num_tokens_to_remove: + window_len = min(len(ids), stride + num_tokens_to_remove) + if self.truncation_side == "left": + overflowing_tokens = ids[:window_len] + ids = ids[num_tokens_to_remove:] + elif self.truncation_side == "right": + overflowing_tokens = ids[-window_len:] + ids = ids[:-num_tokens_to_remove] + else: + raise ValueError(f"invalid truncation strategy: {self.truncation_side}, use 'left' or 'right'.") + + else: + error_msg = ( + f"We need to remove {num_tokens_to_remove} to truncate the input " + f"but the first sequence has a length {len(ids)}. " + ) + if truncation_strategy == TruncationStrategy.ONLY_FIRST: + error_msg = ( + error_msg + "Please select another truncation strategy than " + f"{truncation_strategy}, for instance 'longest_first' or 'only_second'." + ) + logger.error(error_msg) + elif truncation_strategy == TruncationStrategy.LONGEST_FIRST: + logger.warning( + "Be aware, overflowing tokens are not returned for the setting you have chosen," + f" i.e. sequence pairs with the '{TruncationStrategy.LONGEST_FIRST.value}' " + "truncation strategy. So the returned list will always be empty even if some " + "tokens have been removed." + ) + len_pair_ids = len(pair_ids) if pair_ids is not None else 0 + len_ids = len(ids) + first_remove = min(abs(len_pair_ids - len_ids), num_tokens_to_remove) + second_remove = num_tokens_to_remove - first_remove + if len_ids > len_pair_ids: + ids_to_move = first_remove + second_remove // 2 + pair_ids_to_move = second_remove - second_remove // 2 + else: + ids_to_move = second_remove // 2 + pair_ids_to_move = first_remove + second_remove - (second_remove // 2) + + if self.truncation_side == "right": + ids = ids[:-ids_to_move] if ids_to_move > 0 else ids + pair_ids = pair_ids[:-pair_ids_to_move] if pair_ids is not None and pair_ids_to_move > 0 else pair_ids + elif self.truncation_side == "left": + ids = ids[ids_to_move:] + pair_ids = pair_ids[pair_ids_to_move:] if pair_ids is not None else None + else: + raise ValueError(f"invalid truncation strategy:{self.truncation_side}") + + elif truncation_strategy == TruncationStrategy.ONLY_SECOND and pair_ids is not None: + if len(pair_ids) > num_tokens_to_remove: + window_len = min(len(pair_ids), stride + num_tokens_to_remove) + if self.truncation_side == "right": + overflowing_tokens = pair_ids[-window_len:] + pair_ids = pair_ids[:-num_tokens_to_remove] + elif self.truncation_side == "left": + overflowing_tokens = pair_ids[:window_len] + pair_ids = pair_ids[num_tokens_to_remove:] + else: + raise ValueError(f"invalid truncation strategy:{self.truncation_side}") + else: + logger.error( + f"We need to remove {num_tokens_to_remove} to truncate the input " + f"but the second sequence has a length {len(pair_ids)}. " + f"Please select another truncation strategy than {truncation_strategy}, " + "for instance 'longest_first' or 'only_first'." + ) + + return (ids, pair_ids, overflowing_tokens) + + def _pad( + self, + encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], + max_length: Optional[int] = None, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_attention_mask: Optional[bool] = None, + ) -> dict: + """ + Pad encoded inputs (on left/right and up to predefined length or max length in the batch) + + Args: + encoded_inputs: + Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`). + max_length: maximum length of the returned list and optionally padding length (see below). + Will truncate by taking into account the special tokens. + padding_strategy: PaddingStrategy to use for padding. + + - PaddingStrategy.LONGEST Pad to the longest sequence in the batch + - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) + - PaddingStrategy.DO_NOT_PAD: Do not pad + The tokenizer padding sides are defined in `padding_side` argument: + + - 'left': pads on the left of the sequences + - 'right': pads on the right of the sequences + pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. + This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability + `>= 7.5` (Volta). + padding_side: + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + return_attention_mask: + (optional) Set to False to avoid returning attention mask (default: set to model specifics) + """ + # Load from model defaults + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + required_input = encoded_inputs[self.model_input_names[0]] + + if padding_strategy == PaddingStrategy.LONGEST: + max_length = len(required_input) + + if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0): + max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of + + needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length + + # Initialize attention mask if not present. + if return_attention_mask and "attention_mask" not in encoded_inputs: + encoded_inputs["attention_mask"] = [1] * len(required_input) + + if needs_to_be_padded: + difference = max_length - len(required_input) + padding_side = padding_side if padding_side is not None else self.padding_side + + if padding_side == "right": + if return_attention_mask: + encoded_inputs["attention_mask"] = encoded_inputs["attention_mask"] + [0] * difference + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = ( + encoded_inputs["token_type_ids"] + [self.pad_token_type_id] * difference + ) + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = encoded_inputs["special_tokens_mask"] + [1] * difference + encoded_inputs[self.model_input_names[0]] = required_input + [self.pad_token_id] * difference + elif padding_side == "left": + if return_attention_mask: + encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"] + if "token_type_ids" in encoded_inputs: + encoded_inputs["token_type_ids"] = [self.pad_token_type_id] * difference + encoded_inputs[ + "token_type_ids" + ] + if "special_tokens_mask" in encoded_inputs: + encoded_inputs["special_tokens_mask"] = [1] * difference + encoded_inputs["special_tokens_mask"] + encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input + else: + raise ValueError(f"Invalid padding strategy:{padding_side}") + + return encoded_inputs + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + """ + Converts a sequence of tokens in a single string. The most simple way to do it is `" ".join(tokens)` but we + often want to remove sub-word tokenization artifacts at the same time. + + Args: + tokens (`List[str]`): The token to join in a string. + + Returns: + `str`: The joined tokens. + """ + raise NotImplementedError + + def batch_decode( + self, + sequences: Union[List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = None, + **kwargs, + ) -> List[str]: + """ + Convert a list of lists of token ids into a list of strings by calling decode. + + Args: + sequences (`Union[List[int], List[List[int]], np.ndarray, torch.Tensor, tf.Tensor]`): + List of tokenized input ids. Can be obtained using the `__call__` method. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether or not to clean up the tokenization spaces. If `None`, will default to + `self.clean_up_tokenization_spaces`. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific decode method. + + Returns: + `List[str]`: The list of decoded sentences. + """ + return [ + self.decode( + seq, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + for seq in sequences + ] + + def decode( + self, + token_ids: Union[int, List[int], "np.ndarray", "torch.Tensor", "tf.Tensor"], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = None, + **kwargs, + ) -> str: + """ + Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special + tokens and clean up tokenization spaces. + + Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`. + + Args: + token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`): + List of tokenized input ids. Can be obtained using the `__call__` method. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + clean_up_tokenization_spaces (`bool`, *optional*): + Whether or not to clean up the tokenization spaces. If `None`, will default to + `self.clean_up_tokenization_spaces`. + kwargs (additional keyword arguments, *optional*): + Will be passed to the underlying model specific decode method. + + Returns: + `str`: The decoded sentence. + """ + # Convert inputs to python lists + token_ids = to_py_obj(token_ids) + + return self._decode( + token_ids=token_ids, + skip_special_tokens=skip_special_tokens, + clean_up_tokenization_spaces=clean_up_tokenization_spaces, + **kwargs, + ) + + def _decode( + self, + token_ids: Union[int, List[int]], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = None, + **kwargs, + ) -> str: + raise NotImplementedError + + def get_special_tokens_mask( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False + ) -> List[int]: + """ + Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding + special tokens using the tokenizer `prepare_for_model` or `encode_plus` methods. + + Args: + token_ids_0 (`List[int]`): + List of ids of the first sequence. + token_ids_1 (`List[int]`, *optional*): + List of ids of the second sequence. + already_has_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not the token list is already formatted with special tokens for the model. + + Returns: + A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token. + """ + assert already_has_special_tokens and token_ids_1 is None, ( + "You cannot use ``already_has_special_tokens=False`` with this tokenizer. " + "Please use a slow (full python) tokenizer to activate this argument. " + "Or set `return_special_tokens_mask=True` when calling the encoding method " + "to get the special tokens mask in any tokenizer. " + ) + + all_special_ids = self.all_special_ids # cache the property + + special_tokens_mask = [1 if token in all_special_ids else 0 for token in token_ids_0] + + return special_tokens_mask + + @staticmethod + def clean_up_tokenization(out_string: str) -> str: + """ + Clean up a list of simple English tokenization artifacts like spaces before punctuations and abbreviated forms. + + Args: + out_string (`str`): The text to clean up. + + Returns: + `str`: The cleaned-up string. + """ + out_string = ( + out_string.replace(" .", ".") + .replace(" ?", "?") + .replace(" !", "!") + .replace(" ,", ",") + .replace(" ' ", "'") + .replace(" n't", "n't") + .replace(" 'm", "'m") + .replace(" 's", "'s") + .replace(" 've", "'ve") + .replace(" 're", "'re") + ) + return out_string + + def _eventual_warn_about_too_long_sequence(self, ids: List[int], max_length: Optional[int], verbose: bool): + """ + Depending on the input and internal state we might trigger a warning about a sequence that is too long for its + corresponding model + + Args: + ids (`List[str]`): The ids produced by the tokenization + max_length (`int`, *optional*): The max_length desired (does not trigger a warning if it is set) + verbose (`bool`): Whether or not to print more information and warnings. + + """ + if max_length is None and len(ids) > self.model_max_length and verbose and self.model_max_length != 0: + if not self.deprecation_warnings.get("sequence-length-is-longer-than-the-specified-maximum", False): + logger.warning( + "Token indices sequence length is longer than the specified maximum sequence length " + f"for this model ({len(ids)} > {self.model_max_length}). Running this sequence through the model " + "will result in indexing errors" + ) + self.deprecation_warnings["sequence-length-is-longer-than-the-specified-maximum"] = True + + def _switch_to_input_mode(self): + """ + Private method to put the tokenizer in input mode (when it has different modes for input/outputs) + """ + pass + + def _switch_to_target_mode(self): + """ + Private method to put the tokenizer in target mode (when it has different modes for input/outputs) + """ + pass + + @contextmanager + def as_target_tokenizer(self): + """ + Temporarily sets the tokenizer for encoding the targets. Useful for tokenizer associated to + sequence-to-sequence models that need a slightly different processing for the labels. + """ + warnings.warn( + "`as_target_tokenizer` is deprecated and will be removed in v5 of Transformers. You can tokenize your " + "labels by using the argument `text_target` of the regular `__call__` method (either in the same call as " + "your input texts if you use the same keyword arguments, or in a separate call." + ) + self._switch_to_target_mode() + self._in_target_context_manager = True + yield + self._in_target_context_manager = False + self._switch_to_input_mode() + + @classmethod + def register_for_auto_class(cls, auto_class="AutoTokenizer"): + """ + Register this class with a given auto class. This should only be used for custom tokenizers as the ones in the + library are already mapped with `AutoTokenizer`. + + + + This API is experimental and may have some slight breaking changes in the next releases. + + + + Args: + auto_class (`str` or `type`, *optional*, defaults to `"AutoTokenizer"`): + The auto class to register this new tokenizer with. + """ + if not isinstance(auto_class, str): + auto_class = auto_class.__name__ + + import transformers.models.auto as auto_module + + if not hasattr(auto_module, auto_class): + raise ValueError(f"{auto_class} is not a valid auto class.") + + cls._auto_class = auto_class + + def prepare_seq2seq_batch( + self, + src_texts: List[str], + tgt_texts: Optional[List[str]] = None, + max_length: Optional[int] = None, + max_target_length: Optional[int] = None, + padding: str = "longest", + return_tensors: Optional[str] = None, + truncation: bool = True, + **kwargs, + ) -> BatchEncoding: + """ + Prepare model inputs for translation. For best performance, translate one sentence at a time. + + Arguments: + src_texts (`List[str]`): + List of documents to summarize or source language texts. + tgt_texts (`list`, *optional*): + List of summaries or target language texts. + max_length (`int`, *optional*): + Controls the maximum length for encoder inputs (documents to summarize or source language texts) If + left unset or set to `None`, this will use the predefined model maximum length if a maximum length is + required by one of the truncation/padding parameters. If the model has no specific maximum input length + (like XLNet) truncation/padding to a maximum length will be deactivated. + max_target_length (`int`, *optional*): + Controls the maximum length of decoder inputs (target language texts or summaries) If left unset or set + to `None`, this will use the max_length value. + padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`): + Activates and controls padding. Accepts the following values: + + - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single + sequence if provided). + - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum + acceptable input length for the model if that argument is not provided. + - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different + lengths). + return_tensors (`str` or [`~utils.TensorType`], *optional*): + If set, will return tensors instead of list of python integers. Acceptable values are: + + - `'tf'`: Return TensorFlow `tf.constant` objects. + - `'pt'`: Return PyTorch `torch.Tensor` objects. + - `'np'`: Return Numpy `np.ndarray` objects. + truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `True`): + Activates and controls truncation. Accepts the following values: + + - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or + to the maximum acceptable input length for the model if that argument is not provided. This will + truncate token by token, removing a token from the longest sequence in the pair if a pair of + sequences (or a batch of pairs) is provided. + - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the + maximum acceptable input length for the model if that argument is not provided. This will only + truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided. + - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths + greater than the model maximum admissible input size). + **kwargs: + Additional keyword arguments passed along to `self.__call__`. + + Return: + [`BatchEncoding`]: A [`BatchEncoding`] with the following fields: + + - **input_ids** -- List of token ids to be fed to the encoder. + - **attention_mask** -- List of indices specifying which tokens should be attended to by the model. + - **labels** -- List of token ids for tgt_texts. + + The full set of keys `[input_ids, attention_mask, labels]`, will only be returned if tgt_texts is passed. + Otherwise, input_ids, attention_mask will be the only keys. + """ + # docstyle-ignore + formatted_warning = """ +`prepare_seq2seq_batch` is deprecated and will be removed in version 5 of HuggingFace Transformers. Use the regular +`__call__` method to prepare your inputs and targets. + +Here is a short example: + +model_inputs = tokenizer(src_texts, text_target=tgt_texts, ...) + +If you either need to use different keyword arguments for the source and target texts, you should do two calls like +this: + +model_inputs = tokenizer(src_texts, ...) +labels = tokenizer(text_target=tgt_texts, ...) +model_inputs["labels"] = labels["input_ids"] + +See the documentation of your specific tokenizer for more details on the specific arguments to the tokenizer of choice. +For a more complete example, see the implementation of `prepare_seq2seq_batch`. +""" + warnings.warn(formatted_warning, FutureWarning) + # mBART-specific kwargs that should be ignored by other models. + kwargs.pop("src_lang", None) + kwargs.pop("tgt_lang", None) + if max_length is None: + max_length = self.model_max_length + model_inputs = self( + src_texts, + add_special_tokens=True, + return_tensors=return_tensors, + max_length=max_length, + padding=padding, + truncation=truncation, + **kwargs, + ) + if tgt_texts is None: + return model_inputs + # Process tgt_texts + if max_target_length is None: + max_target_length = max_length + with self.as_target_tokenizer(): + labels = self( + tgt_texts, + add_special_tokens=True, + return_tensors=return_tensors, + padding=padding, + max_length=max_target_length, + truncation=truncation, + **kwargs, + ) + model_inputs["labels"] = labels["input_ids"] + return model_inputs + + +def get_fast_tokenizer_file(tokenization_files: List[str]) -> str: + """ + Get the tokenization file to use for this version of transformers. + + Args: + tokenization_files (`List[str]`): The list of available configuration files. + + Returns: + `str`: The tokenization file to use. + """ + tokenizer_files_map = {} + for file_name in tokenization_files: + search = _re_tokenizer_file.search(file_name) + if search is not None: + v = search.groups()[0] + tokenizer_files_map[v] = file_name + available_versions = sorted(tokenizer_files_map.keys()) + + # Defaults to FULL_TOKENIZER_FILE and then try to look at some newer versions. + tokenizer_file = FULL_TOKENIZER_FILE + transformers_version = version.parse(__version__) + for v in available_versions: + if version.parse(v) <= transformers_version: + tokenizer_file = tokenizer_files_map[v] + else: + # No point going further since the versions are sorted. + break + + return tokenizer_file + + +# To update the docstring, we need to copy the method, otherwise we change the original docstring. +PreTrainedTokenizerBase.push_to_hub = copy_func(PreTrainedTokenizerBase.push_to_hub) +if PreTrainedTokenizerBase.push_to_hub.__doc__ is not None: + PreTrainedTokenizerBase.push_to_hub.__doc__ = PreTrainedTokenizerBase.push_to_hub.__doc__.format( + object="tokenizer", object_class="AutoTokenizer", object_files="tokenizer files" + ) diff --git a/lib/python3.12/site-packages/transformers/tokenization_utils_fast.py b/lib/python3.12/site-packages/transformers/tokenization_utils_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..54605dbb02f63aabfc7efb1f957e10462c62e609 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/tokenization_utils_fast.py @@ -0,0 +1,908 @@ +# Copyright 2020 The HuggingFace Inc. team. +# +# 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. +""" +Tokenization classes for fast tokenizers (provided by HuggingFace's tokenizers library). For slow (python) tokenizers +see tokenization_utils.py +""" + +import copy +import json +import os +from collections import defaultdict +from collections.abc import Iterable +from typing import Any, Optional, Union + +import tokenizers.pre_tokenizers as pre_tokenizers_fast +from tokenizers import Encoding as EncodingFast +from tokenizers import Tokenizer as TokenizerFast +from tokenizers.decoders import Decoder as DecoderFast +from tokenizers.trainers import BpeTrainer, UnigramTrainer, WordLevelTrainer, WordPieceTrainer + +from .convert_slow_tokenizer import convert_slow_tokenizer +from .integrations.ggml import convert_gguf_tokenizer +from .modeling_gguf_pytorch_utils import load_gguf_checkpoint +from .tokenization_utils import PreTrainedTokenizer +from .tokenization_utils_base import ( + INIT_TOKENIZER_DOCSTRING, + AddedToken, + BatchEncoding, + PreTokenizedInput, + PreTokenizedInputPair, + PreTrainedTokenizerBase, + SpecialTokensMixin, + TextInput, + TextInputPair, + TruncationStrategy, +) +from .utils import PaddingStrategy, add_end_docstrings, logging + + +logger = logging.get_logger(__name__) + +# Fast tokenizers (provided by HuggingFace tokenizer's library) can be saved in a single file +TOKENIZER_FILE = "tokenizer.json" +SPECIAL_TOKENS_MAP_FILE = "special_tokens_map.json" +TOKENIZER_CONFIG_FILE = "tokenizer_config.json" +TIKTOKEN_VOCAB_FILE = "tokenizer.model" + +# Slow tokenizers have an additional added tokens files +ADDED_TOKENS_FILE = "added_tokens.json" + +INIT_TOKENIZER_DOCSTRING += """ + tokenizer_object ([`tokenizers.Tokenizer`]): + A [`tokenizers.Tokenizer`] object from 🤗 tokenizers to instantiate from. See [Using tokenizers from 🤗 + tokenizers](../fast_tokenizers) for more information. + tokenizer_file ([`str`]): + A path to a local JSON file representing a previously serialized [`tokenizers.Tokenizer`] object from 🤗 + tokenizers. +""" + +MODEL_TO_TRAINER_MAPPING = { + "BPE": BpeTrainer, + "Unigram": UnigramTrainer, + "WordLevel": WordLevelTrainer, + "WordPiece": WordPieceTrainer, +} + +VOCAB_FILES_NAMES = {"tokenizer_file": TOKENIZER_FILE, "vocab_file": TIKTOKEN_VOCAB_FILE} + + +@add_end_docstrings(INIT_TOKENIZER_DOCSTRING) +class PreTrainedTokenizerFast(PreTrainedTokenizerBase): + """ + Base class for all fast tokenizers (wrapping HuggingFace tokenizers library). + + Inherits from [`~tokenization_utils_base.PreTrainedTokenizerBase`]. + + Handles all the shared methods for tokenization and special tokens, as well as methods for + downloading/caching/loading pretrained tokenizers, as well as adding tokens to the vocabulary. + + This class also contains the added tokens in a unified way on top of all tokenizers so we don't have to handle the + specific vocabulary augmentation methods of the various underlying dictionary structures (BPE, sentencepiece...). + """ + + vocab_files_names = VOCAB_FILES_NAMES + slow_tokenizer_class: PreTrainedTokenizer = None + + def __init__(self, *args, **kwargs): + tokenizer_object = kwargs.pop("tokenizer_object", None) + slow_tokenizer = kwargs.pop("__slow_tokenizer", None) + gguf_file = kwargs.pop("gguf_file", None) + fast_tokenizer_file = kwargs.pop("tokenizer_file", None) + from_slow = kwargs.pop("from_slow", False) + added_tokens_decoder = kwargs.pop("added_tokens_decoder", {}) + self.add_prefix_space = kwargs.get("add_prefix_space", False) + + if from_slow and slow_tokenizer is None and self.slow_tokenizer_class is None: + raise ValueError( + "Cannot instantiate this tokenizer from a slow version. If it's based on sentencepiece, make sure you " + "have sentencepiece installed." + ) + + if tokenizer_object is not None: + fast_tokenizer = copy.deepcopy(tokenizer_object) + elif fast_tokenizer_file is not None and not from_slow: + # We have a serialization from tokenizers which let us directly build the backend + fast_tokenizer = TokenizerFast.from_file(fast_tokenizer_file) + elif slow_tokenizer: + # We need to convert a slow tokenizer to build the backend + fast_tokenizer = convert_slow_tokenizer(slow_tokenizer) + elif gguf_file is not None: + # We need to convert a slow tokenizer to build the backend + gguf_param = load_gguf_checkpoint(kwargs.get("vocab_file")) + architecture = gguf_param["config"]["model_type"] + tokenizer_dict = gguf_param["tokenizer"] + tokenizer_config = gguf_param["tokenizer_config"] + fast_tokenizer, additional_kwargs = convert_gguf_tokenizer(architecture, tokenizer_dict) + kwargs.update(tokenizer_config) + if len(additional_kwargs) > 0: + kwargs.update(additional_kwargs) + elif self.slow_tokenizer_class is not None and slow_tokenizer is not False: + # We need to create and convert a slow tokenizer to build the backend + slow_tokenizer = self.slow_tokenizer_class(*args, **kwargs) + fast_tokenizer = convert_slow_tokenizer(slow_tokenizer) + elif not slow_tokenizer: + # We tried loading a slow_tokenizer with spm and failed, try to load with tiktoken + self.vocab_file = kwargs.get("vocab_file", None) + self.additional_special_tokens = kwargs.get("additional_special_tokens", []) + fast_tokenizer = convert_slow_tokenizer(self, from_tiktoken=True) + slow_tokenizer = None + else: + raise ValueError( + "Couldn't instantiate the backend tokenizer from one of: \n" + "(1) a `tokenizers` library serialization file, \n" + "(2) a slow tokenizer instance to convert or \n" + "(3) an equivalent slow tokenizer class to instantiate and convert. \n" + "You need to have sentencepiece or tiktoken installed to convert a slow tokenizer to a fast one." + ) + + self._tokenizer = fast_tokenizer + + if slow_tokenizer is not None: + kwargs.update(slow_tokenizer.init_kwargs) + + self._decode_use_source_tokenizer = False + + _truncation = self._tokenizer.truncation + + if _truncation is not None: + self._tokenizer.enable_truncation(**_truncation) + kwargs.setdefault("max_length", _truncation["max_length"]) + kwargs.setdefault("truncation_side", _truncation["direction"]) + kwargs.setdefault("stride", _truncation["stride"]) + kwargs.setdefault("truncation_strategy", _truncation["strategy"]) + else: + self._tokenizer.no_truncation() + + _padding = self._tokenizer.padding + if _padding is not None: + self._tokenizer.enable_padding(**_padding) + kwargs.setdefault("pad_token", _padding["pad_token"]) + kwargs.setdefault("pad_token_type_id", _padding["pad_type_id"]) + kwargs.setdefault("padding_side", _padding["direction"]) + kwargs.setdefault("max_length", _padding["length"]) + kwargs.setdefault("pad_to_multiple_of", _padding["pad_to_multiple_of"]) + + # We call this after having initialized the backend tokenizer because we update it. + super().__init__(**kwargs) + self._tokenizer.encode_special_tokens = self.split_special_tokens + + added_tokens_decoder_hash = {hash(repr(token)) for token in self.added_tokens_decoder} + tokens_to_add = [ + token + for index, token in sorted(added_tokens_decoder.items(), key=lambda x: x[0]) + if hash(repr(token)) not in added_tokens_decoder_hash + ] + encoder = list(self.added_tokens_encoder.keys()) + [str(token) for token in tokens_to_add] + # if some of the special tokens are strings, we check if we don't already have a token + tokens_to_add += [ + token for token in self.all_special_tokens_extended if token not in encoder and token not in tokens_to_add + ] + + if len(tokens_to_add) > 0: + tokens = [] + special_tokens = self.all_special_tokens + for token in tokens_to_add: + is_special = ( + (token.special or str(token) in special_tokens) + if isinstance(token, AddedToken) + else str(token) in special_tokens + ) + if isinstance(token, str): + token = AddedToken(token, special=is_special) + else: + token.special = is_special + tokens.append(token) + if tokens: + self.add_tokens(tokens) + + try: + pre_tok_state = json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()) + if pre_tok_state.get("add_prefix_space", self.add_prefix_space) != self.add_prefix_space: + pre_tok_class = getattr(pre_tokenizers_fast, pre_tok_state.pop("type")) + pre_tok_state["add_prefix_space"] = self.add_prefix_space + self.backend_tokenizer.pre_tokenizer = pre_tok_class(**pre_tok_state) + except Exception: + # We'll get an error if there is no pre_tokenizer, or if it's a custom pre_tokenizer that can + # not be serialized. In those cases, we just ignore the error as there's no pre_tokenizer + # for which we need to update the `add_prefix_space` attribute. + pass + + @property + def is_fast(self) -> bool: + return True + + @property + def can_save_slow_tokenizer(self) -> bool: + """ + `bool`: Whether or not the slow tokenizer can be saved. Usually for sentencepiece based slow tokenizer, this + can only be `True` if the original `"sentencepiece.model"` was not deleted. + """ + return True + + @property + def vocab_size(self) -> int: + """ + `int`: Size of the base vocabulary (without the added tokens). + """ + return self._tokenizer.get_vocab_size(with_added_tokens=False) + + def get_vocab(self) -> dict[str, int]: + return self._tokenizer.get_vocab(with_added_tokens=True) + + @property + def vocab(self) -> dict[str, int]: + return self.get_vocab() + + @property + def added_tokens_encoder(self) -> dict[str, int]: + """ + Returns the sorted mapping from string to index. The added tokens encoder is cached for performance + optimisation in `self._added_tokens_encoder` for the slow tokenizers. + """ + return {k.content: v for v, k in sorted(self.added_tokens_decoder.items(), key=lambda item: item[0])} + + @property + def added_tokens_decoder(self) -> dict[int, AddedToken]: + """ + Returns the added tokens in the vocabulary as a dictionary of index to AddedToken. + + Returns: + `Dict[str, int]`: The added tokens. + """ + return self._tokenizer.get_added_tokens_decoder() + + def get_added_vocab(self) -> dict[str, int]: + """ + Returns the added tokens in the vocabulary as a dictionary of token to index. + + Returns: + `Dict[str, int]`: The added tokens. + """ + return {k.content: v for v, k in sorted(self.added_tokens_decoder.items(), key=lambda item: item[0])} + + def __len__(self) -> int: + """ + Size of the full vocabulary with the added tokens. + """ + return self._tokenizer.get_vocab_size(with_added_tokens=True) + + @property + def backend_tokenizer(self) -> TokenizerFast: + """ + `tokenizers.implementations.BaseTokenizer`: The Rust tokenizer used as a backend. + """ + return self._tokenizer + + @property + def decoder(self) -> DecoderFast: + """ + `tokenizers.decoders.Decoder`: The Rust decoder for this tokenizer. + """ + return self._tokenizer.decoder + + def _convert_encoding( + self, + encoding: EncodingFast, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + ) -> tuple[dict[str, Any], list[EncodingFast]]: + """ + Convert the encoding representation (from low-level HuggingFace tokenizer output) to a python Dict and a list + of encodings, take care of building a batch from overflowing tokens. + + Overflowing tokens are converted to additional examples (like batches) so the output values of the dict are + lists (overflows) of lists (tokens). + + Output shape: (overflows, sequence length) + """ + if return_token_type_ids is None: + return_token_type_ids = "token_type_ids" in self.model_input_names + if return_attention_mask is None: + return_attention_mask = "attention_mask" in self.model_input_names + + if return_overflowing_tokens and encoding.overflowing is not None: + encodings = [encoding] + encoding.overflowing + else: + encodings = [encoding] + + encoding_dict = defaultdict(list) + for e in encodings: + encoding_dict["input_ids"].append(e.ids) + + if return_token_type_ids: + encoding_dict["token_type_ids"].append(e.type_ids) + if return_attention_mask: + encoding_dict["attention_mask"].append(e.attention_mask) + if return_special_tokens_mask: + encoding_dict["special_tokens_mask"].append(e.special_tokens_mask) + if return_offsets_mapping: + encoding_dict["offset_mapping"].append(e.offsets) + if return_length: + encoding_dict["length"].append(len(e.ids)) + + return encoding_dict, encodings + + def convert_tokens_to_ids(self, tokens: Union[str, Iterable[str]]) -> Union[int, list[int]]: + """ + Converts a token string (or a sequence of tokens) in a single integer id (or a Iterable of ids), using the + vocabulary. + + Args: + tokens (`str` or `Iterable[str]`): One or several token(s) to convert to token id(s). + + Returns: + `int` or `List[int]`: The token id or list of token ids. + """ + if isinstance(tokens, str): + return self._convert_token_to_id_with_added_voc(tokens) + + return [self._convert_token_to_id_with_added_voc(token) for token in tokens] + + def _convert_token_to_id_with_added_voc(self, token: str) -> int: + index = self._tokenizer.token_to_id(token) + if index is None: + return self.unk_token_id + return index + + def _convert_id_to_token(self, index: int) -> Optional[str]: + return self._tokenizer.id_to_token(int(index)) + + def _add_tokens(self, new_tokens: list[Union[str, AddedToken]], special_tokens=False) -> int: + if special_tokens: + return self._tokenizer.add_special_tokens(new_tokens) + + return self._tokenizer.add_tokens(new_tokens) + + def num_special_tokens_to_add(self, pair: bool = False) -> int: + """ + Returns the number of added tokens when encoding a sequence with special tokens. + + + + This encodes a dummy input and checks the number of added tokens, and is therefore not efficient. Do not put + this inside your training loop. + + + + Args: + pair (`bool`, *optional*, defaults to `False`): + Whether the number of added tokens should be computed in the case of a sequence pair or a single + sequence. + + Returns: + `int`: Number of special tokens added to sequences. + """ + return self._tokenizer.num_special_tokens_to_add(pair) + + def convert_ids_to_tokens( + self, ids: Union[int, list[int]], skip_special_tokens: bool = False + ) -> Union[str, list[str]]: + """ + Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and + added tokens. + + Args: + ids (`int` or `List[int]`): + The token id (or token ids) to convert to tokens. + skip_special_tokens (`bool`, *optional*, defaults to `False`): + Whether or not to remove special tokens in the decoding. + + Returns: + `str` or `List[str]`: The decoded token(s). + """ + if isinstance(ids, int): + return self._tokenizer.id_to_token(ids) + tokens = [] + for index in ids: + index = int(index) + if skip_special_tokens and index in self.all_special_ids: + continue + tokens.append(self._tokenizer.id_to_token(index)) + return tokens + + def tokenize(self, text: str, pair: Optional[str] = None, add_special_tokens: bool = False, **kwargs) -> list[str]: + return self.encode_plus(text=text, text_pair=pair, add_special_tokens=add_special_tokens, **kwargs).tokens() + + def set_truncation_and_padding( + self, + padding_strategy: PaddingStrategy, + truncation_strategy: TruncationStrategy, + max_length: int, + stride: int, + pad_to_multiple_of: Optional[int], + padding_side: Optional[str], + ): + """ + Define the truncation and the padding strategies for fast tokenizers (provided by HuggingFace tokenizers + library) and restore the tokenizer settings afterwards. + + The provided tokenizer has no padding / truncation strategy before the managed section. If your tokenizer set a + padding / truncation strategy before, then it will be reset to no padding / truncation when exiting the managed + section. + + Args: + padding_strategy ([`~utils.PaddingStrategy`]): + The kind of padding that will be applied to the input + truncation_strategy ([`~tokenization_utils_base.TruncationStrategy`]): + The kind of truncation that will be applied to the input + max_length (`int`): + The maximum size of a sequence. + stride (`int`): + The stride to use when handling overflow. + pad_to_multiple_of (`int`, *optional*): + If set will pad the sequence to a multiple of the provided value. This is especially useful to enable + the use of Tensor Cores on NVIDIA hardware with compute capability `>= 7.5` (Volta). + padding_side (`str`, *optional*): + The side on which the model should have padding applied. Should be selected between ['right', 'left']. + Default value is picked from the class attribute of the same name. + """ + _truncation = self._tokenizer.truncation + _padding = self._tokenizer.padding + # Set truncation and padding on the backend tokenizer + if truncation_strategy == TruncationStrategy.DO_NOT_TRUNCATE: + if _truncation is not None: + self._tokenizer.no_truncation() + else: + target = { + "max_length": max_length, + "stride": stride, + "strategy": truncation_strategy.value, + "direction": self.truncation_side, + } + + # _truncation might contain more keys that the target `transformers` + # supports. Use only the target keys to trigger `enable_truncation`. + # This should enable this code to works on various `tokenizers` + # targets. + if _truncation is None: + current = None + else: + current = {k: _truncation.get(k, None) for k in target} + + if current != target: + self._tokenizer.enable_truncation(**target) + + if padding_strategy == PaddingStrategy.DO_NOT_PAD: + if _padding is not None: + self._tokenizer.no_padding() + else: + length = max_length if padding_strategy == PaddingStrategy.MAX_LENGTH else None + target = { + "length": length, + "direction": padding_side if padding_side is not None else self.padding_side, + "pad_id": self.pad_token_id, + "pad_token": self.pad_token, + "pad_type_id": self.pad_token_type_id, + "pad_to_multiple_of": pad_to_multiple_of, + } + if _padding != target: + self._tokenizer.enable_padding(**target) + + def _batch_encode_plus( + self, + batch_text_or_text_pairs: Union[ + list[TextInput], list[TextInputPair], list[PreTokenizedInput], list[PreTokenizedInputPair] + ], + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[str] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + ) -> BatchEncoding: + if not isinstance(batch_text_or_text_pairs, (tuple, list)): + raise TypeError( + f"batch_text_or_text_pairs has to be a list or a tuple (got {type(batch_text_or_text_pairs)})" + ) + + # Set the truncation and padding strategy and restore the initial configuration + self.set_truncation_and_padding( + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + ) + + if self._tokenizer.encode_special_tokens != split_special_tokens: + self._tokenizer.encode_special_tokens = split_special_tokens + + encodings = self._tokenizer.encode_batch( + batch_text_or_text_pairs, + add_special_tokens=add_special_tokens, + is_pretokenized=is_split_into_words, + ) + + # Convert encoding to dict + # `Tokens` has type: Tuple[ + # List[Dict[str, List[List[int]]]] or List[Dict[str, 2D-Tensor]], + # List[EncodingFast] + # ] + # with nested dimensions corresponding to batch, overflows, sequence length + tokens_and_encodings = [ + self._convert_encoding( + encoding=encoding, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + ) + for encoding in encodings + ] + + # Convert the output to have dict[list] from list[dict] and remove the additional overflows dimension + # From (variable) shape (batch, overflows, sequence length) to ~ (batch * overflows, sequence length) + # (we say ~ because the number of overflow varies with the example in the batch) + # + # To match each overflowing sample with the original sample in the batch + # we add an overflow_to_sample_mapping array (see below) + sanitized_tokens = {} + for key in tokens_and_encodings[0][0].keys(): + stack = [e for item, _ in tokens_and_encodings for e in item[key]] + sanitized_tokens[key] = stack + sanitized_encodings = [e for _, item in tokens_and_encodings for e in item] + + # If returning overflowing tokens, we need to return a mapping + # from the batch idx to the original sample + if return_overflowing_tokens: + overflow_to_sample_mapping = [] + for i, (toks, _) in enumerate(tokens_and_encodings): + overflow_to_sample_mapping += [i] * len(toks["input_ids"]) + sanitized_tokens["overflow_to_sample_mapping"] = overflow_to_sample_mapping + + for input_ids in sanitized_tokens["input_ids"]: + self._eventual_warn_about_too_long_sequence(input_ids, max_length, verbose) + return BatchEncoding(sanitized_tokens, sanitized_encodings, tensor_type=return_tensors) + + def _encode_plus( + self, + text: Union[TextInput, PreTokenizedInput], + text_pair: Optional[Union[TextInput, PreTokenizedInput]] = None, + add_special_tokens: bool = True, + padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, + truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE, + max_length: Optional[int] = None, + stride: int = 0, + is_split_into_words: bool = False, + pad_to_multiple_of: Optional[int] = None, + padding_side: Optional[str] = None, + return_tensors: Optional[bool] = None, + return_token_type_ids: Optional[bool] = None, + return_attention_mask: Optional[bool] = None, + return_overflowing_tokens: bool = False, + return_special_tokens_mask: bool = False, + return_offsets_mapping: bool = False, + return_length: bool = False, + verbose: bool = True, + split_special_tokens: bool = False, + **kwargs, + ) -> BatchEncoding: + batched_input = [(text, text_pair)] if text_pair else [text] + batched_output = self._batch_encode_plus( + batched_input, + is_split_into_words=is_split_into_words, + add_special_tokens=add_special_tokens, + padding_strategy=padding_strategy, + truncation_strategy=truncation_strategy, + max_length=max_length, + stride=stride, + pad_to_multiple_of=pad_to_multiple_of, + padding_side=padding_side, + return_tensors=return_tensors, + return_token_type_ids=return_token_type_ids, + return_attention_mask=return_attention_mask, + return_overflowing_tokens=return_overflowing_tokens, + return_special_tokens_mask=return_special_tokens_mask, + return_offsets_mapping=return_offsets_mapping, + return_length=return_length, + verbose=verbose, + split_special_tokens=split_special_tokens, + **kwargs, + ) + + # Return tensor is None, then we can remove the leading batch axis + # Overflowing tokens are returned as a batch of output so we keep them in this case + if return_tensors is None and not return_overflowing_tokens: + batched_output = BatchEncoding( + { + key: (value[0] if len(value) > 0 and isinstance(value[0], list) else value) + for key, value in batched_output.items() + }, + batched_output.encodings, + ) + + self._eventual_warn_about_too_long_sequence(batched_output["input_ids"], max_length, verbose) + + return batched_output + + def convert_tokens_to_string(self, tokens: list[str]) -> str: + return ( + self.backend_tokenizer.decoder.decode(tokens) + if self.backend_tokenizer.decoder is not None + else " ".join(tokens) + ) + + def _decode( + self, + token_ids: Union[int, list[int]], + skip_special_tokens: bool = False, + clean_up_tokenization_spaces: Optional[bool] = None, + **kwargs, + ) -> str: + self._decode_use_source_tokenizer = kwargs.pop("use_source_tokenizer", False) + + if isinstance(token_ids, int): + token_ids = [token_ids] + text = self._tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens) + + clean_up_tokenization_spaces = ( + clean_up_tokenization_spaces + if clean_up_tokenization_spaces is not None + else self.clean_up_tokenization_spaces + ) + if clean_up_tokenization_spaces: + clean_text = self.clean_up_tokenization(text) + return clean_text + else: + return text + + def _save_pretrained( + self, + save_directory: Union[str, os.PathLike], + file_names: tuple[str], + legacy_format: Optional[bool] = None, + filename_prefix: Optional[str] = None, + ) -> tuple[str]: + """ + Save a tokenizer using the slow-tokenizer/legacy format: vocabulary + added tokens as well as in a unique JSON + file containing {config + vocab + added-tokens}. + """ + save_directory = str(save_directory) + + if self.slow_tokenizer_class is None and legacy_format is True: + raise ValueError( + "Your tokenizer does not have a legacy version defined and therefore cannot register this version. You" + " might consider leaving the legacy_format at `None` or setting it to `False`." + ) + + save_slow = ( + (legacy_format is None or legacy_format is True) + and self.slow_tokenizer_class is not None + and self.can_save_slow_tokenizer + ) + save_fast = legacy_format is None or legacy_format is False + + if save_slow: + added_tokens_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + ADDED_TOKENS_FILE + ) + # make sure to be forward compatible + added_vocab = {tok: index for tok, index in self.added_tokens_encoder.items() if index >= self.vocab_size} + if added_vocab: + with open(added_tokens_file, "w", encoding="utf-8") as f: + out_str = json.dumps(added_vocab, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + f.write(out_str) + + vocab_files = self.save_vocabulary(save_directory, filename_prefix=filename_prefix) + file_names = file_names + vocab_files + (added_tokens_file,) + + if save_fast: + tokenizer_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + TOKENIZER_FILE + ) + self.backend_tokenizer.save(tokenizer_file) + file_names = file_names + (tokenizer_file,) + + return file_names + + def train_new_from_iterator( + self, + text_iterator, + vocab_size, + length=None, + new_special_tokens=None, + special_tokens_map=None, + **kwargs, + ): + """ + Trains a tokenizer on a new corpus with the same defaults (in terms of special tokens or tokenization pipeline) + as the current one. + + Args: + text_iterator (generator of `List[str]`): + The training corpus. Should be a generator of batches of texts, for instance a list of lists of texts + if you have everything in memory. + vocab_size (`int`): + The size of the vocabulary you want for your tokenizer. + length (`int`, *optional*): + The total number of sequences in the iterator. This is used to provide meaningful progress tracking + new_special_tokens (list of `str` or `AddedToken`, *optional*): + A list of new special tokens to add to the tokenizer you are training. + special_tokens_map (`Dict[str, str]`, *optional*): + If you want to rename some of the special tokens this tokenizer uses, pass along a mapping old special + token name to new special token name in this argument. + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments passed along to the trainer from the 🤗 Tokenizers library. + + Returns: + [`PreTrainedTokenizerFast`]: A new tokenizer of the same type as the original one, trained on + `text_iterator`. + + """ + tokenizer_json = json.loads(self._tokenizer.to_str()) + # Remove added tokens for now (uses IDs of tokens) + added_tokens = tokenizer_json.pop("added_tokens") + # Remove post processor for now (uses IDs of tokens) + post_processor = tokenizer_json.pop("post_processor") + + unk_token = None + # Remove vocab + if tokenizer_json["model"]["type"] == "BPE": + tokenizer_json["model"]["vocab"] = {} + tokenizer_json["model"]["merges"] = [] + elif tokenizer_json["model"]["type"] == "Unigram": + if tokenizer_json["model"]["unk_id"] is not None: + unk_id = tokenizer_json["model"]["unk_id"] + unk_token = tokenizer_json["model"]["vocab"][unk_id][0] + if special_tokens_map is not None and unk_token in special_tokens_map: + unk_token = special_tokens_map[unk_token] + tokenizer_json["model"]["unk_id"] = 0 + tokenizer_json["model"]["vocab"] = [[unk_token, 0.0]] + elif tokenizer_json["model"]["type"] in ["WordLevel", "WordPiece"]: + tokenizer_json["model"]["vocab"] = {} + else: + raise ValueError( + f"This method does not support this type of tokenizer (found {tokenizer_json['model']['type']}) " + "only BPE, Unigram, WordLevel and WordPiece." + ) + + if ( + special_tokens_map is not None + and "unk_token" in tokenizer_json["model"] + and tokenizer_json["model"]["unk_token"] in special_tokens_map + ): + tokenizer_json["model"]["unk_token"] = special_tokens_map[tokenizer_json["model"]["unk_token"]] + + tokenizer = TokenizerFast.from_str(json.dumps(tokenizer_json)) + + # Get the special tokens from the current tokenizer if none are specified. + special_tokens = [] + for added_token in added_tokens: + special = added_token.pop("special", None) + _ = added_token.pop("id", None) + if tokenizer_json["model"]["type"] != "Unigram" and not special: + continue + if special_tokens_map is not None and added_token["content"] in special_tokens_map: + added_token["content"] = special_tokens_map[added_token["content"]] + special_tokens.append(AddedToken(**added_token)) + + if new_special_tokens is not None: + special_tokens.extend(new_special_tokens) + + # Trainer needs to know the end of word / continuing subword thingies in BPE + if ( + tokenizer_json["model"]["type"] == "BPE" + and "continuing_subword_prefix" not in kwargs + and tokenizer_json["model"]["continuing_subword_prefix"] is not None + ): + kwargs["continuing_subword_prefix"] = tokenizer_json["model"]["continuing_subword_prefix"] + if ( + tokenizer_json["model"]["type"] == "BPE" + and "end_of_word_suffix" not in kwargs + and tokenizer_json["model"]["end_of_word_suffix"] is not None + ): + kwargs["end_of_word_suffix"] = tokenizer_json["model"]["end_of_word_suffix"] + if tokenizer_json["model"]["type"] == "Unigram" and unk_token is not None: + kwargs["unk_token"] = unk_token + if tokenizer_json["pre_tokenizer"] is not None: + if ( + tokenizer_json["pre_tokenizer"]["type"] == "ByteLevel" + or tokenizer_json["pre_tokenizer"]["type"] == "Sequence" + and "pretokenizers" in tokenizer_json["pre_tokenizer"] + and any( + pretokenizer["type"] == "ByteLevel" + for pretokenizer in tokenizer_json["pre_tokenizer"]["pretokenizers"] + ) + ): + kwargs["initial_alphabet"] = pre_tokenizers_fast.ByteLevel.alphabet() + + trainer_class = MODEL_TO_TRAINER_MAPPING[tokenizer_json["model"]["type"]] + trainer = trainer_class(vocab_size=vocab_size, special_tokens=special_tokens, **kwargs) + tokenizer.train_from_iterator(text_iterator, length=length, trainer=trainer) + + if post_processor is not None: + trained_tokenizer_json = json.loads(tokenizer.to_str()) + # Almost done, we just have to adjust the token IDs in the post processor + if "special_tokens" in post_processor: + for key in post_processor["special_tokens"]: + tokens = post_processor["special_tokens"][key]["tokens"] + if special_tokens_map is not None: + tokens = [special_tokens_map.get(token, token) for token in tokens] + post_processor["special_tokens"][key]["tokens"] = tokens + for token in tokens: + token_id = tokenizer.token_to_id(token) + if token_id is None: + raise ValueError( + "Attempted to set a token in the post processor that does not exist in the mapping" + ) + + post_processor["special_tokens"][key]["ids"] = [tokenizer.token_to_id(token) for token in tokens] + + for special_token in ["cls", "sep"]: + if special_token in post_processor: + token, _ = post_processor[special_token] + if special_tokens_map is not None and token in special_tokens_map: + token = special_tokens_map[token] + token_id = tokenizer.token_to_id(token) + if token_id is None: + raise ValueError( + "Attempted to set a token in the post processor that does not exist in the mapping" + ) + post_processor[special_token] = [token, token_id] + + trained_tokenizer_json["post_processor"] = post_processor + tokenizer = TokenizerFast.from_str(json.dumps(trained_tokenizer_json)) + + kwargs = self.init_kwargs.copy() + # Map pad/cls/mask token at the Transformers level + special_tokens_list = SpecialTokensMixin.SPECIAL_TOKENS_ATTRIBUTES.copy() + special_tokens_list.remove("additional_special_tokens") + for token in special_tokens_list: + if getattr(self, token) is not None: + special_token = getattr(self, token) + if special_tokens_map is not None and special_token in special_tokens_map: + special_token = special_tokens_map[special_token] + + special_token_full = self._special_tokens_map.get(token, None) + if isinstance(special_token_full, AddedToken): + # Create an added token with the same parameters except the content + kwargs[token] = AddedToken( + special_token, + single_word=special_token_full.single_word, + lstrip=special_token_full.lstrip, + rstrip=special_token_full.rstrip, + normalized=special_token_full.normalized, + special=True, + ) + else: + kwargs[token] = special_token + + additional_special_tokens = self.additional_special_tokens + if new_special_tokens is not None: + additional_special_tokens.extend(new_special_tokens) + if len(additional_special_tokens) > 0: + kwargs["additional_special_tokens"] = additional_special_tokens + + return self.__class__(tokenizer_object=tokenizer, **kwargs) diff --git a/lib/python3.12/site-packages/transformers/trainer.py b/lib/python3.12/site-packages/transformers/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..727c2ca7fefffa8c5255bc3999bf145e02a557d3 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/trainer.py @@ -0,0 +1,5340 @@ +# Copyright 2020-present the HuggingFace Inc. team. +# +# 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. +""" +The Trainer class, to easily train a 🤗 Transformers from scratch or finetune it on a new task. +""" + +import contextlib +import copy +import functools +import glob +import importlib.metadata +import inspect +import json +import math +import os +import random +import re +import shutil +import sys +import tempfile +import time +import warnings +from collections.abc import Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Optional, Union + + +# Integrations must be imported before ML frameworks: +# isort: off +from .integrations import ( + get_reporting_integration_callbacks, +) + +# isort: on + +import huggingface_hub.utils as hf_hub_utils +import numpy as np +import torch +import torch.distributed as dist +from huggingface_hub import ModelCard, create_repo, upload_folder +from packaging import version +from torch import nn +from torch.utils.data import DataLoader, Dataset, IterableDataset, RandomSampler, SequentialSampler + +from . import __version__ +from .configuration_utils import PretrainedConfig +from .data.data_collator import DataCollator, DataCollatorWithPadding, default_data_collator +from .debug_utils import DebugOption, DebugUnderflowOverflow +from .feature_extraction_sequence_utils import SequenceFeatureExtractor +from .feature_extraction_utils import FeatureExtractionMixin +from .hyperparameter_search import ALL_HYPERPARAMETER_SEARCH_BACKENDS, default_hp_search_backend +from .image_processing_utils import BaseImageProcessor +from .integrations.deepspeed import deepspeed_init, deepspeed_load_checkpoint, is_deepspeed_available +from .integrations.tpu import tpu_spmd_dataloader +from .modelcard import TrainingSummary +from .modeling_utils import PreTrainedModel, load_sharded_checkpoint, unwrap_model +from .models.auto.modeling_auto import ( + MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, + MODEL_MAPPING_NAMES, +) +from .optimization import Adafactor, get_scheduler +from .processing_utils import ProcessorMixin +from .pytorch_utils import ( + ALL_LAYERNORM_LAYERS, + is_torch_greater_or_equal_than_2_3, +) +from .tokenization_utils_base import PreTrainedTokenizerBase +from .trainer_callback import ( + CallbackHandler, + DefaultFlowCallback, + ExportableState, + PrinterCallback, + ProgressCallback, + TrainerCallback, + TrainerControl, + TrainerState, +) +from .trainer_pt_utils import ( + DistributedTensorGatherer, + EvalLoopContainer, + IterableDatasetShard, + LabelSmoother, + LayerWiseDummyOptimizer, + LengthGroupedSampler, + SequentialDistributedSampler, + distributed_broadcast_scalars, + distributed_concat, + find_batch_size, + get_model_param_count, + get_module_class_from_name, + get_parameter_names, + nested_concat, + nested_detach, + nested_numpify, + nested_xla_mesh_reduce, + reissue_pt_warnings, + remove_dummy_checkpoint, + set_rng_state_for_device, +) +from .trainer_utils import ( + PREFIX_CHECKPOINT_DIR, + BestRun, + EvalLoopOutput, + EvalPrediction, + HPSearchBackend, + HubStrategy, + PredictionOutput, + RemoveColumnsCollator, + SaveStrategy, + TrainerMemoryTracker, + TrainOutput, + check_target_module_exists, + default_compute_objective, + denumpify_detensorize, + enable_full_determinism, + find_executable_batch_size, + get_last_checkpoint, + has_length, + neftune_post_forward_hook, + number_of_arguments, + seed_worker, + set_seed, + speed_metrics, +) +from .training_args import OptimizerNames, ParallelMode, TrainingArguments +from .utils import ( + ADAPTER_CONFIG_NAME, + ADAPTER_SAFE_WEIGHTS_NAME, + ADAPTER_WEIGHTS_NAME, + CONFIG_NAME, + SAFE_WEIGHTS_INDEX_NAME, + SAFE_WEIGHTS_NAME, + WEIGHTS_INDEX_NAME, + WEIGHTS_NAME, + XLA_FSDPV2_MIN_VERSION, + PushInProgress, + PushToHubMixin, + can_return_loss, + find_labels, + is_accelerate_available, + is_apex_available, + is_apollo_torch_available, + is_bitsandbytes_available, + is_datasets_available, + is_galore_torch_available, + is_grokadamw_available, + is_in_notebook, + is_ipex_available, + is_liger_kernel_available, + is_lomo_available, + is_peft_available, + is_safetensors_available, + is_sagemaker_dp_enabled, + is_sagemaker_mp_enabled, + is_schedulefree_available, + is_torch_compile_available, + is_torch_hpu_available, + is_torch_mlu_available, + is_torch_mps_available, + is_torch_musa_available, + is_torch_neuroncore_available, + is_torch_npu_available, + is_torch_xla_available, + is_torch_xpu_available, + is_torchao_available, + logging, + strtobool, +) +from .utils.deprecation import deprecate_kwarg +from .utils.quantization_config import QuantizationMethod + + +DEFAULT_CALLBACKS = [DefaultFlowCallback] +DEFAULT_PROGRESS_CALLBACK = ProgressCallback + +if is_in_notebook(): + from .utils.notebook import NotebookProgressCallback + + DEFAULT_PROGRESS_CALLBACK = NotebookProgressCallback + +if is_apex_available(): + from apex import amp + +if is_datasets_available(): + import datasets + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + import torch_xla.debug.metrics as met + from torch_xla import __version__ as XLA_VERSION + + IS_XLA_FSDPV2_POST_2_2 = version.parse(XLA_VERSION) >= version.parse(XLA_FSDPV2_MIN_VERSION) + if IS_XLA_FSDPV2_POST_2_2: + import torch_xla.distributed.spmd as xs + import torch_xla.runtime as xr +else: + IS_XLA_FSDPV2_POST_2_2 = False + + +if is_sagemaker_mp_enabled(): + import smdistributed.modelparallel.torch as smp + from smdistributed.modelparallel import __version__ as SMP_VERSION + + IS_SAGEMAKER_MP_POST_1_10 = version.parse(SMP_VERSION) >= version.parse("1.10") + + from .trainer_pt_utils import smp_forward_backward, smp_forward_only, smp_gather, smp_nested_concat +else: + IS_SAGEMAKER_MP_POST_1_10 = False + + +if is_safetensors_available(): + import safetensors.torch + +if is_peft_available(): + from peft import PeftModel + + +if is_accelerate_available(): + from accelerate import Accelerator, skip_first_batches + from accelerate import __version__ as accelerate_version + from accelerate.state import AcceleratorState + from accelerate.utils import ( + AutocastKwargs, + DistributedDataParallelKwargs, + DistributedType, + load_fsdp_model, + load_fsdp_optimizer, + save_fsdp_model, + save_fsdp_optimizer, + ) + + DATA_SAMPLERS = [RandomSampler] + if version.parse(accelerate_version) > version.parse("1.3.0"): + from accelerate.utils import TorchTensorParallelPlugin + if version.parse(accelerate_version) > version.parse("0.23.0"): + from accelerate.data_loader import SeedableRandomSampler + + DATA_SAMPLERS += [SeedableRandomSampler] + + if is_deepspeed_available(): + from accelerate.utils import DeepSpeedSchedulerWrapper + +if is_accelerate_available("0.28.0"): + from accelerate.utils import DataLoaderConfiguration + + +def _is_peft_model(model): + if is_peft_available(): + classes_to_check = (PeftModel,) if is_peft_available() else () + # Here we also check if the model is an instance of `PeftMixedModel` introduced in peft>=0.7.0: https://github.com/huggingface/transformers/pull/28321 + if version.parse(importlib.metadata.version("peft")) >= version.parse("0.7.0"): + from peft import PeftMixedModel + + classes_to_check = (*classes_to_check, PeftMixedModel) + return isinstance(model, classes_to_check) + return False + + +def _get_fsdp_ckpt_kwargs(): + # TODO: @AjayP13, @younesbelkada replace this check with version check at the next `accelerate` release + if is_accelerate_available() and "adapter_only" in list(inspect.signature(save_fsdp_model).parameters): + return {"adapter_only": True} + else: + return {} + + +def safe_globals(): + # Starting from version 2.4 PyTorch introduces a check for the objects loaded + # with torch.load(weights_only=True). Starting from 2.6 weights_only=True becomes + # a default and requires allowlisting of objects being loaded. + # See: https://github.com/pytorch/pytorch/pull/137602 + # See: https://pytorch.org/docs/stable/notes/serialization.html#torch.serialization.add_safe_globals + # See: https://github.com/huggingface/accelerate/pull/3036 + if version.parse(torch.__version__).release < version.parse("2.6").release: + return contextlib.nullcontext() + + np_core = np._core if version.parse(np.__version__) >= version.parse("2.0.0") else np.core + allowlist = [np_core.multiarray._reconstruct, np.ndarray, np.dtype] + # numpy >1.25 defines numpy.dtypes.UInt32DType, but below works for + # all versions of numpy + allowlist += [type(np.dtype(np.uint32))] + + return torch.serialization.safe_globals(allowlist) + + +if TYPE_CHECKING: + import optuna + + if is_datasets_available(): + import datasets + +logger = logging.get_logger(__name__) + + +# Name of the files used for checkpointing +TRAINING_ARGS_NAME = "training_args.bin" +TRAINER_STATE_NAME = "trainer_state.json" +OPTIMIZER_NAME = "optimizer.pt" +SCALER_NAME = "scaler.pt" +OPTIMIZER_NAME_BIN = "optimizer.bin" +SCHEDULER_NAME = "scheduler.pt" +FSDP_MODEL_NAME = "pytorch_model_fsdp" + + +class Trainer: + """ + Trainer is a simple but feature-complete training and eval loop for PyTorch, optimized for 🤗 Transformers. + + Args: + model ([`PreTrainedModel`] or `torch.nn.Module`, *optional*): + The model to train, evaluate or use for predictions. If not provided, a `model_init` must be passed. + + + + [`Trainer`] is optimized to work with the [`PreTrainedModel`] provided by the library. You can still use + your own models defined as `torch.nn.Module` as long as they work the same way as the 🤗 Transformers + models. + + + + args ([`TrainingArguments`], *optional*): + The arguments to tweak for training. Will default to a basic instance of [`TrainingArguments`] with the + `output_dir` set to a directory named *tmp_trainer* in the current directory if not provided. + data_collator (`DataCollator`, *optional*): + The function to use to form a batch from a list of elements of `train_dataset` or `eval_dataset`. Will + default to [`default_data_collator`] if no `processing_class` is provided, an instance of + [`DataCollatorWithPadding`] otherwise if the processing_class is a feature extractor or tokenizer. + train_dataset (Union[`torch.utils.data.Dataset`, `torch.utils.data.IterableDataset`, `datasets.Dataset`], *optional*): + The dataset to use for training. If it is a [`~datasets.Dataset`], columns not accepted by the + `model.forward()` method are automatically removed. + + Note that if it's a `torch.utils.data.IterableDataset` with some randomization and you are training in a + distributed fashion, your iterable dataset should either use a internal attribute `generator` that is a + `torch.Generator` for the randomization that must be identical on all processes (and the Trainer will + manually set the seed of this `generator` at each epoch) or have a `set_epoch()` method that internally + sets the seed of the RNGs used. + eval_dataset (Union[`torch.utils.data.Dataset`, Dict[str, `torch.utils.data.Dataset`, `datasets.Dataset`]), *optional*): + The dataset to use for evaluation. If it is a [`~datasets.Dataset`], columns not accepted by the + `model.forward()` method are automatically removed. If it is a dictionary, it will evaluate on each + dataset prepending the dictionary key to the metric name. + processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + This supercedes the `tokenizer` argument, which is now deprecated. + model_init (`Callable[[], PreTrainedModel]`, *optional*): + A function that instantiates the model to be used. If provided, each call to [`~Trainer.train`] will start + from a new instance of the model as given by this function. + + The function may have zero argument, or a single one containing the optuna/Ray Tune/SigOpt trial object, to + be able to choose different architectures according to hyper parameters (such as layer count, sizes of + inner layers, dropout probabilities etc). + compute_loss_func (`Callable`, *optional*): + A function that accepts the raw model outputs, labels, and the number of items in the entire accumulated + batch (batch_size * gradient_accumulation_steps) and returns the loss. For example, see the default [loss function](https://github.com/huggingface/transformers/blob/052e652d6d53c2b26ffde87e039b723949a53493/src/transformers/trainer.py#L3618) used by [`Trainer`]. + compute_metrics (`Callable[[EvalPrediction], Dict]`, *optional*): + The function that will be used to compute metrics at evaluation. Must take a [`EvalPrediction`] and return + a dictionary string to metric values. *Note* When passing TrainingArgs with `batch_eval_metrics` set to + `True`, your compute_metrics function must take a boolean `compute_result` argument. This will be triggered + after the last eval batch to signal that the function needs to calculate and return the global summary + statistics rather than accumulating the batch-level statistics + callbacks (List of [`TrainerCallback`], *optional*): + A list of callbacks to customize the training loop. Will add those to the list of default callbacks + detailed in [here](callback). + + If you want to remove one of the default callbacks used, use the [`Trainer.remove_callback`] method. + optimizers (`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + optimizer_cls_and_kwargs (`Tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*): + A tuple containing the optimizer class and keyword arguments to use. + Overrides `optim` and `optim_args` in `args`. Incompatible with the `optimizers` argument. + + Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before initializing the Trainer. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*): + A function that preprocess the logits right before caching them at each evaluation step. Must take two + tensors, the logits and the labels, and return the logits once processed as desired. The modifications made + by this function will be reflected in the predictions received by `compute_metrics`. + + Note that the labels (second parameter) will be `None` if the dataset does not have them. + + Important attributes: + + - **model** -- Always points to the core model. If using a transformers model, it will be a [`PreTrainedModel`] + subclass. + - **model_wrapped** -- Always points to the most external model in case one or more other modules wrap the + original model. This is the model that should be used for the forward pass. For example, under `DeepSpeed`, + the inner model is wrapped in `DeepSpeed` and then again in `torch.nn.DistributedDataParallel`. If the inner + model hasn't been wrapped, then `self.model_wrapped` is the same as `self.model`. + - **is_model_parallel** -- Whether or not a model has been switched to a model parallel mode (different from + data parallelism, this means some of the model layers are split on different GPUs). + - **place_model_on_device** -- Whether or not to automatically place the model on the device - it will be set + to `False` if model parallel or deepspeed is used, or if the default + `TrainingArguments.place_model_on_device` is overridden to return `False` . + - **is_in_train** -- Whether or not a model is currently running `train` (e.g. when `evaluate` is called while + in `train`) + + """ + + # Those are used as methods of the Trainer in examples. + from .trainer_pt_utils import _get_learning_rate, log_metrics, metrics_format, save_metrics, save_state + + @deprecate_kwarg("tokenizer", new_name="processing_class", version="5.0.0", raise_if_both_names=True) + def __init__( + self, + model: Union[PreTrainedModel, nn.Module, None] = None, + args: TrainingArguments = None, + data_collator: Optional[DataCollator] = None, + train_dataset: Optional[Union[Dataset, IterableDataset, "datasets.Dataset"]] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset], "datasets.Dataset"]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + model_init: Optional[Callable[[], PreTrainedModel]] = None, + compute_loss_func: Optional[Callable] = None, + compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None), + optimizer_cls_and_kwargs: Optional[tuple[type[torch.optim.Optimizer], dict[str, Any]]] = None, + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + ): + if args is None: + output_dir = "tmp_trainer" + logger.info(f"No `TrainingArguments` passed, using `output_dir={output_dir}`.") + args = TrainingArguments(output_dir=output_dir) + if args.batch_eval_metrics and compute_metrics is not None: + if "compute_result" not in inspect.signature(compute_metrics).parameters.keys(): + raise ValueError( + "When using `batch_eval_metrics`, your `compute_metrics` function must take a `compute_result`" + " boolean argument which will be triggered after the last batch of the eval set to signal that the" + " summary statistics should be returned by the function." + ) + if args.eval_strategy is not None and args.eval_strategy != "no" and eval_dataset is None: + raise ValueError( + f"You have set `args.eval_strategy` to {args.eval_strategy} but you didn't pass an `eval_dataset` to `Trainer`. Either set `args.eval_strategy` to `no` or pass an `eval_dataset`. " + ) + if args.save_strategy == SaveStrategy.BEST or args.load_best_model_at_end: + if args.metric_for_best_model is None: + raise ValueError( + "`args.metric_for_best_model` must be provided when using 'best' save_strategy or if `args.load_best_model_at_end` is set to `True`." + ) + + self.args = args + self.compute_loss_func = compute_loss_func + # Seed must be set before instantiating the model when using model + enable_full_determinism(self.args.seed) if self.args.full_determinism else set_seed(self.args.seed) + + self.hp_name = None + self.deepspeed = None + self.is_in_train = False + + self.create_accelerator_and_postprocess() + + # memory metrics - must set up as early as possible + self._memory_tracker = TrainerMemoryTracker(self.args.skip_memory_metrics) + self._memory_tracker.start() + + # set the correct log level depending on the node + log_level = args.get_process_log_level() + logging.set_verbosity(log_level) + + # force device and distributed setup init explicitly + args._setup_devices + + if model is None: + if model_init is not None: + self.model_init = model_init + model = self.call_model_init() + else: + raise RuntimeError("`Trainer` requires either a `model` or `model_init` argument") + else: + if model_init is not None: + warnings.warn( + "`Trainer` requires either a `model` or `model_init` argument, but not both. `model_init` will" + " overwrite your model when calling the `train` method. This will become a fatal error in the next" + " release.", + FutureWarning, + ) + self.model_init = model_init + + if model.__class__.__name__ in MODEL_MAPPING_NAMES: + raise ValueError( + f"The model you have picked ({model.__class__.__name__}) cannot be used as is for training: it only " + "computes hidden states and does not accept any labels. You should choose a model with a head " + "suitable for your task like any of the `AutoModelForXxx` listed at " + "https://huggingface.co/docs/transformers/model_doc/auto" + ) + + if getattr(model, "is_parallelizable", False) and getattr(model, "model_parallel", False): + self.is_model_parallel = True + else: + self.is_model_parallel = False + + if getattr(model, "hf_device_map", None) is not None: + devices = [device for device in set(model.hf_device_map.values()) if device not in ["cpu", "disk"]] + if len(devices) > 1: + self.is_model_parallel = True + elif len(devices) == 1: + self.is_model_parallel = self.args.device != torch.device(devices[0]) + else: + self.is_model_parallel = False + + # warn users + if self.is_model_parallel: + logger.info( + "You have loaded a model on multiple GPUs. `is_model_parallel` attribute will be force-set" + " to `True` to avoid any unexpected behavior such as device placement mismatching." + ) + + if self.args.use_liger_kernel: + if is_liger_kernel_available(): + from liger_kernel.transformers import _apply_liger_kernel_to_instance + + if isinstance(model, PreTrainedModel): + # Patch the model with liger kernels. Use the default kernel configurations. + _apply_liger_kernel_to_instance(model=model) + elif hasattr(model, "get_base_model") and isinstance(model.get_base_model(), PreTrainedModel): + # Patch the base model with liger kernels where model is a PeftModel. Use the default kernel configurations. + _apply_liger_kernel_to_instance(model=model.get_base_model()) + else: + logger.warning( + "The model is not an instance of PreTrainedModel. No liger kernels will be applied." + ) + else: + raise ImportError( + "You have set `use_liger_kernel` to `True` but liger-kernel >= 0.3.0 is not available. " + "Please install it with `pip install liger-kernel`" + ) + + _is_quantized_and_base_model = getattr(model, "is_quantized", False) and not getattr( + model, "_hf_peft_config_loaded", False + ) + _quantization_method_supports_training = ( + getattr(model, "hf_quantizer", None) is not None and model.hf_quantizer.is_trainable + ) + + _is_model_quantized_and_qat_trainable = getattr(model, "hf_quantizer", None) is not None and getattr( + model.hf_quantizer, "is_qat_trainable", False + ) + + # Filter out quantized + compiled models + if _is_quantized_and_base_model and hasattr(model, "_orig_mod"): + raise ValueError( + "You cannot fine-tune quantized model with `torch.compile()` make sure to pass a non-compiled model when fine-tuning a quantized model with PEFT" + ) + + # At this stage the model is already loaded + if _is_quantized_and_base_model and not _is_peft_model(model) and not _is_model_quantized_and_qat_trainable: + raise ValueError( + "You cannot perform fine-tuning on purely quantized models. Please attach trainable adapters on top of" + " the quantized model to correctly perform fine-tuning. Please see: https://huggingface.co/docs/transformers/peft" + " for more details" + ) + elif _is_quantized_and_base_model and not _quantization_method_supports_training: + raise ValueError( + f"The model you are trying to fine-tune is quantized with {model.hf_quantizer.quantization_config.quant_method}" + " but that quantization method do not support training. Please open an issue on GitHub: https://github.com/huggingface/transformers" + f" to request the support for training support for {model.hf_quantizer.quantization_config.quant_method}" + ) + + self.is_fsdp_xla_enabled = args.fsdp_config["xla"] + if len(args.fsdp) > 0: + if self.is_deepspeed_enabled: + raise ValueError( + "Using --fsdp xxx together with --deepspeed is not possible, deactivate one of those flags." + ) + if not args.fsdp_config["xla"] and args.parallel_mode != ParallelMode.DISTRIBUTED: + raise ValueError("Using fsdp only works in distributed training.") + + # one place to sort out whether to place the model on device or not + # postpone switching model to cuda when: + # 1. MP - since we are trying to fit a much bigger than 1 gpu model + # 2. fp16-enabled DeepSpeed loads the model in half the size and it doesn't need .to() anyway, + # and we only use deepspeed for training at the moment + # 3. full bf16 or fp16 eval - since the model needs to be cast to the right dtype first + # 4. FSDP - same as MP + self.place_model_on_device = args.place_model_on_device + if ( + self.is_model_parallel + or self.is_deepspeed_enabled + or ((args.fp16_full_eval or args.bf16_full_eval) and not args.do_train) + or self.is_fsdp_xla_enabled + or self.is_fsdp_enabled + ): + self.place_model_on_device = False + + default_collator = ( + DataCollatorWithPadding(processing_class) + if processing_class is not None + and isinstance(processing_class, (PreTrainedTokenizerBase, SequenceFeatureExtractor)) + else default_data_collator + ) + self.data_collator = data_collator if data_collator is not None else default_collator + self.train_dataset = train_dataset + self.eval_dataset = eval_dataset + self.processing_class = processing_class + + # Bnb Quantized models doesn't support `.to` operation. + if ( + self.place_model_on_device + and not getattr(model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES + ): + self._move_model_to_device(model, args.device) + + # Force n_gpu to 1 to avoid DataParallel as MP will manage the GPUs + if self.is_model_parallel: + self.args._n_gpu = 1 + + # later use `self.model is self.model_wrapped` to check if it's wrapped or not + self.model_wrapped = model + self.model = model + + # Just in case the model was wrapped outside of the `Trainer` + unwrapped_model = self.accelerator.unwrap_model(model) + model_forward = ( + unwrapped_model.forward + if not _is_peft_model(unwrapped_model) + else unwrapped_model.get_base_model().forward + ) + forward_params = inspect.signature(model_forward).parameters + + # Check if the model has explicit setup for loss kwargs, + # if not, check if `**kwargs` are in model.forward + if hasattr(model, "accepts_loss_kwargs"): + self.model_accepts_loss_kwargs = model.accepts_loss_kwargs + else: + self.model_accepts_loss_kwargs = any( + k.kind == inspect.Parameter.VAR_KEYWORD for k in forward_params.values() + ) + + self.neftune_noise_alpha = args.neftune_noise_alpha + + self.compute_metrics = compute_metrics + self.preprocess_logits_for_metrics = preprocess_logits_for_metrics + self.optimizer, self.lr_scheduler = optimizers + self.optimizer_cls_and_kwargs = optimizer_cls_and_kwargs + if self.optimizer_cls_and_kwargs is not None and self.optimizer is not None: + raise RuntimeError("Passing both `optimizers` and `optimizer_cls_and_kwargs` arguments is incompatible.") + if model_init is not None and (self.optimizer is not None or self.lr_scheduler is not None): + raise RuntimeError( + "Passing a `model_init` is incompatible with providing the `optimizers` argument. " + "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method." + ) + if is_torch_xla_available() and self.optimizer is not None: + for param in self.model.parameters(): + model_device = param.device + break + for param_group in self.optimizer.param_groups: + if len(param_group["params"]) > 0: + optimizer_device = param_group["params"][0].device + break + if model_device != optimizer_device: + raise ValueError( + "The model and the optimizer parameters are not on the same device, which probably means you" + " created an optimizer around your model **before** putting on the device and passing it to the" + " `Trainer`. Make sure the lines `import torch_xla.core.xla_model as xm` and" + " `model.to(xm.xla_device())` is performed before the optimizer creation in your script." + ) + if (self.is_fsdp_xla_enabled or self.is_fsdp_enabled) and ( + self.optimizer is not None or self.lr_scheduler is not None + ): + raise RuntimeError( + "Passing `optimizers` is not allowed if PyTorch FSDP is enabled. " + "You should subclass `Trainer` and override the `create_optimizer_and_scheduler` method." + ) + default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to) + callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks + self.callback_handler = CallbackHandler( + callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler + ) + self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK) + + # Will be set to True by `self._setup_loggers()` on first call to `self.log()`. + self._loggers_initialized = False + + # Create distant repo and output directory if needed + self.hub_model_id = None + if self.args.push_to_hub: + self.init_hf_repo() + if self.args.should_save: + os.makedirs(self.args.output_dir, exist_ok=True) + + if not callable(self.data_collator) and callable(getattr(self.data_collator, "collate_batch", None)): + raise ValueError("The `data_collator` should be a simple callable (function, class with `__call__`).") + + if args.max_steps > 0 and args.num_train_epochs > 0: + logger.info("max_steps is given, it will override any value given in num_train_epochs") + + if train_dataset is not None and not has_length(train_dataset) and args.max_steps <= 0: + raise ValueError( + "The train_dataset does not implement __len__, max_steps has to be specified. " + "The number of steps needs to be known in advance for the learning rate scheduler." + ) + + if ( + train_dataset is not None + and isinstance(train_dataset, torch.utils.data.IterableDataset) + and args.group_by_length + ): + raise ValueError("the `--group_by_length` option is only available for `Dataset`, not `IterableDataset") + + self._signature_columns = None + + # Mixed precision setup + self.use_apex = False + self.use_cpu_amp = False + + # Mixed precision setup for SageMaker Model Parallel + if is_sagemaker_mp_enabled(): + # BF16 + model parallelism in SageMaker: currently not supported, raise an error + if args.bf16: + raise ValueError("SageMaker Model Parallelism does not support BF16 yet. Please use FP16 instead ") + + if IS_SAGEMAKER_MP_POST_1_10: + # When there's mismatch between SMP config and trainer argument, use SMP config as truth + if args.fp16 != smp.state.cfg.fp16: + logger.warning( + f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16}, " + f"but FP16 provided in trainer argument is {args.fp16}, " + f"setting to {smp.state.cfg.fp16}" + ) + args.fp16 = smp.state.cfg.fp16 + else: + # smp < 1.10 does not support fp16 in trainer. + if hasattr(smp.state.cfg, "fp16"): + logger.warning( + f"FP16 provided in SM_HP_MP_PARAMETERS is {smp.state.cfg.fp16}, " + "but SageMaker Model Parallelism < 1.10 does not support FP16 in trainer." + ) + if (args.fp16 or args.bf16) and args.half_precision_backend == "auto": + if args.device == torch.device("cpu"): + if args.fp16: + if not is_torch_greater_or_equal_than_2_3: + raise ValueError("Tried to use `fp16` but it is not supported on cpu") + else: + args.half_precision_backend = "cpu_amp" + logger.info(f"Using {args.half_precision_backend} half precision backend") + + if (args.fp16 or args.bf16) and not (self.is_deepspeed_enabled or is_sagemaker_mp_enabled()): + # deepspeed and SageMaker Model Parallel manage their own half precision + if args.half_precision_backend == "cpu_amp": + self.use_cpu_amp = True + self.amp_dtype = torch.bfloat16 + elif args.half_precision_backend == "apex": + if not is_apex_available(): + raise ImportError( + "Using FP16 with APEX but APEX is not installed, please refer to" + " https://www.github.com/nvidia/apex." + ) + self.use_apex = True + + # Label smoothing + if self.args.label_smoothing_factor != 0: + self.label_smoother = LabelSmoother(epsilon=self.args.label_smoothing_factor) + else: + self.label_smoother = None + + self.control = TrainerControl() + + self.state = TrainerState( + is_local_process_zero=self.is_local_process_zero(), + is_world_process_zero=self.is_world_process_zero(), + stateful_callbacks=[ + cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) + ], + ) + # Internal variable to count flos in each process, will be accumulated in `self.state.total_flos` then + # returned to 0 every time flos need to be logged + self.current_flos = 0 + self.hp_search_backend = None + if _is_peft_model(self.model) and self.args.label_names is None: + logger.warning( + f"No label_names provided for model class `{self.model.__class__.__name__}`." + " Since `PeftModel` hides base models input arguments, if label_names is not given, label_names can't be set automatically within `Trainer`." + " Note that empty label_names list will be used instead." + ) + default_label_names = find_labels(self.model.__class__) + self.label_names = default_label_names if self.args.label_names is None else self.args.label_names + self.can_return_loss = can_return_loss(self.model.__class__) + self.control = self.callback_handler.on_init_end(self.args, self.state, self.control) + + # Internal variables to help with automatic batch size reduction + self._train_batch_size = args.train_batch_size + self._created_lr_scheduler = False + + # very last + self._memory_tracker.stop_and_update_metrics() + + # torch.compile + if args.torch_compile and not is_torch_compile_available(): + raise RuntimeError("Using torch.compile requires PyTorch 2.0 or higher.") + + self.is_fsdp_xla_v2_enabled = args.fsdp_config.get("xla_fsdp_v2", False) + if self.is_fsdp_xla_v2_enabled: + if not IS_XLA_FSDPV2_POST_2_2: + raise ValueError("FSDPv2 requires `torch_xla` 2.2 or higher.") + # Prepare the SPMD mesh that is going to be used by the data loader and the FSDPv2 wrapper. + # Tensor axis is just a placeholder where it will not be used in FSDPv2. + num_devices = xr.global_runtime_device_count() + xs.set_global_mesh(xs.Mesh(np.array(range(num_devices)), (num_devices, 1), axis_names=("fsdp", "tensor"))) + self.is_fsdp_xla_v1_enabled = self.is_fsdp_xla_enabled and not self.is_fsdp_xla_v2_enabled + + @property + def tokenizer(self) -> Optional[PreTrainedTokenizerBase]: + logger.warning("Trainer.tokenizer is now deprecated. You should use Trainer.processing_class instead.") + return self.processing_class + + @tokenizer.setter + def tokenizer(self, processing_class) -> None: + logger.warning( + "Trainer.tokenizer is now deprecated. You should use `Trainer.processing_class = processing_class` instead." + ) + self.processing_class = processing_class + + def _activate_neftune(self, model): + r""" + Activates the neftune as presented in this code: https://github.com/neelsjain/NEFTune and paper: + https://arxiv.org/abs/2310.05914 + """ + unwrapped_model = self.accelerator.unwrap_model(model) + + if _is_peft_model(unwrapped_model): + embeddings = unwrapped_model.base_model.model.get_input_embeddings() + else: + embeddings = unwrapped_model.get_input_embeddings() + + del unwrapped_model + + embeddings.neftune_noise_alpha = self.neftune_noise_alpha + hook_handle = embeddings.register_forward_hook(neftune_post_forward_hook) + self.neftune_hook_handle = hook_handle + return model + + def _deactivate_neftune(self, model): + """ + Deactivates the neftune method. Make sure to call `_activate_neftune` first. + """ + if not hasattr(self, "neftune_hook_handle"): + raise ValueError("Neftune is not activated make sure to call `trainer._activate_neftune()` first") + + unwrapped_model = self.accelerator.unwrap_model(model) + + if _is_peft_model(unwrapped_model): + embeddings = unwrapped_model.base_model.model.get_input_embeddings() + else: + embeddings = unwrapped_model.get_input_embeddings() + + self.neftune_hook_handle.remove() + del embeddings.neftune_noise_alpha, unwrapped_model + + def add_callback(self, callback): + """ + Add a callback to the current list of [`~transformers.TrainerCallback`]. + + Args: + callback (`type` or [`~transformers.TrainerCallback]`): + A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the + first case, will instantiate a member of that class. + """ + self.callback_handler.add_callback(callback) + + def pop_callback(self, callback): + """ + Remove a callback from the current list of [`~transformers.TrainerCallback`] and returns it. + + If the callback is not found, returns `None` (and no error is raised). + + Args: + callback (`type` or [`~transformers.TrainerCallback]`): + A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the + first case, will pop the first member of that class found in the list of callbacks. + + Returns: + [`~transformers.TrainerCallback`]: The callback removed, if found. + """ + return self.callback_handler.pop_callback(callback) + + def remove_callback(self, callback): + """ + Remove a callback from the current list of [`~transformers.TrainerCallback`]. + + Args: + callback (`type` or [`~transformers.TrainerCallback]`): + A [`~transformers.TrainerCallback`] class or an instance of a [`~transformers.TrainerCallback`]. In the + first case, will remove the first member of that class found in the list of callbacks. + """ + self.callback_handler.remove_callback(callback) + + def _move_model_to_device(self, model, device): + model = model.to(device) + # Moving a model to an XLA device disconnects the tied weights, so we have to retie them. + if self.args.parallel_mode == ParallelMode.TPU and hasattr(model, "tie_weights"): + model.tie_weights() + + def _set_signature_columns_if_needed(self): + if self._signature_columns is None: + # Inspect model forward signature to keep only the arguments it accepts. + model_to_inspect = self.model + if _is_peft_model(self.model): + if hasattr(self.model, "get_base_model"): + model_to_inspect = self.model.get_base_model() + else: + # PeftMixedModel do not provide a `get_base_model` method + model_to_inspect = self.model.base_model.model + signature = inspect.signature(model_to_inspect.forward) + self._signature_columns = list(signature.parameters.keys()) + # Labels may be named label or label_ids, the default data collator handles that. + self._signature_columns += list(set(["label", "label_ids"] + self.label_names)) + + def _remove_unused_columns(self, dataset: "datasets.Dataset", description: Optional[str] = None): + if not self.args.remove_unused_columns: + return dataset + self._set_signature_columns_if_needed() + signature_columns = self._signature_columns + + ignored_columns = list(set(dataset.column_names) - set(signature_columns)) + if len(ignored_columns) > 0: + dset_description = "" if description is None else f"in the {description} set" + logger.info( + f"The following columns {dset_description} don't have a corresponding argument in " + f"`{self.model.__class__.__name__}.forward` and have been ignored: {', '.join(ignored_columns)}." + f" If {', '.join(ignored_columns)} are not expected by `{self.model.__class__.__name__}.forward`, " + " you can safely ignore this message." + ) + + columns = [k for k in signature_columns if k in dataset.column_names] + if len(columns) == 0: + raise ValueError( + "No columns in the dataset match the model's forward method signature. " + f"The following columns have been ignored: [{', '.join(ignored_columns)}]. " + "Please check the dataset and model. You may need to set `remove_unused_columns=False` in `TrainingArguments`." + ) + + if version.parse(datasets.__version__) < version.parse("1.4.0"): + dataset.set_format( + type=dataset.format["type"], columns=columns, format_kwargs=dataset.format["format_kwargs"] + ) + return dataset + else: + return dataset.remove_columns(ignored_columns) + + def _get_collator_with_removed_columns( + self, data_collator: Callable, description: Optional[str] = None + ) -> Callable: + """Wrap the data collator in a callable removing unused columns.""" + if not self.args.remove_unused_columns: + return data_collator + self._set_signature_columns_if_needed() + signature_columns = self._signature_columns + + remove_columns_collator = RemoveColumnsCollator( + data_collator=data_collator, + signature_columns=signature_columns, + logger=logger, + description=description, + model_name=self.model.__class__.__name__, + ) + return remove_columns_collator + + def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]: + if self.train_dataset is None or not has_length(self.train_dataset): + return None + + # Build the sampler. + if self.args.group_by_length: + if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset): + lengths = ( + self.train_dataset[self.args.length_column_name] + if self.args.length_column_name in self.train_dataset.column_names + else None + ) + else: + lengths = None + model_input_name = ( + self.processing_class.model_input_names[0] if self.processing_class is not None else None + ) + return LengthGroupedSampler( + self.args.train_batch_size * self.args.gradient_accumulation_steps, + dataset=self.train_dataset, + lengths=lengths, + model_input_name=model_input_name, + ) + + else: + return RandomSampler(self.train_dataset) + + def get_train_dataloader(self) -> DataLoader: + """ + Returns the training [`~torch.utils.data.DataLoader`]. + + Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed + training if necessary) otherwise. + + Subclass and override this method if you want to inject some custom behavior. + """ + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + data_collator = self.data_collator + if is_datasets_available() and isinstance(train_dataset, datasets.Dataset): + train_dataset = self._remove_unused_columns(train_dataset, description="training") + else: + data_collator = self._get_collator_with_removed_columns(data_collator, description="training") + + dataloader_params = { + "batch_size": self._train_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(train_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_train_sampler() + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["worker_init_fn"] = seed_worker + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) + + def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.Sampler]: + if eval_dataset is None or not has_length(eval_dataset): + return None + # Build the sampler. + + # Deprecated code + if self.args.use_legacy_prediction_loop: + if is_torch_xla_available(): + return SequentialDistributedSampler( + eval_dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal() + ) + elif is_sagemaker_mp_enabled(): + return SequentialDistributedSampler( + eval_dataset, + num_replicas=smp.dp_size(), + rank=smp.dp_rank(), + batch_size=self.args.per_device_eval_batch_size, + ) + else: + return SequentialSampler(eval_dataset) + + if self.args.group_by_length: + if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset): + lengths = ( + eval_dataset[self.args.length_column_name] + if self.args.length_column_name in eval_dataset.column_names + else None + ) + else: + lengths = None + model_input_name = ( + self.processing_class.model_input_names[0] if self.processing_class is not None else None + ) + return LengthGroupedSampler( + self.args.eval_batch_size, + dataset=eval_dataset, + lengths=lengths, + model_input_name=model_input_name, + ) + + if self.args.world_size <= 1: + return SequentialSampler(eval_dataset) + else: + return None + + def get_eval_dataloader(self, eval_dataset: Optional[Union[str, Dataset]] = None) -> DataLoader: + """ + Returns the evaluation [`~torch.utils.data.DataLoader`]. + + Subclass and override this method if you want to inject some custom behavior. + + Args: + eval_dataset (`str` or `torch.utils.data.Dataset`, *optional*): + If a `str`, will use `self.eval_dataset[eval_dataset]` as the evaluation dataset. If a `Dataset`, will override `self.eval_dataset` and must implement `__len__`. If it is a [`~datasets.Dataset`], columns not accepted by the `model.forward()` method are automatically removed. + """ + if eval_dataset is None and self.eval_dataset is None: + raise ValueError("Trainer: evaluation requires an eval_dataset.") + + # If we have persistent workers, don't do a fork bomb especially as eval datasets + # don't change during training + dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval" + if ( + hasattr(self, "_eval_dataloaders") + and dataloader_key in self._eval_dataloaders + and self.args.dataloader_persistent_workers + ): + return self.accelerator.prepare(self._eval_dataloaders[dataloader_key]) + + eval_dataset = ( + self.eval_dataset[eval_dataset] + if isinstance(eval_dataset, str) + else eval_dataset + if eval_dataset is not None + else self.eval_dataset + ) + data_collator = self.data_collator + + if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset): + eval_dataset = self._remove_unused_columns(eval_dataset, description="evaluation") + else: + data_collator = self._get_collator_with_removed_columns(data_collator, description="evaluation") + + dataloader_params = { + "batch_size": self.args.eval_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(eval_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_eval_sampler(eval_dataset) + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + # accelerator.free_memory() will destroy the references, so + # we need to store the non-prepared version + eval_dataloader = DataLoader(eval_dataset, **dataloader_params) + if self.args.dataloader_persistent_workers: + if hasattr(self, "_eval_dataloaders"): + self._eval_dataloaders[dataloader_key] = eval_dataloader + else: + self._eval_dataloaders = {dataloader_key: eval_dataloader} + + return self.accelerator.prepare(eval_dataloader) + + def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader: + """ + Returns the test [`~torch.utils.data.DataLoader`]. + + Subclass and override this method if you want to inject some custom behavior. + + Args: + test_dataset (`torch.utils.data.Dataset`, *optional*): + The test dataset to use. If it is a [`~datasets.Dataset`], columns not accepted by the + `model.forward()` method are automatically removed. It must implement `__len__`. + """ + data_collator = self.data_collator + + if is_datasets_available() and isinstance(test_dataset, datasets.Dataset): + test_dataset = self._remove_unused_columns(test_dataset, description="test") + else: + data_collator = self._get_collator_with_removed_columns(data_collator, description="test") + + dataloader_params = { + "batch_size": self.args.eval_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(test_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_eval_sampler(test_dataset) + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + # We use the same batch_size as for eval. + return self.accelerator.prepare(DataLoader(test_dataset, **dataloader_params)) + + def create_optimizer_and_scheduler(self, num_training_steps: int): + """ + Setup the optimizer and the learning rate scheduler. + + We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the + Trainer's init through `optimizers`, or subclass and override this method (or `create_optimizer` and/or + `create_scheduler`) in a subclass. + """ + self.create_optimizer() + if IS_SAGEMAKER_MP_POST_1_10 and smp.state.cfg.fp16: + # If smp >= 1.10 and fp16 is enabled, we unwrap the optimizer + optimizer = self.optimizer.optimizer + else: + optimizer = self.optimizer + self.create_scheduler(num_training_steps=num_training_steps, optimizer=optimizer) + + def get_decay_parameter_names(self, model) -> list[str]: + """ + Get all parameter names that weight decay will be applied to. + + This function filters out parameters in two ways: + 1. By layer type (instances of layers specified in ALL_LAYERNORM_LAYERS) + 2. By parameter name patterns (containing 'bias', 'layernorm', or 'rmsnorm') + """ + decay_parameters = get_parameter_names(model, ALL_LAYERNORM_LAYERS, ["bias", "layernorm", "rmsnorm"]) + return decay_parameters + + def create_optimizer(self): + """ + Setup the optimizer. + + We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the + Trainer's init through `optimizers`, or subclass and override this method in a subclass. + """ + opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model + + if self.optimizer is None: + decay_parameters = self.get_decay_parameter_names(opt_model) + optimizer_grouped_parameters = [ + { + "params": [ + p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad) + ], + "weight_decay": self.args.weight_decay, + }, + { + "params": [ + p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad) + ], + "weight_decay": 0.0, + }, + ] + + if self.optimizer_cls_and_kwargs is not None: + optimizer_cls, optimizer_kwargs = self.optimizer_cls_and_kwargs + else: + optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(self.args, opt_model) + + # Overwrite `params` in case it's created by `get_optimizer_cls_and_kwargs` + # e.g. for GaLore optimizer. + if "params" in optimizer_kwargs: + optimizer_grouped_parameters = optimizer_kwargs.pop("params") + + # Overwrite `model` in case it's created by `get_optimizer_cls_and_kwargs` + # e.g. for LOMO optimizer. + if "model" in optimizer_kwargs: + optimizer_grouped_parameters = optimizer_kwargs.pop("model") + + # For layer-wise dummy optimizers we overwrite optimizer_grouped_parameters with `optimizer_dict` + # to avoid arguments conflicts. + if "optimizer_dict" in optimizer_kwargs: + optimizer_grouped_parameters = optimizer_kwargs.pop("optimizer_dict") + + self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs) + + if optimizer_cls.__name__ == "Adam8bit": + import bitsandbytes + + manager = bitsandbytes.optim.GlobalOptimManager.get_instance() + + skipped = 0 + for module in opt_model.modules(): + if isinstance(module, nn.Embedding): + skipped += sum({p.data_ptr(): p.numel() for p in module.parameters()}.values()) + logger.info(f"skipped {module}: {skipped / 2**20}M params") + manager.register_module_override(module, "weight", {"optim_bits": 32}) + logger.debug(f"bitsandbytes: will optimize {module} in fp32") + logger.info(f"skipped: {skipped / 2**20}M params") + + if is_sagemaker_mp_enabled(): + self.optimizer = smp.DistributedOptimizer(self.optimizer) + + return self.optimizer + + def get_num_trainable_parameters(self): + """ + Get the number of trainable parameters. + """ + return sum(p.numel() for p in self.model.parameters() if p.requires_grad) + + def get_learning_rates(self): + """ + Returns the learning rate of each parameter from self.optimizer. + """ + if self.optimizer is None: + raise ValueError("Trainer optimizer is None, please make sure you have setup the optimizer before.") + return [group["lr"] for group in self.optimizer.param_groups] + + def get_optimizer_group(self, param: Optional[Union[str, torch.nn.parameter.Parameter]] = None): + """ + Returns optimizer group for a parameter if given, else returns all optimizer groups for params. + + Args: + param (`str` or `torch.nn.parameter.Parameter`, *optional*): + The parameter for which optimizer group needs to be returned. + """ + if self.optimizer is None: + raise ValueError("Trainer optimizer is None, please make sure you have setup the optimizer before.") + if param is not None: + for group in self.optimizer.param_groups: + if param in group["params"]: + return group + return [group["params"] for group in self.optimizer.param_groups] + + @staticmethod + def get_optimizer_cls_and_kwargs( + args: TrainingArguments, model: Optional[PreTrainedModel] = None + ) -> tuple[Any, Any]: + """ + Returns the optimizer class and optimizer parameters based on the training arguments. + + Args: + args (`transformers.training_args.TrainingArguments`): + The training arguments for the training session. + + """ + + # parse args.optim_args + optim_args = {} + if args.optim_args: + for mapping in args.optim_args.replace(" ", "").split(","): + key, value = mapping.split("=") + optim_args[key] = value + + optimizer_kwargs = {"lr": args.learning_rate} + + adam_kwargs = { + "betas": (args.adam_beta1, args.adam_beta2), + "eps": args.adam_epsilon, + } + + def setup_low_rank_optimizer( + optimizer_name: str, + optimizer_mapping: dict[str, Any], + optim_kwargs: dict[str, Any], + is_layerwise_supported: bool = True, + ) -> tuple[Any, Any]: + """ + Helper function to set up low-rank optimizers like GaLore and Apollo. + + Args: + optimizer_name (str): Name of the optimizer. + optimizer_mapping (dict): Mapping of optimizer names to their classes. + optim_kwargs (dict): Keyword arguments for the optimizer. + is_layerwise_supported (bool): Whether layerwise optimization is supported. + + Returns: + Tuple[Any, Any]: Optimizer class and updated optimizer kwargs. + """ + is_layerwise = optimizer_name.lower().endswith("layerwise") + if is_layerwise and args.parallel_mode == ParallelMode.DISTRIBUTED and is_layerwise_supported: + raise NotImplementedError(f"Layer-wise {optimizer_name} does not support DDP at this time") + + optimizer_cls = optimizer_mapping[optimizer_name] + + if args.optim_target_modules is None: + raise ValueError(f"You need to define `optim_target_modules` to use {optimizer_name} optimizers") + + if not isinstance(args.optim_target_modules, (list, str)): + raise ValueError( + f"`optim_target_modules` must be a list of strings, a regex string, or 'all-linear'. Got: {args.optim_target_modules}" + ) + + if model is None: + raise ValueError(f"You need to pass a model to initialize {optimizer_name} optimizer.") + + all_linear = ( + isinstance(args.optim_target_modules, str) + and args.optim_target_modules.replace("_", "-") == "all-linear" + ) + + target_params = [] + target_params_names = [] + for module_name, module in model.named_modules(): + target_module_exists, is_regex = check_target_module_exists( + args.optim_target_modules, module_name, return_is_regex=True + ) + + if not isinstance(module, nn.Linear): + if target_module_exists and not is_regex: + logger.warning( + f"{module_name} matched but ignored. {optimizer_name} only supports linear layers." + ) + continue + + if not target_module_exists and not all_linear: + continue + + target_params.append(module.weight) + target_params_names.append(module_name + ".weight") + + if len(target_params) == 0: + raise ValueError(f"No target modules found for {optimizer_name} ({args.optim_target_modules}).") + + non_target_params = [p for n, p in model.named_parameters() if n not in target_params_names] + optim_kwargs.update(optim_args) + + param_groups = [ + {"params": non_target_params}, + {"params": target_params, **optim_kwargs}, + ] + + if is_layerwise: + if args.gradient_accumulation_steps != 1: + raise ValueError(f"Layerwise {optimizer_name} does not support gradient accumulation!") + + optimizer_dict = {} + for param in non_target_params: + optimizer_dict[param] = optimizer_cls([{"params": [param]}], **optimizer_kwargs) + for param in target_params: + optimizer_dict[param] = optimizer_cls([{"params": [param], **optim_kwargs}], **optimizer_kwargs) + + def optimizer_hook(param): + if param.grad is not None: + optimizer_dict[param].step() + optimizer_dict[param].zero_grad() + + for param in model.parameters(): + if param.requires_grad: + param.register_post_accumulate_grad_hook(optimizer_hook) + + optimizer_cls = LayerWiseDummyOptimizer + optimizer_kwargs.update({"optimizer_dict": optimizer_dict}) + + optimizer_kwargs.update({"params": param_groups}) + return optimizer_cls, optimizer_kwargs + + if args.optim == OptimizerNames.ADAFACTOR: + optimizer_cls = Adafactor + optimizer_kwargs.update({"scale_parameter": False, "relative_step": False}) + elif args.optim in [OptimizerNames.ADAMW_TORCH, OptimizerNames.ADAMW_TORCH_FUSED]: + from torch.optim import AdamW + + optimizer_cls = AdamW + optimizer_kwargs.update(adam_kwargs) + if args.optim == OptimizerNames.ADAMW_TORCH_FUSED: + optimizer_kwargs.update({"fused": True}) + elif args.optim == OptimizerNames.ADAMW_TORCH_XLA: + try: + from torch_xla.amp.syncfree import AdamW + + optimizer_cls = AdamW + optimizer_kwargs.update(adam_kwargs) + except ImportError: + raise ValueError("Trainer failed to import syncfree AdamW from torch_xla.") + elif args.optim == OptimizerNames.ADAMW_TORCH_NPU_FUSED: + try: + from torch_npu.optim import NpuFusedAdamW + + optimizer_cls = NpuFusedAdamW + optimizer_kwargs.update(adam_kwargs) + except ImportError: + raise ValueError("Trainer failed to import FusedAdamW from torch_npu.") + elif args.optim == OptimizerNames.ADAMW_APEX_FUSED: + try: + from apex.optimizers import FusedAdam + + optimizer_cls = FusedAdam + optimizer_kwargs.update(adam_kwargs) + except ImportError: + raise ValueError("Trainer tried to instantiate apex FusedAdam but apex is not installed!") + elif args.optim in [ + OptimizerNames.ADAMW_BNB, + OptimizerNames.ADAMW_8BIT, + OptimizerNames.PAGED_ADAMW, + OptimizerNames.PAGED_ADAMW_8BIT, + OptimizerNames.ADEMAMIX, + OptimizerNames.ADEMAMIX_8BIT, + OptimizerNames.PAGED_ADEMAMIX, + OptimizerNames.PAGED_ADEMAMIX_8BIT, + OptimizerNames.LION, + OptimizerNames.LION_8BIT, + OptimizerNames.PAGED_LION, + OptimizerNames.PAGED_LION_8BIT, + OptimizerNames.RMSPROP_BNB, + OptimizerNames.RMSPROP_8BIT, + OptimizerNames.RMSPROP_32BIT, + ]: + try: + from bitsandbytes.optim import AdamW, Lion, RMSprop + + is_paged = False + optim_bits = 32 + optimizer_cls = None + additional_optim_kwargs = adam_kwargs + if "paged" in args.optim: + is_paged = True + if "8bit" in args.optim: + optim_bits = 8 + if "adam" in args.optim: + optimizer_cls = AdamW + elif "lion" in args.optim: + optimizer_cls = Lion + additional_optim_kwargs = {"betas": (args.adam_beta1, args.adam_beta2)} + elif "rmsprop" in args.optim: + optimizer_cls = RMSprop + # Above we pass all `adam_kwargs` to the optimizer, here + # we only pass `optim_args` which can be passed by the user. + additional_optim_kwargs = optim_args + elif "ademamix" in args.optim: + if is_bitsandbytes_available() and version.parse( + importlib.metadata.version("bitsandbytes") + ) < version.parse("0.44.0"): + raise ValueError( + "The AdEMAMix optimizer is not supported by your current version of `bitsandbytes`. " + "Please install `bitsandbytes` >= 0.44.0." + ) + + from bitsandbytes.optim import AdEMAMix + + optimizer_cls = AdEMAMix + additional_optim_kwargs = { + "betas": ( + float(optim_args.get("beta1", args.adam_beta1)), + float(optim_args.get("beta2", args.adam_beta2)), + float(optim_args.get("beta3", 0.9999)), + ), + "alpha": float(optim_args.get("alpha", 5.0)), + "eps": float(optim_args.get("eps", args.adam_epsilon)), + } + + if "t_alpha" in optim_args: + additional_optim_kwargs["t_alpha"] = int(optim_args["t_alpha"]) + + if "t_beta3" in optim_args: + additional_optim_kwargs["t_beta3"] = int(optim_args["t_beta3"]) + + bnb_kwargs = {"optim_bits": optim_bits} + if "rmsprop" not in args.optim: + bnb_kwargs["is_paged"] = is_paged + + optimizer_kwargs.update(additional_optim_kwargs) + optimizer_kwargs.update(bnb_kwargs) + except ImportError: + raise ValueError("Trainer tried to instantiate bnb optimizer but `bitsandbytes` is not installed!") + if is_bitsandbytes_available() and version.parse( + importlib.metadata.version("bitsandbytes") + ) < version.parse("0.41.1"): + logger.warning( + "You are using 8-bit optimizers with a version of `bitsandbytes` < 0.41.1. " + "It is recommended to update your version as a major bug has been fixed in 8-bit optimizers." + ) + elif args.optim == OptimizerNames.ADAMW_ANYPRECISION: + try: + from torchdistx.optimizers import AnyPrecisionAdamW + + optimizer_cls = AnyPrecisionAdamW + optimizer_kwargs.update(adam_kwargs) + + # TODO Change dtypes back to M=FP32, Var = BF16, Kahan = False once they can be cast together in torchdistx. + optimizer_kwargs.update( + { + "use_kahan_summation": strtobool(optim_args.get("use_kahan_summation", "False")), + "momentum_dtype": getattr(torch, optim_args.get("momentum_dtype", "float32")), + "variance_dtype": getattr(torch, optim_args.get("variance_dtype", "float32")), + "compensation_buffer_dtype": getattr( + torch, optim_args.get("compensation_buffer_dtype", "bfloat16") + ), + } + ) + except ImportError: + raise ValueError("Please install https://github.com/pytorch/torchdistx") + elif args.optim == OptimizerNames.SGD: + optimizer_cls = torch.optim.SGD + elif args.optim == OptimizerNames.ADAGRAD: + optimizer_cls = torch.optim.Adagrad + elif args.optim == OptimizerNames.RMSPROP: + optimizer_cls = torch.optim.RMSprop + elif args.optim in [ + OptimizerNames.GALORE_ADAMW, + OptimizerNames.GALORE_ADAMW_8BIT, + OptimizerNames.GALORE_ADAFACTOR, + OptimizerNames.GALORE_ADAMW_LAYERWISE, + OptimizerNames.GALORE_ADAMW_8BIT_LAYERWISE, + OptimizerNames.GALORE_ADAFACTOR_LAYERWISE, + ]: + if not is_galore_torch_available(): + raise ImportError( + "You need to install `galore_torch` in order to use GaLore optimizers" + " install it with `pip install git+https://github.com/jiaweizzhao/GaLore`" + ) + from galore_torch import GaLoreAdafactor, GaLoreAdamW, GaLoreAdamW8bit + + optimizer_mapping = { + OptimizerNames.GALORE_ADAMW: GaLoreAdamW, + OptimizerNames.GALORE_ADAMW_8BIT: GaLoreAdamW8bit, + OptimizerNames.GALORE_ADAFACTOR: GaLoreAdafactor, + OptimizerNames.GALORE_ADAMW_LAYERWISE: GaLoreAdamW, + OptimizerNames.GALORE_ADAMW_8BIT_LAYERWISE: GaLoreAdamW8bit, + OptimizerNames.GALORE_ADAFACTOR_LAYERWISE: GaLoreAdafactor, + } + + galore_optim_kwargs = { + "rank": int(optim_args.pop("rank", 128)), + "update_proj_gap": int(optim_args.pop("update_proj_gap", 200)), + "scale": float(optim_args.pop("scale", 0.25)), + "proj_type": optim_args.pop("proj_type", "std"), + } + + optimizer_cls, optimizer_kwargs = setup_low_rank_optimizer( + args.optim, optimizer_mapping, galore_optim_kwargs + ) + if args.optim == OptimizerNames.GALORE_ADAFACTOR: + optimizer_kwargs.update({"scale_parameter": False, "relative_step": False}) + elif args.optim in [ + OptimizerNames.APOLLO_ADAMW, + OptimizerNames.APOLLO_ADAMW_LAYERWISE, + ]: + if not is_apollo_torch_available(): + raise ImportError( + "You need to install `apollo_torch` in order to use APOLLO optimizers" + " install it with `pip install git+https://github.com/zhuhanqing/APOLLO`" + ) + from apollo_torch import APOLLOAdamW + + optimizer_mapping = { + OptimizerNames.APOLLO_ADAMW: APOLLOAdamW, + OptimizerNames.APOLLO_ADAMW_LAYERWISE: APOLLOAdamW, + } + + apollo_optim_kwargs = { + "rank": int(optim_args.pop("rank", 128)), + "proj": optim_args.pop("proj", "random"), + "scale_type": optim_args.pop("scale_type", "channel"), + "update_proj_gap": int(optim_args.pop("update_proj_gap", 200)), + "scale": float(optim_args.pop("scale", 1.0)), + "proj_type": optim_args.pop("proj_type", "std"), + } + + optimizer_cls, optimizer_kwargs = setup_low_rank_optimizer( + args.optim, optimizer_mapping, apollo_optim_kwargs + ) + elif args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: + if not is_lomo_available(): + raise ImportError( + "You need to install `lomo_optim` in order to use LOMO optimizers" + " install it with `pip install lomo-optim`" + ) + if not is_accelerate_available("0.30.0"): + raise ImportError("You need to have `accelerate>=0.30.0` to be able to use LOMO optimizers") + + if model is None: + raise ValueError("You need to pass a `model` in order to correctly initialize a LOMO optimizer.") + + from lomo_optim import AdaLomo, Lomo + + if "ada" in args.optim: + optimizer_cls = AdaLomo + else: + optimizer_cls = Lomo + + optimizer_kwargs.update({"model": model}) + elif args.optim == OptimizerNames.GROKADAMW: + if not is_grokadamw_available(): + raise ValueError("Please install grokadamw with `pip install grokadamw`") + + from grokadamw import GrokAdamW + + optimizer_cls = GrokAdamW + optimizer_kwargs.update( + { + "alpha_init": float(optim_args.get("alpha_init", 0.98)), + "lamb": float(optim_args.get("lamb", 2.0)), + "gamma": float(optim_args.get("gamma", 0.1)), + "grokking_signal_decay_rate": float(optim_args.get("grokking_signal_decay_rate", 0.1)), + "gradient_clipping": float(optim_args.get("gradient_clipping", 1.0)), + } + ) + elif args.optim in [ + OptimizerNames.ADAMW_TORCH_4BIT, + OptimizerNames.ADAMW_TORCH_8BIT, + ]: + if not is_torchao_available() or version.parse(importlib.metadata.version("torchao")) < version.parse( + "0.4.0" + ): + raise ImportError( + "You need to have `torchao>=0.4.0` in order to use torch 4-bit optimizers." + "Install it with `pip install torchao` or follow the instructions here: https://github.com/pytorch/ao" + ) + if version.parse(importlib.metadata.version("torch")) <= version.parse("2.4"): + raise ImportError( + "You need to have `torch>2.4` in order to use torch 4-bit optimizers. " + "Install it with `pip install --upgrade torch` it is available on pipy. Otherwise, you need to install torch nightly." + ) + from torchao.prototype.low_bit_optim import AdamW4bit, AdamW8bit + + if args.optim == OptimizerNames.ADAMW_TORCH_4BIT: + optimizer_cls = AdamW4bit + elif args.optim == OptimizerNames.ADAMW_TORCH_8BIT: + optimizer_cls = AdamW8bit + else: + raise ValueError("Invalid optimizer") + optimizer_kwargs.update(adam_kwargs) + elif args.optim in [ + OptimizerNames.SCHEDULE_FREE_RADAM, + OptimizerNames.SCHEDULE_FREE_ADAMW, + OptimizerNames.SCHEDULE_FREE_SGD, + ]: + if not is_schedulefree_available(): + raise ImportError( + "You need to install `schedulefree` in order to use schedulefree optimizers. " + "Install it with `pip install schedulefree.`" + ) + if not is_accelerate_available("0.30.0"): + raise ImportError("You need to have `accelerate>=0.30.0` to be able to use schedulefree optimizers") + from schedulefree import AdamWScheduleFree, SGDScheduleFree + + additional_optim_kwargs = {} + require_warmup = True + + if args.optim == OptimizerNames.SCHEDULE_FREE_RADAM: + if not is_schedulefree_available("1.4.0"): + raise ImportError( + "You need to install `schedulefree>=1.4.0` in order to use RAdamScheduleFree optimizer. " + "Install it with `pip install schedulefree.`" + ) + from schedulefree import RAdamScheduleFree + + optimizer_cls = RAdamScheduleFree + additional_optim_kwargs = adam_kwargs + require_warmup = False + elif args.optim == OptimizerNames.SCHEDULE_FREE_ADAMW: + optimizer_cls = AdamWScheduleFree + additional_optim_kwargs = adam_kwargs + elif args.optim == OptimizerNames.SCHEDULE_FREE_SGD: + optimizer_cls = SGDScheduleFree + else: + raise ValueError("Invalid schedulefree optimizer") + + additional_optim_kwargs["weight_decay"] = args.weight_decay + if require_warmup: + additional_optim_kwargs["warmup_steps"] = args.warmup_steps + additional_optim_kwargs.update( + { + "weight_lr_power": float(optim_args.get("weight_lr_power", 2.0)), + "r": float(optim_args.get("r", 0.0)), + } + ) + optimizer_kwargs.update(additional_optim_kwargs) + else: + raise ValueError(f"Trainer cannot instantiate unsupported optimizer: {args.optim}") + return optimizer_cls, optimizer_kwargs + + def create_scheduler(self, num_training_steps: int, optimizer: torch.optim.Optimizer = None): + """ + Setup the scheduler. The optimizer of the trainer must have been set up either before this method is called or + passed as an argument. + + Args: + num_training_steps (int): The number of training steps to do. + """ + if self.lr_scheduler is None: + self.lr_scheduler = get_scheduler( + self.args.lr_scheduler_type, + optimizer=self.optimizer if optimizer is None else optimizer, + num_warmup_steps=self.args.get_warmup_steps(num_training_steps), + num_training_steps=num_training_steps, + scheduler_specific_kwargs=self.args.lr_scheduler_kwargs, + ) + self._created_lr_scheduler = True + return self.lr_scheduler + + def num_examples(self, dataloader: DataLoader) -> int: + """ + Helper to get number of samples in a [`~torch.utils.data.DataLoader`] by accessing its dataset. When + dataloader.dataset does not exist or has no length, estimates as best it can + """ + try: + dataset = dataloader.dataset + # Special case for IterableDatasetShard, we need to dig deeper + if isinstance(dataset, IterableDatasetShard): + return len(dataloader.dataset.dataset) + return len(dataloader.dataset) + except (NameError, AttributeError, TypeError): # no dataset or length, estimate by length of dataloader + return len(dataloader) * self.args.per_device_train_batch_size + + @staticmethod + def num_tokens(train_dl: DataLoader, max_steps: Optional[int] = None) -> int: + """ + Helper to get number of tokens in a [`~torch.utils.data.DataLoader`] by enumerating dataloader. + """ + train_tokens = 0 + try: + for batch in train_dl: + tokens = batch["input_ids"].numel() + if max_steps is not None: + return tokens * max_steps + train_tokens += tokens + except KeyError: + logger.warning("Cannot get num_tokens from dataloader") + return train_tokens + + def _hp_search_setup(self, trial: Union["optuna.Trial", dict[str, Any]]): + """HP search setup code""" + self._trial = trial + + if self.hp_search_backend is None or trial is None: + return + if self.hp_search_backend == HPSearchBackend.OPTUNA: + params = self.hp_space(trial) + elif self.hp_search_backend == HPSearchBackend.RAY: + params = trial + params.pop("wandb", None) + elif self.hp_search_backend == HPSearchBackend.SIGOPT: + params = {k: int(v) if isinstance(v, str) else v for k, v in trial.assignments.items()} + elif self.hp_search_backend == HPSearchBackend.WANDB: + params = trial + + for key, value in params.items(): + if not hasattr(self.args, key): + logger.warning( + f"Trying to set {key} in the hyperparameter search but there is no corresponding field in" + " `TrainingArguments`." + ) + continue + old_attr = getattr(self.args, key, None) + # Casting value to the proper type + if old_attr is not None: + value = type(old_attr)(value) + + setattr(self.args, key, value) + if self.hp_search_backend == HPSearchBackend.OPTUNA: + logger.info(f"Trial: {trial.params}") + if self.hp_search_backend == HPSearchBackend.SIGOPT: + logger.info(f"SigOpt Assignments: {trial.assignments}") + if self.hp_search_backend == HPSearchBackend.WANDB: + logger.info(f"W&B Sweep parameters: {trial}") + if self.is_deepspeed_enabled: + if self.args.deepspeed is None: + raise ValueError("For sweeps with deepspeed, `args.deepspeed` must be set") + + self.accelerator.free_memory() + + # Rebuild the deepspeed config to reflect the updated training parameters + from accelerate.utils import DeepSpeedPlugin + + from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig + + self.args.hf_deepspeed_config = HfTrainerDeepSpeedConfig(self.args.deepspeed) + self.args.hf_deepspeed_config.trainer_config_process(self.args) + self.args.deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.args.hf_deepspeed_config) + + # From 1.0 on, we need to fully wipe the DS plugin when doing sweeps. + # Simply calling `_reset_state` is enough and doesn't need a version pin. + AcceleratorState()._reset_state() + + self.create_accelerator_and_postprocess() + + def _report_to_hp_search(self, trial: Union["optuna.Trial", dict[str, Any]], step: int, metrics: dict[str, float]): + if self.hp_search_backend is None or trial is None: + return + metrics = metrics.copy() + self.objective = self.compute_objective(metrics) + if self.hp_search_backend == HPSearchBackend.OPTUNA: + import optuna + + if hasattr(trial, "study") and not trial.study._is_multi_objective(): + trial.report(self.objective, step) + if trial.should_prune(): + self.callback_handler.on_train_end(self.args, self.state, self.control) + raise optuna.TrialPruned() + elif self.hp_search_backend == HPSearchBackend.RAY: + import ray.train + + with tempfile.TemporaryDirectory() as temp_checkpoint_dir: + checkpoint = None + if self.control.should_save: + self._tune_save_checkpoint(checkpoint_dir=temp_checkpoint_dir) + checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir) + metrics["objective"] = self.objective + ray.train.report(metrics, checkpoint=checkpoint) + + def _tune_save_checkpoint(self, checkpoint_dir: str): + output_dir = os.path.join(checkpoint_dir, f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}") + self.save_model(output_dir, _internal_call=True) + if self.args.should_save: + # Update the `TrainerControl` state to where we are currently + self.state.stateful_callbacks["TrainerControl"] = self.control.state() + self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME)) + torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME)) + torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME)) + + def call_model_init(self, trial=None): + model_init_argcount = number_of_arguments(self.model_init) + if model_init_argcount == 0: + model = self.model_init() + elif model_init_argcount == 1: + model = self.model_init(trial) + else: + raise RuntimeError("model_init should have 0 or 1 argument.") + + if model is None: + raise RuntimeError("model_init should not return None.") + + return model + + def torch_jit_model_eval(self, model, dataloader, training=False): + if not training: + if dataloader is None: + logger.warning("failed to use PyTorch jit mode due to current dataloader is none.") + return model + example_batch = next(iter(dataloader)) + example_batch = self._prepare_inputs(example_batch) + try: + jit_model = copy.copy(model) + jit_model.eval() + original_forward = jit_model.__dict__.pop("_original_forward", None) + # remove mixed precision hooks from the model + if original_forward: + jit_model.forward = original_forward + autocast_handler = AutocastKwargs(cache_enabled=False) + with self.accelerator.autocast(autocast_handler=autocast_handler), torch.no_grad(): + if version.parse(version.parse(torch.__version__).base_version) >= version.parse("2.0.0"): + if isinstance(example_batch, dict): + jit_model = torch.jit.trace(jit_model, example_kwarg_inputs=example_batch, strict=False) + else: + jit_model = torch.jit.trace( + jit_model, + example_kwarg_inputs={key: example_batch[key] for key in example_batch}, + strict=False, + ) + else: + jit_inputs = [] + for key in example_batch: + example_tensor = torch.ones_like(example_batch[key]) + jit_inputs.append(example_tensor) + jit_inputs = tuple(jit_inputs) + jit_model = torch.jit.trace(jit_model, jit_inputs, strict=False) + jit_model = torch.jit.freeze(jit_model) + with torch.no_grad(): + jit_model(**example_batch) + jit_model(**example_batch) + model = jit_model + self.use_cpu_amp = False + except (RuntimeError, TypeError, ValueError, NameError, IndexError) as e: + logger.warning(f"failed to use PyTorch jit mode due to: {e}.") + + return model + + def ipex_optimize_model(self, model, training=False, dtype=torch.float32): + if not is_ipex_available(): + raise ImportError( + "Using IPEX but IPEX is not installed or IPEX's version does not match current PyTorch, please refer" + " to https://github.com/intel/intel-extension-for-pytorch." + ) + + import intel_extension_for_pytorch as ipex + + if not training: + model.eval() + dtype = torch.bfloat16 if not self.is_in_train and self.args.bf16_full_eval else dtype + # conv_bn_folding is disabled as it fails in symbolic tracing, resulting in ipex warnings + model = ipex.optimize(model, dtype=dtype, level="O1", conv_bn_folding=False, inplace=not self.is_in_train) + else: + if not model.training: + model.train() + model, self.optimizer = ipex.optimize( + model, dtype=dtype, optimizer=self.optimizer, inplace=True, level="O1" + ) + + return model + + def compare_trainer_and_checkpoint_args(self, training_args, trainer_state): + attributes_map = { + "logging_steps": "logging_steps", + "eval_steps": "eval_steps", + "save_steps": "save_steps", + } + + has_warning = False + warning_str = "Warning: The following arguments do not match the ones in the `trainer_state.json` within the checkpoint directory: " + for arg_attr, state_attr in attributes_map.items(): + arg_value = getattr(training_args, arg_attr, None) + state_value = getattr(trainer_state, state_attr, None) + + if arg_value is not None and state_value is not None and arg_value != state_value: + warning_str += f"\n\t{arg_attr}: {arg_value} (from args) != {state_value} (from trainer_state.json)" + has_warning = True + + # train bs is special as we need to account for multi-GPU + train_bs_args = training_args.per_device_train_batch_size + train_bs_state = trainer_state.train_batch_size // max(1, training_args.n_gpu) + + if train_bs_args != train_bs_state: + warning_str += f"\n\tper_device_train_batch_size: {train_bs_args} (from args) != {train_bs_state} (from trainer_state.json)" + has_warning = True + + if has_warning: + logger.warning_once(warning_str) + + def _wrap_model(self, model, training=True, dataloader=None): + if self.args.use_ipex: + dtype = torch.bfloat16 if self.use_cpu_amp else torch.float32 + model = self.ipex_optimize_model(model, training, dtype=dtype) + + if is_sagemaker_mp_enabled(): + # Wrapping the base model twice in a DistributedModel will raise an error. + if isinstance(self.model_wrapped, smp.model.DistributedModel): + return self.model_wrapped + return smp.DistributedModel(model, backward_passes_per_step=self.args.gradient_accumulation_steps) + + # train/eval could be run multiple-times - if already wrapped, don't re-wrap it again + if self.accelerator.unwrap_model(model) is not model: + return model + + # Mixed precision training with apex (torch < 1.6) + if self.use_apex and training: + model, self.optimizer = amp.initialize(model, self.optimizer, opt_level=self.args.fp16_opt_level) + + # Multi-gpu training (should be after apex fp16 initialization) / 8bit models does not support DDP + if self.args.n_gpu > 1 and not getattr(model, "is_loaded_in_8bit", False): + model = nn.DataParallel(model) + + if self.args.jit_mode_eval: + start_time = time.time() + model = self.torch_jit_model_eval(model, dataloader, training) + self.jit_compilation_time = round(time.time() - start_time, 4) + + # Note: in torch.distributed mode, there's no point in wrapping the model + # inside a DistributedDataParallel as we'll be under `no_grad` anyways. + if not training: + return model + + # Distributed training (should be after apex fp16 initialization) + # Distributed training using PyTorch FSDP + if self.is_fsdp_xla_enabled: + try: + from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as FSDP + from torch_xla.distributed.fsdp import checkpoint_module + from torch_xla.distributed.fsdp.wrap import ( + size_based_auto_wrap_policy, + transformer_auto_wrap_policy, + ) + + if self.is_fsdp_xla_v2_enabled: + from torch_xla.experimental.spmd_fully_sharded_data_parallel import ( + SpmdFullyShardedDataParallel as FSDPv2, + ) + except ImportError: + raise ImportError("Missing XLA FSDP related module; please make sure to use torch-xla >= 2.0.") + auto_wrap_policy = None + auto_wrapper_callable = None + default_transformer_cls_names_to_wrap = getattr(model, "_no_split_modules", None) + fsdp_transformer_layer_cls_to_wrap = self.args.fsdp_config.get( + "transformer_layer_cls_to_wrap", default_transformer_cls_names_to_wrap + ) + + if self.args.fsdp_config["min_num_params"] > 0: + auto_wrap_policy = functools.partial( + size_based_auto_wrap_policy, min_num_params=self.args.fsdp_config["min_num_params"] + ) + elif fsdp_transformer_layer_cls_to_wrap is not None: + transformer_cls_to_wrap = set() + for layer_class in fsdp_transformer_layer_cls_to_wrap: + transformer_cls = get_module_class_from_name(model, layer_class) + if transformer_cls is None: + raise Exception("Could not find the transformer layer class to wrap in the model.") + else: + transformer_cls_to_wrap.add(transformer_cls) + + auto_wrap_policy = functools.partial( + transformer_auto_wrap_policy, + # Transformer layer class to wrap + transformer_layer_cls=transformer_cls_to_wrap, + ) + fsdp_kwargs = self.args.xla_fsdp_config + if self.args.fsdp_config["xla_fsdp_grad_ckpt"]: + if model.config.use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + model.config.use_cache = False + + # Apply gradient checkpointing to auto-wrapped sub-modules if specified + def auto_wrapper_callable(m, *args, **kwargs): + target_cls = FSDP if not self.is_fsdp_xla_v2_enabled else FSDPv2 + return target_cls(checkpoint_module(m), *args, **kwargs) + + # Wrap the base model with an outer FSDP wrapper + if self.is_fsdp_xla_v2_enabled: + + def shard_output(output, mesh): + from .modeling_outputs import CausalLMOutputWithPast + + real_output = None + if isinstance(output, torch.Tensor): + real_output = output + elif isinstance(output, tuple): + real_output = output[0] + elif isinstance(output, CausalLMOutputWithPast): + real_output = output.logits + + if real_output is None: + raise ValueError("Something went wrong, the output of the model shouldn't be `None`") + xs.mark_sharding(real_output, mesh, ("fsdp", None, None)) + + self.model = model = FSDPv2( + model, + shard_output=shard_output, + auto_wrap_policy=auto_wrap_policy, + auto_wrapper_callable=auto_wrapper_callable, + ) + else: + self.model = model = FSDP( + model, + auto_wrap_policy=auto_wrap_policy, + auto_wrapper_callable=auto_wrapper_callable, + **fsdp_kwargs, + ) + + # Patch `xm.optimizer_step` should not reduce gradients in this case, + # as FSDP does not need gradient reduction over sharded parameters. + def patched_optimizer_step(optimizer, barrier=False, optimizer_args={}): + loss = optimizer.step(**optimizer_args) + if barrier: + xm.mark_step() + return loss + + xm.optimizer_step = patched_optimizer_step + elif is_sagemaker_dp_enabled(): + model = nn.parallel.DistributedDataParallel( + model, device_ids=[int(os.getenv("SMDATAPARALLEL_LOCAL_RANK"))] + ) + elif self.args.parallel_mode == ParallelMode.DISTRIBUTED: + if is_torch_neuroncore_available(): + return model + kwargs = {} + if self.args.ddp_find_unused_parameters is not None: + kwargs["find_unused_parameters"] = self.args.ddp_find_unused_parameters + elif isinstance(model, PreTrainedModel): + # find_unused_parameters breaks checkpointing as per + # https://github.com/huggingface/transformers/pull/4659#issuecomment-643356021 + kwargs["find_unused_parameters"] = not model.is_gradient_checkpointing + else: + kwargs["find_unused_parameters"] = True + + if self.args.ddp_bucket_cap_mb is not None: + kwargs["bucket_cap_mb"] = self.args.ddp_bucket_cap_mb + + if self.args.ddp_broadcast_buffers is not None: + kwargs["broadcast_buffers"] = self.args.ddp_broadcast_buffers + + self.accelerator.ddp_handler = DistributedDataParallelKwargs(**kwargs) + + return model + + def train( + self, + resume_from_checkpoint: Optional[Union[str, bool]] = None, + trial: Union["optuna.Trial", dict[str, Any], None] = None, + ignore_keys_for_eval: Optional[list[str]] = None, + **kwargs, + ): + """ + Main training entry point. + + Args: + resume_from_checkpoint (`str` or `bool`, *optional*): + If a `str`, local path to a saved checkpoint as saved by a previous instance of [`Trainer`]. If a + `bool` and equals `True`, load the last checkpoint in *args.output_dir* as saved by a previous instance + of [`Trainer`]. If present, training will resume from the model/optimizer/scheduler states loaded here. + trial (`optuna.Trial` or `Dict[str, Any]`, *optional*): + The trial run or the hyperparameter dictionary for hyperparameter search. + ignore_keys_for_eval (`List[str]`, *optional*) + A list of keys in the output of your model (if it is a dictionary) that should be ignored when + gathering predictions for evaluation during the training. + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments used to hide deprecated arguments + """ + if resume_from_checkpoint is False: + resume_from_checkpoint = None + + # memory metrics - must set up as early as possible + self._memory_tracker.start() + + args = self.args + + self.is_in_train = True + + # Attach NEFTune hooks if necessary + if self.neftune_noise_alpha is not None: + self.model = self._activate_neftune(self.model) + + # do_train is not a reliable argument, as it might not be set and .train() still called, so + # the following is a workaround: + if ( + (args.fp16_full_eval or args.bf16_full_eval) + and not args.do_train + and not self.is_model_parallel + and self.model_init is None + ): + self._move_model_to_device(self.model, args.device) + + if "model_path" in kwargs: + resume_from_checkpoint = kwargs.pop("model_path") + warnings.warn( + "`model_path` is deprecated and will be removed in a future version. Use `resume_from_checkpoint` " + "instead.", + FutureWarning, + ) + if len(kwargs) > 0: + raise TypeError(f"train() got unexpected keyword arguments: {', '.join(list(kwargs.keys()))}.") + # This might change the seed so needs to run first. + self._hp_search_setup(trial) + self._train_batch_size = self.args.train_batch_size + + # Model re-init + model_reloaded = False + if self.model_init is not None: + # Seed must be set before instantiating the model when using model_init. + enable_full_determinism(self.args.seed) if self.args.full_determinism else set_seed(self.args.seed) + self.model = self.call_model_init(trial) + model_reloaded = True + # Reinitializes optimizer and scheduler + self.optimizer, self.lr_scheduler = None, None + + # Load potential model checkpoint + if isinstance(resume_from_checkpoint, bool) and resume_from_checkpoint: + resume_from_checkpoint = get_last_checkpoint(args.output_dir) + if resume_from_checkpoint is None: + raise ValueError(f"No valid checkpoint found in output directory ({args.output_dir})") + + if resume_from_checkpoint is not None: + if not is_sagemaker_mp_enabled() and not self.is_deepspeed_enabled and not self.is_fsdp_enabled: + self._load_from_checkpoint(resume_from_checkpoint) + # In case of repeating the find_executable_batch_size, set `self._train_batch_size` properly + state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)) + if state.train_batch_size is not None: + self._train_batch_size = state.train_batch_size + + # If model was re-initialized, put it on the right device and update self.model_wrapped + if model_reloaded: + if self.place_model_on_device: + self._move_model_to_device(self.model, args.device) + self.model_wrapped = self.model + + inner_training_loop = find_executable_batch_size( + self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size + ) + if args.push_to_hub: + try: + # Disable progress bars when uploading models during checkpoints to avoid polluting stdout + hf_hub_utils.disable_progress_bars() + return inner_training_loop( + args=args, + resume_from_checkpoint=resume_from_checkpoint, + trial=trial, + ignore_keys_for_eval=ignore_keys_for_eval, + ) + finally: + hf_hub_utils.enable_progress_bars() + else: + return inner_training_loop( + args=args, + resume_from_checkpoint=resume_from_checkpoint, + trial=trial, + ignore_keys_for_eval=ignore_keys_for_eval, + ) + + def _inner_training_loop( + self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None + ): + self.accelerator.free_memory() + self._train_batch_size = batch_size + if self.args.auto_find_batch_size: + if self.state.train_batch_size != self._train_batch_size: + from accelerate.utils import release_memory + + (self.model_wrapped,) = release_memory(self.model_wrapped) + self.model_wrapped = self.model + + # Check for DeepSpeed *after* the initial pass and modify the config + if self.is_deepspeed_enabled: + # Temporarily unset `self.args.train_batch_size` + original_bs = self.args.per_device_train_batch_size + self.args.per_device_train_batch_size = self._train_batch_size // max(1, self.args.n_gpu) + self.propagate_args_to_deepspeed(True) + self.args.per_device_train_batch_size = original_bs + self.state.train_batch_size = self._train_batch_size + logger.debug(f"Currently training with a batch size of: {self._train_batch_size}") + # Data loader and number of training steps + train_dataloader = self.get_train_dataloader() + if self.is_fsdp_xla_v2_enabled: + train_dataloader = tpu_spmd_dataloader(train_dataloader) + + # Setting up training control variables: + # number of training epochs: num_train_epochs + # number of training steps per epoch: num_update_steps_per_epoch + # total number of training steps to execute: max_steps + total_train_batch_size = self._train_batch_size * args.gradient_accumulation_steps * args.world_size + ( + num_train_epochs, + num_update_steps_per_epoch, + num_examples, + num_train_samples, + epoch_based, + len_dataloader, + max_steps, + ) = self.set_initial_training_values(args, train_dataloader, total_train_batch_size) + + num_train_tokens = None + if self.args.include_tokens_per_second: + num_train_tokens = self.num_tokens(train_dataloader, None if epoch_based else max_steps) + # If going by epochs, multiply tokens linearly + if len_dataloader is not None and epoch_based: + num_train_tokens *= args.num_train_epochs + # Otherwise since its steps, we just multiply by grad accum + else: + num_train_tokens *= args.gradient_accumulation_steps + + if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug: + if self.args.n_gpu > 1: + # nn.DataParallel(model) replicates the model, creating new variables and module + # references registered here no longer work on other gpus, breaking the module + raise ValueError( + "Currently --debug underflow_overflow is not supported under DP. Please use DDP" + " (torchrun or torch.distributed.launch (deprecated))." + ) + else: + debug_overflow = DebugUnderflowOverflow(self.model) # noqa + + delay_optimizer_creation = is_sagemaker_mp_enabled() or self.is_fsdp_xla_enabled or self.is_fsdp_enabled + + # Can't delay optimizer creation when using FSDP2: https://github.com/huggingface/accelerate/blob/3f636d626063ffcf9a337c7d3624d61b7d187d59/src/accelerate/accelerator.py#L1404 + is_fsdp2 = self.is_fsdp_enabled and (getattr(self.accelerator.state.fsdp_plugin, "fsdp_version", 1) == 2) + if is_fsdp2: + delay_optimizer_creation = False + + # We need to reset the scheduler, as its parameters may be different on subsequent calls + if self._created_lr_scheduler: + self.lr_scheduler = None + self._created_lr_scheduler = False + + if self.is_deepspeed_enabled: + self.optimizer, self.lr_scheduler = deepspeed_init(self, num_training_steps=max_steps) + + if not delay_optimizer_creation: + self.create_optimizer_and_scheduler(num_training_steps=max_steps) + + self.state = TrainerState( + stateful_callbacks=[ + cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) + ] + ) + self.state.is_hyper_param_search = trial is not None + self.state.train_batch_size = self._train_batch_size + + # Compute absolute values for logging, eval, and save if given as ratio + self.state.compute_steps(args, max_steps) + + # Activate gradient checkpointing if needed + if args.gradient_checkpointing: + self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=args.gradient_checkpointing_kwargs) + + model = self._wrap_model(self.model_wrapped) + + # as the model is wrapped, don't use `accelerator.prepare` + # this is for unhandled cases such as + # FSDP-XLA, SageMaker MP/DP, DataParallel, IPEX + use_accelerator_prepare = True if model is self.model else False + + if use_accelerator_prepare and self.is_fsdp_enabled: + # In case of auto_find_batch_size=True + # Remove FSDP wrapping from sub-models. + self.model = unwrap_model(self.model, recursive=True) + + if delay_optimizer_creation: + if use_accelerator_prepare: + # configure fsdp plugin for qlora if any + self._fsdp_qlora_plugin_updates() + if self.accelerator.mixed_precision != "fp8": + self.model = self.accelerator.prepare(self.model) + self.create_optimizer_and_scheduler(num_training_steps=max_steps) + + # prepare using `accelerator` prepare + if use_accelerator_prepare: + self.model.train() + if hasattr(self.lr_scheduler, "step"): + if self.use_apex: + model = self.accelerator.prepare(self.model) + else: + model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer) + else: + # to handle cases wherein we pass "DummyScheduler" such as when it is specified in DeepSpeed config. + model, self.optimizer, self.lr_scheduler = self.accelerator.prepare( + self.model, self.optimizer, self.lr_scheduler + ) + elif self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: + # In this case we are in DDP + LOMO, which should be supported + self.optimizer = self.accelerator.prepare(self.optimizer) + + if self.is_fsdp_enabled: + self.model = self.model_wrapped = model + + # for the rest of this function `model` is the outside model, whether it was wrapped or not + if model is not self.model: + self.model_wrapped = model + + # backward compatibility + if self.is_deepspeed_enabled: + self.deepspeed = self.model_wrapped + + # ckpt loading + if resume_from_checkpoint is not None: + if self.is_deepspeed_enabled: + deepspeed_load_checkpoint( + self.model_wrapped, resume_from_checkpoint, load_module_strict=not _is_peft_model(self.model) + ) + elif is_sagemaker_mp_enabled() or self.is_fsdp_enabled: + self._load_from_checkpoint(resume_from_checkpoint, self.model_wrapped) + + # Check if saved optimizer or scheduler states exist + self._load_optimizer_and_scheduler(resume_from_checkpoint) + self._load_scaler(resume_from_checkpoint) + + # important: at this point: + # self.model is the Transformers Model + # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model), + # FSDP(Transformers Model), Dynamo Optimized Module(Transformers Model) etc. + + # Train! + logger.info("***** Running training *****") + logger.info(f" Num examples = {num_examples:,}") + logger.info(f" Num Epochs = {num_train_epochs:,}") + logger.info(f" Instantaneous batch size per device = {self.args.per_device_train_batch_size:,}") + if self.args.per_device_train_batch_size != self._train_batch_size: + logger.info(f" Training with DataParallel so batch size has been adjusted to: {self._train_batch_size:,}") + logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size:,}") + logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") + logger.info(f" Total optimization steps = {max_steps:,}") + logger.info(f" Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}") + + self.state.epoch = 0 + start_time = time.time() + epochs_trained = 0 + steps_trained_in_current_epoch = 0 + steps_trained_progress_bar = None + + # Check if continuing training from a checkpoint + if resume_from_checkpoint is not None and os.path.isfile( + os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME) + ): + self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)) + self.compare_trainer_and_checkpoint_args(self.args, self.state) + self._load_callback_state() + epochs_trained = int(self.state.global_step // num_update_steps_per_epoch) + if not args.ignore_data_skip: + steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch) + steps_trained_in_current_epoch *= args.gradient_accumulation_steps + else: + steps_trained_in_current_epoch = 0 + + logger.info(" Continuing training from checkpoint, will skip to saved global_step") + logger.info(f" Continuing training from epoch {epochs_trained}") + logger.info(f" Continuing training from global step {self.state.global_step}") + if not args.ignore_data_skip: + logger.info( + f" Will skip the first {epochs_trained} epochs then the first" + f" {steps_trained_in_current_epoch} batches in the first epoch." + ) + + # Update the references + for attr in ("model", "optimizer", "lr_scheduler"): + setattr(self.callback_handler, attr, getattr(self, attr)) + self.callback_handler.train_dataloader = train_dataloader + + self.state.init_training_references(self, max_steps, num_train_epochs, trial) + + # tr_loss is a tensor to avoid synchronization of TPUs through .item() + tr_loss = torch.tensor(0.0, device=args.device) + # _total_loss_scalar is updated everytime .item() has to be called on tr_loss and stores the sum of all losses + self._total_loss_scalar = 0.0 + self._globalstep_last_logged = self.state.global_step + model.zero_grad() + grad_norm: Optional[float] = None + learning_rate = None + self.control = self.callback_handler.on_train_begin(args, self.state, self.control) + + if args.eval_on_start: + self._evaluate(trial, ignore_keys_for_eval, skip_scheduler=True) + + for epoch in range(epochs_trained, num_train_epochs): + epoch_dataloader = train_dataloader + if hasattr(epoch_dataloader, "set_epoch"): + epoch_dataloader.set_epoch(epoch) + + # Reset the past mems state at the beginning of each epoch if necessary. + if args.past_index >= 0: + self._past = None + + steps_in_epoch = ( + len(epoch_dataloader) + if len_dataloader is not None + else args.max_steps * args.gradient_accumulation_steps + ) + self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control) + + if epoch == epochs_trained and resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0: + self._load_rng_state(resume_from_checkpoint) + + rng_to_sync = False + steps_skipped = 0 + if steps_trained_in_current_epoch > 0: + epoch_dataloader = skip_first_batches(epoch_dataloader, steps_trained_in_current_epoch) + steps_skipped = steps_trained_in_current_epoch + steps_trained_in_current_epoch = 0 + rng_to_sync = True + + step = -1 + epoch_iterator = iter(epoch_dataloader) + # We chunkify the epoch iterator into gradient accumulation steps `n` batches + remainder = num_examples % args.gradient_accumulation_steps + if remainder == 0: + remainder = args.gradient_accumulation_steps + update_step = -1 + total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 + if args.gradient_accumulation_steps == 1: + total_updates -= 1 + for _ in range(total_updates): + update_step += 1 + num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder + batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches, args.device) + for i, inputs in enumerate(batch_samples): + step += 1 + do_sync_step = (step + 1) % args.gradient_accumulation_steps == 0 or (step + 1) == steps_in_epoch + # Since we perform prefetching, we need to manually set sync_gradients + self.accelerator.gradient_state._set_sync_gradients(do_sync_step) + + if self.args.include_num_input_tokens_seen: + main_input_name = getattr(self.model, "main_input_name", "input_ids") + if main_input_name not in inputs: + logger.warning( + "Tried to track the number of tokens seen, however the current model is " + "not configured properly to know what item is the input. To fix this, add " + "a `main_input_name` attribute to the model class you are using." + ) + else: + input_tokens = inputs[main_input_name].numel() + input_tokens = torch.tensor(input_tokens, device=self.args.device, dtype=torch.int64) + self.state.num_input_tokens_seen += self.accelerator.gather(input_tokens).sum().item() + if rng_to_sync: + self._load_rng_state(resume_from_checkpoint) + rng_to_sync = False + + # Skip past any already trained steps if resuming training + if steps_trained_in_current_epoch > 0: + steps_trained_in_current_epoch -= 1 + if steps_trained_progress_bar is not None: + steps_trained_progress_bar.update(1) + if steps_trained_in_current_epoch == 0: + self._load_rng_state(resume_from_checkpoint) + continue + elif steps_trained_progress_bar is not None: + steps_trained_progress_bar.close() + steps_trained_progress_bar = None + + if step % args.gradient_accumulation_steps == 0: + self.control = self.callback_handler.on_step_begin(args, self.state, self.control) + + # We explicitly want to avoid relying on `accelerator.accumulate` for generation training + context = ( + functools.partial(self.accelerator.no_sync, model=model) + if i != len(batch_samples) - 1 + and self.accelerator.distributed_type != DistributedType.DEEPSPEED + else contextlib.nullcontext + ) + with context(): + tr_loss_step = self.training_step(model, inputs, num_items_in_batch) + + if ( + args.logging_nan_inf_filter + and not is_torch_xla_available() + and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step)) + ): + # if loss is nan or inf simply add the average of previous logged losses + tr_loss = tr_loss + tr_loss / (1 + self.state.global_step - self._globalstep_last_logged) + else: + if tr_loss.device != tr_loss_step.device: + raise ValueError( + f"Calculated loss must be on the original device: {tr_loss.device} but device in use is {tr_loss_step.device}" + ) + tr_loss = tr_loss + tr_loss_step + + self.current_flos += float(self.floating_point_ops(inputs)) + + if do_sync_step: + # Since we perform prefetching, we need to manually set sync_gradients to True + self.accelerator.gradient_state._set_sync_gradients(True) + + # Gradient clipping + if args.max_grad_norm is not None and args.max_grad_norm > 0: + if is_sagemaker_mp_enabled() and args.fp16: + _grad_norm = self.optimizer.clip_master_grads(args.max_grad_norm) + elif self.use_apex: + # Revert to normal clipping otherwise, handling Apex or full precision + _grad_norm = nn.utils.clip_grad_norm_( + amp.master_params(self.optimizer), + args.max_grad_norm, + ) + else: + _grad_norm = self.accelerator.clip_grad_norm_( + model.parameters(), + args.max_grad_norm, + ) + + if ( + is_accelerate_available() + and self.accelerator.distributed_type == DistributedType.DEEPSPEED + ): + grad_norm = model.get_global_grad_norm() + # In some cases the grad norm may not return a float + if hasattr(grad_norm, "item"): + grad_norm = grad_norm.item() + else: + grad_norm = _grad_norm + + self.control = self.callback_handler.on_pre_optimizer_step(args, self.state, self.control) + + self.optimizer.step() + + self.control = self.callback_handler.on_optimizer_step(args, self.state, self.control) + + # get leaning rate before update + learning_rate = self._get_learning_rate() + + if not self.accelerator.optimizer_step_was_skipped: + # Delay optimizer scheduling until metrics are generated + if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): + self.lr_scheduler.step() + + model.zero_grad() + self.state.global_step += 1 + self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch + self.control = self.callback_handler.on_step_end(args, self.state, self.control) + self._maybe_log_save_evaluate( + tr_loss, + grad_norm, + model, + trial, + epoch, + ignore_keys_for_eval, + start_time, + learning_rate=learning_rate, + ) + else: + self.control = self.callback_handler.on_substep_end(args, self.state, self.control) + + # PyTorch/XLA relies on the data loader to insert the mark_step for + # each step. Since we are breaking the loop early, we need to manually + # insert the mark_step here. + if self.control.should_epoch_stop or self.control.should_training_stop: + if is_torch_xla_available(): + xm.mark_step() + break + # We also need to break out of the nested loop + if self.control.should_epoch_stop or self.control.should_training_stop: + if is_torch_xla_available(): + xm.mark_step() + break + if step < 0: + logger.warning( + "There seems not to be a single sample in your epoch_iterator, stopping training at step" + f" {self.state.global_step}! This is expected if you're using an IterableDataset and set" + f" num_steps ({max_steps}) higher than the number of available samples." + ) + self.control.should_training_stop = True + + self.control = self.callback_handler.on_epoch_end(args, self.state, self.control) + self._maybe_log_save_evaluate( + tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time, learning_rate=learning_rate + ) + + if DebugOption.TPU_METRICS_DEBUG in self.args.debug: + if is_torch_xla_available(): + # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.) + xm.master_print(met.metrics_report()) + else: + logger.warning( + "You enabled PyTorch/XLA debug metrics but you don't have a TPU " + "configured. Check your training configuration if this is unexpected." + ) + if self.control.should_training_stop: + break + + if args.past_index and hasattr(self, "_past"): + # Clean the state at the end of training + delattr(self, "_past") + + logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n") + if args.load_best_model_at_end and self.state.best_model_checkpoint is not None: + # Wait for everyone to get here so we are sure the model has been saved by process 0. + if is_torch_xla_available(): + xm.rendezvous("load_best_model_at_end") + elif args.parallel_mode == ParallelMode.DISTRIBUTED: + dist.barrier() + elif is_sagemaker_mp_enabled(): + smp.barrier() + + self._load_best_model() + + # add remaining tr_loss + self._total_loss_scalar += tr_loss.item() + effective_global_step = max(self.state.global_step, 0.001) # Avoid ZeroDivisionError + train_loss = self._total_loss_scalar / effective_global_step + + metrics = speed_metrics( + "train", + start_time, + num_samples=num_train_samples, + num_steps=self.state.max_steps, + num_tokens=num_train_tokens, + ) + self.store_flos() + metrics["total_flos"] = self.state.total_flos + metrics["train_loss"] = train_loss + + self.is_in_train = False + + self._memory_tracker.stop_and_update_metrics(metrics) + + self.log(metrics) + + run_dir = self._get_output_dir(trial) + checkpoints_sorted = self._sorted_checkpoints(use_mtime=False, output_dir=run_dir) + + # Delete the last checkpoint when save_total_limit=1 if it's different from the best checkpoint and process allowed to save. + if self.args.should_save and self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1: + for checkpoint in checkpoints_sorted: + if not os.path.samefile(checkpoint, self.state.best_model_checkpoint): + logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit") + shutil.rmtree(checkpoint, ignore_errors=True) + + self.control = self.callback_handler.on_train_end(args, self.state, self.control) + + # Wait for the checkpoint to be uploaded. + self._finish_current_push() + + # After training we make sure to retrieve back the original forward pass method + # for the embedding layer by removing the forward post hook. + if self.neftune_noise_alpha is not None: + self._deactivate_neftune(self.model) + + return TrainOutput(self.state.global_step, train_loss, metrics) + + def _get_output_dir(self, trial): + if self.hp_search_backend is not None and trial is not None: + if self.hp_search_backend == HPSearchBackend.OPTUNA: + run_id = trial.number + elif self.hp_search_backend == HPSearchBackend.RAY: + import ray.train + + run_id = ray.train.get_context().get_trial_id() + elif self.hp_search_backend == HPSearchBackend.SIGOPT: + run_id = trial.id + elif self.hp_search_backend == HPSearchBackend.WANDB: + import wandb + + run_id = wandb.run.id + run_name = self.hp_name(trial) if self.hp_name is not None else f"run-{run_id}" + run_dir = os.path.join(self.args.output_dir, run_name) + else: + run_dir = self.args.output_dir + return run_dir + + def _load_from_checkpoint(self, resume_from_checkpoint, model=None): + if model is None: + model = self.model + + config_file = os.path.join(resume_from_checkpoint, CONFIG_NAME) + adapter_weights_file = os.path.join(resume_from_checkpoint, ADAPTER_WEIGHTS_NAME) + adapter_safe_weights_file = os.path.join(resume_from_checkpoint, ADAPTER_SAFE_WEIGHTS_NAME) + weights_file = os.path.join(resume_from_checkpoint, WEIGHTS_NAME) + weights_index_file = os.path.join(resume_from_checkpoint, WEIGHTS_INDEX_NAME) + safe_weights_file = os.path.join(resume_from_checkpoint, SAFE_WEIGHTS_NAME) + safe_weights_index_file = os.path.join(resume_from_checkpoint, SAFE_WEIGHTS_INDEX_NAME) + is_fsdp_ckpt = os.path.isdir(resume_from_checkpoint) and ( + # this checks the FSDP state dict when `SHARDED_STATE_DICT` is used + any( + FSDP_MODEL_NAME in folder_name + for folder_name in os.listdir(resume_from_checkpoint) + if os.path.isdir(os.path.join(resume_from_checkpoint, folder_name)) + ) + # this checks the FSDP state dict when `FULL_STATE_DICT` is used + or os.path.isfile(os.path.join(resume_from_checkpoint, f"{FSDP_MODEL_NAME}.bin")) + ) + # if multiple adapters exist, they get saved in sub directories + adapter_subdirs = ( + [ + folder_name + for folder_name in os.listdir(resume_from_checkpoint) + if os.path.isdir(os.path.join(resume_from_checkpoint, folder_name)) + and ( + os.path.isfile(os.path.join(resume_from_checkpoint, folder_name, ADAPTER_WEIGHTS_NAME)) + or os.path.isfile(os.path.join(resume_from_checkpoint, folder_name, ADAPTER_SAFE_WEIGHTS_NAME)) + ) + ] + if os.path.isdir(resume_from_checkpoint) + else [] + ) + + if is_fsdp_ckpt and not self.is_fsdp_enabled: + raise ValueError(f"Checkpoint found at {resume_from_checkpoint} is only supported when using PyTorch FSDP") + + if not ( + any( + os.path.isfile(f) + for f in [ + weights_file, + safe_weights_file, + weights_index_file, + safe_weights_index_file, + adapter_weights_file, + adapter_safe_weights_file, + ] + ) + or is_fsdp_ckpt + or adapter_subdirs + ): + raise ValueError(f"Can't find a valid checkpoint at {resume_from_checkpoint}") + + logger.info(f"Loading model from {resume_from_checkpoint}.") + + if os.path.isfile(config_file): + config = PretrainedConfig.from_json_file(config_file) + checkpoint_version = config.transformers_version + if checkpoint_version is not None and checkpoint_version != __version__: + logger.warning( + f"You are resuming training from a checkpoint trained with {checkpoint_version} of " + f"Transformers but your current version is {__version__}. This is not recommended and could " + "yield to errors or unwanted behaviors." + ) + + if os.path.isfile(weights_file) or os.path.isfile(safe_weights_file) or is_fsdp_ckpt: + # If the model is on the GPU, it still works! + if is_sagemaker_mp_enabled(): + if os.path.isfile(os.path.join(resume_from_checkpoint, "user_content.pt")): + # If the 'user_content.pt' file exists, load with the new smp api. + # Checkpoint must have been saved with the new smp api. + smp.resume_from_checkpoint( + path=resume_from_checkpoint, tag=WEIGHTS_NAME, partial=False, load_optimizer=False + ) + else: + # If the 'user_content.pt' file does NOT exist, load with the old smp api. + # Checkpoint must have been saved with the old smp api. + if hasattr(self.args, "fp16") and self.args.fp16 is True: + logger.warning( + "Enabling FP16 and loading from smp < 1.10 checkpoint together is not supported." + ) + state_dict = torch.load(weights_file, map_location="cpu", weights_only=True) + # Required for smp to not auto-translate state_dict from hf to smp (is already smp). + state_dict["_smp_is_partial"] = False + load_result = model.load_state_dict(state_dict, strict=True) + # release memory + del state_dict + elif self.is_fsdp_enabled: + load_fsdp_model( + self.accelerator.state.fsdp_plugin, + self.accelerator, + model, + resume_from_checkpoint, + **_get_fsdp_ckpt_kwargs(), + ) + else: + # We load the model state dict on the CPU to avoid an OOM error. + if self.args.save_safetensors and os.path.isfile(safe_weights_file): + state_dict = safetensors.torch.load_file(safe_weights_file, device="cpu") + else: + state_dict = torch.load(weights_file, map_location="cpu", weights_only=True) + + # workaround for FSDP bug https://github.com/pytorch/pytorch/issues/82963 + # which takes *args instead of **kwargs + load_result = model.load_state_dict(state_dict, False) + # release memory + del state_dict + self._issue_warnings_after_load(load_result) + + # Load adapters following PR # 24096 + elif _is_peft_model(model): + # If train a model using PEFT & LoRA, assume that adapter have been saved properly. + # TODO: in the future support only specific min PEFT versions + if (hasattr(model, "active_adapter") or hasattr(model, "active_adapters")) and hasattr( + model, "load_adapter" + ): + if os.path.exists(resume_from_checkpoint): + # For BC for older PEFT versions + if hasattr(model, "active_adapters"): + active_adapters = model.active_adapters + if len(active_adapters) > 1: + logger.warning("Multiple active adapters detected will only consider the first adapter") + active_adapter = active_adapters[0] + else: + active_adapter = model.active_adapter + + if adapter_subdirs: + for subdir_name in adapter_subdirs: + peft_id = os.path.join(resume_from_checkpoint, subdir_name) + model.load_adapter(peft_id, subdir_name, is_trainable=(subdir_name == active_adapter)) + model.set_adapter(active_adapter) + else: + model.load_adapter(resume_from_checkpoint, active_adapter, is_trainable=True) + else: + logger.warning( + "The intermediate checkpoints of PEFT may not be saved correctly, " + f"consider using a custom callback to save {ADAPTER_WEIGHTS_NAME} in corresponding saving folders. " + "Check some examples here: https://github.com/huggingface/peft/issues/96" + ) + else: + logger.warning("Could not load adapter model, make sure to have `peft>=0.3.0` installed") + else: + # We load the sharded checkpoint + load_result = load_sharded_checkpoint( + model, resume_from_checkpoint, strict=is_sagemaker_mp_enabled(), prefer_safe=self.args.save_safetensors + ) + if not is_sagemaker_mp_enabled(): + self._issue_warnings_after_load(load_result) + + def _load_best_model(self): + logger.info(f"Loading best model from {self.state.best_model_checkpoint} (score: {self.state.best_metric}).") + best_model_path = os.path.join(self.state.best_model_checkpoint, WEIGHTS_NAME) + best_safe_model_path = os.path.join(self.state.best_model_checkpoint, SAFE_WEIGHTS_NAME) + best_adapter_model_path = os.path.join(self.state.best_model_checkpoint, ADAPTER_WEIGHTS_NAME) + best_safe_adapter_model_path = os.path.join(self.state.best_model_checkpoint, ADAPTER_SAFE_WEIGHTS_NAME) + + model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model + if self.is_deepspeed_enabled: + deepspeed_load_checkpoint( + self.model_wrapped, + self.state.best_model_checkpoint, + load_module_strict=not _is_peft_model(self.model), + ) + elif self.is_fsdp_enabled: + load_result = load_fsdp_model( + self.accelerator.state.fsdp_plugin, + self.accelerator, + model, + self.state.best_model_checkpoint, + **_get_fsdp_ckpt_kwargs(), + ) + elif ( + os.path.exists(best_model_path) + or os.path.exists(best_safe_model_path) + or os.path.exists(best_adapter_model_path) + or os.path.exists(best_safe_adapter_model_path) + ): + has_been_loaded = True + if is_sagemaker_mp_enabled(): + if os.path.isfile(os.path.join(self.state.best_model_checkpoint, "user_content.pt")): + # If the 'user_content.pt' file exists, load with the new smp api. + # Checkpoint must have been saved with the new smp api. + smp.resume_from_checkpoint( + path=self.state.best_model_checkpoint, + tag=WEIGHTS_NAME, + partial=False, + load_optimizer=False, + ) + else: + # If the 'user_content.pt' file does NOT exist, load with the old smp api. + # Checkpoint must have been saved with the old smp api. + if self.args.save_safetensors and os.path.isfile(best_safe_model_path): + state_dict = safetensors.torch.load_file(best_safe_model_path, device="cpu") + else: + state_dict = torch.load(best_model_path, map_location="cpu", weights_only=True) + + state_dict["_smp_is_partial"] = False + load_result = model.load_state_dict(state_dict, strict=True) + else: + if _is_peft_model(model): + # If train a model using PEFT & LoRA, assume that adapter have been saved properly. + # TODO: in the future support only specific min PEFT versions + if (hasattr(model, "active_adapter") or hasattr(model, "active_adapters")) and hasattr( + model, "load_adapter" + ): + # For BC for older PEFT versions + if hasattr(model, "active_adapters"): + active_adapter = model.active_adapters[0] + if len(model.active_adapters) > 1: + logger.warning("Detected multiple active adapters, will only consider the first one") + else: + active_adapter = model.active_adapter + + if os.path.exists(best_adapter_model_path) or os.path.exists(best_safe_adapter_model_path): + try: + model.load_adapter(self.state.best_model_checkpoint, active_adapter) + except RuntimeError as exc: + if model.peft_config[active_adapter].is_prompt_learning: + # for context: https://github.com/huggingface/peft/issues/2256 + msg = ( + "When using prompt learning PEFT methods such as " + f"{model.peft_config[active_adapter].peft_type.value}, setting " + "load_best_model_at_end=True can lead to errors, it is recommended " + "to set this to False and to load the model manually from the checkpoint " + "directory using PeftModel.from_pretrained(base_model, ) after training " + "has finished." + ) + raise RuntimeError(msg) from exc + else: + raise + # Load_adapter has no return value present, modify it when appropriate. + from torch.nn.modules.module import _IncompatibleKeys + + load_result = _IncompatibleKeys([], []) + else: + logger.warning( + "The intermediate checkpoints of PEFT may not be saved correctly, " + f"consider using a custom callback to save {ADAPTER_WEIGHTS_NAME} in corresponding saving folders. " + "Check some examples here: https://github.com/huggingface/peft/issues/96" + ) + has_been_loaded = False + else: + logger.warning("Could not load adapter model, make sure to have `peft>=0.3.0` installed") + has_been_loaded = False + else: + # We load the model state dict on the CPU to avoid an OOM error. + if self.args.save_safetensors and os.path.isfile(best_safe_model_path): + state_dict = safetensors.torch.load_file(best_safe_model_path, device="cpu") + else: + state_dict = torch.load(best_model_path, map_location="cpu", weights_only=True) + + # If the model is on the GPU, it still works! + # workaround for FSDP bug https://github.com/pytorch/pytorch/issues/82963 + # which takes *args instead of **kwargs + load_result = model.load_state_dict(state_dict, False) + if not is_sagemaker_mp_enabled() and has_been_loaded: + self._issue_warnings_after_load(load_result) + elif os.path.exists(os.path.join(self.state.best_model_checkpoint, SAFE_WEIGHTS_INDEX_NAME)) or os.path.exists( + os.path.join(self.state.best_model_checkpoint, WEIGHTS_INDEX_NAME) + ): + load_result = load_sharded_checkpoint( + model, self.state.best_model_checkpoint, strict=is_sagemaker_mp_enabled() + ) + if not is_sagemaker_mp_enabled(): + self._issue_warnings_after_load(load_result) + else: + logger.warning( + f"Could not locate the best model at {best_model_path}, if you are running a distributed training " + "on multiple nodes, you should activate `--save_on_each_node`." + ) + + def _issue_warnings_after_load(self, load_result): + if len(load_result.missing_keys) != 0: + if self.model._keys_to_ignore_on_save is not None and set(load_result.missing_keys) == set( + self.model._keys_to_ignore_on_save + ): + self.model.tie_weights() + else: + logger.warning(f"There were missing keys in the checkpoint model loaded: {load_result.missing_keys}.") + if len(load_result.unexpected_keys) != 0: + logger.warning( + f"There were unexpected keys in the checkpoint model loaded: {load_result.unexpected_keys}." + ) + + def _evaluate(self, trial, ignore_keys_for_eval, skip_scheduler=False): + metrics = self.evaluate(ignore_keys=ignore_keys_for_eval) + self._report_to_hp_search(trial, self.state.global_step, metrics) + + # Run delayed LR scheduler now that metrics are populated + if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau) and not skip_scheduler: + metric_to_check = self.args.metric_for_best_model + if not metric_to_check.startswith("eval_"): + metric_to_check = f"eval_{metric_to_check}" + try: + self.lr_scheduler.step(metrics[metric_to_check]) + except KeyError as exc: + raise KeyError( + f"The `metric_for_best_model` training argument is set to '{metric_to_check}', " + f"which is not found in the evaluation metrics. " + f"The available evaluation metrics are: {list(metrics.keys())}. " + f"Please ensure that the `compute_metrics` function returns a dictionary that includes '{metric_to_check}' or " + f"consider changing the `metric_for_best_model` via the TrainingArguments." + ) from exc + return metrics + + def _maybe_log_save_evaluate( + self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time, learning_rate=None + ): + if self.control.should_log and self.state.global_step > self._globalstep_last_logged: + if is_torch_xla_available(): + xm.mark_step() + + logs: dict[str, float] = {} + + # all_gather + mean() to get average loss over all processes + tr_loss_scalar = self._nested_gather(tr_loss).mean().item() + + # reset tr_loss to zero + tr_loss -= tr_loss + + logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4) + if grad_norm is not None: + logs["grad_norm"] = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm + if learning_rate is not None: + logs["learning_rate"] = learning_rate + else: + logs["learning_rate"] = self._get_learning_rate() + + self._total_loss_scalar += tr_loss_scalar + self._globalstep_last_logged = self.state.global_step + self.store_flos() + + self.log(logs, start_time) + + metrics = None + if self.control.should_evaluate: + metrics = self._evaluate(trial, ignore_keys_for_eval) + is_new_best_metric = self._determine_best_metric(metrics=metrics, trial=trial) + + if self.args.save_strategy == SaveStrategy.BEST: + self.control.should_save = is_new_best_metric + + if self.control.should_save: + self._save_checkpoint(model, trial) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + + def _load_rng_state(self, checkpoint): + # Load RNG states from `checkpoint` + if checkpoint is None: + return + + if self.args.world_size > 1: + process_index = self.args.process_index + rng_file = os.path.join(checkpoint, f"rng_state_{process_index}.pth") + if not os.path.isfile(rng_file): + logger.info( + f"Didn't find an RNG file for process {process_index}, if you are resuming a training that " + "wasn't launched in a distributed fashion, reproducibility is not guaranteed." + ) + return + else: + rng_file = os.path.join(checkpoint, "rng_state.pth") + if not os.path.isfile(rng_file): + logger.info( + "Didn't find an RNG file, if you are resuming a training that was launched in a distributed " + "fashion, reproducibility is not guaranteed." + ) + return + + with safe_globals(): + checkpoint_rng_state = torch.load(rng_file, weights_only=True) + random.setstate(checkpoint_rng_state["python"]) + np.random.set_state(checkpoint_rng_state["numpy"]) + torch.random.set_rng_state(checkpoint_rng_state["cpu"]) + if is_torch_xla_available(): + xm.set_rng_state(checkpoint_rng_state["xla"]) + + is_distributed = self.args.parallel_mode == ParallelMode.DISTRIBUTED + if torch.cuda.is_available(): + set_rng_state_for_device("CUDA", torch.cuda, checkpoint_rng_state, is_distributed) + if is_torch_npu_available(): + set_rng_state_for_device("NPU", torch.npu, checkpoint_rng_state, is_distributed) + if is_torch_hpu_available(): + set_rng_state_for_device("HPU", torch.hpu, checkpoint_rng_state, is_distributed) + if is_torch_mlu_available(): + set_rng_state_for_device("MLU", torch.mlu, checkpoint_rng_state, is_distributed) + if is_torch_musa_available(): + set_rng_state_for_device("MUSA", torch.musa, checkpoint_rng_state, is_distributed) + + def _determine_best_metric(self, metrics, trial): + """ + Determine if the model should be saved based on the evaluation metrics. + + Returns: + bool: True if a new best metric was found, else False + """ + is_new_best_metric = False + + if self.args.metric_for_best_model is not None: + metric_to_check = self.args.metric_for_best_model + + if not metric_to_check.startswith("eval_"): + metric_to_check = f"eval_{metric_to_check}" + + try: + metric_value = metrics[metric_to_check] + except KeyError as exc: + raise KeyError( + f"The `metric_for_best_model` training argument is set to '{metric_to_check}', which is not found in the evaluation metrics. " + f"The available evaluation metrics are: {list(metrics.keys())}. Consider changing the `metric_for_best_model` via the TrainingArguments." + ) from exc + + operator = np.greater if self.args.greater_is_better else np.less + + if self.state.best_metric is None: + self.state.best_metric = float("-inf") if self.args.greater_is_better else float("inf") + + if operator(metric_value, self.state.best_metric): + self.state.best_metric = metric_value + + if self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH]: + self.state.best_global_step = self.state.global_step + + is_new_best_metric = True + + return is_new_best_metric + + def _save_checkpoint(self, model, trial): + # In all cases, including ddp/dp/deepspeed, self.model is always a reference to the model we + # want to save except FullyShardedDDP. + # assert unwrap_model(model) is self.model, "internal model should be a reference to self.model" + + # Save model checkpoint + checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}" + + if self.hp_search_backend is None and trial is None: + self.store_flos() + + run_dir = self._get_output_dir(trial=trial) + output_dir = os.path.join(run_dir, checkpoint_folder) + self.save_model(output_dir, _internal_call=True) + + if self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH] and self.state.best_global_step: + best_checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.best_global_step}" + best_checkpoint_dir = os.path.join(run_dir, best_checkpoint_folder) + + if os.path.exists(best_checkpoint_dir): + self.state.best_model_checkpoint = best_checkpoint_dir + + if not self.args.save_only_model: + # Save optimizer and scheduler + self._save_optimizer_and_scheduler(output_dir) + self._save_scaler(output_dir) + # Save RNG state + self._save_rng_state(output_dir) + + # Save the Trainer state + if self.args.should_save: + # Update `ExportableState` callbacks and `TrainerControl` state to where we are currently + for cb in [ + cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) + ]: + cb_name = cb.__class__.__name__ + cb_state = cb.state() + if isinstance(self.state.stateful_callbacks[cb_name], list): + self.state.stateful_callbacks[cb_name].append(cb_state) + else: + self.state.stateful_callbacks[cb_name] = cb_state + self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME)) + + if self.args.push_to_hub: + self._push_from_checkpoint(output_dir) + + # Maybe delete some older checkpoints. + if self.args.should_save: + # Solely rely on numerical checkpoint id for rotation. + # mtime is not reliable especially on some fuse fs in cloud environments. + self._rotate_checkpoints(use_mtime=False, output_dir=run_dir) + + def _save_rng_state(self, output_dir): + # Save RNG state in non-distributed training + rng_states = { + "python": random.getstate(), + "numpy": np.random.get_state(), + "cpu": torch.random.get_rng_state(), + } + if torch.cuda.is_available(): + if self.args.parallel_mode == ParallelMode.DISTRIBUTED: + # In non distributed, we save the global CUDA RNG state (will take care of DataParallel) + rng_states["cuda"] = torch.cuda.random.get_rng_state_all() + else: + rng_states["cuda"] = torch.cuda.random.get_rng_state() + + if is_torch_xla_available(): + rng_states["xla"] = xm.get_rng_state() + + if is_torch_npu_available(): + if self.args.parallel_mode == ParallelMode.DISTRIBUTED: + rng_states["npu"] = torch.npu.random.get_rng_state_all() + else: + rng_states["npu"] = torch.npu.random.get_rng_state() + + if is_torch_hpu_available(): + if self.args.parallel_mode == ParallelMode.DISTRIBUTED: + rng_states["hpu"] = torch.hpu.random.get_rng_state_all() + else: + rng_states["hpu"] = torch.hpu.random.get_rng_state() + + if is_torch_mlu_available(): + if self.args.parallel_mode == ParallelMode.DISTRIBUTED: + rng_states["mlu"] = torch.mlu.random.get_rng_state_all() + else: + rng_states["mlu"] = torch.mlu.random.get_rng_state() + + if is_torch_musa_available(): + if self.args.parallel_mode == ParallelMode.DISTRIBUTED: + rng_states["musa"] = torch.musa.get_rng_state_all() + else: + rng_states["musa"] = torch.musa.get_rng_state() + + # A process can arrive here before the process 0 has a chance to save the model, in which case output_dir may + # not yet exist. + os.makedirs(output_dir, exist_ok=True) + + if self.args.world_size <= 1: + torch.save(rng_states, os.path.join(output_dir, "rng_state.pth")) + else: + torch.save(rng_states, os.path.join(output_dir, f"rng_state_{self.args.process_index}.pth")) + + def _save_optimizer_and_scheduler(self, output_dir): + if is_torch_xla_available(): + xm.rendezvous("saving_optimizer_states") + if self.is_fsdp_xla_v1_enabled: + optm = { + "optimizer": self.optimizer.state_dict(), + "shard_metadata": self.model.get_shard_metadata(), + } + xm.save( + optm, + os.path.join( + output_dir, f"rank{self.args.process_index}-of-{self.args.world_size}-{OPTIMIZER_NAME}" + ), + master_only=False, + ) + else: + xm.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME)) + with warnings.catch_warnings(record=True) as caught_warnings: + xm.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME)) + reissue_pt_warnings(caught_warnings) + elif is_sagemaker_mp_enabled(): + opt_state_dict = self.optimizer.local_state_dict(gather_if_shard=False) + smp.barrier() + if smp.rdp_rank() == 0 or smp.state.cfg.shard_optimizer_state: + smp.save( + opt_state_dict, + os.path.join(output_dir, OPTIMIZER_NAME), + partial=True, + v3=smp.state.cfg.shard_optimizer_state, + ) + elif self.is_deepspeed_enabled: + # under zero3 model file itself doesn't get saved since it's bogus! Unless deepspeed + # config `stage3_gather_16bit_weights_on_model_save` is True + accept_exclude_frozen_parameters = "exclude_frozen_parameters" in set( + inspect.signature(self.model_wrapped.save_checkpoint).parameters.keys() + ) + if accept_exclude_frozen_parameters and _is_peft_model(self.model): + self.model_wrapped.save_checkpoint(output_dir, exclude_frozen_parameters=True) + else: + self.model_wrapped.save_checkpoint(output_dir) + elif self.is_fsdp_enabled: + # save fsdp specific ckpt for resuming from ckpt + save_fsdp_model( + self.accelerator.state.fsdp_plugin, self.accelerator, self.model, output_dir, **_get_fsdp_ckpt_kwargs() + ) + save_fsdp_optimizer( + self.accelerator.state.fsdp_plugin, self.accelerator, self.optimizer, self.model, output_dir + ) + elif self.args.should_save: + # deepspeed.save_checkpoint above saves model/optim/sched + torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME)) + + # Save SCHEDULER & SCALER + is_deepspeed_custom_scheduler = self.is_deepspeed_enabled and not isinstance( + self.lr_scheduler, DeepSpeedSchedulerWrapper + ) + if ( + self.args.should_save + and (not self.is_deepspeed_enabled or is_deepspeed_custom_scheduler) + and not is_torch_xla_available() + ): + with warnings.catch_warnings(record=True) as caught_warnings: + torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME)) + reissue_pt_warnings(caught_warnings) + + def _load_optimizer_and_scheduler(self, checkpoint): + """If optimizer and scheduler states exist, load them.""" + if checkpoint is None: + return + + if self.is_deepspeed_enabled: + # deepspeed loads optimizer/lr_scheduler together with the model in deepspeed_init + if not isinstance(self.lr_scheduler, DeepSpeedSchedulerWrapper): + with warnings.catch_warnings(record=True) as caught_warnings: + self.lr_scheduler.load_state_dict( + torch.load(os.path.join(checkpoint, SCHEDULER_NAME), weights_only=True) + ) + reissue_pt_warnings(caught_warnings) + return + + checkpoint_file_exists = ( + glob.glob(os.path.join(checkpoint, OPTIMIZER_NAME) + "_*") + if is_sagemaker_mp_enabled() + else ( + os.path.isfile(os.path.join(checkpoint, OPTIMIZER_NAME)) + or os.path.isfile(os.path.join(checkpoint, OPTIMIZER_NAME_BIN)) + or ( + os.path.isdir(checkpoint) + and any( + OPTIMIZER_NAME_BIN.split(".")[0] in folder_name + for folder_name in os.listdir(checkpoint) + if os.path.isdir(os.path.join(checkpoint, folder_name)) + ) + ) + ) + ) + checkpoint_file_exists = ( + glob.glob(os.path.join(checkpoint, f"rank*-of-{self.args.world_size}-{OPTIMIZER_NAME}")) + if self.is_fsdp_xla_v1_enabled + else checkpoint_file_exists + ) + if checkpoint_file_exists and os.path.isfile(os.path.join(checkpoint, SCHEDULER_NAME)): + # Load in optimizer and scheduler states + if is_torch_xla_available(): + # On TPU we have to take some extra precautions to properly load the states on the right device. + if self.is_fsdp_xla_v1_enabled: + optimizer_state = torch.load( + os.path.join( + checkpoint, f"rank{self.args.process_index}-of-{self.args.world_size}-{OPTIMIZER_NAME}" + ), + map_location="cpu", + weights_only=True, + ) + # We only need `optimizer` when resuming from checkpoint + optimizer_state = optimizer_state["optimizer"] + else: + optimizer_state = torch.load( + os.path.join(checkpoint, OPTIMIZER_NAME), map_location="cpu", weights_only=True + ) + with warnings.catch_warnings(record=True) as caught_warnings: + lr_scheduler_state = torch.load( + os.path.join(checkpoint, SCHEDULER_NAME), map_location="cpu", weights_only=True + ) + reissue_pt_warnings(caught_warnings) + + xm.send_cpu_data_to_device(optimizer_state, self.args.device) + xm.send_cpu_data_to_device(lr_scheduler_state, self.args.device) + + self.optimizer.load_state_dict(optimizer_state) + self.lr_scheduler.load_state_dict(lr_scheduler_state) + else: + if is_sagemaker_mp_enabled(): + if os.path.isfile(os.path.join(checkpoint, "user_content.pt")): + # Optimizer checkpoint was saved with smp >= 1.10 + def opt_load_hook(mod, opt): + opt.load_state_dict(smp.load(os.path.join(checkpoint, OPTIMIZER_NAME), partial=True)) + + else: + # Optimizer checkpoint was saved with smp < 1.10 + def opt_load_hook(mod, opt): + if IS_SAGEMAKER_MP_POST_1_10: + opt.load_state_dict( + smp.load(os.path.join(checkpoint, OPTIMIZER_NAME), partial=True, back_compat=True) + ) + else: + opt.load_state_dict(smp.load(os.path.join(checkpoint, OPTIMIZER_NAME), partial=True)) + + self.model_wrapped.register_post_step_hook(opt_load_hook) + else: + # We use the CPU when training on one GPU to avoid OOM for GPU RAM when training big models. + # In distributed training however, we load directly on each GPU and risk the GPU OOM as it's more + # likely to get OOM on CPU (since we load num_gpu times the optimizer state + map_location = self.args.device if self.args.world_size > 1 else "cpu" + if self.is_fsdp_enabled: + load_fsdp_optimizer( + self.accelerator.state.fsdp_plugin, + self.accelerator, + self.optimizer, + self.model, + checkpoint, + **_get_fsdp_ckpt_kwargs(), + ) + else: + self.optimizer.load_state_dict( + torch.load( + os.path.join(checkpoint, OPTIMIZER_NAME), map_location=map_location, weights_only=True + ) + ) + with warnings.catch_warnings(record=True) as caught_warnings: + self.lr_scheduler.load_state_dict( + torch.load(os.path.join(checkpoint, SCHEDULER_NAME), weights_only=True) + ) + reissue_pt_warnings(caught_warnings) + + def _save_scaler(self, output_dir): + # See if there is a scaler attribute + try: + scaler = self.accelerator.scaler + except AttributeError: + return + if scaler is None: + return + if is_torch_xla_available(): + xm.rendezvous("saving_scaler_state") + with warnings.catch_warnings(record=True) as caught_warnings: + xm.save(self.accelerator.scaler.state_dict(), os.path.join(output_dir, SCALER_NAME)) + reissue_pt_warnings(caught_warnings) + + # Save SCALER + if self.args.should_save and not is_torch_xla_available(): + with warnings.catch_warnings(record=True) as caught_warnings: + torch.save(self.accelerator.scaler.state_dict(), os.path.join(output_dir, SCALER_NAME)) + reissue_pt_warnings(caught_warnings) + + def _load_scaler(self, checkpoint): + """If scaler state exists, load it.""" + if checkpoint is None: + return + + checkpoint_file_exists = os.path.isfile(os.path.join(checkpoint, SCALER_NAME)) + + if checkpoint_file_exists: + # On TPU we have to take some extra precautions to properly load the states on the right device. + # Load in scaler states + if is_torch_xla_available(): + with warnings.catch_warnings(record=True) as caught_warnings: + scaler_state = torch.load( + os.path.join(checkpoint, SCALER_NAME), map_location="cpu", weights_only=True + ) + reissue_pt_warnings(caught_warnings) + xm.send_cpu_data_to_device(scaler_state, self.args.device) + self.accelerator.scaler.load_state_dict(scaler_state) + else: + with warnings.catch_warnings(record=True) as caught_warnings: + self.accelerator.scaler.load_state_dict( + torch.load(os.path.join(checkpoint, SCALER_NAME), weights_only=True) + ) + reissue_pt_warnings(caught_warnings) + + def _load_callback_state(self): + """If callback states exist and were passed in, restore their states if enabled""" + if not self.args.restore_callback_states_from_checkpoint: + return + # Callback states are stored in stateful_callbacks + not_found = [] + new_callbacks = [] + original_callbacks = self.callback_handler.callbacks + [self.control] + for stored_callback, data in self.state.stateful_callbacks.items(): + if not isinstance(data, list): + data = [data] + if any(callback.__class__.__name__ == stored_callback for callback in original_callbacks): + # We can load/restore from multiple callbacks of the same type. + duplicates = [ + callback for callback in original_callbacks if callback.__class__.__name__ == stored_callback + ] + for callback, callback_data in zip(duplicates, data): + args = callback_data.get("args", {}) + attributes = callback_data.get("attributes", {}) + new_callback = type(callback)(**args) + for attribute, value in attributes.items(): + setattr(new_callback, attribute, value) + if isinstance(callback, TrainerControl): + # Specifically for restoring the `control` state + self.control = new_callback + else: + new_callbacks.append(new_callback) + # We remove the existing callback and add it to the list of new callbacks + self.callback_handler.remove_callback(type(new_callback)) + logger.info("Continuing training from checkpoint, restoring any callbacks that were passed in") + else: + not_found.append(stored_callback) + if len(not_found) > 0: + logger.warning( + f"Checkpoint included callbacks not included in current configuration. Ignoring. ({', '.join(not_found)})" + ) + for callback in new_callbacks: + self.callback_handler.add_callback(callback) + + def hyperparameter_search( + self, + hp_space: Optional[Callable[["optuna.Trial"], dict[str, float]]] = None, + compute_objective: Optional[Callable[[dict[str, float]], float]] = None, + n_trials: int = 20, + direction: Union[str, list[str]] = "minimize", + backend: Optional[Union["str", HPSearchBackend]] = None, + hp_name: Optional[Callable[["optuna.Trial"], str]] = None, + **kwargs, + ) -> Union[BestRun, list[BestRun]]: + """ + Launch an hyperparameter search using `optuna` or `Ray Tune` or `SigOpt`. The optimized quantity is determined + by `compute_objective`, which defaults to a function returning the evaluation loss when no metric is provided, + the sum of all metrics otherwise. + + + + To use this method, you need to have provided a `model_init` when initializing your [`Trainer`]: we need to + reinitialize the model at each new run. This is incompatible with the `optimizers` argument, so you need to + subclass [`Trainer`] and override the method [`~Trainer.create_optimizer_and_scheduler`] for custom + optimizer/scheduler. + + + + Args: + hp_space (`Callable[["optuna.Trial"], Dict[str, float]]`, *optional*): + A function that defines the hyperparameter search space. Will default to + [`~trainer_utils.default_hp_space_optuna`] or [`~trainer_utils.default_hp_space_ray`] or + [`~trainer_utils.default_hp_space_sigopt`] depending on your backend. + compute_objective (`Callable[[Dict[str, float]], float]`, *optional*): + A function computing the objective to minimize or maximize from the metrics returned by the `evaluate` + method. Will default to [`~trainer_utils.default_compute_objective`]. + n_trials (`int`, *optional*, defaults to 100): + The number of trial runs to test. + direction (`str` or `List[str]`, *optional*, defaults to `"minimize"`): + If it's single objective optimization, direction is `str`, can be `"minimize"` or `"maximize"`, you + should pick `"minimize"` when optimizing the validation loss, `"maximize"` when optimizing one or + several metrics. If it's multi objectives optimization, direction is `List[str]`, can be List of + `"minimize"` and `"maximize"`, you should pick `"minimize"` when optimizing the validation loss, + `"maximize"` when optimizing one or several metrics. + backend (`str` or [`~training_utils.HPSearchBackend`], *optional*): + The backend to use for hyperparameter search. Will default to optuna or Ray Tune or SigOpt, depending + on which one is installed. If all are installed, will default to optuna. + hp_name (`Callable[["optuna.Trial"], str]]`, *optional*): + A function that defines the trial/run name. Will default to None. + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments for each backend: + + - `optuna`: parameters from + [optuna.study.create_study](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html) + and also the parameters `timeout`, `n_jobs` and `gc_after_trial` from + [optuna.study.Study.optimize](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.Study.html#optuna.study.Study.optimize) + - `ray`: parameters from [tune.run](https://docs.ray.io/en/latest/tune/api_docs/execution.html#tune-run). + If `resources_per_trial` is not set in the `kwargs`, it defaults to 1 CPU core and 1 GPU (if available). + If `progress_reporter` is not set in the `kwargs`, + [ray.tune.CLIReporter](https://docs.ray.io/en/latest/tune/api/doc/ray.tune.CLIReporter.html) is used. + - `sigopt`: the parameter `proxies` from + [sigopt.Connection.set_proxies](https://docs.sigopt.com/support/faq#how-do-i-use-sigopt-with-a-proxy). + + Returns: + [`trainer_utils.BestRun` or `List[trainer_utils.BestRun]`]: All the information about the best run or best + runs for multi-objective optimization. Experiment summary can be found in `run_summary` attribute for Ray + backend. + """ + if backend is None: + backend = default_hp_search_backend() + backend = HPSearchBackend(backend) + backend_obj = ALL_HYPERPARAMETER_SEARCH_BACKENDS[backend]() + backend_obj.ensure_available() + self.hp_search_backend = backend + if self.model_init is None: + raise RuntimeError( + "To use hyperparameter search, you need to pass your model through a model_init function." + ) + + self.hp_space = backend_obj.default_hp_space if hp_space is None else hp_space + self.hp_name = hp_name + self.compute_objective = default_compute_objective if compute_objective is None else compute_objective + + best_run = backend_obj.run(self, n_trials, direction, **kwargs) + + self.hp_search_backend = None + return best_run + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + """ + Log `logs` on the various objects watching training. + + Subclass and override this method to inject custom behavior. + + Args: + logs (`Dict[str, float]`): + The values to log. + start_time (`Optional[float]`): + The start of training. + """ + if self.state.epoch is not None: + logs["epoch"] = self.state.epoch + if self.args.include_num_input_tokens_seen: + logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen + if start_time is not None: + speed_metrics("train", start_time, num_tokens=self.state.num_input_tokens_seen) + + output = {**logs, **{"step": self.state.global_step}} + self.state.log_history.append(output) + self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs) + + def _prepare_input(self, data: Union[torch.Tensor, Any]) -> Union[torch.Tensor, Any]: + """ + Prepares one `data` before feeding it to the model, be it a tensor or a nested list/dictionary of tensors. + """ + if isinstance(data, Mapping): + return type(data)({k: self._prepare_input(v) for k, v in data.items()}) + elif isinstance(data, (tuple, list)): + return type(data)(self._prepare_input(v) for v in data) + elif isinstance(data, torch.Tensor): + kwargs = {"device": self.args.device} + if self.is_deepspeed_enabled and (torch.is_floating_point(data) or torch.is_complex(data)): + # NLP models inputs are int/uint and those get adjusted to the right dtype of the + # embedding. Other models such as wav2vec2's inputs are already float and thus + # may need special handling to match the dtypes of the model + kwargs.update({"dtype": self.accelerator.state.deepspeed_plugin.hf_ds_config.dtype()}) + return data.to(**kwargs) + return data + + def _prepare_inputs(self, inputs: dict[str, Union[torch.Tensor, Any]]) -> dict[str, Union[torch.Tensor, Any]]: + """ + Prepare `inputs` before feeding them to the model, converting them to tensors if they are not already and + handling potential state. + """ + inputs = self._prepare_input(inputs) + if len(inputs) == 0: + raise ValueError( + "The batch received was empty, your model won't be able to train on it. Double-check that your " + f"training dataset contains keys expected by the model: {','.join(self._signature_columns)}." + ) + if self.args.past_index >= 0 and self._past is not None: + inputs["mems"] = self._past + + return inputs + + def compute_loss_context_manager(self): + """ + A helper wrapper to group together context managers. + """ + return self.autocast_smart_context_manager() + + def autocast_smart_context_manager(self, cache_enabled: Optional[bool] = True): + """ + A helper wrapper that creates an appropriate context manager for `autocast` while feeding it the desired + arguments, depending on the situation. + """ + if self.use_cpu_amp: + ctx_manager = torch.cpu.amp.autocast(cache_enabled=cache_enabled, dtype=self.amp_dtype) + else: + ctx_manager = contextlib.nullcontext() + + return ctx_manager + + def training_step( + self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch=None + ) -> torch.Tensor: + """ + Perform a training step on a batch of inputs. + + Subclass and override to inject custom behavior. + + Args: + model (`nn.Module`): + The model to train. + inputs (`Dict[str, Union[torch.Tensor, Any]]`): + The inputs and targets of the model. + + The dictionary will be unpacked before being fed to the model. Most models expect the targets under the + argument `labels`. Check your model's documentation for all accepted arguments. + + Return: + `torch.Tensor`: The tensor with training loss on this batch. + """ + model.train() + if hasattr(self.optimizer, "train") and callable(self.optimizer.train): + self.optimizer.train() + + inputs = self._prepare_inputs(inputs) + if is_sagemaker_mp_enabled(): + loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps) + return loss_mb.reduce_mean().detach().to(self.args.device) + + with self.compute_loss_context_manager(): + loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch) + + del inputs + if ( + self.args.torch_empty_cache_steps is not None + and self.state.global_step % self.args.torch_empty_cache_steps == 0 + ): + if is_torch_xpu_available(): + torch.xpu.empty_cache() + elif is_torch_mlu_available(): + torch.mlu.empty_cache() + elif is_torch_musa_available(): + torch.musa.empty_cache() + elif is_torch_npu_available(): + torch.npu.empty_cache() + elif is_torch_mps_available(min_version="2.0"): + torch.mps.empty_cache() + elif is_torch_hpu_available(): + logger.warning( + "`torch_empty_cache_steps` is set but HPU device/backend does not support empty_cache()." + ) + else: + torch.cuda.empty_cache() + + kwargs = {} + + # For LOMO optimizers you need to explicitly use the learnign rate + if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: + kwargs["learning_rate"] = self._get_learning_rate() + + if self.args.n_gpu > 1: + loss = loss.mean() # mean() to average on multi-gpu parallel training + + if self.use_apex: + with amp.scale_loss(loss, self.optimizer) as scaled_loss: + scaled_loss.backward() + else: + # Finally we need to normalize the loss for reporting + if not self.model_accepts_loss_kwargs and self.compute_loss_func is None: + loss = loss / self.args.gradient_accumulation_steps + + # Turning off loss scaling w.r.t. gradient accumulation when DeepSpeed is enabled + # https://github.com/huggingface/transformers/pull/35808 + if self.accelerator.distributed_type == DistributedType.DEEPSPEED: + kwargs["scale_wrt_gas"] = False + + self.accelerator.backward(loss, **kwargs) + + return loss.detach() + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + """ + How the loss is computed by Trainer. By default, all models return the loss in the first element. + + Subclass and override for custom behavior. + """ + if (self.label_smoother is not None or self.compute_loss_func is not None) and "labels" in inputs: + labels = inputs.pop("labels") + else: + labels = None + if self.model_accepts_loss_kwargs: + loss_kwargs = {} + if num_items_in_batch is not None: + loss_kwargs["num_items_in_batch"] = num_items_in_batch + inputs = {**inputs, **loss_kwargs} + outputs = model(**inputs) + # Save past state if it exists + # TODO: this needs to be fixed and made cleaner later. + if self.args.past_index >= 0: + self._past = outputs[self.args.past_index] + + if labels is not None: + unwrapped_model = self.accelerator.unwrap_model(model) + if _is_peft_model(unwrapped_model): + model_name = unwrapped_model.base_model.model._get_name() + else: + model_name = unwrapped_model._get_name() + # User-defined compute_loss function + if self.compute_loss_func is not None: + loss = self.compute_loss_func(outputs, labels, num_items_in_batch=num_items_in_batch) + elif model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values(): + loss = self.label_smoother(outputs, labels, shift_labels=True) + else: + loss = self.label_smoother(outputs, labels) + else: + if isinstance(outputs, dict) and "loss" not in outputs: + raise ValueError( + "The model did not return a loss from the inputs, only the following keys: " + f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}." + ) + # We don't use .loss here since the model may return tuples instead of ModelOutput. + loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0] + + if ( + self.args.average_tokens_across_devices + and (self.model_accepts_loss_kwargs or self.compute_loss_func) + and num_items_in_batch is not None + ): + loss *= self.accelerator.num_processes + + return (loss, outputs) if return_outputs else loss + + def is_local_process_zero(self) -> bool: + """ + Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several + machines) main process. + """ + return self.args.local_process_index == 0 + + def is_world_process_zero(self) -> bool: + """ + Whether or not this process is the global main process (when training in a distributed fashion on several + machines, this is only going to be `True` for one process). + """ + # Special case for SageMaker ModelParallel since there process_index is dp_process_index, not the global + # process index. + if is_sagemaker_mp_enabled(): + return smp.rank() == 0 + else: + return self.args.process_index == 0 + + def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False): + """ + Will save the model, so you can reload it using `from_pretrained()`. + + Will only save from the main process. + """ + + if output_dir is None: + output_dir = self.args.output_dir + + if is_torch_xla_available(): + self._save_tpu(output_dir) + elif is_sagemaker_mp_enabled(): + # Calling the state_dict needs to be done on the wrapped model and on all processes. + os.makedirs(output_dir, exist_ok=True) + state_dict = self.model_wrapped.state_dict() + if self.args.should_save: + self._save(output_dir, state_dict=state_dict) + if IS_SAGEMAKER_MP_POST_1_10: + # 'user_content.pt' indicates model state_dict saved with smp >= 1.10 + Path(os.path.join(output_dir, "user_content.pt")).touch() + elif self.is_fsdp_enabled: + if ("FULL_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type)) and ( + version.parse(accelerate_version) > version.parse("0.24.1") + ): + state_dict = self.accelerator.get_state_dict(self.model) + if self.args.should_save: + self._save(output_dir, state_dict=state_dict) + elif self.is_deepspeed_enabled: + try: + state_dict = self.accelerator.get_state_dict(self.deepspeed) + if self.args.should_save: + self._save(output_dir, state_dict=state_dict) + except ValueError: + logger.warning( + " stage3_gather_16bit_weights_on_model_save=false. Saving the full checkpoint instead, use" + " zero_to_fp32.py to recover weights" + ) + if self.args.should_save: + self._save(output_dir, state_dict={}) + # remove the dummy state_dict + remove_dummy_checkpoint(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME]) + self.model_wrapped.save_checkpoint(output_dir) + + elif self.args.should_save: + self._save(output_dir) + + # Push to the Hub when `save_model` is called by the user. + if self.args.push_to_hub and not _internal_call: + self.push_to_hub(commit_message="Model save") + + def _save_tpu(self, output_dir: Optional[str] = None): + output_dir = output_dir if output_dir is not None else self.args.output_dir + + logger.info(f"Saving model checkpoint to {output_dir}") + model = self.model + xm.mark_step() + + if xm.is_master_ordinal(local=False): + os.makedirs(output_dir, exist_ok=True) + torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) + + # Save a trained model and configuration using `save_pretrained()`. + # They can then be reloaded using `from_pretrained()` + supported_classes = (PushToHubMixin,) + xm.rendezvous("saving_checkpoint") + if self.is_fsdp_xla_v1_enabled: + ckpt = { + "model": model.state_dict(), + "shard_metadata": model.get_shard_metadata(), + } + ckpt_path = os.path.join( + output_dir, f"rank{self.args.process_index}-of-{self.args.world_size}-{WEIGHTS_NAME}" + ) + # All ranks save sharded checkpoint + xm.save(ckpt, ckpt_path, master_only=False) + # Make sure all ranks have saved checkpoints + xm.rendezvous("save_full_checkpoints") + # Master save full checkpoint + if self.args.should_save: + from torch_xla.distributed.fsdp import consolidate_sharded_model_checkpoints + + full_state_dict, _ = consolidate_sharded_model_checkpoints( + ckpt_prefix=os.path.join(output_dir, ""), + ckpt_suffix=f"rank*-of-*-{WEIGHTS_NAME}", + save_model=False, + ) + model = model.module.module + unwrapped_model = self.accelerator.unwrap_model(model) + if isinstance(unwrapped_model, supported_classes): + unwrapped_model.save_pretrained( + output_dir, + state_dict=full_state_dict, + save_function=xm.save, + safe_serialization=self.args.save_safetensors, + ) + else: + logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.") + xm.save(full_state_dict, os.path.join(output_dir, WEIGHTS_NAME)) + elif not isinstance(model, supported_classes): + if isinstance(self.accelerator.unwrap_model(model), supported_classes): + self.accelerator.unwrap_model(model).save_pretrained( + output_dir, + is_main_process=self.args.should_save, + state_dict=xm._maybe_convert_to_cpu(model.state_dict()), + save_function=xm.save, + safe_serialization=self.args.save_safetensors, + ) + else: + logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.") + state_dict = xm._maybe_convert_to_cpu(model.state_dict()) + xm.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME)) + else: + model.save_pretrained( + output_dir, + is_main_process=self.args.should_save, + save_function=xm.save, + safe_serialization=self.args.save_safetensors, + state_dict=xm._maybe_convert_to_cpu(model.state_dict()), + ) + if self.processing_class is not None and self.args.should_save: + self.processing_class.save_pretrained(output_dir) + + def _save(self, output_dir: Optional[str] = None, state_dict=None): + # If we are executing this function, we are the process zero, so we don't check for that. + output_dir = output_dir if output_dir is not None else self.args.output_dir + os.makedirs(output_dir, exist_ok=True) + logger.info(f"Saving model checkpoint to {output_dir}") + + supported_classes = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel) + # Save a trained model and configuration using `save_pretrained()`. + # They can then be reloaded using `from_pretrained()` + if not isinstance(self.model, supported_classes): + if state_dict is None: + state_dict = self.model.state_dict() + + if isinstance(self.accelerator.unwrap_model(self.model), supported_classes): + self.accelerator.unwrap_model(self.model).save_pretrained( + output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors + ) + else: + logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.") + if self.args.save_safetensors: + safetensors.torch.save_file( + state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME), metadata={"format": "pt"} + ) + else: + torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME)) + else: + self.model.save_pretrained( + output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors + ) + + if self.processing_class is not None: + self.processing_class.save_pretrained(output_dir) + elif ( + self.data_collator is not None + and hasattr(self.data_collator, "tokenizer") + and self.data_collator.tokenizer is not None + ): + logger.info("Saving Trainer.data_collator.tokenizer by default as Trainer.processing_class is `None`") + self.data_collator.tokenizer.save_pretrained(output_dir) + + # Good practice: save your training arguments together with the trained model + torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) + + def store_flos(self): + # Storing the number of floating-point operations that went into the model + if self.args.parallel_mode == ParallelMode.DISTRIBUTED: + self.state.total_flos += ( + distributed_broadcast_scalars([self.current_flos], device=self.args.device).sum().item() + ) + self.current_flos = 0 + else: + self.state.total_flos += self.current_flos + self.current_flos = 0 + + def _sorted_checkpoints( + self, output_dir=None, checkpoint_prefix=PREFIX_CHECKPOINT_DIR, use_mtime=False + ) -> list[str]: + ordering_and_checkpoint_path = [] + + glob_checkpoints = [str(x) for x in Path(output_dir).glob(f"{checkpoint_prefix}-*") if os.path.isdir(x)] + + for path in glob_checkpoints: + if use_mtime: + ordering_and_checkpoint_path.append((os.path.getmtime(path), path)) + else: + regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path) + if regex_match is not None and regex_match.groups() is not None: + ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path)) + + checkpoints_sorted = sorted(ordering_and_checkpoint_path) + checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted] + # Make sure we don't delete the best model. + if ( + self.state.best_model_checkpoint is not None + and str(Path(self.state.best_model_checkpoint)) in checkpoints_sorted + ): + best_model_index = checkpoints_sorted.index(str(Path(self.state.best_model_checkpoint))) + for i in range(best_model_index, len(checkpoints_sorted) - 2): + checkpoints_sorted[i], checkpoints_sorted[i + 1] = checkpoints_sorted[i + 1], checkpoints_sorted[i] + return checkpoints_sorted + + def _rotate_checkpoints(self, use_mtime=False, output_dir=None) -> None: + if self.args.save_total_limit is None or self.args.save_total_limit <= 0: + return + + # Check if we should delete older checkpoint(s) + checkpoints_sorted = self._sorted_checkpoints(use_mtime=use_mtime, output_dir=output_dir) + if len(checkpoints_sorted) <= self.args.save_total_limit: + return + + # If save_total_limit=1 with load_best_model_at_end=True, we could end up deleting the last checkpoint, which + # we don't do to allow resuming. + save_total_limit = self.args.save_total_limit + if ( + self.state.best_model_checkpoint is not None + and self.args.save_total_limit == 1 + and checkpoints_sorted[-1] != self.state.best_model_checkpoint + ): + save_total_limit = 2 + + number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit) + checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete] + for checkpoint in checkpoints_to_be_deleted: + logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit") + shutil.rmtree(checkpoint, ignore_errors=True) + + def evaluate( + self, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + ignore_keys: Optional[list[str]] = None, + metric_key_prefix: str = "eval", + ) -> dict[str, float]: + """ + Run evaluation and returns metrics. + + The calling script will be responsible for providing a method to compute metrics, as they are task-dependent + (pass it to the init `compute_metrics` argument). + + You can also subclass and override this method to inject custom behavior. + + Args: + eval_dataset (Union[`Dataset`, Dict[str, `Dataset`]), *optional*): + Pass a dataset if you wish to override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns + not accepted by the `model.forward()` method are automatically removed. If it is a dictionary, it will + evaluate on each dataset, prepending the dictionary key to the metric name. Datasets must implement the + `__len__` method. + + + + If you pass a dictionary with names of datasets as keys and datasets as values, evaluate will run + separate evaluations on each dataset. This can be useful to monitor how training affects other + datasets or simply to get a more fine-grained evaluation. + When used with `load_best_model_at_end`, make sure `metric_for_best_model` references exactly one + of the datasets. If you, for example, pass in `{"data1": data1, "data2": data2}` for two datasets + `data1` and `data2`, you could specify `metric_for_best_model="eval_data1_loss"` for using the + loss on `data1` and `metric_for_best_model="eval_data2_loss"` for the loss on `data2`. + + + + ignore_keys (`List[str]`, *optional*): + A list of keys in the output of your model (if it is a dictionary) that should be ignored when + gathering predictions. + metric_key_prefix (`str`, *optional*, defaults to `"eval"`): + An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named + "eval_bleu" if the prefix is "eval" (default) + + Returns: + A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The + dictionary also contains the epoch number which comes from the training state. + """ + # handle multiple eval datasets + override = eval_dataset is not None + eval_dataset = eval_dataset if override else self.eval_dataset + if isinstance(eval_dataset, dict): + metrics = {} + for eval_dataset_name, _eval_dataset in eval_dataset.items(): + dataset_metrics = self.evaluate( + eval_dataset=_eval_dataset if override else eval_dataset_name, + ignore_keys=ignore_keys, + metric_key_prefix=f"{metric_key_prefix}_{eval_dataset_name}", + ) + metrics.update(dataset_metrics) + return metrics + + # memory metrics - must set up as early as possible + self._memory_tracker.start() + + eval_dataloader = self.get_eval_dataloader(eval_dataset) + if self.is_fsdp_xla_v2_enabled: + eval_dataloader = tpu_spmd_dataloader(eval_dataloader) + + start_time = time.time() + + eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop + output = eval_loop( + eval_dataloader, + description="Evaluation", + # No point gathering the predictions if there are no metrics, otherwise we defer to + # self.args.prediction_loss_only + prediction_loss_only=True if self.compute_metrics is None else None, + ignore_keys=ignore_keys, + metric_key_prefix=metric_key_prefix, + ) + + total_batch_size = self.args.eval_batch_size * self.args.world_size + if f"{metric_key_prefix}_jit_compilation_time" in output.metrics: + start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"] + if f"{metric_key_prefix}_model_preparation_time" in output.metrics: + start_time += output.metrics[f"{metric_key_prefix}_model_preparation_time"] + output.metrics.update( + speed_metrics( + metric_key_prefix, + start_time, + num_samples=output.num_samples, + num_steps=math.ceil(output.num_samples / total_batch_size), + ) + ) + + self.log(output.metrics) + + if DebugOption.TPU_METRICS_DEBUG in self.args.debug: + # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.) + xm.master_print(met.metrics_report()) + + self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, output.metrics) + + self._memory_tracker.stop_and_update_metrics(output.metrics) + + return output.metrics + + def predict( + self, test_dataset: Dataset, ignore_keys: Optional[list[str]] = None, metric_key_prefix: str = "test" + ) -> PredictionOutput: + """ + Run prediction and returns predictions and potential metrics. + + Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method + will also return metrics, like in `evaluate()`. + + Args: + test_dataset (`Dataset`): + Dataset to run the predictions on. If it is an `datasets.Dataset`, columns not accepted by the + `model.forward()` method are automatically removed. Has to implement the method `__len__` + ignore_keys (`List[str]`, *optional*): + A list of keys in the output of your model (if it is a dictionary) that should be ignored when + gathering predictions. + metric_key_prefix (`str`, *optional*, defaults to `"test"`): + An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named + "test_bleu" if the prefix is "test" (default) + + + + If your predictions or labels have different sequence length (for instance because you're doing dynamic padding + in a token classification task) the predictions will be padded (on the right) to allow for concatenation into + one array. The padding index is -100. + + + + Returns: *NamedTuple* A namedtuple with the following keys: + + - predictions (`np.ndarray`): The predictions on `test_dataset`. + - label_ids (`np.ndarray`, *optional*): The labels (if the dataset contained some). + - metrics (`Dict[str, float]`, *optional*): The potential dictionary of metrics (if the dataset contained + labels). + """ + # memory metrics - must set up as early as possible + self._memory_tracker.start() + + test_dataloader = self.get_test_dataloader(test_dataset) + start_time = time.time() + + eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop + output = eval_loop( + test_dataloader, description="Prediction", ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix + ) + total_batch_size = self.args.eval_batch_size * self.args.world_size + if f"{metric_key_prefix}_jit_compilation_time" in output.metrics: + start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"] + if f"{metric_key_prefix}_model_preparation_time" in output.metrics: + start_time += output.metrics[f"{metric_key_prefix}_model_preparation_time"] + output.metrics.update( + speed_metrics( + metric_key_prefix, + start_time, + num_samples=output.num_samples, + num_steps=math.ceil(output.num_samples / total_batch_size), + ) + ) + + self.control = self.callback_handler.on_predict(self.args, self.state, self.control, output.metrics) + self._memory_tracker.stop_and_update_metrics(output.metrics) + + return PredictionOutput(predictions=output.predictions, label_ids=output.label_ids, metrics=output.metrics) + + def evaluation_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: Optional[bool] = None, + ignore_keys: Optional[list[str]] = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + args = self.args + + prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only + + # if eval is called w/o train, handle model prep here + if self.is_deepspeed_enabled and self.deepspeed is None: + _, _ = deepspeed_init(self, num_training_steps=0, inference=True) + + model = self._wrap_model(self.model, training=False, dataloader=dataloader) + + if len(self.accelerator._models) == 0 and model is self.model: + start_time = time.time() + model = ( + self.accelerator.prepare(model) + if self.is_deepspeed_enabled or (self.is_fsdp_enabled and self.accelerator.mixed_precision != "fp8") + else self.accelerator.prepare_model(model, evaluation_mode=True) + ) + self.model_preparation_time = round(time.time() - start_time, 4) + + if self.is_fsdp_enabled: + self.model = model + + # for the rest of this function `model` is the outside model, whether it was wrapped or not + if model is not self.model: + self.model_wrapped = model + + # backward compatibility + if self.is_deepspeed_enabled: + self.deepspeed = self.model_wrapped + + # if full fp16 or bf16 eval is wanted and this ``evaluation`` or ``predict`` isn't called + # while ``train`` is running, cast it to the right dtype first and then put on device + if not self.is_in_train: + if args.fp16_full_eval: + model = model.to(dtype=torch.float16, device=args.device) + elif args.bf16_full_eval: + model = model.to(dtype=torch.bfloat16, device=args.device) + + batch_size = self.args.eval_batch_size + + logger.info(f"\n***** Running {description} *****") + if has_length(dataloader): + logger.info(f" Num examples = {self.num_examples(dataloader)}") + else: + logger.info(" Num examples: Unknown") + logger.info(f" Batch size = {batch_size}") + + model.eval() + if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): + self.optimizer.eval() + + self.callback_handler.eval_dataloader = dataloader + # Do this before wrapping. + eval_dataset = getattr(dataloader, "dataset", None) + + if args.past_index >= 0: + self._past = None + + # Initialize containers + all_losses = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100) + all_preds = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100) + all_labels = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100) + all_inputs = EvalLoopContainer(self.args.eval_do_concat_batches, padding_index=-100) + + metrics = None + eval_set_kwargs = {} + + # Will be useful when we have an iterable dataset so don't know its length. + observed_num_examples = 0 + + # Main evaluation loop + for step, inputs in enumerate(dataloader): + # Update the observed num examples + observed_batch_size = find_batch_size(inputs) + if observed_batch_size is not None: + observed_num_examples += observed_batch_size + # For batch samplers, batch_size is not known by the dataloader in advance. + if batch_size is None: + batch_size = observed_batch_size + + # Prediction step + losses, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys) + main_input_name = getattr(self.model, "main_input_name", "input_ids") + inputs_decode = ( + self._prepare_input(inputs[main_input_name]) if "inputs" in args.include_for_metrics else None + ) + + if is_torch_xla_available(): + xm.mark_step() + + # Update containers + if losses is not None: + losses = self.gather_function(losses.repeat(batch_size)) + all_losses.add(losses) + if inputs_decode is not None: + inputs_decode = self.accelerator.pad_across_processes(inputs_decode, dim=1, pad_index=-100) + inputs_decode = self.gather_function(inputs_decode) + if not self.args.batch_eval_metrics or description == "Prediction": + all_inputs.add(inputs_decode) + if labels is not None: + # Pad labels here, preparing for preprocess_logits_for_metrics in next logits block. + labels = self.accelerator.pad_across_processes(labels, dim=1, pad_index=-100) + if logits is not None: + logits = self.accelerator.pad_across_processes(logits, dim=1, pad_index=-100) + if self.preprocess_logits_for_metrics is not None: + logits = self.preprocess_logits_for_metrics(logits, labels) + logits = self.gather_function(logits) + if not self.args.batch_eval_metrics or description == "Prediction": + all_preds.add(logits) + if labels is not None: + labels = self.gather_function(labels) + if not self.args.batch_eval_metrics or description == "Prediction": + all_labels.add(labels) + + self.control = self.callback_handler.on_prediction_step(args, self.state, self.control) + + if self.args.batch_eval_metrics: + if self.compute_metrics is not None and logits is not None and labels is not None: + is_last_step = self.accelerator.gradient_state.end_of_dataloader + batch_kwargs = {} + batch_kwargs["losses"] = losses if "loss" in args.include_for_metrics else None + batch_kwargs["inputs"] = inputs if "inputs" in args.include_for_metrics else None + metrics = self.compute_metrics( + EvalPrediction(predictions=logits, label_ids=labels, **batch_kwargs), + compute_result=is_last_step, + ) + + del losses, logits, labels, inputs + torch.cuda.empty_cache() + + # Gather all tensors and put them back on the CPU if we have done enough accumulation steps. + elif args.eval_accumulation_steps is not None and (step + 1) % args.eval_accumulation_steps == 0: + all_losses.to_cpu_and_numpy() + all_preds.to_cpu_and_numpy() + all_labels.to_cpu_and_numpy() + all_inputs.to_cpu_and_numpy() + + del losses, logits, labels, inputs + torch.cuda.empty_cache() + + # After all calls to `.gather_function`, reset to `gather_for_metrics`: + self.gather_function = self.accelerator.gather_for_metrics + if args.past_index and hasattr(self, "_past"): + # Clean the state at the end of the evaluation loop + delattr(self, "_past") + + # Gather all remaining tensors and put them back on the CPU + all_losses = all_losses.get_arrays() + all_preds = all_preds.get_arrays() + all_labels = all_labels.get_arrays() + all_inputs = all_inputs.get_arrays() + + # Number of samples + if has_length(eval_dataset): + num_samples = len(eval_dataset) + # The instance check is weird and does not actually check for the type, but whether the dataset has the right + # methods. Therefore we need to make sure it also has the attribute. + elif isinstance(eval_dataset, IterableDatasetShard) and getattr(eval_dataset, "num_examples", 0) > 0: + num_samples = eval_dataset.num_examples + else: + if has_length(dataloader): + num_samples = self.num_examples(dataloader) + else: # both len(dataloader.dataset) and len(dataloader) fail + num_samples = observed_num_examples + if num_samples == 0 and observed_num_examples > 0: + num_samples = observed_num_examples + + # Metrics! + if ( + self.compute_metrics is not None + and all_preds is not None + and all_labels is not None + and not self.args.batch_eval_metrics + ): + eval_set_kwargs["losses"] = all_losses if "loss" in args.include_for_metrics else None + eval_set_kwargs["inputs"] = all_inputs if "inputs" in args.include_for_metrics else None + metrics = self.compute_metrics( + EvalPrediction(predictions=all_preds, label_ids=all_labels, **eval_set_kwargs) + ) + elif metrics is None: + metrics = {} + + # To be JSON-serializable, we need to remove numpy types or zero-d tensors + metrics = denumpify_detensorize(metrics) + + if isinstance(all_losses, list) and all_losses: + metrics[f"{metric_key_prefix}_loss"] = np.concatenate(all_losses).mean().item() + elif isinstance(all_losses, np.ndarray): + metrics[f"{metric_key_prefix}_loss"] = all_losses.mean().item() + if hasattr(self, "jit_compilation_time"): + metrics[f"{metric_key_prefix}_jit_compilation_time"] = self.jit_compilation_time + if hasattr(self, "model_preparation_time"): + metrics[f"{metric_key_prefix}_model_preparation_time"] = self.model_preparation_time + + # Prefix all keys with metric_key_prefix + '_' + for key in list(metrics.keys()): + if not key.startswith(f"{metric_key_prefix}_"): + metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key) + + return EvalLoopOutput(predictions=all_preds, label_ids=all_labels, metrics=metrics, num_samples=num_samples) + + def _nested_gather(self, tensors, name=None): + """ + Gather value of `tensors` (tensor or list/tuple of nested tensors) and convert them to numpy before + concatenating them to `gathered` + """ + if tensors is None: + return + if is_torch_xla_available(): + if name is None: + name = "nested_gather" + tensors = nested_xla_mesh_reduce(tensors, name) + elif is_sagemaker_mp_enabled(): + tensors = smp_gather(tensors) + elif (self.args.distributed_state is not None and self.args.distributed_state.distributed_type != "NO") or ( + self.args.distributed_state is None and self.args.local_rank != -1 + ): + tensors = distributed_concat(tensors) + return tensors + + def prediction_step( + self, + model: nn.Module, + inputs: dict[str, Union[torch.Tensor, Any]], + prediction_loss_only: bool, + ignore_keys: Optional[list[str]] = None, + ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + Perform an evaluation step on `model` using `inputs`. + + Subclass and override to inject custom behavior. + + Args: + model (`nn.Module`): + The model to evaluate. + inputs (`Dict[str, Union[torch.Tensor, Any]]`): + The inputs and targets of the model. + + The dictionary will be unpacked before being fed to the model. Most models expect the targets under the + argument `labels`. Check your model's documentation for all accepted arguments. + prediction_loss_only (`bool`): + Whether or not to return the loss only. + ignore_keys (`List[str]`, *optional*): + A list of keys in the output of your model (if it is a dictionary) that should be ignored when + gathering predictions. + + Return: + Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, + logits and labels (each being optional). + """ + has_labels = False if len(self.label_names) == 0 else all(inputs.get(k) is not None for k in self.label_names) + # For CLIP-like models capable of returning loss values. + # If `return_loss` is not specified or being `None` in `inputs`, we check if the default value of `return_loss` + # is `True` in `model.forward`. + return_loss = inputs.get("return_loss", None) + if return_loss is None: + return_loss = self.can_return_loss + loss_without_labels = True if len(self.label_names) == 0 and return_loss else False + + inputs = self._prepare_inputs(inputs) + if ignore_keys is None: + if hasattr(self.model, "config"): + ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", ["past_key_values"]) + else: + ignore_keys = [] + + # labels may be popped when computing the loss (label smoothing for instance) so we grab them first. + if has_labels or loss_without_labels: + labels = nested_detach(tuple(inputs.get(name) for name in self.label_names)) + if len(labels) == 1: + labels = labels[0] + else: + labels = None + + with torch.no_grad(): + if is_sagemaker_mp_enabled(): + raw_outputs = smp_forward_only(model, inputs) + if has_labels or loss_without_labels: + if isinstance(raw_outputs, dict): + loss_mb = raw_outputs["loss"] + logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys + ["loss"]) + else: + loss_mb = raw_outputs[0] + logits_mb = raw_outputs[1:] + + loss = loss_mb.reduce_mean().detach().cpu() + logits = smp_nested_concat(logits_mb) + else: + loss = None + if isinstance(raw_outputs, dict): + logits_mb = tuple(v for k, v in raw_outputs.items() if k not in ignore_keys) + else: + logits_mb = raw_outputs + logits = smp_nested_concat(logits_mb) + else: + if has_labels or loss_without_labels: + with self.compute_loss_context_manager(): + loss, outputs = self.compute_loss(model, inputs, return_outputs=True) + loss = loss.detach().mean() + + if isinstance(outputs, dict): + logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"]) + else: + logits = outputs[1:] + else: + loss = None + with self.compute_loss_context_manager(): + outputs = model(**inputs) + if isinstance(outputs, dict): + logits = tuple(v for k, v in outputs.items() if k not in ignore_keys) + else: + logits = outputs + # TODO: this needs to be fixed and made cleaner later. + if self.args.past_index >= 0: + self._past = outputs[self.args.past_index - 1] + + if prediction_loss_only: + return (loss, None, None) + + logits = nested_detach(logits) + if len(logits) == 1: + logits = logits[0] + + return (loss, logits, labels) + + def floating_point_ops(self, inputs: dict[str, Union[torch.Tensor, Any]]): + """ + For models that inherit from [`PreTrainedModel`], uses that method to compute the number of floating point + operations for every backward + forward pass. If using another model, either implement such a method in the + model or subclass and override this method. + + Args: + inputs (`Dict[str, Union[torch.Tensor, Any]]`): + The inputs and targets of the model. + + Returns: + `int`: The number of floating-point operations. + """ + if hasattr(self.model, "floating_point_ops"): + return self.model.floating_point_ops(inputs) + else: + return 0 + + def init_hf_repo(self, token: Optional[str] = None): + """ + Initializes a git repo in `self.args.hub_model_id`. + """ + # Only on process zero + if not self.is_world_process_zero(): + return + + if self.args.hub_model_id is None: + repo_name = Path(self.args.output_dir).absolute().name + else: + repo_name = self.args.hub_model_id + + token = token if token is not None else self.args.hub_token + repo_url = create_repo(repo_name, token=token, private=self.args.hub_private_repo, exist_ok=True) + self.hub_model_id = repo_url.repo_id + self.push_in_progress = None + + def create_model_card( + self, + language: Optional[str] = None, + license: Optional[str] = None, + tags: Union[str, list[str], None] = None, + model_name: Optional[str] = None, + finetuned_from: Optional[str] = None, + tasks: Union[str, list[str], None] = None, + dataset_tags: Union[str, list[str], None] = None, + dataset: Union[str, list[str], None] = None, + dataset_args: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + language (`str`, *optional*): + The language of the model (if applicable) + license (`str`, *optional*): + The license of the model. Will default to the license of the pretrained model used, if the original + model given to the `Trainer` comes from a repo on the Hub. + tags (`str` or `List[str]`, *optional*): + Some tags to be included in the metadata of the model card. + model_name (`str`, *optional*): + The name of the model. + finetuned_from (`str`, *optional*): + The name of the model used to fine-tune this one (if applicable). Will default to the name of the repo + of the original model given to the `Trainer` (if it comes from the Hub). + tasks (`str` or `List[str]`, *optional*): + One or several task identifiers, to be included in the metadata of the model card. + dataset_tags (`str` or `List[str]`, *optional*): + One or several dataset tags, to be included in the metadata of the model card. + dataset (`str` or `List[str]`, *optional*): + One or several dataset identifiers, to be included in the metadata of the model card. + dataset_args (`str` or `List[str]`, *optional*): + One or several dataset arguments, to be included in the metadata of the model card. + """ + if not self.is_world_process_zero(): + return + + model_card_filepath = os.path.join(self.args.output_dir, "README.md") + is_peft_library = False + if os.path.exists(model_card_filepath): + library_name = ModelCard.load(model_card_filepath).data.get("library_name") + is_peft_library = library_name == "peft" + + # Append existing tags in `tags` + existing_tags = ModelCard.load(model_card_filepath).data.tags + if tags is not None and existing_tags is not None: + if isinstance(tags, str): + tags = [tags] + for tag in existing_tags: + if tag not in tags: + tags.append(tag) + + training_summary = TrainingSummary.from_trainer( + self, + language=language, + license=license, + tags=tags, + model_name=model_name, + finetuned_from=finetuned_from, + tasks=tasks, + dataset_tags=dataset_tags, + dataset=dataset, + dataset_args=dataset_args, + ) + model_card = training_summary.to_model_card() + with open(model_card_filepath, "w") as f: + f.write(model_card) + + if is_peft_library: + self.accelerator.unwrap_model(self.model).create_or_update_model_card(self.args.output_dir) + + def _push_from_checkpoint(self, checkpoint_folder): + # Only push from one node. + if not self.is_world_process_zero() or self.args.hub_strategy == HubStrategy.END: + return + # If we haven't finished the last push, we don't do this one unless args.hub_always_push=True. + if not self.args.hub_always_push and self.push_in_progress is not None and not self.push_in_progress.is_done(): + return + + output_dir = self.args.output_dir + # To avoid a new synchronization of all model weights, we just copy the file from the checkpoint folder + modeling_files = [CONFIG_NAME, WEIGHTS_NAME, SAFE_WEIGHTS_NAME] + # Add sharded checkpoints if we have an index + for index_file in [WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_INDEX_NAME]: + index_path = os.path.join(checkpoint_folder, index_file) + if os.path.isfile(index_path): + modeling_files.append(index_file) + with open(index_path) as f: + index = json.loads(f.read()) + shard_files = list(set(index["weight_map"].values())) + modeling_files.extend(shard_files) + if is_peft_available(): + modeling_files.extend([ADAPTER_CONFIG_NAME, ADAPTER_WEIGHTS_NAME, ADAPTER_SAFE_WEIGHTS_NAME]) + for modeling_file in modeling_files: + if os.path.isfile(os.path.join(checkpoint_folder, modeling_file)): + shutil.copy(os.path.join(checkpoint_folder, modeling_file), os.path.join(output_dir, modeling_file)) + # Saving the processing class is fast and we don't know how many files it may have spawned, so we resave it to be sure. + if self.processing_class is not None: + self.processing_class.save_pretrained(output_dir) + # Same for the training arguments + torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) + + if self.args.save_strategy == SaveStrategy.STEPS: + commit_message = f"Training in progress, step {self.state.global_step}" + else: + commit_message = f"Training in progress, epoch {int(self.state.epoch)}" + + model_push_job = upload_folder( + repo_id=self.hub_model_id, + folder_path=output_dir, + commit_message=commit_message, + token=self.args.hub_token, + run_as_future=True, + ignore_patterns=["_*", f"{PREFIX_CHECKPOINT_DIR}-*"], + ) + + push_jobs = [model_push_job] + + if self.args.hub_strategy in [HubStrategy.CHECKPOINT, HubStrategy.ALL_CHECKPOINTS]: + path_in_repo = ( + "last-checkpoint" if self.args.hub_strategy == HubStrategy.CHECKPOINT else Path(checkpoint_folder).name + ) + checkpoint_push = upload_folder( + repo_id=self.hub_model_id, + folder_path=checkpoint_folder, + path_in_repo=path_in_repo, + commit_message=commit_message + ", checkpoint", + token=self.args.hub_token, + run_as_future=True, + ) + push_jobs.append(checkpoint_push) + + if self.push_in_progress is None or self.push_in_progress.is_done(): + self.push_in_progress = PushInProgress(push_jobs) + else: + self.push_in_progress.jobs.extend(push_jobs) + + def _finish_current_push(self): + if not hasattr(self, "push_in_progress"): + return + if self.push_in_progress is not None and not self.push_in_progress.is_done(): + logger.info("Waiting for the current checkpoint push to be finished, this might take a couple of minutes.") + self.push_in_progress.wait_until_done() + + def push_to_hub( + self, + commit_message: Optional[str] = "End of training", + blocking: bool = True, + token: Optional[str] = None, + revision: Optional[str] = None, + **kwargs, + ) -> str: + """ + Upload `self.model` and `self.processing_class` to the 🤗 model hub on the repo `self.args.hub_model_id`. + + Parameters: + commit_message (`str`, *optional*, defaults to `"End of training"`): + Message to commit while pushing. + blocking (`bool`, *optional*, defaults to `True`): + Whether the function should return only when the `git push` has finished. + token (`str`, *optional*, defaults to `None`): + Token with write permission to overwrite Trainer's original args. + revision (`str`, *optional*): + The git revision to commit from. Defaults to the head of the "main" branch. + kwargs (`Dict[str, Any]`, *optional*): + Additional keyword arguments passed along to [`~Trainer.create_model_card`]. + + Returns: + The URL of the repository where the model was pushed if `blocking=False`, or a `Future` object tracking the + progress of the commit if `blocking=True`. + """ + model_name = kwargs.pop("model_name", None) + if model_name is None and self.args.should_save: + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + token = token if token is not None else self.args.hub_token + + # In case the user calls this method with args.push_to_hub = False + if self.hub_model_id is None: + self.init_hf_repo(token=token) + + # Needs to be executed on all processes for TPU training, but will only save on the processed determined by + # self.args.should_save. + self.save_model(_internal_call=True) + + # Only push from one node. + if not self.is_world_process_zero(): + return + + # Add additional tags in the case the model has already some tags and users pass + # "tags" argument to `push_to_hub` so that trainer automatically handles internal tags + # from all models since Trainer does not call `model.push_to_hub`. + if getattr(self.model, "model_tags", None) is not None: + if "tags" not in kwargs: + kwargs["tags"] = [] + + # If it is a string, convert it to a list + if isinstance(kwargs["tags"], str): + kwargs["tags"] = [kwargs["tags"]] + + for model_tag in self.model.model_tags: + if model_tag not in kwargs["tags"]: + kwargs["tags"].append(model_tag) + + self.create_model_card(model_name=model_name, **kwargs) + + # Wait for the current upload to be finished. + self._finish_current_push() + return upload_folder( + repo_id=self.hub_model_id, + folder_path=self.args.output_dir, + commit_message=commit_message, + token=token, + run_as_future=not blocking, + ignore_patterns=["_*", f"{PREFIX_CHECKPOINT_DIR}-*"], + revision=revision, + ) + + # + # Deprecated code + # + + def prediction_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: Optional[bool] = None, + ignore_keys: Optional[list[str]] = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + args = self.args + + if not has_length(dataloader): + raise ValueError("dataloader must implement a working __len__") + + prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else args.prediction_loss_only + + # if eval is called w/o train, handle model prep here + if self.is_deepspeed_enabled and self.deepspeed is None: + _, _ = deepspeed_init(self, num_training_steps=0, inference=True) + + model = self._wrap_model(self.model, training=False, dataloader=dataloader) + + if len(self.accelerator._models) == 0 and model is self.model: + model = ( + self.accelerator.prepare(model) + if self.is_deepspeed_enabled or self.is_fsdp_enabled + else self.accelerator.prepare_model(model, evaluation_mode=True) + ) + + if self.is_fsdp_enabled: + self.model = model + + # for the rest of this function `model` is the outside model, whether it was wrapped or not + if model is not self.model: + self.model_wrapped = model + + # backward compatibility + if self.is_deepspeed_enabled: + self.deepspeed = self.model_wrapped + + # if full fp16 or bf16 eval is wanted and this ``evaluation`` or ``predict`` isn't called + # while ``train`` is running, cast it to the right dtype first and then put on device + if not self.is_in_train: + if args.fp16_full_eval: + model = model.to(dtype=torch.float16, device=args.device) + elif args.bf16_full_eval: + model = model.to(dtype=torch.bfloat16, device=args.device) + + batch_size = ( + dataloader.total_batch_size + if getattr(dataloader, "_is_accelerate_prepared", False) + else dataloader.batch_size + ) + + if batch_size is None: + raise ValueError( + "Batch size cannot be None. Ensure the dataloader has a valid batch_size or total_batch_size." + ) + + num_examples = self.num_examples(dataloader) + logger.info(f"\n***** Running {description} *****") + logger.info(f" Num examples = {num_examples}") + logger.info(f" Batch size = {batch_size}") + + losses_host: Optional[torch.Tensor] = None + preds_host: Union[torch.Tensor, list[torch.Tensor], None] = None + labels_host: Union[torch.Tensor, list[torch.Tensor], None] = None + inputs_host: Union[torch.Tensor, list[torch.Tensor], None] = None + metrics: Optional[dict] = None + eval_set_kwargs: dict = {} + + world_size = max(1, args.world_size) + + eval_losses_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=batch_size) + if not prediction_loss_only: + # The actual number of eval_sample can be greater than num_examples in distributed settings (when we pass + # a batch size to the sampler) + make_multiple_of = None + if hasattr(dataloader, "sampler") and isinstance(dataloader.sampler, SequentialDistributedSampler): + make_multiple_of = dataloader.sampler.batch_size + preds_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of) + labels_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of) + inputs_gatherer = DistributedTensorGatherer(world_size, num_examples, make_multiple_of=make_multiple_of) + + model.eval() + if hasattr(self.optimizer, "eval") and callable(self.optimizer.eval): + self.optimizer.eval() + + if args.past_index >= 0: + self._past = None + + self.callback_handler.eval_dataloader = dataloader + + for step, inputs in enumerate(dataloader): + loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys) + main_input_name = getattr(self.model, "main_input_name", "input_ids") + inputs_decode = ( + self._prepare_input(inputs[main_input_name]) if "inputs" in args.include_for_metrics else None + ) + + if loss is not None: + losses = loss.repeat(batch_size) + losses_host = losses if losses_host is None else torch.cat((losses_host, losses), dim=0) + if logits is not None: + preds_host = logits if preds_host is None else nested_concat(preds_host, logits, padding_index=-100) + if labels is not None: + labels_host = labels if labels_host is None else nested_concat(labels_host, labels, padding_index=-100) + if inputs_decode is not None: + inputs_host = ( + inputs_decode + if inputs_host is None + else nested_concat(inputs_host, inputs_decode, padding_index=-100) + ) + self.control = self.callback_handler.on_prediction_step(args, self.state, self.control) + + if self.args.batch_eval_metrics: + if self.compute_metrics is not None and preds_host is not None and labels_host is not None: + is_last_step = self.accelerator.gradient_state.end_of_dataloader + batch_kwargs = {} + batch_kwargs["losses"] = losses_host if "loss" in args.include_for_metrics else None + batch_kwargs["inputs"] = inputs_host if "inputs" in args.include_for_metrics else None + metrics = self.compute_metrics( + EvalPrediction(predictions=preds_host, label_ids=labels_host, **batch_kwargs), + compute_result=is_last_step, + ) + + if self.args.batch_eval_metrics or ( + args.eval_accumulation_steps is not None and (step + 1) % args.eval_accumulation_steps == 0 + ): + # Gather all tensors and put them back on the CPU if we have done enough accumulation steps. + eval_losses_gatherer.add_arrays(self._gather_and_numpify(losses_host, "eval_losses")) + if not prediction_loss_only: + preds_gatherer.add_arrays(self._gather_and_numpify(preds_host, "eval_preds")) + labels_gatherer.add_arrays(self._gather_and_numpify(labels_host, "eval_label_ids")) + inputs_gatherer.add_arrays(self._gather_and_numpify(inputs_host, "eval_inputs_ids")) + + # Set back to None to begin a new accumulation + del losses_host, preds_host, labels_host, inputs_host + torch.cuda.empty_cache() + losses_host, preds_host, labels_host, inputs_host = None, None, None, None + + if args.past_index and hasattr(self, "_past"): + # Clean the state at the end of the evaluation loop + delattr(self, "_past") + + # Gather all remaining tensors and put them back on the CPU + eval_losses_gatherer.add_arrays(self._gather_and_numpify(losses_host, "eval_losses")) + if not prediction_loss_only: + preds_gatherer.add_arrays(self._gather_and_numpify(preds_host, "eval_preds")) + labels_gatherer.add_arrays(self._gather_and_numpify(labels_host, "eval_label_ids")) + inputs_gatherer.add_arrays(self._gather_and_numpify(inputs_host, "eval_inputs_ids")) + + eval_loss = eval_losses_gatherer.finalize() + preds = preds_gatherer.finalize() if not prediction_loss_only else None + label_ids = labels_gatherer.finalize() if not prediction_loss_only else None + inputs_ids = inputs_gatherer.finalize() if not prediction_loss_only else None + + if ( + self.compute_metrics is not None + and preds is not None + and label_ids is not None + and not self.args.batch_eval_metrics + ): + eval_set_kwargs["losses"] = eval_loss if "loss" in args.include_for_metrics else None + eval_set_kwargs["inputs"] = inputs_ids if "inputs" in args.include_for_metrics else None + metrics = self.compute_metrics(EvalPrediction(predictions=preds, label_ids=label_ids, **eval_set_kwargs)) + elif metrics is None: + metrics = {} + + # To be JSON-serializable, we need to remove numpy types or zero-d tensors + metrics = denumpify_detensorize(metrics) + + if eval_loss is not None: + metrics[f"{metric_key_prefix}_loss"] = eval_loss.mean().item() + + # Prefix all keys with metric_key_prefix + '_' + for key in list(metrics.keys()): + if not key.startswith(f"{metric_key_prefix}_"): + metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key) + + return EvalLoopOutput(predictions=preds, label_ids=label_ids, metrics=metrics, num_samples=num_examples) + + def _gather_and_numpify(self, tensors, name): + """ + Gather value of `tensors` (tensor or list/tuple of nested tensors) and convert them to numpy before + concatenating them to `gathered` + """ + if tensors is None: + return + if is_torch_xla_available(): + tensors = nested_xla_mesh_reduce(tensors, name) + elif is_sagemaker_mp_enabled(): + tensors = smp_gather(tensors) + elif self.args.parallel_mode == ParallelMode.DISTRIBUTED: + tensors = distributed_concat(tensors) + + return nested_numpify(tensors) + + def _add_sm_patterns_to_gitignore(self) -> None: + """Add SageMaker Checkpointing patterns to .gitignore file.""" + # Make sure we only do this on the main process + if not self.is_world_process_zero(): + return + + patterns = ["*.sagemaker-uploading", "*.sagemaker-uploaded"] + + # Get current .gitignore content + if os.path.exists(os.path.join(self.repo.local_dir, ".gitignore")): + with open(os.path.join(self.repo.local_dir, ".gitignore")) as f: + current_content = f.read() + else: + current_content = "" + + # Add the patterns to .gitignore + content = current_content + for pattern in patterns: + if pattern not in content: + if content.endswith("\n"): + content += pattern + else: + content += f"\n{pattern}" + + # Write the .gitignore file if it has changed + if content != current_content: + with open(os.path.join(self.repo.local_dir, ".gitignore"), "w") as f: + logger.debug(f"Writing .gitignore file. Content: {content}") + f.write(content) + + self.repo.git_add(".gitignore") + + # avoid race condition with git status + time.sleep(0.5) + + if not self.repo.is_repo_clean(): + self.repo.git_commit("Add *.sagemaker patterns to .gitignore.") + self.repo.git_push() + + def create_accelerator_and_postprocess(self): + # We explicitly don't rely on the `Accelerator` to do gradient accumulation + grad_acc_kwargs = {} + if is_accelerate_available("0.28.0") and self.args.accelerator_config.gradient_accumulation_kwargs is not None: + grad_acc_kwargs = self.args.accelerator_config.gradient_accumulation_kwargs + + # check if num_steps is attempted to be passed in gradient_accumulation_kwargs + if "num_steps" in grad_acc_kwargs: + if self.args.gradient_accumulation_steps > 1: + # raise because we do not know which setting is intended. + raise ValueError( + "The `AcceleratorConfig`'s `num_steps` is set but `gradient_accumulation_steps` is greater than 1 in the passed `TrainingArguments`" + "If using the passed `AcceleratorConfig` is desired, do not set the `TrainingArguments` `gradient_accumulation_steps`." + ) + else: + self.args.gradient_accumulation_steps = grad_acc_kwargs["num_steps"] + + accelerator_config = self.args.accelerator_config.to_dict() + + if is_accelerate_available("0.28.0"): + # Extract dataloader config params from accelerator config + dataloader_params = ["split_batches", "dispatch_batches", "even_batches", "use_seedable_sampler"] + dataloader_config = DataLoaderConfiguration( + **{param: accelerator_config.pop(param) for param in dataloader_params} + ) + if is_accelerate_available("1.1.0"): + dataloader_config.data_seed = self.args.data_seed + + non_blocking = accelerator_config.pop("non_blocking") + if not is_accelerate_available("0.30.0"): + if non_blocking: + raise ImportError( + "`non_blocking` is only supported in accelerate v0.30.0 and above. Please upgrade accelerate to use this feature." + ) + else: + if non_blocking and not self.args.dataloader_pin_memory: + logger.warning( + "`non_blocking` is enabled but `dataloader_pin_memory` is not. For the best performance, it's recommended to enable both." + ) + dataloader_config.non_blocking = non_blocking + # this would have been updated above, no need for it anymore + accelerator_config.pop("gradient_accumulation_kwargs") + + args = { + "deepspeed_plugin": self.args.deepspeed_plugin, + } + if is_accelerate_available("0.28.0"): + args["dataloader_config"] = dataloader_config + else: + args.update(accelerator_config) + # tp is initialized at Accelerator init phase so + # args should be prepared here + if self.args.tp_size > 1: + self.is_tp_enabled = True + if version.parse(accelerate_version) > version.parse("1.3.0"): + args["torch_tp_plugin"] = TorchTensorParallelPlugin(tp_size=self.args.tp_size) + else: + raise ValueError("Requires accelerate>1.3.0 to use Tensor Parallelism.") + + # create accelerator object + self.accelerator = Accelerator(**args) + # some Trainer classes need to use `gather` instead of `gather_for_metrics`, thus we store a flag + self.gather_function = self.accelerator.gather_for_metrics + + if "use_gather_object" in inspect.signature(self.gather_function).parameters.keys(): + self.gather_function = functools.partial( + self.gather_function, use_gather_object=self.args.eval_use_gather_object + ) + + # deepspeed and accelerate flags covering both trainer args and accelerate launcher + self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None + self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None + self.is_tp_enabled = getattr(self.accelerator.state, "torch_tp_plugin", None) is not None + # post accelerator creation setup + if self.is_fsdp_enabled: + fsdp_plugin = self.accelerator.state.fsdp_plugin + for param in ["limit_all_gathers", "activation_checkpointing"]: + setattr(fsdp_plugin, param, self.args.fsdp_config.get(param, getattr(fsdp_plugin, param))) + if fsdp_plugin.activation_checkpointing and self.args.gradient_checkpointing: + raise ValueError( + "The activation_checkpointing in FSDP config and the gradient_checkpointing in training arg " + "can't be set to True simultaneously. Please use FSDP's activation_checkpointing logic " + "when using FSDP." + ) + + if self.is_deepspeed_enabled and getattr(self.args, "hf_deepspeed_config", None) is None: + self.propagate_args_to_deepspeed() + + # `save_only_model` can't be used with DeepSpeed/FSDP along with `load_best_model_at_end` + if ( + self.args.save_only_model + and (self.is_deepspeed_enabled or self.is_fsdp_enabled) + and self.args.load_best_model_at_end + ): + wrapper = "DeepSpeed" if self.is_deepspeed_enabled else "FSDP" + raise ValueError(f"{wrapper} can't be used with `save_only_model` along with `load_best_model_at_end`.") + + # `auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3 + if ( + self.is_deepspeed_enabled + and self.accelerator.state.deepspeed_plugin.zero_stage == 3 + and self.args.auto_find_batch_size + ): + raise ValueError( + "`auto_find_batch_size` isn't supported yet with DeepSpeed Zero-3. Please consider using Zero-2, Zero-1, or FSDP" + ) + if ( + self.args.save_only_model + and self.is_fsdp_enabled + and "SHARDED_STATE_DICT" in str(self.accelerator.state.fsdp_plugin.state_dict_type) + ): + raise ValueError("save_only_model option is not compatible with FSDP state dict type 'SHARDED_STATE_DICT'") + + def propagate_args_to_deepspeed(self, auto_find_batch_size=False): + """ + Sets values in the deepspeed plugin based on the Trainer args + """ + from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig + + ds_plugin = self.accelerator.state.deepspeed_plugin + + ds_plugin.hf_ds_config = HfTrainerDeepSpeedConfig(ds_plugin.hf_ds_config.config) + ds_plugin.deepspeed_config = ds_plugin.hf_ds_config.config + ds_plugin.hf_ds_config.trainer_config_process(self.args, auto_find_batch_size) + + def _fsdp_qlora_plugin_updates(self): + if self.is_fsdp_enabled and _is_peft_model(self.model): + from peft import PeftConfig + from peft.utils.other import fsdp_auto_wrap_policy + + if isinstance(self.model.active_peft_config, PeftConfig): + self.accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(self.model) + if ( + getattr(self.model, "quantization_method", None) == QuantizationMethod.BITS_AND_BYTES + and self.model.hf_quantizer.quantization_config.bnb_4bit_quant_storage.is_floating_point + and version.parse(accelerate_version) > version.parse("0.27.0") + ): + self.accelerator.state.fsdp_plugin.set_mixed_precision( + self.model.hf_quantizer.quantization_config.bnb_4bit_quant_storage, override=True + ) + + def get_batch_samples(self, epoch_iterator, num_batches, device): + batch_samples = [] + num_items_in_batch = None + + for _ in range(num_batches): + try: + batch_samples.append(next(epoch_iterator)) + except StopIteration: + break + + count_num_items_in_batch = ( + len(batch_samples) > 0 + and "labels" in batch_samples[0] + and ( + # num_items_in_batch is passed to model forward + # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3757 + self.model_accepts_loss_kwargs + # num_items_in_batch is passed to compute_loss_func + # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3773 + or self.compute_loss_func is not None + # num_items_in_batch is also verified if (self.model_accepts_loss_kwargs or self.compute_loss_func) + # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3790 + ) + ) + + if count_num_items_in_batch: + # For now we don't support object detection + try: + num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples]) + except (TypeError, AttributeError): + pass + + if num_items_in_batch is not None: + if self.args.average_tokens_across_devices: + num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum() + + if torch.is_tensor(num_items_in_batch): + num_items_in_batch = num_items_in_batch.to(device) + + if self.args.n_gpu > 1 and num_items_in_batch.dim() == 0: + # In the DataParallel case, convert the scalar tensor into a 1-dim tensor + num_items_in_batch = num_items_in_batch.unsqueeze(0) + + return batch_samples, num_items_in_batch + + def set_initial_training_values( + self, args: TrainingArguments, dataloader: DataLoader, total_train_batch_size: int + ): + """ + Calculates and returns the following values: + - `num_train_epochs` + - `num_update_steps_per_epoch` + - `num_examples` + - `num_train_samples` + - `epoch_based` + - `len_dataloader` + - `max_steps` + """ + # Case 1: we rely on `args.max_steps` first + max_steps = args.max_steps + # If max_steps is negative, we use the number of epochs to determine the number of total steps later + epoch_based = max_steps < 0 + len_dataloader = len(dataloader) if has_length(dataloader) else None + + # Case 2: We have a dataloader length and can extrapolate + if len_dataloader is not None: + num_update_steps_per_epoch = max(len_dataloader // args.gradient_accumulation_steps, 1) + # Case 3: We have a length but are using epochs, we can extrapolate the number of steps + if epoch_based: + max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch) + + # Now we figure out `num_examples`, `num_train_epochs`, and `train_samples` + if len_dataloader: + num_examples = self.num_examples(dataloader) + if args.max_steps > 0: + num_train_epochs = max_steps // num_update_steps_per_epoch + int( + max_steps % num_update_steps_per_epoch > 0 + ) + # May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's + # the best we can do. + num_train_samples = max_steps * total_train_batch_size + else: + num_train_epochs = math.ceil(args.num_train_epochs) + num_train_samples = self.num_examples(dataloader) * args.num_train_epochs + elif args.max_steps > 0: # Rely on max_steps when dataloader does not have a working size + # Setting a very large number of epochs so we go as many times as necessary over the iterator. + num_train_epochs = sys.maxsize + num_update_steps_per_epoch = max_steps + num_examples = total_train_batch_size * args.max_steps + num_train_samples = args.max_steps * total_train_batch_size + else: + raise ValueError( + "args.max_steps must be set to a positive value if dataloader does not have a length, was" + f" {args.max_steps}" + ) + return ( + num_train_epochs, + num_update_steps_per_epoch, + num_examples, + num_train_samples, + epoch_based, + len_dataloader, + max_steps, + ) diff --git a/lib/python3.12/site-packages/transformers/trainer_callback.py b/lib/python3.12/site-packages/transformers/trainer_callback.py new file mode 100644 index 0000000000000000000000000000000000000000..e4c465dfe558ff6a31accd5fdec2198ce32cd4a8 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/trainer_callback.py @@ -0,0 +1,785 @@ +# Copyright 2020-present the HuggingFace Inc. team. +# +# 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. +""" +Callbacks to use with the Trainer class and customize the training loop. +""" + +import dataclasses +import json +import math +from dataclasses import dataclass +from typing import Optional, Union + +import numpy as np +from tqdm.auto import tqdm + +from .trainer_utils import HPSearchBackend, IntervalStrategy, SaveStrategy, has_length +from .training_args import TrainingArguments +from .utils import logging + + +logger = logging.get_logger(__name__) + + +@dataclass +class TrainerState: + """ + A class containing the [`Trainer`] inner state that will be saved along the model and optimizer when checkpointing + and passed to the [`TrainerCallback`]. + + + + In all this class, one step is to be understood as one update step. When using gradient accumulation, one update + step may require several forward and backward passes: if you use `gradient_accumulation_steps=n`, then one update + step requires going through *n* batches. + + + + Args: + epoch (`float`, *optional*): + Only set during training, will represent the epoch the training is at (the decimal part being the + percentage of the current epoch completed). + global_step (`int`, *optional*, defaults to 0): + During training, represents the number of update steps completed. + max_steps (`int`, *optional*, defaults to 0): + The number of update steps to do during the current training. + logging_steps (`int`, *optional*, defaults to 500): + Log every X updates steps + eval_steps (`int`, *optional*): + Run an evaluation every X steps. + save_steps (`int`, *optional*, defaults to 500): + Save checkpoint every X updates steps. + train_batch_size (`int`, *optional*): + The batch size for the training dataloader. Only needed when + `auto_find_batch_size` has been used. + num_input_tokens_seen (`int`, *optional*, defaults to 0): + When tracking the inputs tokens, the number of tokens seen during training (number of input tokens, not the + number of prediction tokens). + total_flos (`float`, *optional*, defaults to 0): + The total number of floating operations done by the model since the beginning of training (stored as floats + to avoid overflow). + log_history (`List[Dict[str, float]]`, *optional*): + The list of logs done since the beginning of training. + best_metric (`float`, *optional*): + When tracking the best model, the value of the best metric encountered so far. + best_global_step (`int`, *optional*): + When tracking the best model, the step at which the best metric was encountered. + Used for setting `best_model_checkpoint`. + best_model_checkpoint (`str`, *optional*): + When tracking the best model, the value of the name of the checkpoint for the best model encountered so + far. + is_local_process_zero (`bool`, *optional*, defaults to `True`): + Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on + several machines) main process. + is_world_process_zero (`bool`, *optional*, defaults to `True`): + Whether or not this process is the global main process (when training in a distributed fashion on several + machines, this is only going to be `True` for one process). + is_hyper_param_search (`bool`, *optional*, defaults to `False`): + Whether we are in the process of a hyper parameter search using Trainer.hyperparameter_search. This will + impact the way data will be logged in TensorBoard. + stateful_callbacks (`List[StatefulTrainerCallback]`, *optional*): + Callbacks attached to the `Trainer` that should have their states be saved or restored. + Relevant callbacks should implement a `state` and `from_state` function. + """ + + epoch: Optional[float] = None + global_step: int = 0 + max_steps: int = 0 + logging_steps: int = 500 + eval_steps: int = 500 + save_steps: int = 500 + train_batch_size: Optional[int] = None + num_train_epochs: int = 0 + num_input_tokens_seen: int = 0 + total_flos: float = 0 + log_history: list[dict[str, float]] = None + best_metric: Optional[float] = None + best_global_step: Optional[int] = None + best_model_checkpoint: Optional[str] = None + is_local_process_zero: bool = True + is_world_process_zero: bool = True + is_hyper_param_search: bool = False + trial_name: Optional[str] = None + trial_params: dict[str, Union[str, float, int, bool]] = None + stateful_callbacks: list["TrainerCallback"] = None + + def __post_init__(self): + if self.log_history is None: + self.log_history = [] + if self.stateful_callbacks is None: + self.stateful_callbacks = {} + elif isinstance(self.stateful_callbacks, dict): + # We are loading the callbacks in from the state file, no need to process them + pass + else: + # Saveable callbacks get stored as dict of kwargs + stateful_callbacks = {} + for callback in self.stateful_callbacks: + if not isinstance(callback, (ExportableState)): + raise TypeError( + f"All callbacks passed to be saved must inherit `ExportableState`, but received {type(callback)}" + ) + name = callback.__class__.__name__ + if name in stateful_callbacks: + # We can have multiple versions of the same callback + # if so, we store them as a list of states to restore + if not isinstance(stateful_callbacks[name], list): + stateful_callbacks[name] = [stateful_callbacks[name]] + stateful_callbacks[name].append(callback.state()) + else: + stateful_callbacks[name] = callback.state() + self.stateful_callbacks = stateful_callbacks + + def save_to_json(self, json_path: str): + """Save the content of this instance in JSON format inside `json_path`.""" + json_string = json.dumps(dataclasses.asdict(self), indent=2, sort_keys=True) + "\n" + with open(json_path, "w", encoding="utf-8") as f: + f.write(json_string) + + @classmethod + def load_from_json(cls, json_path: str): + """Create an instance from the content of `json_path`.""" + with open(json_path, encoding="utf-8") as f: + text = f.read() + return cls(**json.loads(text)) + + def compute_steps(self, args, max_steps): + """ + Calculates and stores the absolute value for logging, + eval, and save steps based on if it was a proportion + or not. + """ + for step_kind in ("logging", "eval", "save"): + num_steps = getattr(args, f"{step_kind}_steps") + if num_steps is not None: + if num_steps < 1: + num_steps = math.ceil(max_steps * num_steps) + setattr(self, f"{step_kind}_steps", num_steps) + + def init_training_references(self, trainer, max_steps, num_train_epochs, trial): + """ + Stores the initial training references needed in `self` + """ + if trainer.hp_name is not None and trainer._trial is not None: + # use self._trial because the SigOpt/Optuna hpo only call `_hp_search_setup(trial)` instead of passing trial + # parameter to Train when using DDP. + self.trial_name = trainer.hp_name(trainer._trial) + self.trial_params = None + if trial is not None: + from transformers.integrations import hp_params + + assignments = trial.assignments if trainer.hp_search_backend == HPSearchBackend.SIGOPT else trial + self.trial_params = hp_params(assignments) + + self.max_steps = max_steps + self.num_train_epochs = num_train_epochs + self.is_local_process_zero = trainer.is_local_process_zero() + self.is_world_process_zero = trainer.is_world_process_zero() + + +class ExportableState: + """ + A class for objects that include the ability to have its state + be saved during `Trainer._save_checkpoint` and loaded back in during + `Trainer._load_from_checkpoint`. + + These must implement a `state` function that gets called during the respective + Trainer function call. It should only include parameters and attributes needed to + recreate the state at a particular time, to avoid utilizing pickle/maintain standard + file IO writing. + + Example: + + ```python + class EarlyStoppingCallback(TrainerCallback, ExportableState): + def __init__(self, early_stopping_patience: int = 1, early_stopping_threshold: Optional[float] = 0.0): + self.early_stopping_patience = early_stopping_patience + self.early_stopping_threshold = early_stopping_threshold + # early_stopping_patience_counter denotes the number of times validation metrics failed to improve. + self.early_stopping_patience_counter = 0 + + def state(self) -> dict: + return { + "args": { + "early_stopping_patience": self.early_stopping_patience, + "early_stopping_threshold": self.early_stopping_threshold, + }, + "attributes": { + "early_stopping_patience_counter": self.early_stopping_patience_counter, + } + } + ```""" + + def state(self) -> dict: + raise NotImplementedError("You must implement a `state` function to utilize this class.") + + @classmethod + def from_state(cls, state): + instance = cls(**state["args"]) + for k, v in state["attributes"].items(): + setattr(instance, k, v) + return instance + + +@dataclass +class TrainerControl(ExportableState): + """ + A class that handles the [`Trainer`] control flow. This class is used by the [`TrainerCallback`] to activate some + switches in the training loop. + + Args: + should_training_stop (`bool`, *optional*, defaults to `False`): + Whether or not the training should be interrupted. + + If `True`, this variable will not be set back to `False`. The training will just stop. + should_epoch_stop (`bool`, *optional*, defaults to `False`): + Whether or not the current epoch should be interrupted. + + If `True`, this variable will be set back to `False` at the beginning of the next epoch. + should_save (`bool`, *optional*, defaults to `False`): + Whether or not the model should be saved at this step. + + If `True`, this variable will be set back to `False` at the beginning of the next step. + should_evaluate (`bool`, *optional*, defaults to `False`): + Whether or not the model should be evaluated at this step. + + If `True`, this variable will be set back to `False` at the beginning of the next step. + should_log (`bool`, *optional*, defaults to `False`): + Whether or not the logs should be reported at this step. + + If `True`, this variable will be set back to `False` at the beginning of the next step. + """ + + should_training_stop: bool = False + should_epoch_stop: bool = False + should_save: bool = False + should_evaluate: bool = False + should_log: bool = False + + def _new_training(self): + """Internal method that resets the variable for a new training.""" + self.should_training_stop = False + + def _new_epoch(self): + """Internal method that resets the variable for a new epoch.""" + self.should_epoch_stop = False + + def _new_step(self): + """Internal method that resets the variable for a new step.""" + self.should_save = False + self.should_evaluate = False + self.should_log = False + + def state(self) -> dict: + return { + "args": { + "should_training_stop": self.should_training_stop, + "should_epoch_stop": self.should_epoch_stop, + "should_save": self.should_save, + "should_evaluate": self.should_evaluate, + "should_log": self.should_log, + }, + "attributes": {}, + } + + +class TrainerCallback: + # no-format + """ + A class for objects that will inspect the state of the training loop at some events and take some decisions. At + each of those events the following arguments are available: + + Args: + args ([`TrainingArguments`]): + The training arguments used to instantiate the [`Trainer`]. + state ([`TrainerState`]): + The current state of the [`Trainer`]. + control ([`TrainerControl`]): + The object that is returned to the [`Trainer`] and can be used to make some decisions. + model ([`PreTrainedModel`] or `torch.nn.Module`): + The model being trained. + tokenizer ([`PreTrainedTokenizer`]): + The tokenizer used for encoding the data. This is deprecated in favour of `processing_class`. + processing_class ([`PreTrainedTokenizer` or `BaseImageProcessor` or `ProcessorMixin` or `FeatureExtractionMixin`]): + The processing class used for encoding the data. Can be a tokenizer, a processor, an image processor or a feature extractor. + optimizer (`torch.optim.Optimizer`): + The optimizer used for the training steps. + lr_scheduler (`torch.optim.lr_scheduler.LambdaLR`): + The scheduler used for setting the learning rate. + train_dataloader (`torch.utils.data.DataLoader`, *optional*): + The current dataloader used for training. + eval_dataloader (`torch.utils.data.DataLoader`, *optional*): + The current dataloader used for evaluation. + metrics (`Dict[str, float]`): + The metrics computed by the last evaluation phase. + + Those are only accessible in the event `on_evaluate`. + logs (`Dict[str, float]`): + The values to log. + + Those are only accessible in the event `on_log`. + + The `control` object is the only one that can be changed by the callback, in which case the event that changes it + should return the modified version. + + The argument `args`, `state` and `control` are positionals for all events, all the others are grouped in `kwargs`. + You can unpack the ones you need in the signature of the event using them. As an example, see the code of the + simple [`~transformers.PrinterCallback`]. + + Example: + + ```python + class PrinterCallback(TrainerCallback): + def on_log(self, args, state, control, logs=None, **kwargs): + _ = logs.pop("total_flos", None) + if state.is_local_process_zero: + print(logs) + ```""" + + def on_init_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called at the end of the initialization of the [`Trainer`]. + """ + pass + + def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called at the beginning of training. + """ + pass + + def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called at the end of training. + """ + pass + + def on_epoch_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called at the beginning of an epoch. + """ + pass + + def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called at the end of an epoch. + """ + pass + + def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called at the beginning of a training step. If using gradient accumulation, one training step might take + several inputs. + """ + pass + + def on_pre_optimizer_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called before the optimizer step but after gradient clipping. Useful for monitoring gradients. + """ + pass + + def on_optimizer_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called after the optimizer step but before gradients are zeroed out. Useful for monitoring gradients. + """ + pass + + def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called at the end of an substep during gradient accumulation. + """ + pass + + def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called at the end of a training step. If using gradient accumulation, one training step might take + several inputs. + """ + pass + + def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called after an evaluation phase. + """ + pass + + def on_predict(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics, **kwargs): + """ + Event called after a successful prediction. + """ + pass + + def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called after a checkpoint save. + """ + pass + + def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called after logging the last logs. + """ + pass + + def on_prediction_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called after a prediction step. + """ + pass + + +class CallbackHandler(TrainerCallback): + """Internal class that just calls the list of callbacks in order.""" + + def __init__(self, callbacks, model, processing_class, optimizer, lr_scheduler): + self.callbacks = [] + for cb in callbacks: + self.add_callback(cb) + self.model = model + self.processing_class = processing_class + self.optimizer = optimizer + self.lr_scheduler = lr_scheduler + self.train_dataloader = None + self.eval_dataloader = None + + if not any(isinstance(cb, DefaultFlowCallback) for cb in self.callbacks): + logger.warning( + "The Trainer will not work properly if you don't have a `DefaultFlowCallback` in its callbacks. You\n" + + "should add one before training with `trainer.add_callback(DefaultFlowCallback). The current list of" + + "callbacks is\n:" + + self.callback_list + ) + + def add_callback(self, callback): + cb = callback() if isinstance(callback, type) else callback + cb_class = callback if isinstance(callback, type) else callback.__class__ + if cb_class in [c.__class__ for c in self.callbacks]: + logger.warning( + f"You are adding a {cb_class} to the callbacks of this Trainer, but there is already one. The current" + + "list of callbacks is\n:" + + self.callback_list + ) + self.callbacks.append(cb) + + def pop_callback(self, callback): + if isinstance(callback, type): + for cb in self.callbacks: + if isinstance(cb, callback): + self.callbacks.remove(cb) + return cb + else: + for cb in self.callbacks: + if cb == callback: + self.callbacks.remove(cb) + return cb + + def remove_callback(self, callback): + if isinstance(callback, type): + for cb in self.callbacks: + if isinstance(cb, callback): + self.callbacks.remove(cb) + return + else: + self.callbacks.remove(callback) + + @property + def callback_list(self): + return "\n".join(cb.__class__.__name__ for cb in self.callbacks) + + def on_init_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + return self.call_event("on_init_end", args, state, control) + + def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + control.should_training_stop = False + return self.call_event("on_train_begin", args, state, control) + + def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + return self.call_event("on_train_end", args, state, control) + + def on_epoch_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + control.should_epoch_stop = False + return self.call_event("on_epoch_begin", args, state, control) + + def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + return self.call_event("on_epoch_end", args, state, control) + + def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + control.should_log = False + control.should_evaluate = False + control.should_save = False + return self.call_event("on_step_begin", args, state, control) + + def on_pre_optimizer_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + return self.call_event("on_pre_optimizer_step", args, state, control) + + def on_optimizer_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + return self.call_event("on_optimizer_step", args, state, control) + + def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + return self.call_event("on_substep_end", args, state, control) + + def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + return self.call_event("on_step_end", args, state, control) + + def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics): + control.should_evaluate = False + return self.call_event("on_evaluate", args, state, control, metrics=metrics) + + def on_predict(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, metrics): + return self.call_event("on_predict", args, state, control, metrics=metrics) + + def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + control.should_save = False + return self.call_event("on_save", args, state, control) + + def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, logs): + control.should_log = False + return self.call_event("on_log", args, state, control, logs=logs) + + def on_prediction_step(self, args: TrainingArguments, state: TrainerState, control: TrainerControl): + return self.call_event("on_prediction_step", args, state, control) + + def call_event(self, event, args, state, control, **kwargs): + for callback in self.callbacks: + result = getattr(callback, event)( + args, + state, + control, + model=self.model, + processing_class=self.processing_class, + optimizer=self.optimizer, + lr_scheduler=self.lr_scheduler, + train_dataloader=self.train_dataloader, + eval_dataloader=self.eval_dataloader, + **kwargs, + ) + # A Callback can skip the return of `control` if it doesn't change it. + if result is not None: + control = result + return control + + +class DefaultFlowCallback(TrainerCallback): + """ + A [`TrainerCallback`] that handles the default flow of the training loop for logs, evaluation and checkpoints. + """ + + def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + # Log + if state.global_step == 1 and args.logging_first_step: + control.should_log = True + if args.logging_strategy == IntervalStrategy.STEPS and state.global_step % state.logging_steps == 0: + control.should_log = True + + # Evaluate + if ( + args.eval_strategy == IntervalStrategy.STEPS + and state.global_step % state.eval_steps == 0 + and args.eval_delay <= state.global_step + ): + control.should_evaluate = True + + # Save + if ( + args.save_strategy == SaveStrategy.STEPS + and state.save_steps > 0 + and state.global_step % state.save_steps == 0 + ): + control.should_save = True + + # End training + if state.global_step >= state.max_steps: + control.should_training_stop = True + # Save the model at the end if we have a save strategy + if args.save_strategy == SaveStrategy.STEPS: + control.should_save = True + + return control + + def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + # Log + if args.logging_strategy == IntervalStrategy.EPOCH: + control.should_log = True + + # Evaluate + if args.eval_strategy == IntervalStrategy.EPOCH and args.eval_delay <= state.epoch: + control.should_evaluate = True + + # Save + if args.save_strategy == SaveStrategy.EPOCH: + control.should_save = True + + return control + + +class ProgressCallback(TrainerCallback): + """ + A [`TrainerCallback`] that displays the progress of training or evaluation. + You can modify `max_str_len` to control how long strings are truncated when logging. + """ + + def __init__(self, max_str_len: int = 100): + """ + Initialize the callback with optional max_str_len parameter to control string truncation length. + + Args: + max_str_len (`int`): + Maximum length of strings to display in logs. + Longer strings will be truncated with a message. + """ + self.training_bar = None + self.prediction_bar = None + self.max_str_len = max_str_len + + def on_train_begin(self, args, state, control, **kwargs): + if state.is_world_process_zero: + self.training_bar = tqdm(total=state.max_steps, dynamic_ncols=True) + self.current_step = 0 + + def on_step_end(self, args, state, control, **kwargs): + if state.is_world_process_zero: + self.training_bar.update(state.global_step - self.current_step) + self.current_step = state.global_step + + def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs): + if state.is_world_process_zero and has_length(eval_dataloader): + if self.prediction_bar is None: + self.prediction_bar = tqdm( + total=len(eval_dataloader), leave=self.training_bar is None, dynamic_ncols=True + ) + self.prediction_bar.update(1) + + def on_evaluate(self, args, state, control, **kwargs): + if state.is_world_process_zero: + if self.prediction_bar is not None: + self.prediction_bar.close() + self.prediction_bar = None + + def on_predict(self, args, state, control, **kwargs): + if state.is_world_process_zero: + if self.prediction_bar is not None: + self.prediction_bar.close() + self.prediction_bar = None + + def on_log(self, args, state, control, logs=None, **kwargs): + if state.is_world_process_zero and self.training_bar is not None: + # make a shallow copy of logs so we can mutate the fields copied + # but avoid doing any value pickling. + shallow_logs = {} + for k, v in logs.items(): + if isinstance(v, str) and len(v) > self.max_str_len: + shallow_logs[k] = ( + f"[String too long to display, length: {len(v)} > {self.max_str_len}. " + "Consider increasing `max_str_len` if needed.]" + ) + else: + shallow_logs[k] = v + _ = shallow_logs.pop("total_flos", None) + # round numbers so that it looks better in console + if "epoch" in shallow_logs: + shallow_logs["epoch"] = round(shallow_logs["epoch"], 2) + self.training_bar.write(str(shallow_logs)) + + def on_train_end(self, args, state, control, **kwargs): + if state.is_world_process_zero: + self.training_bar.close() + self.training_bar = None + + +class PrinterCallback(TrainerCallback): + """ + A bare [`TrainerCallback`] that just prints the logs. + """ + + def on_log(self, args, state, control, logs=None, **kwargs): + _ = logs.pop("total_flos", None) + if state.is_local_process_zero: + print(logs) + + +class EarlyStoppingCallback(TrainerCallback, ExportableState): + """ + A [`TrainerCallback`] that handles early stopping. + + Args: + early_stopping_patience (`int`): + Use with `metric_for_best_model` to stop training when the specified metric worsens for + `early_stopping_patience` evaluation calls. + early_stopping_threshold(`float`, *optional*): + Use with TrainingArguments `metric_for_best_model` and `early_stopping_patience` to denote how much the + specified metric must improve to satisfy early stopping conditions. ` + + This callback depends on [`TrainingArguments`] argument *load_best_model_at_end* functionality to set best_metric + in [`TrainerState`]. Note that if the [`TrainingArguments`] argument *save_steps* differs from *eval_steps*, the + early stopping will not occur until the next save step. + """ + + def __init__(self, early_stopping_patience: int = 1, early_stopping_threshold: Optional[float] = 0.0): + self.early_stopping_patience = early_stopping_patience + self.early_stopping_threshold = early_stopping_threshold + # early_stopping_patience_counter denotes the number of times validation metrics failed to improve. + self.early_stopping_patience_counter = 0 + + def check_metric_value(self, args, state, control, metric_value): + # best_metric is set by code for load_best_model + operator = np.greater if args.greater_is_better else np.less + if state.best_metric is None or ( + operator(metric_value, state.best_metric) + and abs(metric_value - state.best_metric) > self.early_stopping_threshold + ): + self.early_stopping_patience_counter = 0 + else: + self.early_stopping_patience_counter += 1 + + def on_train_begin(self, args, state, control, **kwargs): + if not args.load_best_model_at_end: + logger.warning( + "Using EarlyStoppingCallback without load_best_model_at_end=True. " + "Once training is finished, the best model will not be loaded automatically." + ) + assert args.metric_for_best_model is not None, ( + "EarlyStoppingCallback requires metric_for_best_model to be defined" + ) + assert args.eval_strategy != IntervalStrategy.NO, ( + "EarlyStoppingCallback requires IntervalStrategy of steps or epoch" + ) + + def on_evaluate(self, args, state, control, metrics, **kwargs): + metric_to_check = args.metric_for_best_model + if not metric_to_check.startswith("eval_"): + metric_to_check = f"eval_{metric_to_check}" + metric_value = metrics.get(metric_to_check) + + if metric_value is None: + logger.warning( + f"early stopping required metric_for_best_model, but did not find {metric_to_check} so early stopping" + " is disabled" + ) + return + + self.check_metric_value(args, state, control, metric_value) + if self.early_stopping_patience_counter >= self.early_stopping_patience: + control.should_training_stop = True + + def state(self) -> dict: + return { + "args": { + "early_stopping_patience": self.early_stopping_patience, + "early_stopping_threshold": self.early_stopping_threshold, + }, + "attributes": { + "early_stopping_patience_counter": self.early_stopping_patience_counter, + }, + } diff --git a/lib/python3.12/site-packages/transformers/trainer_pt_utils.py b/lib/python3.12/site-packages/transformers/trainer_pt_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..30474daea6e435b6567b99d745f12cc8053f6e83 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/trainer_pt_utils.py @@ -0,0 +1,1411 @@ +# Copyright 2020-present the HuggingFace Inc. team. +# +# 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. +""" +Torch utilities for the Trainer class. +""" + +import copy +import datetime +import io +import json +import math +import os +import sys +import warnings +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from itertools import chain +from logging import StreamHandler +from typing import Any, Optional, Union + +import numpy as np +import torch +import torch.distributed as dist +from torch import nn +from torch.utils.data import Dataset, IterableDataset, RandomSampler, Sampler +from torch.utils.data.distributed import DistributedSampler + +from .integrations.deepspeed import is_deepspeed_zero3_enabled +from .tokenization_utils_base import BatchEncoding +from .utils import ( + is_sagemaker_mp_enabled, + is_torch_available, + is_torch_xla_available, + is_training_run_on_sagemaker, + logging, +) + + +if is_training_run_on_sagemaker(): + logging.add_handler(StreamHandler(sys.stdout)) + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + +if is_torch_available(): + from torch.optim.lr_scheduler import LRScheduler + + +logger = logging.get_logger(__name__) + + +def get_dataloader_sampler(dataloader): + if hasattr(dataloader, "batch_sampler") and dataloader.batch_sampler is not None: + return get_dataloader_sampler(dataloader.batch_sampler) + elif hasattr(dataloader, "sampler"): + return dataloader.sampler + + +def atleast_1d(tensor_or_array: Union[torch.Tensor, np.ndarray]): + if isinstance(tensor_or_array, torch.Tensor): + if hasattr(torch, "atleast_1d"): + tensor_or_array = torch.atleast_1d(tensor_or_array) + elif tensor_or_array.ndim < 1: + tensor_or_array = tensor_or_array[None] + else: + tensor_or_array = np.atleast_1d(tensor_or_array) + return tensor_or_array + + +def torch_pad_and_concatenate(tensor1, tensor2, padding_index=-100): + """Concatenates `tensor1` and `tensor2` on first axis, applying padding on the second if necessary.""" + tensor1 = atleast_1d(tensor1) + tensor2 = atleast_1d(tensor2) + + if len(tensor1.shape) == 1 or tensor1.shape[1] == tensor2.shape[1]: + return torch.cat((tensor1, tensor2), dim=0) + + # Let's figure out the new shape + new_shape = (tensor1.shape[0] + tensor2.shape[0], max(tensor1.shape[1], tensor2.shape[1])) + tensor1.shape[2:] + + # Now let's fill the result tensor + result = tensor1.new_full(new_shape, padding_index) + result[: tensor1.shape[0], : tensor1.shape[1]] = tensor1 + result[tensor1.shape[0] :, : tensor2.shape[1]] = tensor2 + return result + + +def numpy_pad_and_concatenate(array1, array2, padding_index=-100): + """Concatenates `array1` and `array2` on first axis, applying padding on the second if necessary.""" + array1 = atleast_1d(array1) + array2 = atleast_1d(array2) + + if len(array1.shape) == 1 or array1.shape[1] == array2.shape[1]: + return np.concatenate((array1, array2), axis=0) + + # Let's figure out the new shape + new_shape = (array1.shape[0] + array2.shape[0], max(array1.shape[1], array2.shape[1])) + array1.shape[2:] + + # Now let's fill the result tensor + result = np.full_like(array1, padding_index, shape=new_shape) + result[: array1.shape[0], : array1.shape[1]] = array1 + result[array1.shape[0] :, : array2.shape[1]] = array2 + return result + + +def nested_concat(tensors, new_tensors, padding_index=-100): + """ + Concat the `new_tensors` to `tensors` on the first dim and pad them on the second if needed. Works for tensors or + nested list/tuples/dict of tensors. + """ + if not (isinstance(tensors, torch.Tensor) and isinstance(new_tensors, torch.Tensor)): + assert type(tensors) is type(new_tensors), ( + f"Expected `tensors` and `new_tensors` to have the same type but found {type(tensors)} and {type(new_tensors)}." + ) + if isinstance(tensors, (list, tuple)): + return type(tensors)(nested_concat(t, n, padding_index=padding_index) for t, n in zip(tensors, new_tensors)) + elif isinstance(tensors, torch.Tensor): + return torch_pad_and_concatenate(tensors, new_tensors, padding_index=padding_index) + elif isinstance(tensors, Mapping): + return type(tensors)( + {k: nested_concat(t, new_tensors[k], padding_index=padding_index) for k, t in tensors.items()} + ) + elif isinstance(tensors, np.ndarray): + return numpy_pad_and_concatenate(tensors, new_tensors, padding_index=padding_index) + else: + raise TypeError(f"Unsupported type for concatenation: got {type(tensors)}") + + +def find_batch_size(tensors): + """ + Find the first dimension of a tensor in a nested list/tuple/dict of tensors. + """ + if isinstance(tensors, (list, tuple)): + for t in tensors: + result = find_batch_size(t) + if result is not None: + return result + elif isinstance(tensors, Mapping): + for key, value in tensors.items(): + result = find_batch_size(value) + if result is not None: + return result + elif isinstance(tensors, torch.Tensor): + return tensors.shape[0] if len(tensors.shape) >= 1 else None + elif isinstance(tensors, np.ndarray): + return tensors.shape[0] if len(tensors.shape) >= 1 else None + + +def nested_numpify(tensors): + "Numpify `tensors` (even if it's a nested list/tuple/dict of tensors)." + if isinstance(tensors, (list, tuple)): + return type(tensors)(nested_numpify(t) for t in tensors) + if isinstance(tensors, Mapping): + return type(tensors)({k: nested_numpify(t) for k, t in tensors.items()}) + + t = tensors.cpu() + if t.dtype == torch.bfloat16: + # As of Numpy 1.21.4, NumPy does not support bfloat16 (see + # https://github.com/numpy/numpy/blob/a47ecdea856986cd60eabbd53265c2ca5916ad5d/doc/source/user/basics.types.rst ). + # Until Numpy adds bfloat16, we must convert float32. + t = t.to(torch.float32) + return t.numpy() + + +def nested_detach(tensors): + "Detach `tensors` (even if it's a nested list/tuple/dict of tensors)." + if isinstance(tensors, (list, tuple)): + return type(tensors)(nested_detach(t) for t in tensors) + elif isinstance(tensors, Mapping): + return type(tensors)({k: nested_detach(t) for k, t in tensors.items()}) + return tensors.detach() if isinstance(tensors, torch.Tensor) else tensors + + +def nested_xla_mesh_reduce(tensors, name): + if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + if isinstance(tensors, (list, tuple)): + return type(tensors)(nested_xla_mesh_reduce(t, f"{name}_{i}") for i, t in enumerate(tensors)) + if isinstance(tensors, Mapping): + return type(tensors)( + {k: nested_xla_mesh_reduce(t, f"{name}_{i}") for i, (k, t) in enumerate(tensors.items())} + ) + + tensors = atleast_1d(tensors) + return xm.mesh_reduce(name, tensors, torch.cat) + else: + raise ImportError("Torch xla must be installed to use `nested_xla_mesh_reduce`") + + +def distributed_concat(tensor: Any, num_total_examples: Optional[int] = None) -> Any: + try: + if isinstance(tensor, (tuple, list)): + return type(tensor)(distributed_concat(t, num_total_examples) for t in tensor) + if isinstance(tensor, Mapping): + return type(tensor)({k: distributed_concat(t, num_total_examples) for k, t in tensor.items()}) + tensor = atleast_1d(tensor).contiguous() + output_tensors = [tensor.clone() for _ in range(dist.get_world_size())] + dist.all_gather(output_tensors, tensor) + concat = torch.cat(output_tensors, dim=0) + + # truncate the dummy elements added by SequentialDistributedSampler + if num_total_examples is not None: + concat = concat[:num_total_examples] + return concat + except AssertionError: + raise AssertionError("Not currently using distributed training") + + +def distributed_broadcast_scalars( + scalars: list[Union[int, float]], + num_total_examples: Optional[int] = None, + device: Optional[torch.device] = torch.device("cuda"), +) -> torch.Tensor: + try: + tensorized_scalar = torch.tensor(scalars, device=device) + output_tensors = [tensorized_scalar.clone() for _ in range(dist.get_world_size())] + dist.all_gather(output_tensors, tensorized_scalar) + concat = torch.cat(output_tensors, dim=0) + + # truncate the dummy elements added by SequentialDistributedSampler + if num_total_examples is not None: + concat = concat[:num_total_examples] + return concat + except AssertionError: + raise AssertionError("Not currently using distributed training") + + +def reissue_pt_warnings(caught_warnings): + # Reissue warnings + if len(caught_warnings) > 1: + for w in caught_warnings: + if w.category is not UserWarning: + warnings.warn(w.message, w.category) + + +@contextmanager +def torch_distributed_zero_first(local_rank: int): + """ + Decorator to make all processes in distributed training wait for each local_master to do something. + + Args: + local_rank (`int`): The rank of the local process. + """ + if local_rank not in [-1, 0]: + dist.barrier() + yield + if local_rank == 0: + dist.barrier() + + +class DistributedSamplerWithLoop(DistributedSampler): + """ + Like a torch.utils.data.distributed.DistributedSampler` but loops at the end back to the beginning of the shuffled + samples to make each process have a round multiple of batch_size samples. + + Args: + dataset (`torch.utils.data.Dataset`): + Dataset used for sampling. + batch_size (`int`): + The batch size used with this sampler + kwargs (`Dict[str, Any]`, *optional*): + All other keyword arguments passed to `DistributedSampler`. + """ + + def __init__(self, dataset, batch_size, **kwargs): + super().__init__(dataset, **kwargs) + self.batch_size = batch_size + + def __iter__(self): + indices = list(super().__iter__()) + remainder = 0 if len(indices) % self.batch_size == 0 else self.batch_size - len(indices) % self.batch_size + # DistributedSampler already added samples from the beginning to make the number of samples a round multiple + # of the world size, so we skip those. + start_remainder = 1 if self.rank < len(self.dataset) % self.num_replicas else 0 + indices += indices[start_remainder : start_remainder + remainder] + return iter(indices) + + +class EvalLoopContainer: + """ + Container to store intermediate results of evaluation loop. + + Args: + do_nested_concat (`bool`, *optional*, defaults to `True`): + If set to `True`, each iteration will recursively concatenate a new object containing tensors to + the existing stored tensors, provided that the structure of the existing object and the new one + are identical. If set to `False`, all newly added tensors will be stored in a list. + padding_index (`int`, *optional*, defaults to -100): + Value used to pad tensors of different shapes when `do_nested_concat=True`. + """ + + def __init__(self, do_nested_concat: bool = True, padding_index: int = -100): + self.do_nested_concat = do_nested_concat + self.padding_index = padding_index + self.tensors = None + self.arrays = None + + def add(self, tensors) -> None: + """Add tensors to the stored objects. If `do_nested_concat=True`, the tensors will be concatenated recursively.""" + if self.tensors is None: + self.tensors = tensors if self.do_nested_concat else [tensors] + elif self.do_nested_concat: + self.tensors = nested_concat(self.tensors, tensors, padding_index=self.padding_index) + else: + self.tensors.append(tensors) + + def to_cpu_and_numpy(self) -> None: + """Move tensors in stored objects to CPU and convert them to numpy arrays.""" + + # Check if we have something to add, if not just return + if self.tensors is None: + return + + new_arrays = nested_numpify(self.tensors) + if self.arrays is None: + self.arrays = new_arrays + elif self.do_nested_concat: + self.arrays = nested_concat(self.arrays, new_arrays, padding_index=self.padding_index) + else: + self.arrays.extend(new_arrays) + + # reset device tensors after adding to cpu + self.tensors = None + + def get_arrays(self): + """Returns the numpified and moved to CPU stored objects.""" + self.to_cpu_and_numpy() + return self.arrays + + +class SequentialDistributedSampler(Sampler): + """ + Distributed Sampler that subsamples indices sequentially, making it easier to collate all results at the end. + + Even though we only use this sampler for eval and predict (no training), which means that the model params won't + have to be synced (i.e. will not hang for synchronization even if varied number of forward passes), we still add + extra samples to the sampler to make it evenly divisible (like in `DistributedSampler`) to make it easy to `gather` + or `reduce` resulting tensors at the end of the loop. + """ + + def __init__(self, dataset, num_replicas=None, rank=None, batch_size=None): + warnings.warn( + "SequentialDistributedSampler is deprecated and will be removed in v5 of Transformers.", + FutureWarning, + ) + if num_replicas is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + num_replicas = dist.get_world_size() + if rank is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + rank = dist.get_rank() + self.dataset = dataset + self.num_replicas = num_replicas + self.rank = rank + num_samples = len(self.dataset) + # Add extra samples to make num_samples a multiple of batch_size if passed + if batch_size is not None: + self.num_samples = int(math.ceil(num_samples / (batch_size * num_replicas))) * batch_size + else: + self.num_samples = int(math.ceil(num_samples / num_replicas)) + self.total_size = self.num_samples * self.num_replicas + self.batch_size = batch_size + + def __iter__(self): + indices = list(range(len(self.dataset))) + + # add extra samples to make it evenly divisible + indices += indices[: (self.total_size - len(indices))] + assert len(indices) == self.total_size, ( + f"Indices length {len(indices)} and total size {self.total_size} mismatched" + ) + + # subsample + indices = indices[self.rank * self.num_samples : (self.rank + 1) * self.num_samples] + assert len(indices) == self.num_samples, ( + f"Indices length {len(indices)} and sample number {self.num_samples} mismatched" + ) + + return iter(indices) + + def __len__(self): + return self.num_samples + + +def get_tpu_sampler(dataset: torch.utils.data.Dataset, batch_size: int): + if xm.xrt_world_size() <= 1: + return RandomSampler(dataset) + return DistributedSampler(dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal()) + + +def nested_new_like(arrays, num_samples, padding_index=-100): + """Create the same nested structure as `arrays` with a first dimension always at `num_samples`.""" + if isinstance(arrays, (list, tuple)): + return type(arrays)(nested_new_like(x, num_samples) for x in arrays) + return np.full_like(arrays, padding_index, shape=(num_samples, *arrays.shape[1:])) + + +def expand_like(arrays, new_seq_length, padding_index=-100): + """Expand the `arrays` so that the second dimension grows to `new_seq_length`. Uses `padding_index` for padding.""" + result = np.full_like(arrays, padding_index, shape=(arrays.shape[0], new_seq_length) + arrays.shape[2:]) + result[:, : arrays.shape[1]] = arrays + return result + + +def nested_truncate(tensors, limit): + "Truncate `tensors` at `limit` (even if it's a nested list/tuple/dict of tensors)." + if isinstance(tensors, (list, tuple)): + return type(tensors)(nested_truncate(t, limit) for t in tensors) + if isinstance(tensors, Mapping): + return type(tensors)({k: nested_truncate(t, limit) for k, t in tensors.items()}) + + return tensors[:limit] + + +class DistributedTensorGatherer: + """ + A class responsible for properly gathering tensors (or nested list/tuple of tensors) on the CPU by chunks. + + If our dataset has 16 samples with a batch size of 2 on 3 processes and we gather then transfer on CPU at every + step, our sampler will generate the following indices: + + `[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 0, 1]` + + to get something of size a multiple of 3 (so that each process gets the same dataset length). Then process 0, 1 and + 2 will be responsible of making predictions for the following samples: + + - P0: `[0, 1, 2, 3, 4, 5]` + - P1: `[6, 7, 8, 9, 10, 11]` + - P2: `[12, 13, 14, 15, 0, 1]` + + The first batch treated on each process will be: + + - P0: `[0, 1]` + - P1: `[6, 7]` + - P2: `[12, 13]` + + So if we gather at the end of the first batch, we will get a tensor (nested list/tuple of tensor) corresponding to + the following indices: + + `[0, 1, 6, 7, 12, 13]` + + If we directly concatenate our results without taking any precautions, the user will then get the predictions for + the indices in this order at the end of the prediction loop: + + `[0, 1, 6, 7, 12, 13, 2, 3, 8, 9, 14, 15, 4, 5, 10, 11, 0, 1]` + + For some reason, that's not going to roll their boat. This class is there to solve that problem. + + Args: + world_size (`int`): + The number of processes used in the distributed training. + num_samples (`int`): + The number of samples in our dataset. + make_multiple_of (`int`, *optional*): + If passed, the class assumes the datasets passed to each process are made to be a multiple of this argument + (by adding samples). + padding_index (`int`, *optional*, defaults to -100): + The padding index to use if the arrays don't all have the same sequence length. + """ + + def __init__(self, world_size, num_samples, make_multiple_of=None, padding_index=-100): + warnings.warn( + "DistributedTensorGatherer is deprecated and will be removed in v5 of Transformers.", + FutureWarning, + ) + self.world_size = world_size + self.num_samples = num_samples + total_size = world_size if make_multiple_of is None else world_size * make_multiple_of + self.total_samples = int(np.ceil(num_samples / total_size)) * total_size + self.process_length = self.total_samples // world_size + self._storage = None + self._offsets = None + self.padding_index = padding_index + + def add_arrays(self, arrays): + """ + Add `arrays` to the internal storage, Will initialize the storage to the full size at the first arrays passed + so that if we're bound to get an OOM, it happens at the beginning. + """ + if arrays is None: + return + if self._storage is None: + self._storage = nested_new_like(arrays, self.total_samples, padding_index=self.padding_index) + self._offsets = list(range(0, self.total_samples, self.process_length)) + + slice_len, self._storage = self._nested_set_tensors(self._storage, arrays) + for i in range(self.world_size): + self._offsets[i] += slice_len + + def _nested_set_tensors(self, storage, arrays): + if isinstance(arrays, (list, tuple)): + result = [self._nested_set_tensors(x, y) for x, y in zip(storage, arrays)] + return result[0][0], type(arrays)(r[1] for r in result) + assert arrays.shape[0] % self.world_size == 0, ( + f"Arrays passed should all have a first dimension multiple of {self.world_size}, found {arrays.shape[0]}." + ) + + slice_len = arrays.shape[0] // self.world_size + for i in range(self.world_size): + if len(arrays.shape) == 1: + storage[self._offsets[i] : self._offsets[i] + slice_len] = arrays[i * slice_len : (i + 1) * slice_len] + else: + # Expand the array on the fly if needed. + if len(storage.shape) > 1 and storage.shape[1] < arrays.shape[1]: + storage = expand_like(storage, arrays.shape[1], padding_index=self.padding_index) + storage[self._offsets[i] : self._offsets[i] + slice_len, : arrays.shape[1]] = arrays[ + i * slice_len : (i + 1) * slice_len + ] + return slice_len, storage + + def finalize(self): + """ + Return the properly gathered arrays and truncate to the number of samples (since the sampler added some extras + to get each process a dataset of the same length). + """ + if self._storage is None: + return + if self._offsets[0] != self.process_length: + logger.warning("Not all data has been set. Are you sure you passed all values?") + return nested_truncate(self._storage, self.num_samples) + + +@dataclass +class LabelSmoother: + """ + Adds label-smoothing on a pre-computed output from a Transformers model. + + Args: + epsilon (`float`, *optional*, defaults to 0.1): + The label smoothing factor. + ignore_index (`int`, *optional*, defaults to -100): + The index in the labels to ignore when computing the loss. + """ + + epsilon: float = 0.1 + ignore_index: int = -100 + + def __call__(self, model_output, labels, shift_labels=False): + logits = model_output["logits"] if isinstance(model_output, dict) else model_output[0] + if shift_labels: + logits = logits[..., :-1, :].contiguous() + labels = labels[..., 1:].contiguous() + + log_probs = -nn.functional.log_softmax(logits, dim=-1) + if labels.dim() == log_probs.dim() - 1: + labels = labels.unsqueeze(-1) + + padding_mask = labels.eq(self.ignore_index) + # In case the ignore_index is -100, the gather will fail, so we replace labels by 0. The padding_mask + # will ignore them in any case. + labels = torch.clamp(labels, min=0) + nll_loss = log_probs.gather(dim=-1, index=labels) + # works for fp16 input tensor too, by internally upcasting it to fp32 + smoothed_loss = log_probs.sum(dim=-1, keepdim=True, dtype=torch.float32) + + nll_loss.masked_fill_(padding_mask, 0.0) + smoothed_loss.masked_fill_(padding_mask, 0.0) + + # Take the mean over the label dimensions, then divide by the number of active elements (i.e. not-padded): + num_active_elements = padding_mask.numel() - padding_mask.long().sum() + nll_loss = nll_loss.sum() / num_active_elements + smoothed_loss = smoothed_loss.sum() / (num_active_elements * log_probs.shape[-1]) + return (1 - self.epsilon) * nll_loss + self.epsilon * smoothed_loss + + +def get_length_grouped_indices(lengths, batch_size, mega_batch_mult=None, generator=None): + """ + Return a list of indices so that each slice of `batch_size` consecutive indices correspond to elements of similar + lengths. To do this, the indices are: + + - randomly permuted + - grouped in mega-batches of size `mega_batch_mult * batch_size` + - sorted by length in each mega-batch + + The result is the concatenation of all mega-batches, with the batch of `batch_size` containing the element of + maximum length placed first, so that an OOM happens sooner rather than later. + """ + # Default for mega_batch_mult: 50 or the number to get 4 megabatches, whichever is smaller. + if mega_batch_mult is None: + mega_batch_mult = min(len(lengths) // (batch_size * 4), 50) + # Just in case, for tiny datasets + if mega_batch_mult == 0: + mega_batch_mult = 1 + + # We need to use torch for the random part as a distributed sampler will set the random seed for torch. + indices = torch.randperm(len(lengths), generator=generator) + megabatch_size = mega_batch_mult * batch_size + megabatches = [indices[i : i + megabatch_size].tolist() for i in range(0, len(lengths), megabatch_size)] + megabatches = [sorted(megabatch, key=lambda i: lengths[i], reverse=True) for megabatch in megabatches] + + # The rest is to get the biggest batch first. + # Since each megabatch is sorted by descending length, the longest element is the first + megabatch_maximums = [lengths[megabatch[0]] for megabatch in megabatches] + max_idx = torch.argmax(torch.tensor(megabatch_maximums)).item() + # Switch to put the longest element in first position + megabatches[0][0], megabatches[max_idx][0] = megabatches[max_idx][0], megabatches[0][0] + + return [i for megabatch in megabatches for i in megabatch] + + +class LengthGroupedSampler(Sampler): + r""" + Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while + keeping a bit of randomness. + """ + + def __init__( + self, + batch_size: int, + dataset: Optional[Dataset] = None, + lengths: Optional[list[int]] = None, + model_input_name: Optional[str] = None, + generator=None, + ): + if dataset is None and lengths is None: + raise ValueError("One of dataset and lengths must be provided.") + + self.batch_size = batch_size + if lengths is None: + model_input_name = model_input_name if model_input_name is not None else "input_ids" + if ( + not (isinstance(dataset[0], dict) or isinstance(dataset[0], BatchEncoding)) + or model_input_name not in dataset[0] + ): + raise ValueError( + "Can only automatically infer lengths for datasets whose items are dictionaries with an " + f"'{model_input_name}' key." + ) + lengths = [len(feature[model_input_name]) for feature in dataset] + elif isinstance(lengths, torch.Tensor): + logger.info( + "If lengths is a torch.Tensor, LengthGroupedSampler will be slow. Converting lengths to List[int]..." + ) + lengths = lengths.tolist() + + self.lengths = lengths + self.generator = generator + + def __len__(self): + return len(self.lengths) + + def __iter__(self): + indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=self.generator) + return iter(indices) + + +class DistributedLengthGroupedSampler(DistributedSampler): + r""" + Distributed Sampler that samples indices in a way that groups together features of the dataset of roughly the same + length while keeping a bit of randomness. + """ + + # Copied and adapted from PyTorch DistributedSampler. + def __init__( + self, + batch_size: int, + dataset: Optional[Dataset] = None, + num_replicas: Optional[int] = None, + rank: Optional[int] = None, + seed: int = 0, + drop_last: bool = False, + lengths: Optional[list[int]] = None, + model_input_name: Optional[str] = None, + ): + if dataset is None and lengths is None: + raise ValueError("One of dataset and lengths must be provided.") + if num_replicas is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + num_replicas = dist.get_world_size() + if rank is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + rank = dist.get_rank() + + self.batch_size = batch_size + self.num_replicas = num_replicas + self.rank = rank + self.epoch = 0 + self.drop_last = drop_last + + if lengths is None: + model_input_name = model_input_name if model_input_name is not None else "input_ids" + if ( + not (isinstance(dataset[0], dict) or isinstance(dataset[0], BatchEncoding)) + or model_input_name not in dataset[0] + ): + raise ValueError( + "Can only automatically infer lengths for datasets whose items are dictionaries with an " + f"'{model_input_name}' key." + ) + lengths = [len(feature[model_input_name]) for feature in dataset] + elif isinstance(lengths, torch.Tensor): + logger.info( + "If lengths is a torch.Tensor, DistributedLengthGroupedSampler will be slow. Converting lengths to" + " List[int]..." + ) + lengths = lengths.tolist() + + self.lengths = lengths + + # If the dataset length is evenly divisible by # of replicas, then there + # is no need to drop any data, since the dataset will be split equally. + if self.drop_last and len(self.lengths) % self.num_replicas != 0: + # Split to nearest available length that is evenly divisible. + # This is to ensure each rank receives the same amount of data when + # using this Sampler. + self.num_samples = math.ceil((len(self.lengths) - self.num_replicas) / self.num_replicas) + else: + self.num_samples = math.ceil(len(self.lengths) / self.num_replicas) + self.total_size = self.num_samples * self.num_replicas + self.seed = seed + + def __iter__(self) -> Iterator: + # Deterministically shuffle based on epoch and seed + g = torch.Generator() + g.manual_seed(self.seed + self.epoch) + indices = get_length_grouped_indices(self.lengths, self.batch_size, generator=g) + + if not self.drop_last: + # add extra samples to make it evenly divisible + indices += indices[: (self.total_size - len(indices))] + else: + # remove tail of data to make it evenly divisible + indices = indices[: self.total_size] + assert len(indices) == self.total_size + + # subsample + indices = indices[self.rank : self.total_size : self.num_replicas] + assert len(indices) == self.num_samples + + return iter(indices) + + +class ShardSampler(Sampler): + """ + Sampler that shards batches between several processes. Dispatches indices batch by batch: on 2 processes with batch + size 4, the first two batches are `[0, 1, 2, 3, 4, 5, 6, 7]` and `[8, 9, 10, 11, 12, 13, 14, 15]`, which shard into + `[0, 1, 2, 3]` and `[8, 9, 10, 11]` for GPU-0 and `[4, 5, 6, 7]` and `[12, 13, 14, 15]` for GPU-1. + + The sampler thus yields `[0, 1, 2, 3, 8, 9, 10, 11]` on GPU-0 and `[4, 5, 6, 7, 12, 13, 14, 15]` on GPU-1. + """ + + def __init__( + self, + dataset: Dataset, + batch_size: int = 1, + drop_last: bool = False, + num_processes: int = 1, + process_index: int = 0, + ): + self.dataset = dataset + self.batch_size = batch_size + self.drop_last = drop_last + self.num_processes = num_processes + self.process_index = process_index + + self.total_batch_size = total_batch_size = batch_size * num_processes + + num_batches = len(dataset) // total_batch_size if drop_last else math.ceil(len(dataset) / total_batch_size) + self.total_num_samples = num_batches * total_batch_size + + def __iter__(self): + indices = list(range(len(self.dataset))) + + # Add extra samples to make it evenly divisible. While loop is there in the edge case we have a tiny dataset + # and it needs to be done several times. + while len(indices) < self.total_num_samples: + indices += indices[: (self.total_num_samples - len(indices))] + + result = [] + for batch_start in range(self.batch_size * self.process_index, self.total_num_samples, self.total_batch_size): + result += indices[batch_start : batch_start + self.batch_size] + + return iter(result) + + def __len__(self): + # Each shard only sees a fraction of total_num_samples. + return self.total_num_samples // self.num_processes + + +class IterableDatasetShard(IterableDataset): + """ + Wraps a PyTorch `IterableDataset` to generate samples for one of the processes only. Instances of this class will + always yield a number of samples that is a round multiple of the actual batch size (which is `batch_size x + num_processes`). Depending on the value of the `drop_last` attribute, it will either stop the iteration at the + first batch that would be too small or loop with indices from the beginning. + + On two processes with an iterable dataset yielding of `[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]` with a batch size of + 2: + + - the shard on process 0 will yield `[0, 1, 4, 5, 8, 9]` so will see batches `[0, 1]`, `[4, 5]`, `[8, 9]` + - the shard on process 1 will yield `[2, 3, 6, 7, 10, 11]` so will see batches `[2, 3]`, `[6, 7]`, `[10, 11]` + + + + If your IterableDataset implements some randomization that needs to be applied the same way on all processes + (for instance, a shuffling), you should use a `torch.Generator` in a `generator` attribute of the `dataset` to + generate your random numbers and call the [`~trainer_pt_utils.IterableDatasetShard.set_epoch`] method of this + object. It will set the seed of this `generator` to `seed + epoch` on all processes before starting the + iteration. Alternatively, you can also implement a `set_epoch()` method in your iterable dataset to deal with + this. + + + + Args: + dataset (`torch.utils.data.IterableDataset`): + The batch sampler to split in several shards. + batch_size (`int`, *optional*, defaults to 1): + The size of the batches per shard. + drop_last (`bool`, *optional*, defaults to `False`): + Whether or not to drop the last incomplete batch or complete the last batches by using the samples from the + beginning. + num_processes (`int`, *optional*, defaults to 1): + The number of processes running concurrently. + process_index (`int`, *optional*, defaults to 0): + The index of the current process. + seed (`int`, *optional*, defaults to 0): + A random seed that will be used for the random number generation in + [`~trainer_pt_utils.IterableDatasetShard.set_epoch`]. + """ + + def __init__( + self, + dataset: IterableDataset, + batch_size: int = 1, + drop_last: bool = False, + num_processes: int = 1, + process_index: int = 0, + seed: int = 0, + ): + self.dataset = dataset + self.batch_size = batch_size + self.drop_last = drop_last + self.num_processes = num_processes + self.process_index = process_index + self.seed = seed + self.epoch = 0 + self.num_examples = 0 + + def set_epoch(self, epoch): + self.epoch = epoch + if hasattr(self.dataset, "set_epoch"): + self.dataset.set_epoch(epoch) + + def __iter__(self): + self.num_examples = 0 + if ( + not hasattr(self.dataset, "set_epoch") + and hasattr(self.dataset, "generator") + and isinstance(self.dataset.generator, torch.Generator) + ): + self.dataset.generator.manual_seed(self.seed + self.epoch) + real_batch_size = self.batch_size * self.num_processes + process_slice = range(self.process_index * self.batch_size, (self.process_index + 1) * self.batch_size) + + first_batch = None + current_batch = [] + for element in self.dataset: + self.num_examples += 1 + current_batch.append(element) + # Wait to have a full batch before yielding elements. + if len(current_batch) == real_batch_size: + for i in process_slice: + yield current_batch[i] + if first_batch is None: + first_batch = current_batch.copy() + current_batch = [] + + # Finished if drop_last is True, otherwise complete the last batch with elements from the beginning. + if not self.drop_last and len(current_batch) > 0: + if first_batch is None: + first_batch = current_batch.copy() + while len(current_batch) < real_batch_size: + current_batch += first_batch + for i in process_slice: + yield current_batch[i] + + def __len__(self): + # Will raise an error if the underlying dataset is not sized. + if self.drop_last: + return (len(self.dataset) // (self.batch_size * self.num_processes)) * self.batch_size + else: + return math.ceil(len(self.dataset) / (self.batch_size * self.num_processes)) * self.batch_size + + +# In order to keep `trainer.py` compact and easy to understand, place any secondary PT Trainer +# helper methods here + + +def _get_learning_rate(self): + if self.is_deepspeed_enabled: + # with deepspeed's fp16 and dynamic loss scale enabled the optimizer/scheduler steps may + # not run for the first few dozen steps while loss scale is too large, and thus during + # that time `get_last_lr` will fail if called during that warm up stage, so work around it: + try: + last_lr = self.lr_scheduler.get_last_lr()[0] + except AssertionError as e: + if "need to call step" in str(e): + logger.warning("tried to get lr value before scheduler/optimizer started stepping, returning lr=0") + last_lr = 0 + else: + raise + else: + if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): + last_lr = self.optimizer.param_groups[0]["lr"] + else: + last_lr = self.lr_scheduler.get_last_lr()[0] + if torch.is_tensor(last_lr): + last_lr = last_lr.item() + return last_lr + + +def _secs2timedelta(secs): + """ + Convert seconds to hh:mm:ss.msec, msecs rounded to 2 decimal places. + """ + + msec = int(abs(secs - int(secs)) * 100) + return f"{datetime.timedelta(seconds=int(secs))}.{msec:02d}" + + +def metrics_format(self, metrics: dict[str, float]) -> dict[str, float]: + """ + Reformat Trainer metrics values to a human-readable format. + + Args: + metrics (`Dict[str, float]`): + The metrics returned from train/evaluate/predict + + Returns: + metrics (`Dict[str, float]`): The reformatted metrics + """ + + metrics_copy = metrics.copy() + for k, v in metrics_copy.items(): + if "_mem_" in k: + metrics_copy[k] = f"{v >> 20}MB" + elif "_runtime" in k: + metrics_copy[k] = _secs2timedelta(v) + elif k == "total_flos": + metrics_copy[k] = f"{int(v) >> 30}GF" + elif isinstance(metrics_copy[k], float): + metrics_copy[k] = round(v, 4) + + return metrics_copy + + +def log_metrics(self, split, metrics): + """ + Log metrics in a specially formatted way. + + Under distributed environment this is done only for a process with rank 0. + + Args: + split (`str`): + Mode/split name: one of `train`, `eval`, `test` + metrics (`Dict[str, float]`): + The metrics returned from train/evaluate/predictmetrics: metrics dict + + Notes on memory reports: + + In order to get memory usage report you need to install `psutil`. You can do that with `pip install psutil`. + + Now when this method is run, you will see a report that will include: + + ``` + init_mem_cpu_alloc_delta = 1301MB + init_mem_cpu_peaked_delta = 154MB + init_mem_gpu_alloc_delta = 230MB + init_mem_gpu_peaked_delta = 0MB + train_mem_cpu_alloc_delta = 1345MB + train_mem_cpu_peaked_delta = 0MB + train_mem_gpu_alloc_delta = 693MB + train_mem_gpu_peaked_delta = 7MB + ``` + + **Understanding the reports:** + + - the first segment, e.g., `train__`, tells you which stage the metrics are for. Reports starting with `init_` + will be added to the first stage that gets run. So that if only evaluation is run, the memory usage for the + `__init__` will be reported along with the `eval_` metrics. + - the third segment, is either `cpu` or `gpu`, tells you whether it's the general RAM or the gpu0 memory + metric. + - `*_alloc_delta` - is the difference in the used/allocated memory counter between the end and the start of the + stage - it can be negative if a function released more memory than it allocated. + - `*_peaked_delta` - is any extra memory that was consumed and then freed - relative to the current allocated + memory counter - it is never negative. When you look at the metrics of any stage you add up `alloc_delta` + + `peaked_delta` and you know how much memory was needed to complete that stage. + + The reporting happens only for process of rank 0 and gpu 0 (if there is a gpu). Typically this is enough since the + main process does the bulk of work, but it could be not quite so if model parallel is used and then other GPUs may + use a different amount of gpu memory. This is also not the same under DataParallel where gpu0 may require much more + memory than the rest since it stores the gradient and optimizer states for all participating GPUs. Perhaps in the + future these reports will evolve to measure those too. + + The CPU RAM metric measures RSS (Resident Set Size) includes both the memory which is unique to the process and the + memory shared with other processes. It is important to note that it does not include swapped out memory, so the + reports could be imprecise. + + The CPU peak memory is measured using a sampling thread. Due to python's GIL it may miss some of the peak memory if + that thread didn't get a chance to run when the highest memory was used. Therefore this report can be less than + reality. Using `tracemalloc` would have reported the exact peak memory, but it doesn't report memory allocations + outside of python. So if some C++ CUDA extension allocated its own memory it won't be reported. And therefore it + was dropped in favor of the memory sampling approach, which reads the current process memory usage. + + The GPU allocated and peak memory reporting is done with `torch.cuda.memory_allocated()` and + `torch.cuda.max_memory_allocated()`. This metric reports only "deltas" for pytorch-specific allocations, as + `torch.cuda` memory management system doesn't track any memory allocated outside of pytorch. For example, the very + first cuda call typically loads CUDA kernels, which may take from 0.5 to 2GB of GPU memory. + + Note that this tracker doesn't account for memory allocations outside of [`Trainer`]'s `__init__`, `train`, + `evaluate` and `predict` calls. + + Because `evaluation` calls may happen during `train`, we can't handle nested invocations because + `torch.cuda.max_memory_allocated` is a single counter, so if it gets reset by a nested eval call, `train`'s tracker + will report incorrect info. If this [pytorch issue](https://github.com/pytorch/pytorch/issues/16266) gets resolved + it will be possible to change this class to be re-entrant. Until then we will only track the outer level of + `train`, `evaluate` and `predict` methods. Which means that if `eval` is called during `train`, it's the latter + that will account for its memory usage and that of the former. + + This also means that if any other tool that is used along the [`Trainer`] calls + `torch.cuda.reset_peak_memory_stats`, the gpu peak memory stats could be invalid. And the [`Trainer`] will disrupt + the normal behavior of any such tools that rely on calling `torch.cuda.reset_peak_memory_stats` themselves. + + For best performance you may want to consider turning the memory profiling off for production runs. + """ + if not self.is_world_process_zero(): + return + + print(f"***** {split} metrics *****") + metrics_formatted = self.metrics_format(metrics) + k_width = max(len(str(x)) for x in metrics_formatted.keys()) + v_width = max(len(str(x)) for x in metrics_formatted.values()) + for key in sorted(metrics_formatted.keys()): + print(f" {key: <{k_width}} = {metrics_formatted[key]:>{v_width}}") + + +def save_metrics(self, split, metrics, combined=True): + """ + Save metrics into a json file for that split, e.g. `train_results.json`. + + Under distributed environment this is done only for a process with rank 0. + + Args: + split (`str`): + Mode/split name: one of `train`, `eval`, `test`, `all` + metrics (`Dict[str, float]`): + The metrics returned from train/evaluate/predict + combined (`bool`, *optional*, defaults to `True`): + Creates combined metrics by updating `all_results.json` with metrics of this call + + To understand the metrics please read the docstring of [`~Trainer.log_metrics`]. The only difference is that raw + unformatted numbers are saved in the current method. + + """ + if not self.is_world_process_zero(): + return + + path = os.path.join(self.args.output_dir, f"{split}_results.json") + with open(path, "w") as f: + json.dump(metrics, f, indent=4, sort_keys=True) + + if combined: + path = os.path.join(self.args.output_dir, "all_results.json") + if os.path.exists(path): + with open(path) as f: + all_metrics = json.load(f) + else: + all_metrics = {} + + all_metrics.update(metrics) + with open(path, "w") as f: + json.dump(all_metrics, f, indent=4, sort_keys=True) + + +def save_state(self): + """ + Saves the Trainer state, since Trainer.save_model saves only the tokenizer with the model. + + Under distributed environment this is done only for a process with rank 0. + """ + if not self.is_world_process_zero(): + return + + path = os.path.join(self.args.output_dir, "trainer_state.json") + self.state.save_to_json(path) + + +def get_model_param_count(model, trainable_only=False): + """ + Calculate model's total param count. If trainable_only is True then count only those requiring grads. + """ + if is_deepspeed_zero3_enabled(): + + def numel(p): + return p.ds_numel if hasattr(p, "ds_numel") else p.numel() + + else: + + def numel(p): + return p.numel() + + return sum(numel(p) for p in model.parameters() if not trainable_only or p.requires_grad) + + +def get_parameter_names(model, forbidden_layer_types, forbidden_layer_names=None): + """ + Returns the names of the model parameters that are not inside a forbidden layer. + """ + if forbidden_layer_names is None: + forbidden_layer_names = [] + result = [] + for name, child in model.named_children(): + child_params = get_parameter_names(child, forbidden_layer_types, forbidden_layer_names) + result += [ + f"{name}.{n}" + for n in child_params + if not isinstance(child, tuple(forbidden_layer_types)) + and not any(forbidden in f"{name}.{n}".lower() for forbidden in forbidden_layer_names) + ] + # Add model specific parameters that are not in any child + result += [ + k for k in model._parameters.keys() if not any(forbidden in k.lower() for forbidden in forbidden_layer_names) + ] + return result + + +def get_module_class_from_name(module, name): + """ + Gets a class from a module by its name. + + Args: + module (`torch.nn.Module`): The module to get the class from. + name (`str`): The name of the class. + """ + modules_children = list(module.children()) + if module.__class__.__name__ == name: + return module.__class__ + elif len(modules_children) == 0: + return + else: + for child_module in modules_children: + module_class = get_module_class_from_name(child_module, name) + if module_class is not None: + return module_class + + +def remove_dummy_checkpoint(is_main_process, output_dir, filenames): + if is_main_process: + for filename in filenames: + file = os.path.join(output_dir, filename) + if os.path.isfile(file): + os.remove(file) + + +if is_sagemaker_mp_enabled(): + import smdistributed.modelparallel.torch as smp + + @smp.step() + def smp_forward_backward(model, inputs, gradient_accumulation_steps=1): + outputs = model(**inputs) + loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0] + loss /= gradient_accumulation_steps + model.backward(loss) + return loss + + @smp.step() + def smp_forward_only(model, inputs): + return model(**inputs) + + def smp_gather(tensor): + if isinstance(tensor, (list, tuple)): + return type(tensor)(smp_gather(t) for t in tensor) + elif isinstance(tensor, dict): + return type(tensor)({k: smp_gather(v) for k, v in tensor.items()}) + elif not isinstance(tensor, torch.Tensor): + raise TypeError( + f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors." + ) + all_tensors = smp.allgather(tensor, smp.CommGroup.DP_GROUP) + all_tensors = [atleast_1d(t) for t in all_tensors] + return torch.cat([t.cpu() for t in all_tensors], dim=0) + + def smp_nested_concat(tensor): + if isinstance(tensor, (list, tuple)): + return type(tensor)(smp_nested_concat(t) for t in tensor) + elif isinstance(tensor, dict): + return type(tensor)({k: smp_nested_concat(v) for k, v in tensor.items()}) + # It doesn't seem possible to check here if `tensor` is a StepOutput because StepOutput lives in `smp.step` + # which is also the name of the decorator so Python is confused. + return tensor.detach().concat().cpu() + + +@dataclass +class AcceleratorConfig: + """ + A subset of arguments relating to the underlying [`accelerate.Accelerator`] + implementation utilized in the `Trainer` that can be customized. + Mostly relating to data. + + Parameters: + split_batches (`bool`, *optional*, defaults to `False`): + Whether or not the accelerator should split the batches yielded by the dataloaders across the devices. If + `True` the actual batch size used will be the same on any kind of distributed processes, but it must be a + round multiple of the `num_processes` you are using. If `False`, actual batch size used will be the one set + in your script multiplied by the number of processes. + dispatch_batches (`bool`, *optional*): + If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process + and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose + underlying dataset is an `IterableDataset`, `False` otherwise. + even_batches (`bool`, *optional*, defaults to `True`): + If set to `True`, in cases where the total batch size across all processes does not exactly divide the + dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among + all workers. + use_seedable_sampler (`bool`, *optional*, defaults to `True`): + Whether or not use a fully seedable random sampler ([`accelerate.data_loader.SeedableRandomSampler`]). Ensures + training results are fully reproducible using a different sampling technique. While seed-to-seed results + may differ, on average the differences are negligible when using multiple different seeds to compare. Should + also be ran with [`~utils.set_seed`] for the best results. + gradient_accumulation_kwargs (`dict`, *optional*): + Additional kwargs to configure gradient accumulation, see [`accelerate.utils.GradientAccumulationPlugin`]. + Any of the following (optional) keys are acceptable: + num_steps (`int`): Will take precedence over [`~.TrainingArguments.gradient_accumulation_steps`] if + the latter is set to 1, otherwise an exception will be raised. + adjust_scheduler (`bool`): Whether to adjust the scheduler steps to account for [`~.TrainingArguments.gradient_accumulation_steps`]. + The [`accelerate.utils.GradientAccumulationPlugin`] default is `True`. + sync_each_batch (`bool`): Whether to synchronize the gradients at each data batch. + The [`accelerate.utils.GradientAccumulationPlugin`] default is `False`. + non_blocking (`bool`, *optional*, defaults to `False`): + Whether to use non-blocking CUDA calls to help minimize synchronization during + distributed training with prepared `DataLoader` inputs being moved to device. + Best if used with `pin_memory=True` in the `TrainingArguments`. + use_configured_state (`bool*, *optional*, defaults to `False`): + Whether or not to use a pre-configured `AcceleratorState` or `PartialState` defined + before calling `TrainingArguments`. If `True`, an `Accelerator` or `PartialState` + must be initialized. May lead to issues using sweeps or hyperparameter tuning. + + """ + + # Data related arguments + split_batches: bool = field( + default=False, + metadata={ + "help": "Whether or not the accelerator should split the batches yielded by the dataloaders across the devices. If" + " `True` the actual batch size used will be the same on any kind of distributed processes, but it must be a" + " round multiple of the `num_processes` you are using. If `False`, actual batch size used will be the one set" + " in your script multiplied by the number of processes." + }, + ) + dispatch_batches: Optional[bool] = field( + default=None, + metadata={ + "help": "If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process" + " and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose" + " underlying dataset is an `IterableDataslet`, `False` otherwise." + }, + ) + even_batches: bool = field( + default=True, + metadata={ + "help": "If set to `True`, in cases where the total batch size across all processes does not exactly divide the" + " dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among" + " all workers." + }, + ) + use_seedable_sampler: bool = field( + default=True, + metadata={ + "help": "Whether or not use a fully seedable random sampler ([`accelerate.data_loader.SeedableRandomSampler`])." + "Ensures training results are fully reproducible using a different sampling technique. " + "While seed-to-seed results may differ, on average the differences are negligible when using" + "multiple different seeds to compare. Should also be ran with [`~utils.set_seed`] for the best results." + }, + ) + + non_blocking: Optional[bool] = field( + default=False, + metadata={ + "help": "Whether to use non-blocking CUDA calls to help minimize synchronization during " + "distributed training with prepared `DataLoader` inputs being moved to device. " + "Best if used with `pin_memory=True` in the `TrainingArguments`. Requires accelerate " + "v0.30.0." + }, + ) + + gradient_accumulation_kwargs: Optional[dict] = field( + default=None, + metadata={ + "help": "Additional kwargs to configure gradient accumulation, see [`accelerate.utils.GradientAccumulationPlugin`]. " + "Any of the following (optional) keys are acceptable: " + " num_steps (`int`): Will take precedence over [`~.TrainingArguments.gradient_accumulation_steps`] if " + " the latter is set to 1, otherwise an exception will be raised. " + " adjust_scheduler (`bool`): Whether to adjust the scheduler steps to account for [`~.TrainingArguments.gradient_accumulation_steps`]. " + " The [`accelerate.utils.GradientAccumulationPlugin`] default is `True`. " + " sync_each_batch (`bool`): Whether to synchronize the gradients at each data batch. " + " The [`accelerate.utils.GradientAccumulationPlugin`] default is `False`." + }, + ) + use_configured_state: bool = field( + default=False, + metadata={ + "help": "Whether or not to use a pre-configured `AcceleratorState` or `PartialState` defined before calling `TrainingArguments`." + "If `True`, an `Accelerator` or `PartialState` must be initialized. May lead to issues using sweeps or hyperparameter tuning." + }, + ) + + @classmethod + def from_json_file(cls, json_file): + # Check if exists + open_file = io.open if os.path.exists(json_file) else open + with open_file(json_file, "r", encoding="utf-8") as f: + config_dict = json.load(f) + # Check for keys and load sensible defaults + extra_keys = sorted(key for key in config_dict.keys() if key not in cls.__dataclass_fields__.keys()) + if len(extra_keys) > 0: + raise ValueError( + f"The config file at {json_file} had unknown keys ({extra_keys}), please try upgrading your `transformers`" + " version or fix (and potentially remove these keys) from your config file." + ) + return cls(**config_dict) + + def to_dict(self): + return copy.deepcopy(self.__dict__) + + def pop(self, key, default=None): + return self.__dict__.pop(key, default) + + +class LayerWiseDummyOptimizer(torch.optim.Optimizer): + """ + For Layer-wise optimizers such as GaLoRE optimizer, the optimization + step is already done through the post gradient hooks. Therefore + the trick is to create a dummy optimizer that can take arbitrary + args and kwargs and return a no-op during training. + + Initial idea from @hiyouga in LLaMA-Factory: + https://github.com/hiyouga/LLaMA-Factory/commit/8664262cde3919e10eaecbd66e8c5d356856362e#diff-ebe08ab14496dfb9e06075f0fdd36799ef6d1535cc4dd4715b74c4e3e06fe3ba + """ + + def __init__(self, optimizer_dict=None, *args, **kwargs): + dummy_tensor = torch.randn(1, 1) + self.optimizer_dict = optimizer_dict + super().__init__([dummy_tensor], {"lr": kwargs.get("lr", 1e-03)}) + + def zero_grad(self, set_to_none: bool = True) -> None: + pass + + def step(self, closure=None) -> Optional[float]: + pass + + +class LayerWiseDummyScheduler(LRScheduler): + """ + For Layer-wise optimizers such as GaLoRE optimizer, the optimization and scheduling step + are already done through the post gradient hooks. Therefore + the trick is to create a dummy scheduler that can take arbitrary + args and kwargs and return a no-op during training. + """ + + def __init__(self, *args, **kwargs): + self.default_lr = kwargs["lr"] + optimizer = LayerWiseDummyOptimizer(**kwargs) + last_epoch = -1 + verbose = False + super().__init__(optimizer, last_epoch, verbose) + + def get_lr(self): + # default value + lrs = [self.default_lr] + + # we take each lr in the parameters if they exist, assumes the optimizer to be the `LayerWiseDummyOptimizer` + if self.optimizer is not None: + param_wise_lrs = [ + [group["lr"] for group in optim.param_groups] for optim in self.optimizer.optimizer_dict.values() + ] + lrs = list(chain(*param_wise_lrs)) + + return lrs + + def _get_closed_form_lr(self): + return self.base_lrs + + +def set_rng_state_for_device(device_name, device_module, checkpoint_rng_state, is_distributed): + """Helper to set RNG state for a specific device type (CUDA, NPU, MLU, MUSA)""" + device_state_key = device_name.lower() + err_template = "Didn't manage to set back the RNG states of the {backend} because of the following error:\n {exception}\nThis won't yield the same results as if the training had not been interrupted." + try: + if is_distributed: + device_module.random.set_rng_state_all(checkpoint_rng_state[device_state_key]) + else: + device_module.random.set_rng_state(checkpoint_rng_state[device_state_key]) + except Exception as e: + # Log error if setting RNG state fails + logger.error(err_template.format(backend=device_name, exception=e)) diff --git a/lib/python3.12/site-packages/transformers/trainer_seq2seq.py b/lib/python3.12/site-packages/transformers/trainer_seq2seq.py new file mode 100644 index 0000000000000000000000000000000000000000..e9fa797f0628e37fc2478e793129c8355cb3ebf4 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/trainer_seq2seq.py @@ -0,0 +1,392 @@ +# Copyright 2020 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. + +import contextlib +import warnings +from copy import deepcopy +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Optional, Union + +import torch +from torch import nn +from torch.distributed.fsdp import FullyShardedDataParallel +from torch.utils.data import Dataset + +from .generation.configuration_utils import GenerationConfig +from .integrations.deepspeed import is_deepspeed_zero3_enabled +from .integrations.fsdp import is_fsdp_managed_module +from .trainer import Trainer +from .utils import is_datasets_available, logging +from .utils.deprecation import deprecate_kwarg + + +if is_datasets_available(): + import datasets + +if TYPE_CHECKING: + from torch.utils.data import IterableDataset + + from .data.data_collator import DataCollator + from .feature_extraction_utils import FeatureExtractionMixin + from .image_processing_utils import BaseImageProcessor + from .modeling_utils import PreTrainedModel + from .processing_utils import ProcessorMixin + from .tokenization_utils_base import PreTrainedTokenizerBase + from .trainer_callback import TrainerCallback + from .trainer_utils import EvalPrediction, PredictionOutput + from .training_args import TrainingArguments + + +logger = logging.get_logger(__name__) + + +class Seq2SeqTrainer(Trainer): + @deprecate_kwarg("tokenizer", new_name="processing_class", version="5.0.0", raise_if_both_names=True) + def __init__( + self, + model: Union["PreTrainedModel", nn.Module] = None, + args: "TrainingArguments" = None, + data_collator: Optional["DataCollator"] = None, + train_dataset: Optional[Union[Dataset, "IterableDataset", "datasets.Dataset"]] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union["PreTrainedTokenizerBase", "BaseImageProcessor", "FeatureExtractionMixin", "ProcessorMixin"] + ] = None, + model_init: Optional[Callable[[], "PreTrainedModel"]] = None, + compute_loss_func: Optional[Callable] = None, + compute_metrics: Optional[Callable[["EvalPrediction"], dict]] = None, + callbacks: Optional[list["TrainerCallback"]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + ): + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_loss_func=compute_loss_func, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Override self.model.generation_config if a GenerationConfig is specified in args. + # Priority: args.generation_config > model.generation_config > default GenerationConfig. + if self.args.generation_config is not None: + gen_config = self.load_generation_config(self.args.generation_config) + self.model.generation_config = gen_config + + @staticmethod + def load_generation_config(gen_config_arg: Union[str, GenerationConfig]) -> GenerationConfig: + """ + Loads a `~generation.GenerationConfig` from the `Seq2SeqTrainingArguments.generation_config` arguments. + + Args: + gen_config_arg (`str` or [`~generation.GenerationConfig]`): + `Seq2SeqTrainingArguments.generation_config` argument. + + Returns: + A `~generation.GenerationConfig`. + """ + + # GenerationConfig provided, nothing to do + if isinstance(gen_config_arg, GenerationConfig): + gen_config = deepcopy(gen_config_arg) + else: + # str or Path + pretrained_model_name = Path(gen_config_arg) if isinstance(gen_config_arg, str) else gen_config_arg + config_file_name = None + + # Figuring if it is path pointing to a file, pointing to a directory or else a model id or URL + # This step is required in order to determine config_file_name + if pretrained_model_name.is_file(): + config_file_name = pretrained_model_name.name + pretrained_model_name = pretrained_model_name.parent + # dir path + elif pretrained_model_name.is_dir(): + pass + # model id or URL + else: + pretrained_model_name = gen_config_arg + + gen_config = GenerationConfig.from_pretrained(pretrained_model_name, config_file_name) + + # Strict validation to fail early. `GenerationConfig.save_pretrained()`, run at the end of training, throws + # an exception if there are warnings at validation time. + try: + with warnings.catch_warnings(record=True) as caught_warnings: + gen_config.validate() + if len(caught_warnings) > 0: + raise ValueError(str([w.message for w in caught_warnings])) + except ValueError as exc: + raise ValueError( + "The loaded generation config instance is invalid -- `GenerationConfig.validate()` throws warnings " + "and/or exceptions. Fix these issues to train your model.\n\nThrown during validation:\n" + str(exc) + ) + return gen_config + + def evaluate( + self, + eval_dataset: Optional[Dataset] = None, + ignore_keys: Optional[list[str]] = None, + metric_key_prefix: str = "eval", + **gen_kwargs, + ) -> dict[str, float]: + """ + Run evaluation and returns metrics. + + The calling script will be responsible for providing a method to compute metrics, as they are task-dependent + (pass it to the init `compute_metrics` argument). + + You can also subclass and override this method to inject custom behavior. + + Args: + eval_dataset (`Dataset`, *optional*): + Pass a dataset if you wish to override `self.eval_dataset`. If it is an [`~datasets.Dataset`], columns + not accepted by the `model.forward()` method are automatically removed. It must implement the `__len__` + method. + ignore_keys (`List[str]`, *optional*): + A list of keys in the output of your model (if it is a dictionary) that should be ignored when + gathering predictions. + metric_key_prefix (`str`, *optional*, defaults to `"eval"`): + An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named + "eval_bleu" if the prefix is `"eval"` (default) + max_length (`int`, *optional*): + The maximum target length to use when predicting with the generate method. + num_beams (`int`, *optional*): + Number of beams for beam search that will be used when predicting with the generate method. 1 means no + beam search. + gen_kwargs: + Additional `generate` specific kwargs. + + Returns: + A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The + dictionary also contains the epoch number which comes from the training state. + """ + + gen_kwargs = gen_kwargs.copy() + + # Use legacy argument setting if a) the option is not explicitly passed; and b) the argument is set in the + # training args + if ( + gen_kwargs.get("max_length") is None + and gen_kwargs.get("max_new_tokens") is None + and self.args.generation_max_length is not None + ): + gen_kwargs["max_length"] = self.args.generation_max_length + if gen_kwargs.get("num_beams") is None and self.args.generation_num_beams is not None: + gen_kwargs["num_beams"] = self.args.generation_num_beams + # We don't want to drop samples in general + self.gather_function = self.accelerator.gather + self._gen_kwargs = gen_kwargs + return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) + + def predict( + self, + test_dataset: Dataset, + ignore_keys: Optional[list[str]] = None, + metric_key_prefix: str = "test", + **gen_kwargs, + ) -> "PredictionOutput": + """ + Run prediction and returns predictions and potential metrics. + + Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method + will also return metrics, like in `evaluate()`. + + Args: + test_dataset (`Dataset`): + Dataset to run the predictions on. If it is a [`~datasets.Dataset`], columns not accepted by the + `model.forward()` method are automatically removed. Has to implement the method `__len__` + ignore_keys (`List[str]`, *optional*): + A list of keys in the output of your model (if it is a dictionary) that should be ignored when + gathering predictions. + metric_key_prefix (`str`, *optional*, defaults to `"eval"`): + An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named + "eval_bleu" if the prefix is `"eval"` (default) + max_length (`int`, *optional*): + The maximum target length to use when predicting with the generate method. + num_beams (`int`, *optional*): + Number of beams for beam search that will be used when predicting with the generate method. 1 means no + beam search. + gen_kwargs: + Additional `generate` specific kwargs. + + + + If your predictions or labels have different sequence lengths (for instance because you're doing dynamic + padding in a token classification task) the predictions will be padded (on the right) to allow for + concatenation into one array. The padding index is -100. + + + + Returns: *NamedTuple* A namedtuple with the following keys: + + - predictions (`np.ndarray`): The predictions on `test_dataset`. + - label_ids (`np.ndarray`, *optional*): The labels (if the dataset contained some). + - metrics (`Dict[str, float]`, *optional*): The potential dictionary of metrics (if the dataset contained + labels). + """ + + gen_kwargs = gen_kwargs.copy() + + # Use legacy argument setting if a) the option is not explicitly passed; and b) the argument is set in the + # training args + if ( + gen_kwargs.get("max_length") is None + and gen_kwargs.get("max_new_tokens") is None + and self.args.generation_max_length is not None + ): + gen_kwargs["max_length"] = self.args.generation_max_length + if gen_kwargs.get("num_beams") is None and self.args.generation_num_beams is not None: + gen_kwargs["num_beams"] = self.args.generation_num_beams + self.gather_function = self.accelerator.gather + self._gen_kwargs = gen_kwargs + + return super().predict(test_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) + + def prediction_step( + self, + model: nn.Module, + inputs: dict[str, Union[torch.Tensor, Any]], + prediction_loss_only: bool, + ignore_keys: Optional[list[str]] = None, + **gen_kwargs, + ) -> tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + Perform an evaluation step on `model` using `inputs`. + + Subclass and override to inject custom behavior. + + Args: + model (`nn.Module`): + The model to evaluate. + inputs (`Dict[str, Union[torch.Tensor, Any]]`): + The inputs and targets of the model. + + The dictionary will be unpacked before being fed to the model. Most models expect the targets under the + argument `labels`. Check your model's documentation for all accepted arguments. + prediction_loss_only (`bool`): + Whether or not to return the loss only. + gen_kwargs: + Additional `generate` specific kwargs. + + Return: + Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and + labels (each being optional). + """ + + if not self.args.predict_with_generate or prediction_loss_only: + return super().prediction_step( + model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys + ) + + has_labels = "labels" in inputs + inputs = self._prepare_inputs(inputs) + + # Priority (handled in generate): + # non-`None` gen_kwargs > model.generation_config > default GenerationConfig() + if len(gen_kwargs) == 0 and hasattr(self, "_gen_kwargs"): + gen_kwargs = self._gen_kwargs.copy() + if "num_beams" in gen_kwargs and gen_kwargs["num_beams"] is None: + gen_kwargs.pop("num_beams") + if "max_length" in gen_kwargs and gen_kwargs["max_length"] is None: + gen_kwargs.pop("max_length") + + default_synced_gpus = is_deepspeed_zero3_enabled() or is_fsdp_managed_module(self.model) + gen_kwargs["synced_gpus"] = gen_kwargs.get("synced_gpus", default_synced_gpus) + + generation_inputs = inputs.copy() + # If the `decoder_input_ids` was created from `labels`, evict the former, so that the model can freely generate + # (otherwise, it would continue generating from the padded `decoder_input_ids`) + if ( + "labels" in generation_inputs + and "decoder_input_ids" in generation_inputs + and generation_inputs["labels"].shape == generation_inputs["decoder_input_ids"].shape + ): + generation_inputs = { + k: v for k, v in inputs.items() if k not in ("decoder_input_ids", "decoder_attention_mask") + } + + summon_full_params_context = ( + FullyShardedDataParallel.summon_full_params(self.model) + if isinstance(self.model, FullyShardedDataParallel) + else contextlib.nullcontext() + ) + + with summon_full_params_context: + generated_tokens = self.model.generate(**generation_inputs, **gen_kwargs) + + # Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop + # TODO: remove this hack when the legacy code that initializes generation_config from a model config is + # removed in https://github.com/huggingface/transformers/blob/98d88b23f54e5a23e741833f1e973fdf600cc2c5/src/transformers/generation/utils.py#L1183 + if self.model.generation_config._from_model_config: + self.model.generation_config._from_model_config = False + + # Retrieves GenerationConfig from model.generation_config + gen_config = self.model.generation_config + # in case the batch is shorter than max length, the output should be padded + if generated_tokens.shape[-1] < gen_config.max_length: + generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_config.max_length) + elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1: + generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_config.max_new_tokens + 1) + + with torch.no_grad(): + if has_labels: + with self.compute_loss_context_manager(): + outputs = model(**inputs) + if self.label_smoother is not None: + loss = self.label_smoother(outputs, inputs["labels"]).detach().mean() + else: + loss = (outputs["loss"] if isinstance(outputs, dict) else outputs[0]).detach().mean() + else: + loss = None + + if self.args.prediction_loss_only: + return loss, None, None + + if has_labels: + labels = inputs["labels"] + if labels.shape[-1] < gen_config.max_length: + labels = self._pad_tensors_to_max_len(labels, gen_config.max_length) + elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1: + labels = self._pad_tensors_to_max_len(labels, gen_config.max_new_tokens + 1) + else: + labels = None + + return loss, generated_tokens, labels + + def _pad_tensors_to_max_len(self, tensor, max_length): + if self.processing_class is not None and hasattr(self.processing_class, "pad_token_id"): + # If PAD token is not defined at least EOS token has to be defined + pad_token_id = ( + self.processing_class.pad_token_id + if self.processing_class.pad_token_id is not None + else self.processing_class.eos_token_id + ) + else: + if self.model.config.pad_token_id is not None: + pad_token_id = self.model.config.pad_token_id + else: + raise ValueError("Pad_token_id must be set in the configuration of the model, in order to pad tensors") + + padded_tensor = pad_token_id * torch.ones( + (tensor.shape[0], max_length), dtype=tensor.dtype, device=tensor.device + ) + padded_tensor[:, : tensor.shape[-1]] = tensor + return padded_tensor diff --git a/lib/python3.12/site-packages/transformers/trainer_utils.py b/lib/python3.12/site-packages/transformers/trainer_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4d3dd6d6bb17fe0ca850c70bd58ffb534e99263b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/trainer_utils.py @@ -0,0 +1,910 @@ +# Copyright 2020-present the HuggingFace Inc. team. +# +# 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. +""" +PyTorch-independent utilities for the Trainer class. +""" + +import copy +import functools +import gc +import inspect +import os +import random +import re +import threading +import time +from typing import Any, NamedTuple, Optional, Union + +import numpy as np + +from .utils import ( + ExplicitEnum, + is_psutil_available, + is_tf_available, + is_torch_available, + is_torch_cuda_available, + is_torch_hpu_available, + is_torch_mlu_available, + is_torch_mps_available, + is_torch_musa_available, + is_torch_npu_available, + is_torch_xla_available, + is_torch_xpu_available, + requires_backends, +) + + +if is_torch_available(): + import torch + + +def seed_worker(_): + """ + Helper function to set worker seed during Dataloader initialization. + """ + worker_seed = torch.initial_seed() % 2**32 + set_seed(worker_seed) + + +def enable_full_determinism(seed: int, warn_only: bool = False): + """ + Helper function for reproducible behavior during distributed training. See + - https://pytorch.org/docs/stable/notes/randomness.html for pytorch + - https://www.tensorflow.org/api_docs/python/tf/config/experimental/enable_op_determinism for tensorflow + """ + # set seed first + set_seed(seed) + + if is_torch_available(): + # Enable PyTorch deterministic mode. This potentially requires either the environment + # variable 'CUDA_LAUNCH_BLOCKING' or 'CUBLAS_WORKSPACE_CONFIG' to be set, + # depending on the CUDA version, so we set them both here + os.environ["CUDA_LAUNCH_BLOCKING"] = "1" + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" + # The environment variable required to enable deterministic mode on Ascend NPUs. + os.environ["ASCEND_LAUNCH_BLOCKING"] = "1" + os.environ["HCCL_DETERMINISTIC"] = "1" + + os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1" + torch.use_deterministic_algorithms(True, warn_only=warn_only) + + # Enable CUDNN deterministic mode + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + if is_tf_available(): + import tensorflow as tf + + tf.config.experimental.enable_op_determinism() + + +def set_seed(seed: int, deterministic: bool = False): + """ + Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch` and/or `tf` (if installed). + + Args: + seed (`int`): + The seed to set. + deterministic (`bool`, *optional*, defaults to `False`): + Whether to use deterministic algorithms where available. Can slow down training. + """ + random.seed(seed) + np.random.seed(seed) + if is_torch_available(): + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + # ^^ safe to call this function even if cuda is not available + if deterministic: + torch.use_deterministic_algorithms(True) + if is_torch_mlu_available(): + torch.mlu.manual_seed_all(seed) + if is_torch_musa_available(): + torch.musa.manual_seed_all(seed) + if is_torch_npu_available(): + torch.npu.manual_seed_all(seed) + if is_torch_hpu_available(): + torch.hpu.manual_seed_all(seed) + if is_torch_xpu_available(): + torch.xpu.manual_seed_all(seed) + if is_tf_available(): + import tensorflow as tf + + tf.random.set_seed(seed) + if deterministic: + tf.config.experimental.enable_op_determinism() + + +def neftune_post_forward_hook(module, input, output): + """ + Implements the NEFTune forward pass for the model using forward hooks. Note this works only for torch.nn.Embedding + layers. This method is slightly adapted from the original source code that can be found here: + https://github.com/neelsjain/NEFTune Simply add it to your model as follows: + ```python + model = ... + model.embed_tokens.neftune_noise_alpha = 0.1 + model.embed_tokens.register_forward_hook(neftune_post_forward_hook) + ``` + Args: + module (`torch.nn.Module`): + The embedding module where the hook is attached. Note that you need to set `module.neftune_noise_alpha` to + the desired noise alpha value. + input (`torch.Tensor`): + The input tensor to the model. + output (`torch.Tensor`): + The output tensor of the model (i.e. the embeddings). + """ + if module.training: + dims = torch.tensor(output.size(1) * output.size(2)) + mag_norm = module.neftune_noise_alpha / torch.sqrt(dims) + output = output + torch.zeros_like(output).uniform_(-mag_norm, mag_norm) + return output + + +class EvalPrediction: + """ + Evaluation output (always contains labels), to be used to compute metrics. + + Parameters: + predictions (`np.ndarray`): Predictions of the model. + label_ids (`np.ndarray`): Targets to be matched. + inputs (`np.ndarray`, *optional*): Input data passed to the model. + losses (`np.ndarray`, *optional*): Loss values computed during evaluation. + """ + + def __init__( + self, + predictions: Union[np.ndarray, tuple[np.ndarray]], + label_ids: Union[np.ndarray, tuple[np.ndarray]], + inputs: Optional[Union[np.ndarray, tuple[np.ndarray]]] = None, + losses: Optional[Union[np.ndarray, tuple[np.ndarray]]] = None, + ): + self.predictions = predictions + self.label_ids = label_ids + self.inputs = inputs + self.losses = losses + self.elements = (self.predictions, self.label_ids) + if self.inputs is not None: + self.elements += (self.inputs,) + if self.losses is not None: + self.elements += (self.losses,) + + def __iter__(self): + return iter(self.elements) + + def __getitem__(self, idx): + if idx < 0 or idx >= len(self.elements): + raise IndexError("tuple index out of range") + return self.elements[idx] + + +class EvalLoopOutput(NamedTuple): + predictions: Union[np.ndarray, tuple[np.ndarray]] + label_ids: Optional[Union[np.ndarray, tuple[np.ndarray]]] + metrics: Optional[dict[str, float]] + num_samples: Optional[int] + + +class PredictionOutput(NamedTuple): + predictions: Union[np.ndarray, tuple[np.ndarray]] + label_ids: Optional[Union[np.ndarray, tuple[np.ndarray]]] + metrics: Optional[dict[str, float]] + + +class TrainOutput(NamedTuple): + global_step: int + training_loss: float + metrics: dict[str, float] + + +PREFIX_CHECKPOINT_DIR = "checkpoint" +_re_checkpoint = re.compile(r"^" + PREFIX_CHECKPOINT_DIR + r"\-(\d+)$") + + +def get_last_checkpoint(folder): + content = os.listdir(folder) + checkpoints = [ + path + for path in content + if _re_checkpoint.search(path) is not None and os.path.isdir(os.path.join(folder, path)) + ] + if len(checkpoints) == 0: + return + return os.path.join(folder, max(checkpoints, key=lambda x: int(_re_checkpoint.search(x).groups()[0]))) + + +class IntervalStrategy(ExplicitEnum): + NO = "no" + STEPS = "steps" + EPOCH = "epoch" + + +class SaveStrategy(ExplicitEnum): + NO = "no" + STEPS = "steps" + EPOCH = "epoch" + BEST = "best" + + +class EvaluationStrategy(ExplicitEnum): + NO = "no" + STEPS = "steps" + EPOCH = "epoch" + + +class HubStrategy(ExplicitEnum): + END = "end" + EVERY_SAVE = "every_save" + CHECKPOINT = "checkpoint" + ALL_CHECKPOINTS = "all_checkpoints" + + +class BestRun(NamedTuple): + """ + The best run found by a hyperparameter search (see [`~Trainer.hyperparameter_search`]). + + Parameters: + run_id (`str`): + The id of the best run (if models were saved, the corresponding checkpoint will be in the folder ending + with run-{run_id}). + objective (`float`): + The objective that was obtained for this run. + hyperparameters (`Dict[str, Any]`): + The hyperparameters picked to get this run. + run_summary (`Optional[Any]`): + A summary of tuning experiments. `ray.tune.ExperimentAnalysis` object for Ray backend. + """ + + run_id: str + objective: Union[float, list[float]] + hyperparameters: dict[str, Any] + run_summary: Optional[Any] = None + + +def default_compute_objective(metrics: dict[str, float]) -> float: + """ + The default objective to maximize/minimize when doing an hyperparameter search. It is the evaluation loss if no + metrics are provided to the [`Trainer`], the sum of all metrics otherwise. + + Args: + metrics (`Dict[str, float]`): The metrics returned by the evaluate method. + + Return: + `float`: The objective to minimize or maximize + """ + metrics = copy.deepcopy(metrics) + loss = metrics.pop("eval_loss", None) + _ = metrics.pop("epoch", None) + # Remove speed metrics + speed_metrics = [ + m + for m in metrics.keys() + if m.endswith("_runtime") or m.endswith("_per_second") or m.endswith("_compilation_time") + ] + for sm in speed_metrics: + _ = metrics.pop(sm, None) + return loss if len(metrics) == 0 else sum(metrics.values()) + + +def default_hp_space_optuna(trial) -> dict[str, float]: + from .integrations import is_optuna_available + + assert is_optuna_available(), "This function needs Optuna installed: `pip install optuna`" + return { + "learning_rate": trial.suggest_float("learning_rate", 1e-6, 1e-4, log=True), + "num_train_epochs": trial.suggest_int("num_train_epochs", 1, 5), + "seed": trial.suggest_int("seed", 1, 40), + "per_device_train_batch_size": trial.suggest_categorical("per_device_train_batch_size", [4, 8, 16, 32, 64]), + } + + +def default_hp_space_ray(trial) -> dict[str, float]: + from .integrations import is_ray_tune_available + + assert is_ray_tune_available(), "This function needs ray installed: `pip install ray[tune]`" + from ray import tune + + return { + "learning_rate": tune.loguniform(1e-6, 1e-4), + "num_train_epochs": tune.choice(list(range(1, 6))), + "seed": tune.uniform(1, 40), + "per_device_train_batch_size": tune.choice([4, 8, 16, 32, 64]), + } + + +def default_hp_space_sigopt(trial): + return [ + {"bounds": {"min": 1e-6, "max": 1e-4}, "name": "learning_rate", "type": "double", "transformation": "log"}, + {"bounds": {"min": 1, "max": 6}, "name": "num_train_epochs", "type": "int"}, + {"bounds": {"min": 1, "max": 40}, "name": "seed", "type": "int"}, + { + "categorical_values": ["4", "8", "16", "32", "64"], + "name": "per_device_train_batch_size", + "type": "categorical", + }, + ] + + +def default_hp_space_wandb(trial) -> dict[str, float]: + from .integrations import is_wandb_available + + if not is_wandb_available(): + raise ImportError("This function needs wandb installed: `pip install wandb`") + + return { + "method": "random", + "metric": {"name": "objective", "goal": "minimize"}, + "parameters": { + "learning_rate": {"distribution": "uniform", "min": 1e-6, "max": 1e-4}, + "num_train_epochs": {"distribution": "int_uniform", "min": 1, "max": 6}, + "seed": {"distribution": "int_uniform", "min": 1, "max": 40}, + "per_device_train_batch_size": {"values": [4, 8, 16, 32, 64]}, + }, + } + + +class HPSearchBackend(ExplicitEnum): + OPTUNA = "optuna" + RAY = "ray" + SIGOPT = "sigopt" + WANDB = "wandb" + + +def is_main_process(local_rank): + """ + Whether or not the current process is the local process, based on `xm.get_ordinal()` (for TPUs) first, then on + `local_rank`. + """ + if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + return xm.get_ordinal() == 0 + return local_rank in [-1, 0] + + +def total_processes_number(local_rank): + """ + Return the number of processes launched in parallel. Works with `torch.distributed` and TPUs. + """ + if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + return xm.xrt_world_size() + elif local_rank != -1 and is_torch_available(): + import torch + + return torch.distributed.get_world_size() + return 1 + + +def speed_metrics(split, start_time, num_samples=None, num_steps=None, num_tokens=None): + """ + Measure and return speed performance metrics. + + This function requires a time snapshot `start_time` before the operation to be measured starts and this function + should be run immediately after the operation to be measured has completed. + + Args: + - split: name to prefix metric (like train, eval, test...) + - start_time: operation start time + - num_samples: number of samples processed + - num_steps: number of steps processed + - num_tokens: number of tokens processed + """ + runtime = time.time() - start_time + result = {f"{split}_runtime": round(runtime, 4)} + if runtime == 0: + return result + if num_samples is not None: + samples_per_second = num_samples / runtime + result[f"{split}_samples_per_second"] = round(samples_per_second, 3) + if num_steps is not None: + steps_per_second = num_steps / runtime + result[f"{split}_steps_per_second"] = round(steps_per_second, 3) + if num_tokens is not None: + tokens_per_second = num_tokens / runtime + result[f"{split}_tokens_per_second"] = round(tokens_per_second, 3) + return result + + +class SchedulerType(ExplicitEnum): + """ + Scheduler names for the parameter `lr_scheduler_type` in [`TrainingArguments`]. + By default, it uses "linear". Internally, this retrieves `get_linear_schedule_with_warmup` scheduler from [`Trainer`]. + Scheduler types: + - "linear" = get_linear_schedule_with_warmup + - "cosine" = get_cosine_schedule_with_warmup + - "cosine_with_restarts" = get_cosine_with_hard_restarts_schedule_with_warmup + - "polynomial" = get_polynomial_decay_schedule_with_warmup + - "constant" = get_constant_schedule + - "constant_with_warmup" = get_constant_schedule_with_warmup + - "inverse_sqrt" = get_inverse_sqrt_schedule + - "reduce_lr_on_plateau" = get_reduce_on_plateau_schedule + - "cosine_with_min_lr" = get_cosine_with_min_lr_schedule_with_warmup + - "warmup_stable_decay" = get_wsd_schedule + """ + + LINEAR = "linear" + COSINE = "cosine" + COSINE_WITH_RESTARTS = "cosine_with_restarts" + POLYNOMIAL = "polynomial" + CONSTANT = "constant" + CONSTANT_WITH_WARMUP = "constant_with_warmup" + INVERSE_SQRT = "inverse_sqrt" + REDUCE_ON_PLATEAU = "reduce_lr_on_plateau" + COSINE_WITH_MIN_LR = "cosine_with_min_lr" + WARMUP_STABLE_DECAY = "warmup_stable_decay" + + +class TrainerMemoryTracker: + """ + A helper class that tracks cpu and gpu memory. + + This class will silently skip unless `psutil` is available. Install with `pip install psutil`. + + When a stage completes, it can pass metrics dict to update with the memory metrics gathered during this stage. + + Example : + + ```python + self._memory_tracker = TrainerMemoryTracker(self.args.skip_memory_metrics) + self._memory_tracker.start() + # code ... + metrics = {"train_runtime": 10.5} + self._memory_tracker.stop_and_update_metrics(metrics) + ``` + + At the moment GPU tracking is only for `pytorch`, but can be extended to support `tensorflow`. + + To understand this class' intricacies please read the documentation of [`~Trainer.log_metrics`]. + """ + + # map trainer methods to metrics prefix + stages = { + "__init__": "init", + "train": "train", + "_inner_training_loop": "train", + "evaluate": "eval", + "predict": "test", + } + + def __init__(self, skip_memory_metrics=False): + self.skip_memory_metrics = skip_memory_metrics + + if not is_psutil_available(): + # soft dependency on psutil + self.skip_memory_metrics = True + + if self.skip_memory_metrics: + return + + import psutil # noqa + + if is_torch_cuda_available() or is_torch_mlu_available() or is_torch_musa_available(): + import torch + + self.torch = torch + self.gpu = {} + elif is_torch_mps_available(): + import torch + + self.torch = torch + self.gpu = {} + elif is_torch_xpu_available(): + import torch + + self.torch = torch + self.gpu = {} + elif is_torch_npu_available(): + import torch + + self.torch = torch + self.gpu = {} + elif is_torch_hpu_available(): + import torch + + self.torch = torch + self.gpu = {} + else: + self.torch = None + + self.process = psutil.Process() + + self.cur_stage = None + self.cpu = {} + self.init_reported = False + + def derive_stage(self): + """derives the stage/caller name automatically""" + caller = inspect.currentframe().f_back.f_back.f_code.co_name + if caller in self.stages: + return self.stages[caller] + else: + raise ValueError( + f"was called from {caller}, but only expect to be called from one of {self.stages.keys()}" + ) + + def cpu_mem_used(self): + """get resident set size memory for the current process""" + return self.process.memory_info().rss + + def peak_monitor_func(self): + self.cpu_mem_used_peak = -1 + + while True: + self.cpu_mem_used_peak = max(self.cpu_mem_used(), self.cpu_mem_used_peak) + + # can't sleep or will not catch the peak right (this comment is here on purpose) + # time.sleep(0.001) # 1msec + + if not self.peak_monitoring: + break + + def start(self): + """start tracking for the caller's stage""" + if self.skip_memory_metrics: + return + + stage = self.derive_stage() + # deal with nested calls of eval during train - simply ignore those + if self.cur_stage is not None and self.cur_stage != stage: + return + + self.cur_stage = stage + + gc.collect() + + if self.torch is not None: + if torch.cuda.is_available(): + self.torch.cuda.reset_peak_memory_stats() + self.torch.cuda.empty_cache() + elif is_torch_mlu_available(): + self.torch.mlu.reset_peak_memory_stats() + self.torch.mlu.empty_cache() + elif is_torch_musa_available(): + self.torch.musa.reset_peak_memory_stats() + self.torch.musa.empty_cache() + elif is_torch_xpu_available(): + self.torch.xpu.reset_peak_memory_stats() + self.torch.xpu.empty_cache() + elif is_torch_npu_available(): + self.torch.npu.reset_peak_memory_stats() + self.torch.npu.empty_cache() + elif is_torch_hpu_available(): + self.torch.hpu.reset_peak_memory_stats() + # not available on hpu as it reserves all device memory for the current process + # self.torch.hpu.empty_cache() + elif is_torch_mps_available(): + self.torch.mps.empty_cache() + + # gpu + if self.torch is not None: + if torch.cuda.is_available(): + self.gpu_mem_used_at_start = self.torch.cuda.memory_allocated() + elif is_torch_mlu_available(): + self.gpu_mem_used_at_start = self.torch.mlu.memory_allocated() + elif is_torch_musa_available(): + self.gpu_mem_used_at_start = self.torch.musa.memory_allocated() + elif is_torch_xpu_available(): + self.gpu_mem_used_at_start = self.torch.xpu.memory_allocated() + elif is_torch_npu_available(): + self.gpu_mem_used_at_start = self.torch.npu.memory_allocated() + elif is_torch_hpu_available(): + self.gpu_mem_used_at_start = self.torch.hpu.memory_allocated() + elif is_torch_mps_available(): + self.gpu_mem_used_at_start = self.torch.mps.current_allocated_memory() + + # cpu + self.cpu_mem_used_at_start = self.cpu_mem_used() + + self.peak_monitoring = True + peak_monitor_thread = threading.Thread(target=self.peak_monitor_func) + peak_monitor_thread.daemon = True + peak_monitor_thread.start() + + def stop(self, stage): + """stop tracking for the passed stage""" + + # deal with nested calls of eval during train - simply ignore those + if self.cur_stage is not None and self.cur_stage != stage: + return + + # this sends a signal to peak_monitor_func to complete its loop + self.peak_monitoring = False + + # first ensure all objects get collected and their memory is freed + gc.collect() + + if self.torch is not None: + if torch.cuda.is_available(): + self.torch.cuda.empty_cache() + elif is_torch_mlu_available(): + self.torch.mlu.empty_cache() + elif is_torch_musa_available(): + self.torch.musa.empty_cache() + elif is_torch_xpu_available(): + self.torch.xpu.empty_cache() + elif is_torch_npu_available(): + self.torch.npu.empty_cache() + elif is_torch_hpu_available(): + # not available on hpu as it reserves all device memory for the current process + # self.torch.npu.empty_cache() + pass + elif is_torch_mps_available(): + self.torch.mps.empty_cache() + + # concepts: + # - alloc_delta: the difference of allocated memory between the end and the start + # - peaked_delta: the difference between the peak memory and the current memory + # in order to know how much memory the measured code consumed one needs to sum these two + + # gpu + if self.torch is not None: + if torch.cuda.is_available(): + self.gpu_mem_used_now = self.torch.cuda.memory_allocated() + self.gpu_mem_used_peak = self.torch.cuda.max_memory_allocated() + elif is_torch_mlu_available(): + self.gpu_mem_used_now = self.torch.mlu.memory_allocated() + self.gpu_mem_used_peak = self.torch.mlu.max_memory_allocated() + elif is_torch_musa_available(): + self.gpu_mem_used_now = self.torch.musa.memory_allocated() + self.gpu_mem_used_peak = self.torch.musa.max_memory_allocated() + elif is_torch_xpu_available(): + self.gpu_mem_used_now = self.torch.xpu.memory_allocated() + self.gpu_mem_used_peak = self.torch.xpu.max_memory_allocated() + elif is_torch_npu_available(): + self.gpu_mem_used_now = self.torch.npu.memory_allocated() + self.gpu_mem_used_peak = self.torch.npu.max_memory_allocated() + elif is_torch_hpu_available(): + self.gpu_mem_used_now = self.torch.hpu.memory_allocated() + self.gpu_mem_used_peak = self.torch.hpu.max_memory_allocated() + elif is_torch_mps_available(): + self.gpu_mem_used_now = self.torch.mps.current_allocated_memory() + # self.torch.mps.max_memory_allocated() does not exist yet + self.gpu_mem_used_peak = None + + else: + raise ValueError("No available GPU device found!") + + self.gpu[self.cur_stage] = { + "begin": self.gpu_mem_used_at_start, + "end": self.gpu_mem_used_now, + "alloc": (self.gpu_mem_used_now - self.gpu_mem_used_at_start), + } + if self.gpu_mem_used_peak is not None: + self.gpu[self.cur_stage]["peaked"] = max(0, self.gpu_mem_used_peak - self.gpu_mem_used_now) + else: + self.gpu[self.cur_stage]["peaked"] = "Not available" + + # cpu + self.cpu_mem_used_now = self.cpu_mem_used() + self.cpu[self.cur_stage] = { + "begin": self.cpu_mem_used_at_start, + "end": self.cpu_mem_used_now, + "alloc": (self.cpu_mem_used_now - self.cpu_mem_used_at_start), + "peaked": max(0, self.cpu_mem_used_peak - self.cpu_mem_used_now), + } + + # reset - cycle finished + self.cur_stage = None + + def update_metrics(self, stage, metrics): + """updates the metrics""" + if self.skip_memory_metrics: + return + + # deal with nested calls of eval during train - simply ignore those + if self.cur_stage is not None and self.cur_stage != stage: + return + + # since we don't have a way to return init metrics, we push them into the first of train/val/predict + stages = [stage] + if not self.init_reported: + stages.insert(0, "init") + self.init_reported = True + + for stage in stages: + for t in ["alloc", "peaked"]: + if stage in self.cpu and t in self.cpu[stage]: + metrics[f"{stage}_mem_cpu_{t}_delta"] = self.cpu[stage][t] + if self.torch is not None and stage in self.gpu and t in self.gpu[stage]: + metrics[f"{stage}_mem_gpu_{t}_delta"] = self.gpu[stage][t] + # if we need additional debug info, enable the following + # for t in ["begin", "end"]: + # if stage in self.cpu and t in self.cpu[stage]: + # metrics[f"{stage}_mem_cpu_{t}"] = self.cpu[stage][t] + # if self.torch is not None and stage in self.gpu and t in self.gpu[stage]: + # metrics[f"{stage}_mem_gpu_{t}"] = self.gpu[stage][t] + + # since memory can be allocated before init, and it might be difficult to track overall + # memory usage, in particular for GPU, let's report memory usage at the point init was called + if stages[0] == "init": + metrics["before_init_mem_cpu"] = self.cpu["init"]["begin"] + if self.torch is not None: + metrics["before_init_mem_gpu"] = self.gpu["init"]["begin"] + # if we also wanted to report any additional memory allocations in between init and + # whatever the next stage was we could also report this: + # if self.cpu["init"]["end"] != self.cpu[stage]["begin"]: + # metrics[f"after_init_mem_cpu_delta"] = self.cpu[stage]["begin"] - self.cpu["init"]["end"] + # if self.torch is not None and self.gpu["init"]["end"] != self.gpu[stage]["begin"]: + # metrics[f"after_init_mem_gpu_delta"] = self.gpu[stage]["begin"] - self.gpu["init"]["end"] + + def stop_and_update_metrics(self, metrics=None): + """combine stop and metrics update in one call for simpler code""" + if self.skip_memory_metrics: + return + + stage = self.derive_stage() + self.stop(stage) + + # init doesn't have metrics to update so we just save that data for later stages to retrieve + if metrics is not None: + self.update_metrics(stage, metrics) + + +def has_length(dataset): + """ + Checks if the dataset implements __len__() and it doesn't raise an error + """ + try: + return len(dataset) is not None + except TypeError: + # TypeError: len() of unsized object + return False + except AttributeError: + # Ray DataSets raises an AttributeError: https://github.com/ray-project/ray/blob/master/python/ray/data/dataset.py#L5616 + return False + + +def denumpify_detensorize(metrics): + """ + Recursively calls `.item()` on the element of the dictionary passed + """ + if isinstance(metrics, (list, tuple)): + return type(metrics)(denumpify_detensorize(m) for m in metrics) + elif isinstance(metrics, dict): + return type(metrics)({k: denumpify_detensorize(v) for k, v in metrics.items()}) + elif isinstance(metrics, np.generic): + return metrics.item() + elif is_torch_available() and isinstance(metrics, torch.Tensor) and metrics.numel() == 1: + return metrics.item() + return metrics + + +def number_of_arguments(func): + """ + Return the number of arguments of the passed function, even if it's a partial function. + """ + if isinstance(func, functools.partial): + total_args = len(inspect.signature(func.func).parameters) + return total_args - len(func.args) - len(func.keywords) + return len(inspect.signature(func).parameters) + + +def find_executable_batch_size( + function: callable = None, starting_batch_size: int = 128, auto_find_batch_size: bool = False +): + """ + Args: + A basic decorator that will try to execute `function`. If it fails from exceptions related to out-of-memory or + CUDNN, the batch size is cut in half and passed to `function`. `function` must take in a `batch_size` parameter as + its first argument. + function (`callable`, *optional*) + A function to wrap + starting_batch_size (`int`, *optional*) + The batch size to try and fit into memory + auto_find_batch_size (`bool`, *optional*) + If False, will just execute `function` + """ + if function is None: + return functools.partial( + find_executable_batch_size, + starting_batch_size=starting_batch_size, + auto_find_batch_size=auto_find_batch_size, + ) + + if auto_find_batch_size: + requires_backends(find_executable_batch_size, "accelerate") + from accelerate.utils import find_executable_batch_size as accelerate_find_executable_batch_size + + return accelerate_find_executable_batch_size(function=function, starting_batch_size=starting_batch_size) + + return functools.partial(function, batch_size=starting_batch_size) + + +class FSDPOption(ExplicitEnum): + FULL_SHARD = "full_shard" + SHARD_GRAD_OP = "shard_grad_op" + NO_SHARD = "no_shard" + HYBRID_SHARD = "hybrid_shard" + HYBRID_SHARD_ZERO2 = "hybrid_shard_zero2" + OFFLOAD = "offload" + AUTO_WRAP = "auto_wrap" + + +class RemoveColumnsCollator: + """Wrap the data collator to remove unused columns before they are passed to the collator.""" + + def __init__( + self, + data_collator, + signature_columns, + logger=None, + model_name: Optional[str] = None, + description: Optional[str] = None, + ): + self.data_collator = data_collator + self.signature_columns = signature_columns + self.logger = logger + self.description = description + self.model_name = model_name + self.message_logged = False + + def _remove_columns(self, feature: dict) -> dict: + if not isinstance(feature, dict): + return feature + if not self.message_logged and self.logger and self.model_name: + ignored_columns = list(set(feature.keys()) - set(self.signature_columns)) + if len(ignored_columns) > 0: + dset_description = "" if self.description is None else f"in the {self.description} set" + self.logger.info( + f"The following columns {dset_description} don't have a corresponding argument in " + f"`{self.model_name}.forward` and have been ignored: {', '.join(ignored_columns)}." + f" If {', '.join(ignored_columns)} are not expected by `{self.model_name}.forward`, " + " you can safely ignore this message." + ) + self.message_logged = True + return {k: v for k, v in feature.items() if k in self.signature_columns} + + def __call__(self, features: list[dict]): + features = [self._remove_columns(feature) for feature in features] + return self.data_collator(features) + + +def check_target_module_exists(optim_target_modules, key: str, return_is_regex: bool = False): + """A helper method to check if the passed module's key name matches any of the target modules in the optim_target_modules. + + Args: + optim_target_modules (`Union[str, List[str]]`): + A list of strings to try to match. Can be also a full string. + key (`str`): + A key to search any matches in optim_target_modules + return_is_regex (`bool`): + If set to `True`, the method will return whether the passed `optim_target_modules` + is a regex or not. + + Returns: + `bool` : True of match object if key matches any target modules from config, False or + None if no match found + `bool` : If the matched target module is a regex to silence out the warnings in Trainer + for extra modules being found (only if `target_module_found=True` for an array of regex). + """ + target_module_found = False + is_regex = False + + if isinstance(optim_target_modules, str): + target_module_found = bool(re.fullmatch(optim_target_modules, key)) + is_regex = True if not optim_target_modules == key else False + elif key in optim_target_modules: # from here, target_module_found must be a list of str + # this module is specified directly in target_modules + target_module_found = True + elif any(target_key in key for target_key in optim_target_modules): + target_module_found = True + elif any(bool(re.fullmatch(optim_target_module, key)) for optim_target_module in optim_target_modules): + target_module_found = True + is_regex = True + + if return_is_regex: + return target_module_found, is_regex + + return target_module_found diff --git a/lib/python3.12/site-packages/transformers/training_args.py b/lib/python3.12/site-packages/transformers/training_args.py new file mode 100644 index 0000000000000000000000000000000000000000..0e153df47a289ff24afd5f2c200c8f13f5178bc6 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/training_args.py @@ -0,0 +1,3097 @@ +# Copyright 2020 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. + +import contextlib +import json +import math +import os +import warnings +from dataclasses import asdict, dataclass, field, fields +from datetime import timedelta +from enum import Enum +from pathlib import Path +from typing import Any, Optional, Union + +from huggingface_hub import get_full_repo_name +from packaging import version + +from .debug_utils import DebugOption +from .trainer_utils import ( + EvaluationStrategy, + FSDPOption, + HubStrategy, + IntervalStrategy, + SaveStrategy, + SchedulerType, +) +from .utils import ( + ACCELERATE_MIN_VERSION, + ExplicitEnum, + cached_property, + is_accelerate_available, + is_ipex_available, + is_safetensors_available, + is_sagemaker_dp_enabled, + is_sagemaker_mp_enabled, + is_torch_available, + is_torch_bf16_cpu_available, + is_torch_bf16_gpu_available, + is_torch_hpu_available, + is_torch_mlu_available, + is_torch_mps_available, + is_torch_musa_available, + is_torch_neuroncore_available, + is_torch_npu_available, + is_torch_tf32_available, + is_torch_xla_available, + is_torch_xpu_available, + logging, + requires_backends, +) +from .utils.generic import strtobool +from .utils.import_utils import is_optimum_neuron_available + + +logger = logging.get_logger(__name__) +log_levels = logging.get_log_levels_dict().copy() +trainer_log_levels = dict(**log_levels, passive=-1) + +if is_torch_available(): + import torch + import torch.distributed as dist + +if is_accelerate_available(): + from accelerate.state import AcceleratorState, PartialState + from accelerate.utils import DistributedType + + from .trainer_pt_utils import AcceleratorConfig + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + +if is_torch_neuroncore_available(check_device=False): + # torchrun support + # https://github.com/pytorch/xla/pull/3609 + if os.environ.get("TORCHELASTIC_RUN_ID"): + if is_optimum_neuron_available(): + logger.info( + "Make sure that you are performing the training with the NeuronTrainer from optimum[neuron], this " + "will fail otherwise." + ) + else: + logger.warning( + "Please use the NeuronTrainer from optimum[neuron] instead of the Transformers library to perform " + "training on AWS Trainium instances. More information here: " + "https://github.com/huggingface/optimum-neuron" + ) + import torch_xla.distributed.xla_backend as xbn + + if not isinstance(dist.group.WORLD, xbn.ProcessGroupXla): + dist.init_process_group(backend="xla") + if not isinstance(dist.group.WORLD, xbn.ProcessGroupXla): + raise AssertionError("Failed to initialize torch.distributed process group using XLA backend.") + + +if is_sagemaker_mp_enabled(): + import smdistributed.modelparallel.torch as smp + + smp.init() + + +def default_logdir() -> str: + """ + Same default as PyTorch + """ + import socket + from datetime import datetime + + current_time = datetime.now().strftime("%b%d_%H-%M-%S") + return os.path.join("runs", current_time + "_" + socket.gethostname()) + + +def get_int_from_env(env_keys, default): + """Returns the first positive env value found in the `env_keys` list or the default.""" + for e in env_keys: + val = int(os.environ.get(e, -1)) + if val >= 0: + return val + return default + + +def get_xla_device_type(device: "torch.device") -> Optional[str]: + """ + Returns the xla device type (CPU|GPU|TPU) or None if the device is a non-xla device. + """ + if is_torch_xla_available(): + if device.type == "cpu": + return "CPU" + return xm.xla_real_devices([device])[0].split(":")[0] + return None + + +class OptimizerNames(ExplicitEnum): + """ + Stores the acceptable string identifiers for optimizers. + """ + + ADAMW_TORCH = "adamw_torch" + ADAMW_TORCH_FUSED = "adamw_torch_fused" + ADAMW_TORCH_XLA = "adamw_torch_xla" + ADAMW_TORCH_NPU_FUSED = "adamw_torch_npu_fused" + ADAMW_APEX_FUSED = "adamw_apex_fused" + ADAFACTOR = "adafactor" + ADAMW_ANYPRECISION = "adamw_anyprecision" + ADAMW_TORCH_4BIT = "adamw_torch_4bit" + ADAMW_TORCH_8BIT = "adamw_torch_8bit" + ADEMAMIX = "ademamix" + SGD = "sgd" + ADAGRAD = "adagrad" + ADAMW_BNB = "adamw_bnb_8bit" + ADAMW_8BIT = "adamw_8bit" # just an alias for adamw_bnb_8bit + ADEMAMIX_8BIT = "ademamix_8bit" + LION_8BIT = "lion_8bit" + LION = "lion_32bit" + PAGED_ADAMW = "paged_adamw_32bit" + PAGED_ADAMW_8BIT = "paged_adamw_8bit" + PAGED_ADEMAMIX = "paged_ademamix_32bit" + PAGED_ADEMAMIX_8BIT = "paged_ademamix_8bit" + PAGED_LION = "paged_lion_32bit" + PAGED_LION_8BIT = "paged_lion_8bit" + RMSPROP = "rmsprop" + RMSPROP_BNB = "rmsprop_bnb" + RMSPROP_8BIT = "rmsprop_bnb_8bit" + RMSPROP_32BIT = "rmsprop_bnb_32bit" + GALORE_ADAMW = "galore_adamw" + GALORE_ADAMW_8BIT = "galore_adamw_8bit" + GALORE_ADAFACTOR = "galore_adafactor" + GALORE_ADAMW_LAYERWISE = "galore_adamw_layerwise" + GALORE_ADAMW_8BIT_LAYERWISE = "galore_adamw_8bit_layerwise" + GALORE_ADAFACTOR_LAYERWISE = "galore_adafactor_layerwise" + LOMO = "lomo" + ADALOMO = "adalomo" + GROKADAMW = "grokadamw" + SCHEDULE_FREE_RADAM = "schedule_free_radam" + SCHEDULE_FREE_ADAMW = "schedule_free_adamw" + SCHEDULE_FREE_SGD = "schedule_free_sgd" + APOLLO_ADAMW = "apollo_adamw" + APOLLO_ADAMW_LAYERWISE = "apollo_adamw_layerwise" + + +def _convert_str_dict(passed_value: dict): + "Safely checks that a passed value is a dictionary and converts any string values to their appropriate types." + for key, value in passed_value.items(): + if isinstance(value, dict): + passed_value[key] = _convert_str_dict(value) + elif isinstance(value, str): + # First check for bool and convert + if value.lower() in ("true", "false"): + passed_value[key] = value.lower() == "true" + # Check for digit + elif value.isdigit(): + passed_value[key] = int(value) + elif value.replace(".", "", 1).isdigit(): + passed_value[key] = float(value) + + return passed_value + + +# TODO: `TrainingArguments` users rely on it being fully mutable. In the future see if we can narrow this to a few keys: https://github.com/huggingface/transformers/pull/25903 +@dataclass +class TrainingArguments: + """ + TrainingArguments is the subset of the arguments we use in our example scripts **which relate to the training loop + itself**. + + Using [`HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + output_dir (`str`, *optional*, defaults to `"trainer_output"`): + The output directory where the model predictions and checkpoints will be written. + overwrite_output_dir (`bool`, *optional*, defaults to `False`): + If `True`, overwrite the content of the output directory. Use this to continue training if `output_dir` + points to a checkpoint directory. + do_train (`bool`, *optional*, defaults to `False`): + Whether to run training or not. This argument is not directly used by [`Trainer`], it's intended to be used + by your training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + do_eval (`bool`, *optional*): + Whether to run evaluation on the validation set or not. Will be set to `True` if `eval_strategy` is + different from `"no"`. This argument is not directly used by [`Trainer`], it's intended to be used by your + training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + do_predict (`bool`, *optional*, defaults to `False`): + Whether to run predictions on the test set or not. This argument is not directly used by [`Trainer`], it's + intended to be used by your training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + eval_strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"no"`): + The evaluation strategy to adopt during training. Possible values are: + + - `"no"`: No evaluation is done during training. + - `"steps"`: Evaluation is done (and logged) every `eval_steps`. + - `"epoch"`: Evaluation is done at the end of each epoch. + + prediction_loss_only (`bool`, *optional*, defaults to `False`): + When performing evaluation and generating predictions, only returns the loss. + per_device_train_batch_size (`int`, *optional*, defaults to 8): + The batch size per device accelerator core/CPU for training. + per_device_eval_batch_size (`int`, *optional*, defaults to 8): + The batch size per device accelerator core/CPU for evaluation. + gradient_accumulation_steps (`int`, *optional*, defaults to 1): + Number of updates steps to accumulate the gradients for, before performing a backward/update pass. + + + + When using gradient accumulation, one step is counted as one step with backward pass. Therefore, logging, + evaluation, save will be conducted every `gradient_accumulation_steps * xxx_step` training examples. + + + + eval_accumulation_steps (`int`, *optional*): + Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU. If + left unset, the whole predictions are accumulated on the device accelerator before being moved to the CPU (faster but + requires more memory). + eval_delay (`float`, *optional*): + Number of epochs or steps to wait for before the first evaluation can be performed, depending on the + eval_strategy. + torch_empty_cache_steps (`int`, *optional*): + Number of steps to wait before calling `torch..empty_cache()`. If left unset or set to None, cache will not be emptied. + + + + This can help avoid CUDA out-of-memory errors by lowering peak VRAM usage at a cost of about [10% slower performance](https://github.com/huggingface/transformers/issues/31372). + + + + learning_rate (`float`, *optional*, defaults to 5e-5): + The initial learning rate for [`AdamW`] optimizer. + weight_decay (`float`, *optional*, defaults to 0): + The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights in [`AdamW`] + optimizer. + adam_beta1 (`float`, *optional*, defaults to 0.9): + The beta1 hyperparameter for the [`AdamW`] optimizer. + adam_beta2 (`float`, *optional*, defaults to 0.999): + The beta2 hyperparameter for the [`AdamW`] optimizer. + adam_epsilon (`float`, *optional*, defaults to 1e-8): + The epsilon hyperparameter for the [`AdamW`] optimizer. + max_grad_norm (`float`, *optional*, defaults to 1.0): + Maximum gradient norm (for gradient clipping). + num_train_epochs(`float`, *optional*, defaults to 3.0): + Total number of training epochs to perform (if not an integer, will perform the decimal part percents of + the last epoch before stopping training). + max_steps (`int`, *optional*, defaults to -1): + If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`. + For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until + `max_steps` is reached. + lr_scheduler_type (`str` or [`SchedulerType`], *optional*, defaults to `"linear"`): + The scheduler type to use. See the documentation of [`SchedulerType`] for all possible values. + lr_scheduler_kwargs ('dict', *optional*, defaults to {}): + The extra arguments for the lr_scheduler. See the documentation of each scheduler for possible values. + warmup_ratio (`float`, *optional*, defaults to 0.0): + Ratio of total training steps used for a linear warmup from 0 to `learning_rate`. + warmup_steps (`int`, *optional*, defaults to 0): + Number of steps used for a linear warmup from 0 to `learning_rate`. Overrides any effect of `warmup_ratio`. + log_level (`str`, *optional*, defaults to `passive`): + Logger log level to use on the main process. Possible choices are the log levels as strings: 'debug', + 'info', 'warning', 'error' and 'critical', plus a 'passive' level which doesn't set anything and keeps the + current log level for the Transformers library (which will be `"warning"` by default). + log_level_replica (`str`, *optional*, defaults to `"warning"`): + Logger log level to use on replicas. Same choices as `log_level`" + log_on_each_node (`bool`, *optional*, defaults to `True`): + In multinode distributed training, whether to log using `log_level` once per node, or only on the main + node. + logging_dir (`str`, *optional*): + [TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to + *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***. + logging_strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`): + The logging strategy to adopt during training. Possible values are: + + - `"no"`: No logging is done during training. + - `"epoch"`: Logging is done at the end of each epoch. + - `"steps"`: Logging is done every `logging_steps`. + + logging_first_step (`bool`, *optional*, defaults to `False`): + Whether to log the first `global_step` or not. + logging_steps (`int` or `float`, *optional*, defaults to 500): + Number of update steps between two logs if `logging_strategy="steps"`. Should be an integer or a float in + range `[0,1)`. If smaller than 1, will be interpreted as ratio of total training steps. + logging_nan_inf_filter (`bool`, *optional*, defaults to `True`): + Whether to filter `nan` and `inf` losses for logging. If set to `True` the loss of every step that is `nan` + or `inf` is filtered and the average loss of the current logging window is taken instead. + + + + `logging_nan_inf_filter` only influences the logging of loss values, it does not change the behavior the + gradient is computed or applied to the model. + + + + save_strategy (`str` or [`~trainer_utils.SaveStrategy`], *optional*, defaults to `"steps"`): + The checkpoint save strategy to adopt during training. Possible values are: + + - `"no"`: No save is done during training. + - `"epoch"`: Save is done at the end of each epoch. + - `"steps"`: Save is done every `save_steps`. + - `"best"`: Save is done whenever a new `best_metric` is achieved. + + If `"epoch"` or `"steps"` is chosen, saving will also be performed at the + very end of training, always. + save_steps (`int` or `float`, *optional*, defaults to 500): + Number of updates steps before two checkpoint saves if `save_strategy="steps"`. Should be an integer or a + float in range `[0,1)`. If smaller than 1, will be interpreted as ratio of total training steps. + save_total_limit (`int`, *optional*): + If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in + `output_dir`. When `load_best_model_at_end` is enabled, the "best" checkpoint according to + `metric_for_best_model` will always be retained in addition to the most recent ones. For example, for + `save_total_limit=5` and `load_best_model_at_end`, the four last checkpoints will always be retained + alongside the best model. When `save_total_limit=1` and `load_best_model_at_end`, it is possible that two + checkpoints are saved: the last one and the best one (if they are different). + save_safetensors (`bool`, *optional*, defaults to `True`): + Use [safetensors](https://huggingface.co/docs/safetensors) saving and loading for state dicts instead of + default `torch.load` and `torch.save`. + save_on_each_node (`bool`, *optional*, defaults to `False`): + When doing multi-node distributed training, whether to save models and checkpoints on each node, or only on + the main one. + + This should not be activated when the different nodes use the same storage as the files will be saved with + the same names for each node. + save_only_model (`bool`, *optional*, defaults to `False`): + When checkpointing, whether to only save the model, or also the optimizer, scheduler & rng state. + Note that when this is true, you won't be able to resume training from checkpoint. + This enables you to save storage by not storing the optimizer, scheduler & rng state. + You can only load the model using `from_pretrained` with this option set to `True`. + restore_callback_states_from_checkpoint (`bool`, *optional*, defaults to `False`): + Whether to restore the callback states from the checkpoint. If `True`, will override + callbacks passed to the `Trainer` if they exist in the checkpoint." + use_cpu (`bool`, *optional*, defaults to `False`): + Whether or not to use cpu. If set to False, we will use cuda or mps device if available. + seed (`int`, *optional*, defaults to 42): + Random seed that will be set at the beginning of training. To ensure reproducibility across runs, use the + [`~Trainer.model_init`] function to instantiate the model if it has some randomly initialized parameters. + data_seed (`int`, *optional*): + Random seed to be used with data samplers. If not set, random generators for data sampling will use the + same seed as `seed`. This can be used to ensure reproducibility of data sampling, independent of the model + seed. + jit_mode_eval (`bool`, *optional*, defaults to `False`): + Whether or not to use PyTorch jit trace for inference. + use_ipex (`bool`, *optional*, defaults to `False`): + Use Intel extension for PyTorch when it is available. [IPEX + installation](https://github.com/intel/intel-extension-for-pytorch). + bf16 (`bool`, *optional*, defaults to `False`): + Whether to use bf16 16-bit (mixed) precision training instead of 32-bit training. Requires Ampere or higher + NVIDIA architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change. + fp16 (`bool`, *optional*, defaults to `False`): + Whether to use fp16 16-bit (mixed) precision training instead of 32-bit training. + fp16_opt_level (`str`, *optional*, defaults to 'O1'): + For `fp16` training, Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']. See details on + the [Apex documentation](https://nvidia.github.io/apex/amp). + fp16_backend (`str`, *optional*, defaults to `"auto"`): + This argument is deprecated. Use `half_precision_backend` instead. + half_precision_backend (`str`, *optional*, defaults to `"auto"`): + The backend to use for mixed precision training. Must be one of `"auto", "apex", "cpu_amp"`. `"auto"` will + use CPU/CUDA AMP or APEX depending on the PyTorch version detected, while the other choices will force the + requested backend. + bf16_full_eval (`bool`, *optional*, defaults to `False`): + Whether to use full bfloat16 evaluation instead of 32-bit. This will be faster and save memory but can harm + metric values. This is an experimental API and it may change. + fp16_full_eval (`bool`, *optional*, defaults to `False`): + Whether to use full float16 evaluation instead of 32-bit. This will be faster and save memory but can harm + metric values. + tf32 (`bool`, *optional*): + Whether to enable the TF32 mode, available in Ampere and newer GPU architectures. The default value depends + on PyTorch's version default of `torch.backends.cuda.matmul.allow_tf32`. For more details please refer to + the [TF32](https://huggingface.co/docs/transformers/perf_train_gpu_one#tf32) documentation. This is an + experimental API and it may change. + local_rank (`int`, *optional*, defaults to -1): + Rank of the process during distributed training. + ddp_backend (`str`, *optional*): + The backend to use for distributed training. Must be one of `"nccl"`, `"mpi"`, `"ccl"`, `"gloo"`, `"hccl"`. + tpu_num_cores (`int`, *optional*): + When training on TPU, the number of TPU cores (automatically passed by launcher script). + dataloader_drop_last (`bool`, *optional*, defaults to `False`): + Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size) + or not. + eval_steps (`int` or `float`, *optional*): + Number of update steps between two evaluations if `eval_strategy="steps"`. Will default to the same + value as `logging_steps` if not set. Should be an integer or a float in range `[0,1)`. If smaller than 1, + will be interpreted as ratio of total training steps. + dataloader_num_workers (`int`, *optional*, defaults to 0): + Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in the + main process. + past_index (`int`, *optional*, defaults to -1): + Some models like [TransformerXL](../model_doc/transformerxl) or [XLNet](../model_doc/xlnet) can make use of + the past hidden states for their predictions. If this argument is set to a positive int, the `Trainer` will + use the corresponding output (usually index 2) as the past state and feed it to the model at the next + training step under the keyword argument `mems`. + run_name (`str`, *optional*, defaults to `output_dir`): + A descriptor for the run. Typically used for [wandb](https://www.wandb.com/), + [mlflow](https://www.mlflow.org/), [comet](https://www.comet.com/site) and [swanlab](https://swanlab.cn) + logging. If not specified, will be the same as `output_dir`. + disable_tqdm (`bool`, *optional*): + Whether or not to disable the tqdm progress bars and table of metrics produced by + [`~notebook.NotebookTrainingTracker`] in Jupyter Notebooks. Will default to `True` if the logging level is + set to warn or lower (default), `False` otherwise. + remove_unused_columns (`bool`, *optional*, defaults to `True`): + Whether or not to automatically remove the columns unused by the model forward method. + label_names (`List[str]`, *optional*): + The list of keys in your dictionary of inputs that correspond to the labels. + + Will eventually default to the list of argument names accepted by the model that contain the word "label", + except if the model used is one of the `XxxForQuestionAnswering` in which case it will also include the + `["start_positions", "end_positions"]` keys. + load_best_model_at_end (`bool`, *optional*, defaults to `False`): + Whether or not to load the best model found during training at the end of training. When this option is + enabled, the best checkpoint will always be saved. See + [`save_total_limit`](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.save_total_limit) + for more. + + + + When set to `True`, the parameters `save_strategy` needs to be the same as `eval_strategy`, and in + the case it is "steps", `save_steps` must be a round multiple of `eval_steps`. + + + + metric_for_best_model (`str`, *optional*): + Use in conjunction with `load_best_model_at_end` to specify the metric to use to compare two different + models. Must be the name of a metric returned by the evaluation with or without the prefix `"eval_"`. + + If not specified, this will default to `"loss"` when either `load_best_model_at_end == True` + or `lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU` (to use the evaluation loss). + + If you set this value, `greater_is_better` will default to `True` unless the name ends with "loss". + Don't forget to set it to `False` if your metric is better when lower. + greater_is_better (`bool`, *optional*): + Use in conjunction with `load_best_model_at_end` and `metric_for_best_model` to specify if better models + should have a greater metric or not. Will default to: + + - `True` if `metric_for_best_model` is set to a value that doesn't end in `"loss"`. + - `False` if `metric_for_best_model` is not set, or set to a value that ends in `"loss"`. + ignore_data_skip (`bool`, *optional*, defaults to `False`): + When resuming training, whether or not to skip the epochs and batches to get the data loading at the same + stage as in the previous training. If set to `True`, the training will begin faster (as that skipping step + can take a long time) but will not yield the same results as the interrupted training would have. + fsdp (`bool`, `str` or list of [`~trainer_utils.FSDPOption`], *optional*, defaults to `''`): + Use PyTorch Distributed Parallel Training (in distributed training only). + + A list of options along the following: + + - `"full_shard"`: Shard parameters, gradients and optimizer states. + - `"shard_grad_op"`: Shard optimizer states and gradients. + - `"hybrid_shard"`: Apply `FULL_SHARD` within a node, and replicate parameters across nodes. + - `"hybrid_shard_zero2"`: Apply `SHARD_GRAD_OP` within a node, and replicate parameters across nodes. + - `"offload"`: Offload parameters and gradients to CPUs (only compatible with `"full_shard"` and + `"shard_grad_op"`). + - `"auto_wrap"`: Automatically recursively wrap layers with FSDP using `default_auto_wrap_policy`. + fsdp_config (`str` or `dict`, *optional*): + Config to be used with fsdp (Pytorch Distributed Parallel Training). The value is either a location of + fsdp json config file (e.g., `fsdp_config.json`) or an already loaded json file as `dict`. + + A List of config and its options: + - min_num_params (`int`, *optional*, defaults to `0`): + FSDP's minimum number of parameters for Default Auto Wrapping. (useful only when `fsdp` field is + passed). + - transformer_layer_cls_to_wrap (`List[str]`, *optional*): + List of transformer layer class names (case-sensitive) to wrap, e.g, `BertLayer`, `GPTJBlock`, + `T5Block` .... (useful only when `fsdp` flag is passed). + - backward_prefetch (`str`, *optional*) + FSDP's backward prefetch mode. Controls when to prefetch next set of parameters (useful only when + `fsdp` field is passed). + + A list of options along the following: + + - `"backward_pre"` : Prefetches the next set of parameters before the current set of parameter's + gradient + computation. + - `"backward_post"` : This prefetches the next set of parameters after the current set of + parameter’s + gradient computation. + - forward_prefetch (`bool`, *optional*, defaults to `False`) + FSDP's forward prefetch mode (useful only when `fsdp` field is passed). + If `"True"`, then FSDP explicitly prefetches the next upcoming all-gather while executing in the + forward pass. + - limit_all_gathers (`bool`, *optional*, defaults to `False`) + FSDP's limit_all_gathers (useful only when `fsdp` field is passed). + If `"True"`, FSDP explicitly synchronizes the CPU thread to prevent too many in-flight + all-gathers. + - use_orig_params (`bool`, *optional*, defaults to `True`) + If `"True"`, allows non-uniform `requires_grad` during init, which means support for interspersed + frozen and trainable parameters. Useful in cases such as parameter-efficient fine-tuning. Please + refer this + [blog](https://dev-discuss.pytorch.org/t/rethinking-pytorch-fully-sharded-data-parallel-fsdp-from-first-principles/1019 + - sync_module_states (`bool`, *optional*, defaults to `True`) + If `"True"`, each individually wrapped FSDP unit will broadcast module parameters from rank 0 to + ensure they are the same across all ranks after initialization + - cpu_ram_efficient_loading (`bool`, *optional*, defaults to `False`) + If `"True"`, only the first process loads the pretrained model checkpoint while all other processes + have empty weights. When this setting as `"True"`, `sync_module_states` also must to be `"True"`, + otherwise all the processes except the main process would have random weights leading to unexpected + behaviour during training. + - activation_checkpointing (`bool`, *optional*, defaults to `False`): + If `"True"`, activation checkpointing is a technique to reduce memory usage by clearing activations of + certain layers and recomputing them during a backward pass. Effectively, this trades extra + computation time for reduced memory usage. + - xla (`bool`, *optional*, defaults to `False`): + Whether to use PyTorch/XLA Fully Sharded Data Parallel Training. This is an experimental feature + and its API may evolve in the future. + - xla_fsdp_settings (`dict`, *optional*) + The value is a dictionary which stores the XLA FSDP wrapping parameters. + + For a complete list of options, please see [here]( + https://github.com/pytorch/xla/blob/master/torch_xla/distributed/fsdp/xla_fully_sharded_data_parallel.py). + - xla_fsdp_grad_ckpt (`bool`, *optional*, defaults to `False`): + Will use gradient checkpointing over each nested XLA FSDP wrapped layer. This setting can only be + used when the xla flag is set to true, and an auto wrapping policy is specified through + fsdp_min_num_params or fsdp_transformer_layer_cls_to_wrap. + tp_size (`int`, *optional*): + Use tp_size to enable PyTorch tensor parallelism. Tensor parallelism support is only available to models having `base_tp_plan` + in their respective config classes. + Set a value greater than 1 to activate TP. The same is used to prepare device mesh internally. Requires accelerate>1.3.0. + deepspeed (`str` or `dict`, *optional*): + Use [Deepspeed](https://github.com/deepspeedai/DeepSpeed). This is an experimental feature and its API may + evolve in the future. The value is either the location of DeepSpeed json config file (e.g., + `ds_config.json`) or an already loaded json file as a `dict`" + + + If enabling any Zero-init, make sure that your model is not initialized until + *after* initializing the `TrainingArguments`, else it will not be applied. + + + accelerator_config (`str`, `dict`, or `AcceleratorConfig`, *optional*): + Config to be used with the internal `Accelerator` implementation. The value is either a location of + accelerator json config file (e.g., `accelerator_config.json`), an already loaded json file as `dict`, + or an instance of [`~trainer_pt_utils.AcceleratorConfig`]. + + A list of config and its options: + - split_batches (`bool`, *optional*, defaults to `False`): + Whether or not the accelerator should split the batches yielded by the dataloaders across the devices. If + `True` the actual batch size used will be the same on any kind of distributed processes, but it must be a + round multiple of the `num_processes` you are using. If `False`, actual batch size used will be the one set + in your script multiplied by the number of processes. + - dispatch_batches (`bool`, *optional*): + If set to `True`, the dataloader prepared by the Accelerator is only iterated through on the main process + and then the batches are split and broadcast to each process. Will default to `True` for `DataLoader` whose + underlying dataset is an `IterableDataset`, `False` otherwise. + - even_batches (`bool`, *optional*, defaults to `True`): + If set to `True`, in cases where the total batch size across all processes does not exactly divide the + dataset, samples at the start of the dataset will be duplicated so the batch can be divided equally among + all workers. + - use_seedable_sampler (`bool`, *optional*, defaults to `True`): + Whether or not use a fully seedable random sampler ([`accelerate.data_loader.SeedableRandomSampler`]). Ensures + training results are fully reproducible using a different sampling technique. While seed-to-seed results + may differ, on average the differences are negligible when using multiple different seeds to compare. Should + also be ran with [`~utils.set_seed`] for the best results. + - use_configured_state (`bool`, *optional*, defaults to `False`): + Whether or not to use a pre-configured `AcceleratorState` or `PartialState` defined before calling `TrainingArguments`. + If `True`, an `Accelerator` or `PartialState` must be initialized. Note that by doing so, this could lead to issues + with hyperparameter tuning. + + label_smoothing_factor (`float`, *optional*, defaults to 0.0): + The label smoothing factor to use. Zero means no label smoothing, otherwise the underlying onehot-encoded + labels are changed from 0s and 1s to `label_smoothing_factor/num_labels` and `1 - label_smoothing_factor + + label_smoothing_factor/num_labels` respectively. + debug (`str` or list of [`~debug_utils.DebugOption`], *optional*, defaults to `""`): + Enable one or more debug features. This is an experimental feature. + + Possible options are: + + - `"underflow_overflow"`: detects overflow in model's input/outputs and reports the last frames that led to + the event + - `"tpu_metrics_debug"`: print debug metrics on TPU + + The options should be separated by whitespaces. + optim (`str` or [`training_args.OptimizerNames`], *optional*, defaults to `"adamw_torch"`): + The optimizer to use, such as "adamw_torch", "adamw_torch_fused", "adamw_apex_fused", "adamw_anyprecision", + "adafactor". See `OptimizerNames` in [training_args.py](https://github.com/huggingface/transformers/blob/main/src/transformers/training_args.py) + for a full list of optimizers. + optim_args (`str`, *optional*): + Optional arguments that are supplied to optimizers such as AnyPrecisionAdamW, AdEMAMix, and GaLore. + group_by_length (`bool`, *optional*, defaults to `False`): + Whether or not to group together samples of roughly the same length in the training dataset (to minimize + padding applied and be more efficient). Only useful if applying dynamic padding. + length_column_name (`str`, *optional*, defaults to `"length"`): + Column name for precomputed lengths. If the column exists, grouping by length will use these values rather + than computing them on train startup. Ignored unless `group_by_length` is `True` and the dataset is an + instance of `Dataset`. + report_to (`str` or `List[str]`, *optional*, defaults to `"all"`): + The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`, + `"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`, `"neptune"`, + `"swanlab"`, `"tensorboard"`, and `"wandb"`. Use `"all"` to report to all integrations installed, `"none"` + for no integrations. + ddp_find_unused_parameters (`bool`, *optional*): + When using distributed training, the value of the flag `find_unused_parameters` passed to + `DistributedDataParallel`. Will default to `False` if gradient checkpointing is used, `True` otherwise. + ddp_bucket_cap_mb (`int`, *optional*): + When using distributed training, the value of the flag `bucket_cap_mb` passed to `DistributedDataParallel`. + ddp_broadcast_buffers (`bool`, *optional*): + When using distributed training, the value of the flag `broadcast_buffers` passed to + `DistributedDataParallel`. Will default to `False` if gradient checkpointing is used, `True` otherwise. + dataloader_pin_memory (`bool`, *optional*, defaults to `True`): + Whether you want to pin memory in data loaders or not. Will default to `True`. + dataloader_persistent_workers (`bool`, *optional*, defaults to `False`): + If True, the data loader will not shut down the worker processes after a dataset has been consumed once. + This allows to maintain the workers Dataset instances alive. Can potentially speed up training, but will + increase RAM usage. Will default to `False`. + dataloader_prefetch_factor (`int`, *optional*): + Number of batches loaded in advance by each worker. + 2 means there will be a total of 2 * num_workers batches prefetched across all workers. + skip_memory_metrics (`bool`, *optional*, defaults to `True`): + Whether to skip adding of memory profiler reports to metrics. This is skipped by default because it slows + down the training and evaluation speed. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether or not to push the model to the Hub every time the model is saved. If this is activated, + `output_dir` will begin a git directory synced with the repo (determined by `hub_model_id`) and the content + will be pushed each time a save is triggered (depending on your `save_strategy`). Calling + [`~Trainer.save_model`] will also trigger a push. + + + + If `output_dir` exists, it needs to be a local clone of the repository to which the [`Trainer`] will be + pushed. + + + + resume_from_checkpoint (`str`, *optional*): + The path to a folder with a valid checkpoint for your model. This argument is not directly used by + [`Trainer`], it's intended to be used by your training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + hub_model_id (`str`, *optional*): + The name of the repository to keep in sync with the local *output_dir*. It can be a simple model ID in + which case the model will be pushed in your namespace. Otherwise it should be the whole repository name, + for instance `"user_name/model"`, which allows you to push to an organization you are a member of with + `"organization_name/model"`. Will default to `user_name/output_dir_name` with *output_dir_name* being the + name of `output_dir`. + + Will default to the name of `output_dir`. + hub_strategy (`str` or [`~trainer_utils.HubStrategy`], *optional*, defaults to `"every_save"`): + Defines the scope of what is pushed to the Hub and when. Possible values are: + + - `"end"`: push the model, its configuration, the processing class e.g. tokenizer (if passed along to the [`Trainer`]) and a + draft of a model card when the [`~Trainer.save_model`] method is called. + - `"every_save"`: push the model, its configuration, the processing class e.g. tokenizer (if passed along to the [`Trainer`]) and + a draft of a model card each time there is a model save. The pushes are asynchronous to not block + training, and in case the save are very frequent, a new push is only attempted if the previous one is + finished. A last push is made with the final model at the end of training. + - `"checkpoint"`: like `"every_save"` but the latest checkpoint is also pushed in a subfolder named + last-checkpoint, allowing you to resume training easily with + `trainer.train(resume_from_checkpoint="last-checkpoint")`. + - `"all_checkpoints"`: like `"checkpoint"` but all checkpoints are pushed like they appear in the output + folder (so you will get one checkpoint folder per folder in your final repository) + + hub_token (`str`, *optional*): + The token to use to push the model to the Hub. Will default to the token in the cache folder obtained with + `huggingface-cli login`. + hub_private_repo (`bool`, *optional*): + Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. + hub_always_push (`bool`, *optional*, defaults to `False`): + Unless this is `True`, the `Trainer` will skip pushing a checkpoint when the previous push is not finished. + gradient_checkpointing (`bool`, *optional*, defaults to `False`): + If True, use gradient checkpointing to save memory at the expense of slower backward pass. + gradient_checkpointing_kwargs (`dict`, *optional*, defaults to `None`): + Key word arguments to be passed to the `gradient_checkpointing_enable` method. + include_inputs_for_metrics (`bool`, *optional*, defaults to `False`): + This argument is deprecated. Use `include_for_metrics` instead, e.g, `include_for_metrics = ["inputs"]`. + include_for_metrics (`List[str]`, *optional*, defaults to `[]`): + Include additional data in the `compute_metrics` function if needed for metrics computation. + Possible options to add to `include_for_metrics` list: + - `"inputs"`: Input data passed to the model, intended for calculating input dependent metrics. + - `"loss"`: Loss values computed during evaluation, intended for calculating loss dependent metrics. + eval_do_concat_batches (`bool`, *optional*, defaults to `True`): + Whether to recursively concat inputs/losses/labels/predictions across batches. If `False`, + will instead store them as lists, with each batch kept separate. + auto_find_batch_size (`bool`, *optional*, defaults to `False`) + Whether to find a batch size that will fit into memory automatically through exponential decay, avoiding + CUDA Out-of-Memory errors. Requires accelerate to be installed (`pip install accelerate`) + full_determinism (`bool`, *optional*, defaults to `False`) + If `True`, [`enable_full_determinism`] is called instead of [`set_seed`] to ensure reproducible results in + distributed training. Important: this will negatively impact the performance, so only use it for debugging. + torchdynamo (`str`, *optional*): + If set, the backend compiler for TorchDynamo. Possible choices are `"eager"`, `"aot_eager"`, `"inductor"`, + `"nvfuser"`, `"aot_nvfuser"`, `"aot_cudagraphs"`, `"ofi"`, `"fx2trt"`, `"onnxrt"` and `"ipex"`. + ray_scope (`str`, *optional*, defaults to `"last"`): + The scope to use when doing hyperparameter search with Ray. By default, `"last"` will be used. Ray will + then use the last checkpoint of all trials, compare those, and select the best one. However, other options + are also available. See the [Ray documentation]( + https://docs.ray.io/en/latest/tune/api_docs/analysis.html#ray.tune.ExperimentAnalysis.get_best_trial) for + more options. + ddp_timeout (`int`, *optional*, defaults to 1800): + The timeout for `torch.distributed.init_process_group` calls, used to avoid GPU socket timeouts when + performing slow operations in distributed runnings. Please refer the [PyTorch documentation] + (https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group) for more + information. + use_mps_device (`bool`, *optional*, defaults to `False`): + This argument is deprecated.`mps` device will be used if it is available similar to `cuda` device. + torch_compile (`bool`, *optional*, defaults to `False`): + Whether or not to compile the model using PyTorch 2.0 + [`torch.compile`](https://pytorch.org/get-started/pytorch-2.0/). + + This will use the best defaults for the [`torch.compile` + API](https://pytorch.org/docs/stable/generated/torch.compile.html?highlight=torch+compile#torch.compile). + You can customize the defaults with the argument `torch_compile_backend` and `torch_compile_mode` but we + don't guarantee any of them will work as the support is progressively rolled in in PyTorch. + + This flag and the whole compile API is experimental and subject to change in future releases. + torch_compile_backend (`str`, *optional*): + The backend to use in `torch.compile`. If set to any value, `torch_compile` will be set to `True`. + + Refer to the PyTorch doc for possible values and note that they may change across PyTorch versions. + + This flag is experimental and subject to change in future releases. + torch_compile_mode (`str`, *optional*): + The mode to use in `torch.compile`. If set to any value, `torch_compile` will be set to `True`. + + Refer to the PyTorch doc for possible values and note that they may change across PyTorch versions. + + This flag is experimental and subject to change in future releases. + include_tokens_per_second (`bool`, *optional*): + Whether or not to compute the number of tokens per second per device for training speed metrics. + + This will iterate over the entire training dataloader once beforehand, + + and will slow down the entire process. + + include_num_input_tokens_seen (`bool`, *optional*): + Whether or not to track the number of input tokens seen throughout training. + + May be slower in distributed training as gather operations must be called. + + neftune_noise_alpha (`Optional[float]`): + If not `None`, this will activate NEFTune noise embeddings. This can drastically improve model performance + for instruction fine-tuning. Check out the [original paper](https://arxiv.org/abs/2310.05914) and the + [original code](https://github.com/neelsjain/NEFTune). Support transformers `PreTrainedModel` and also + `PeftModel` from peft. The original paper used values in the range [5.0, 15.0]. + optim_target_modules (`Union[str, List[str]]`, *optional*): + The target modules to optimize, i.e. the module names that you would like to train. + Currently used for the GaLore algorithm (https://arxiv.org/abs/2403.03507) and APOLLO algorithm (https://arxiv.org/abs/2412.05270). + See GaLore implementation (https://github.com/jiaweizzhao/GaLore) and APOLLO implementation (https://github.com/zhuhanqing/APOLLO) for more details. + You need to make sure to pass a valid GaLore or APOLLO optimizer, e.g., one of: "apollo_adamw", "galore_adamw", "galore_adamw_8bit", "galore_adafactor" and make sure that the target modules are `nn.Linear` modules only. + + batch_eval_metrics (`Optional[bool]`, defaults to `False`): + If set to `True`, evaluation will call compute_metrics at the end of each batch to accumulate statistics + rather than saving all eval logits in memory. When set to `True`, you must pass a compute_metrics function + that takes a boolean argument `compute_result`, which when passed `True`, will trigger the final global + summary statistics from the batch-level summary statistics you've accumulated over the evaluation set. + + eval_on_start (`bool`, *optional*, defaults to `False`): + Whether to perform a evaluation step (sanity check) before the training to ensure the validation steps works correctly. + + eval_use_gather_object (`bool`, *optional*, defaults to `False`): + Whether to run recursively gather object in a nested list/tuple/dictionary of objects from all devices. This should only be enabled if users are not just returning tensors, and this is actively discouraged by PyTorch. + + use_liger_kernel (`bool`, *optional*, defaults to `False`): + Whether enable [Liger](https://github.com/linkedin/Liger-Kernel) Kernel for LLM model training. + It can effectively increase multi-GPU training throughput by ~20% and reduces memory usage by ~60%, works out of the box with + flash attention, PyTorch FSDP, and Microsoft DeepSpeed. Currently, it supports llama, mistral, mixtral and gemma models. + + average_tokens_across_devices (`bool`, *optional*, defaults to `False`): + Whether or not to average tokens across devices. If enabled, will use all_reduce to synchronize + num_tokens_in_batch for precise loss calculation. Reference: + https://github.com/huggingface/transformers/issues/34242 + """ + + # Sometimes users will pass in a `str` repr of a dict in the CLI + # We need to track what fields those can be. Each time a new arg + # has a dict type, it must be added to this list. + # Important: These should be typed with Optional[Union[dict,str,...]] + _VALID_DICT_FIELDS = [ + "accelerator_config", + "fsdp_config", + "deepspeed", + "gradient_checkpointing_kwargs", + "lr_scheduler_kwargs", + ] + + framework = "pt" + output_dir: Optional[str] = field( + default=None, + metadata={ + "help": "The output directory where the model predictions and checkpoints will be written. Defaults to 'trainer_output' if not provided." + }, + ) + overwrite_output_dir: bool = field( + default=False, + metadata={ + "help": ( + "Overwrite the content of the output directory. " + "Use this to continue training if output_dir points to a checkpoint directory." + ) + }, + ) + + do_train: bool = field(default=False, metadata={"help": "Whether to run training."}) + do_eval: bool = field(default=False, metadata={"help": "Whether to run eval on the dev set."}) + do_predict: bool = field(default=False, metadata={"help": "Whether to run predictions on the test set."}) + eval_strategy: Union[IntervalStrategy, str] = field( + default="no", + metadata={"help": "The evaluation strategy to use."}, + ) + prediction_loss_only: bool = field( + default=False, + metadata={"help": "When performing evaluation and predictions, only returns the loss."}, + ) + + per_device_train_batch_size: int = field( + default=8, metadata={"help": "Batch size per device accelerator core/CPU for training."} + ) + per_device_eval_batch_size: int = field( + default=8, metadata={"help": "Batch size per device accelerator core/CPU for evaluation."} + ) + + per_gpu_train_batch_size: Optional[int] = field( + default=None, + metadata={ + "help": ( + "Deprecated, the use of `--per_device_train_batch_size` is preferred. " + "Batch size per GPU/TPU core/CPU for training." + ) + }, + ) + per_gpu_eval_batch_size: Optional[int] = field( + default=None, + metadata={ + "help": ( + "Deprecated, the use of `--per_device_eval_batch_size` is preferred. " + "Batch size per GPU/TPU core/CPU for evaluation." + ) + }, + ) + + gradient_accumulation_steps: int = field( + default=1, + metadata={"help": "Number of updates steps to accumulate before performing a backward/update pass."}, + ) + eval_accumulation_steps: Optional[int] = field( + default=None, + metadata={"help": "Number of predictions steps to accumulate before moving the tensors to the CPU."}, + ) + + eval_delay: Optional[float] = field( + default=0, + metadata={ + "help": ( + "Number of epochs or steps to wait for before the first evaluation can be performed, depending on the" + " eval_strategy." + ) + }, + ) + + torch_empty_cache_steps: Optional[int] = field( + default=None, + metadata={ + "help": "Number of steps to wait before calling `torch..empty_cache()`." + "This can help avoid CUDA out-of-memory errors by lowering peak VRAM usage at a cost of about [10% slower performance](https://github.com/huggingface/transformers/issues/31372)." + "If left unset or set to None, cache will not be emptied." + }, + ) + + learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for AdamW."}) + weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."}) + adam_beta1: float = field(default=0.9, metadata={"help": "Beta1 for AdamW optimizer"}) + adam_beta2: float = field(default=0.999, metadata={"help": "Beta2 for AdamW optimizer"}) + adam_epsilon: float = field(default=1e-8, metadata={"help": "Epsilon for AdamW optimizer."}) + max_grad_norm: float = field(default=1.0, metadata={"help": "Max gradient norm."}) + + num_train_epochs: float = field(default=3.0, metadata={"help": "Total number of training epochs to perform."}) + max_steps: int = field( + default=-1, + metadata={"help": "If > 0: set total number of training steps to perform. Override num_train_epochs."}, + ) + lr_scheduler_type: Union[SchedulerType, str] = field( + default="linear", + metadata={"help": "The scheduler type to use."}, + ) + lr_scheduler_kwargs: Optional[Union[dict, str]] = field( + default_factory=dict, + metadata={ + "help": ( + "Extra parameters for the lr_scheduler such as {'num_cycles': 1} for the cosine with hard restarts." + ) + }, + ) + warmup_ratio: float = field( + default=0.0, metadata={"help": "Linear warmup over warmup_ratio fraction of total steps."} + ) + warmup_steps: int = field(default=0, metadata={"help": "Linear warmup over warmup_steps."}) + + log_level: Optional[str] = field( + default="passive", + metadata={ + "help": ( + "Logger log level to use on the main node. Possible choices are the log levels as strings: 'debug'," + " 'info', 'warning', 'error' and 'critical', plus a 'passive' level which doesn't set anything and" + " lets the application set the level. Defaults to 'passive'." + ), + "choices": trainer_log_levels.keys(), + }, + ) + log_level_replica: Optional[str] = field( + default="warning", + metadata={ + "help": "Logger log level to use on replica nodes. Same choices and defaults as ``log_level``", + "choices": trainer_log_levels.keys(), + }, + ) + log_on_each_node: bool = field( + default=True, + metadata={ + "help": ( + "When doing a multinode distributed training, whether to log once per node or just once on the main" + " node." + ) + }, + ) + logging_dir: Optional[str] = field(default=None, metadata={"help": "Tensorboard log dir."}) + logging_strategy: Union[IntervalStrategy, str] = field( + default="steps", + metadata={"help": "The logging strategy to use."}, + ) + logging_first_step: bool = field(default=False, metadata={"help": "Log the first global_step"}) + logging_steps: float = field( + default=500, + metadata={ + "help": ( + "Log every X updates steps. Should be an integer or a float in range `[0,1)`. " + "If smaller than 1, will be interpreted as ratio of total training steps." + ) + }, + ) + logging_nan_inf_filter: bool = field(default=True, metadata={"help": "Filter nan and inf losses for logging."}) + save_strategy: Union[SaveStrategy, str] = field( + default="steps", + metadata={"help": "The checkpoint save strategy to use."}, + ) + save_steps: float = field( + default=500, + metadata={ + "help": ( + "Save checkpoint every X updates steps. Should be an integer or a float in range `[0,1)`. " + "If smaller than 1, will be interpreted as ratio of total training steps." + ) + }, + ) + save_total_limit: Optional[int] = field( + default=None, + metadata={ + "help": ( + "If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in" + " `output_dir`. When `load_best_model_at_end` is enabled, the 'best' checkpoint according to" + " `metric_for_best_model` will always be retained in addition to the most recent ones. For example," + " for `save_total_limit=5` and `load_best_model_at_end=True`, the four last checkpoints will always be" + " retained alongside the best model. When `save_total_limit=1` and `load_best_model_at_end=True`," + " it is possible that two checkpoints are saved: the last one and the best one (if they are different)." + " Default is unlimited checkpoints" + ) + }, + ) + save_safetensors: Optional[bool] = field( + default=True, + metadata={ + "help": "Use safetensors saving and loading for state dicts instead of default torch.load and torch.save." + }, + ) + save_on_each_node: bool = field( + default=False, + metadata={ + "help": ( + "When doing multi-node distributed training, whether to save models and checkpoints on each node, or" + " only on the main one" + ) + }, + ) + save_only_model: bool = field( + default=False, + metadata={ + "help": ( + "When checkpointing, whether to only save the model, or also the optimizer, scheduler & rng state." + "Note that when this is true, you won't be able to resume training from checkpoint." + "This enables you to save storage by not storing the optimizer, scheduler & rng state." + "You can only load the model using from_pretrained with this option set to True." + ) + }, + ) + restore_callback_states_from_checkpoint: bool = field( + default=False, + metadata={ + "help": "Whether to restore the callback states from the checkpoint. If `True`, will override callbacks passed to the `Trainer` if they exist in the checkpoint." + }, + ) + no_cuda: bool = field( + default=False, + metadata={"help": "This argument is deprecated. It will be removed in version 5.0 of 🤗 Transformers."}, + ) + use_cpu: bool = field( + default=False, + metadata={ + "help": "Whether or not to use cpu. If left to False, we will use the available torch device/backend (cuda/mps/xpu/hpu etc.)" + }, + ) + use_mps_device: bool = field( + default=False, + metadata={ + "help": "This argument is deprecated. `mps` device will be used if available similar to `cuda` device." + " It will be removed in version 5.0 of 🤗 Transformers" + }, + ) + seed: int = field(default=42, metadata={"help": "Random seed that will be set at the beginning of training."}) + data_seed: Optional[int] = field(default=None, metadata={"help": "Random seed to be used with data samplers."}) + jit_mode_eval: bool = field( + default=False, metadata={"help": "Whether or not to use PyTorch jit trace for inference"} + ) + use_ipex: bool = field( + default=False, + metadata={ + "help": ( + "Use Intel extension for PyTorch when it is available, installation:" + " 'https://github.com/intel/intel-extension-for-pytorch'" + ) + }, + ) + bf16: bool = field( + default=False, + metadata={ + "help": ( + "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA" + " architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change." + ) + }, + ) + fp16: bool = field( + default=False, + metadata={"help": "Whether to use fp16 (mixed) precision instead of 32-bit"}, + ) + fp16_opt_level: str = field( + default="O1", + metadata={ + "help": ( + "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']. " + "See details at https://nvidia.github.io/apex/amp.html" + ) + }, + ) + half_precision_backend: str = field( + default="auto", + metadata={ + "help": "The backend to be used for half precision.", + "choices": ["auto", "apex", "cpu_amp"], + }, + ) + bf16_full_eval: bool = field( + default=False, + metadata={ + "help": ( + "Whether to use full bfloat16 evaluation instead of 32-bit. This is an experimental API and it may" + " change." + ) + }, + ) + fp16_full_eval: bool = field( + default=False, + metadata={"help": "Whether to use full float16 evaluation instead of 32-bit"}, + ) + tf32: Optional[bool] = field( + default=None, + metadata={ + "help": ( + "Whether to enable tf32 mode, available in Ampere and newer GPU architectures. This is an experimental" + " API and it may change." + ) + }, + ) + local_rank: int = field(default=-1, metadata={"help": "For distributed training: local_rank"}) + ddp_backend: Optional[str] = field( + default=None, + metadata={ + "help": "The backend to be used for distributed training", + "choices": ["nccl", "gloo", "mpi", "ccl", "hccl", "cncl", "mccl"], + }, + ) + tpu_num_cores: Optional[int] = field( + default=None, metadata={"help": "TPU: Number of TPU cores (automatically passed by launcher script)"} + ) + tpu_metrics_debug: bool = field( + default=False, + metadata={ + "help": ( + "Deprecated, the use of `--debug tpu_metrics_debug` is preferred. TPU: Whether to print debug metrics" + ) + }, + ) + debug: Union[str, list[DebugOption]] = field( + default="", + metadata={ + "help": ( + "Whether or not to enable debug mode. Current options: " + "`underflow_overflow` (Detect underflow and overflow in activations and weights), " + "`tpu_metrics_debug` (print debug metrics on TPU)." + ) + }, + ) + + dataloader_drop_last: bool = field( + default=False, metadata={"help": "Drop the last incomplete batch if it is not divisible by the batch size."} + ) + eval_steps: Optional[float] = field( + default=None, + metadata={ + "help": ( + "Run an evaluation every X steps. Should be an integer or a float in range `[0,1)`. " + "If smaller than 1, will be interpreted as ratio of total training steps." + ) + }, + ) + dataloader_num_workers: int = field( + default=0, + metadata={ + "help": ( + "Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded" + " in the main process." + ) + }, + ) + dataloader_prefetch_factor: Optional[int] = field( + default=None, + metadata={ + "help": ( + "Number of batches loaded in advance by each worker. " + "2 means there will be a total of 2 * num_workers batches prefetched across all workers. " + "Default is 2 for PyTorch < 2.0.0 and otherwise None." + ) + }, + ) + past_index: int = field( + default=-1, + metadata={"help": "If >=0, uses the corresponding part of the output as the past state for next step."}, + ) + + run_name: Optional[str] = field( + default=None, + metadata={ + "help": "An optional descriptor for the run. Notably used for wandb, mlflow comet and swanlab logging." + }, + ) + disable_tqdm: Optional[bool] = field( + default=None, metadata={"help": "Whether or not to disable the tqdm progress bars."} + ) + + remove_unused_columns: Optional[bool] = field( + default=True, metadata={"help": "Remove columns not required by the model when using an nlp.Dataset."} + ) + label_names: Optional[list[str]] = field( + default=None, metadata={"help": "The list of keys in your dictionary of inputs that correspond to the labels."} + ) + load_best_model_at_end: Optional[bool] = field( + default=False, + metadata={ + "help": ( + "Whether or not to load the best model found during training at the end of training. When this option" + " is enabled, the best checkpoint will always be saved. See `save_total_limit` for more." + ) + }, + ) + metric_for_best_model: Optional[str] = field( + default=None, metadata={"help": "The metric to use to compare two different models."} + ) + greater_is_better: Optional[bool] = field( + default=None, metadata={"help": "Whether the `metric_for_best_model` should be maximized or not."} + ) + ignore_data_skip: bool = field( + default=False, + metadata={ + "help": ( + "When resuming training, whether or not to skip the first epochs and batches to get to the same" + " training data." + ) + }, + ) + fsdp: Optional[Union[list[FSDPOption], str]] = field( + default="", + metadata={ + "help": ( + "Whether or not to use PyTorch Fully Sharded Data Parallel (FSDP) training (in distributed training" + " only). The base option should be `full_shard`, `shard_grad_op` or `no_shard` and you can add" + " CPU-offload to `full_shard` or `shard_grad_op` like this: full_shard offload` or `shard_grad_op" + " offload`. You can add auto-wrap to `full_shard` or `shard_grad_op` with the same syntax: full_shard" + " auto_wrap` or `shard_grad_op auto_wrap`." + ), + }, + ) + fsdp_min_num_params: int = field( + default=0, + metadata={ + "help": ( + "This parameter is deprecated. FSDP's minimum number of parameters for Default Auto Wrapping. (useful" + " only when `fsdp` field is passed)." + ) + }, + ) + fsdp_config: Optional[Union[dict, str]] = field( + default=None, + metadata={ + "help": ( + "Config to be used with FSDP (Pytorch Fully Sharded Data Parallel). The value is either a " + "fsdp json config file (e.g., `fsdp_config.json`) or an already loaded json file as `dict`." + ) + }, + ) + tp_size: Optional[int] = field( + default=0, + metadata={ + "help": ( + "Use tp_size to enable pytorch tensor parallelism." + "Tensor parallelism support is only available to models having `base_tp_plan` in their respective config classes." + "Set a value greater than 1 to activate TP." + "The same is used to prepare device mesh internally." + "Requires accelerate>1.3.0." + ) + }, + ) + fsdp_transformer_layer_cls_to_wrap: Optional[str] = field( + default=None, + metadata={ + "help": ( + "This parameter is deprecated. Transformer layer class name (case-sensitive) to wrap, e.g," + " `BertLayer`, `GPTJBlock`, `T5Block` .... (useful only when `fsdp` flag is passed)." + ) + }, + ) + accelerator_config: Optional[Union[dict, str]] = field( + default=None, + metadata={ + "help": ( + "Config to be used with the internal Accelerator object initialization. The value is either a " + "accelerator json config file (e.g., `accelerator_config.json`) or an already loaded json file as `dict`." + ) + }, + ) + deepspeed: Optional[Union[dict, str]] = field( + default=None, + metadata={ + "help": ( + "Enable deepspeed and pass the path to deepspeed json config file (e.g. `ds_config.json`) or an already" + " loaded json file as a dict" + ) + }, + ) + label_smoothing_factor: float = field( + default=0.0, metadata={"help": "The label smoothing epsilon to apply (zero means no label smoothing)."} + ) + + default_optim = "adamw_torch" + # XXX: enable when pytorch==2.0.1 comes out - we want to give it time to get all the bugs sorted out + # if is_torch_available() and version.parse(version.parse(torch.__version__).base_version) >= version.parse("2.1.0"): + # default_optim = "adamw_torch_fused" + # and update the doc above to: + # optim (`str` or [`training_args.OptimizerNames`], *optional*, defaults to `"adamw_torch_fused"` (for torch<2.1.0 `"adamw_torch"`): + optim: Union[OptimizerNames, str] = field( + default=default_optim, + metadata={"help": "The optimizer to use."}, + ) + optim_args: Optional[str] = field(default=None, metadata={"help": "Optional arguments to supply to optimizer."}) + adafactor: bool = field(default=False, metadata={"help": "Whether or not to replace AdamW by Adafactor."}) + group_by_length: bool = field( + default=False, + metadata={"help": "Whether or not to group samples of roughly the same length together when batching."}, + ) + length_column_name: Optional[str] = field( + default="length", + metadata={"help": "Column name with precomputed lengths to use when grouping by length."}, + ) + report_to: Union[None, str, list[str]] = field( + default=None, metadata={"help": "The list of integrations to report the results and logs to."} + ) + ddp_find_unused_parameters: Optional[bool] = field( + default=None, + metadata={ + "help": ( + "When using distributed training, the value of the flag `find_unused_parameters` passed to " + "`DistributedDataParallel`." + ) + }, + ) + ddp_bucket_cap_mb: Optional[int] = field( + default=None, + metadata={ + "help": ( + "When using distributed training, the value of the flag `bucket_cap_mb` passed to " + "`DistributedDataParallel`." + ) + }, + ) + ddp_broadcast_buffers: Optional[bool] = field( + default=None, + metadata={ + "help": ( + "When using distributed training, the value of the flag `broadcast_buffers` passed to " + "`DistributedDataParallel`." + ) + }, + ) + dataloader_pin_memory: bool = field( + default=True, metadata={"help": "Whether or not to pin memory for DataLoader."} + ) + dataloader_persistent_workers: bool = field( + default=False, + metadata={ + "help": "If True, the data loader will not shut down the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training, but will increase RAM usage." + }, + ) + skip_memory_metrics: bool = field( + default=True, metadata={"help": "Whether or not to skip adding of memory profiler reports to metrics."} + ) + use_legacy_prediction_loop: bool = field( + default=False, metadata={"help": "Whether or not to use the legacy prediction_loop in the Trainer."} + ) + push_to_hub: bool = field( + default=False, metadata={"help": "Whether or not to upload the trained model to the model hub after training."} + ) + resume_from_checkpoint: Optional[str] = field( + default=None, + metadata={"help": "The path to a folder with a valid checkpoint for your model."}, + ) + hub_model_id: Optional[str] = field( + default=None, metadata={"help": "The name of the repository to keep in sync with the local `output_dir`."} + ) + hub_strategy: Union[HubStrategy, str] = field( + default="every_save", + metadata={"help": "The hub strategy to use when `--push_to_hub` is activated."}, + ) + hub_token: Optional[str] = field(default=None, metadata={"help": "The token to use to push to the Model Hub."}) + hub_private_repo: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists." + }, + ) + hub_always_push: bool = field( + default=False, + metadata={"help": "Unless `True`, the Trainer will skip pushes if the previous one wasn't finished yet."}, + ) + gradient_checkpointing: bool = field( + default=False, + metadata={ + "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass." + }, + ) + gradient_checkpointing_kwargs: Optional[Union[dict, str]] = field( + default=None, + metadata={ + "help": "Gradient checkpointing key word arguments such as `use_reentrant`. Will be passed to `torch.utils.checkpoint.checkpoint` through `model.gradient_checkpointing_enable`." + }, + ) + include_inputs_for_metrics: bool = field( + default=False, + metadata={ + "help": "This argument is deprecated and will be removed in version 5 of 🤗 Transformers. Use `include_for_metrics` instead." + }, + ) + include_for_metrics: list[str] = field( + default_factory=list, + metadata={ + "help": "List of strings to specify additional data to include in the `compute_metrics` function." + "Options: 'inputs', 'loss'." + }, + ) + eval_do_concat_batches: bool = field( + default=True, + metadata={ + "help": "Whether to recursively concat inputs/losses/labels/predictions across batches. If `False`, will instead store them as lists, with each batch kept separate." + }, + ) + # Deprecated arguments + fp16_backend: str = field( + default="auto", + metadata={ + "help": "Deprecated. Use half_precision_backend instead", + "choices": ["auto", "apex", "cpu_amp"], + }, + ) + push_to_hub_model_id: Optional[str] = field( + default=None, metadata={"help": "The name of the repository to which push the `Trainer`."} + ) + push_to_hub_organization: Optional[str] = field( + default=None, metadata={"help": "The name of the organization in with to which push the `Trainer`."} + ) + push_to_hub_token: Optional[str] = field( + default=None, metadata={"help": "The token to use to push to the Model Hub."} + ) + _n_gpu: int = field(init=False, repr=False, default=-1) + mp_parameters: str = field( + default="", + metadata={"help": "Used by the SageMaker launcher to send mp-specific args. Ignored in Trainer"}, + ) + + auto_find_batch_size: bool = field( + default=False, + metadata={ + "help": ( + "Whether to automatically decrease the batch size in half and rerun the training loop again each time" + " a CUDA Out-of-Memory was reached" + ) + }, + ) + full_determinism: bool = field( + default=False, + metadata={ + "help": ( + "Whether to call enable_full_determinism instead of set_seed for reproducibility in distributed" + " training. Important: this will negatively impact the performance, so only use it for debugging." + ) + }, + ) + torchdynamo: Optional[str] = field( + default=None, + metadata={ + "help": "This argument is deprecated, use `--torch_compile_backend` instead.", + }, + ) + ray_scope: Optional[str] = field( + default="last", + metadata={ + "help": ( + 'The scope to use when doing hyperparameter search with Ray. By default, `"last"` will be used. Ray' + " will then use the last checkpoint of all trials, compare those, and select the best one. However," + " other options are also available. See the Ray documentation" + " (https://docs.ray.io/en/latest/tune/api_docs/analysis.html" + "#ray.tune.ExperimentAnalysis.get_best_trial)" + " for more options." + ) + }, + ) + ddp_timeout: Optional[int] = field( + default=1800, + metadata={ + "help": "Overrides the default timeout for distributed training (value should be given in seconds)." + }, + ) + torch_compile: bool = field( + default=False, metadata={"help": "If set to `True`, the model will be wrapped in `torch.compile`."} + ) + torch_compile_backend: Optional[str] = field( + default=None, + metadata={ + "help": "Which backend to use with `torch.compile`, passing one will trigger a model compilation.", + }, + ) + torch_compile_mode: Optional[str] = field( + default=None, + metadata={ + "help": "Which mode to use with `torch.compile`, passing one will trigger a model compilation.", + }, + ) + + include_tokens_per_second: Optional[bool] = field( + default=False, + metadata={"help": "If set to `True`, the speed metrics will include `tgs` (tokens per second per device)."}, + ) + + include_num_input_tokens_seen: Optional[bool] = field( + default=False, + metadata={ + "help": "If set to `True`, will track the number of input tokens seen throughout training. (May be slower in distributed training)" + }, + ) + + neftune_noise_alpha: Optional[float] = field( + default=None, + metadata={ + "help": "Activates neftune noise embeddings into the model. NEFTune has been proven to drastically improve model performances for instruction fine-tuning. Check out the original paper here: https://arxiv.org/abs/2310.05914 and the original code here: https://github.com/neelsjain/NEFTune. Only supported for `PreTrainedModel` and `PeftModel` classes." + }, + ) + + optim_target_modules: Union[None, str, list[str]] = field( + default=None, + metadata={ + "help": "Target modules for the optimizer defined in the `optim` argument. Only used for the GaLore optimizer at the moment." + }, + ) + + batch_eval_metrics: bool = field( + default=False, + metadata={"help": "Break eval metrics calculation into batches to save memory."}, + ) + + eval_on_start: bool = field( + default=False, + metadata={ + "help": "Whether to run through the entire `evaluation` step at the very beginning of training as a sanity check." + }, + ) + + use_liger_kernel: Optional[bool] = field( + default=False, + metadata={"help": "Whether or not to enable the Liger Kernel for model training."}, + ) + + eval_use_gather_object: Optional[bool] = field( + default=False, + metadata={ + "help": "Whether to run recursively gather object in a nested list/tuple/dictionary of objects from all devices." + }, + ) + + average_tokens_across_devices: Optional[bool] = field( + default=False, + metadata={ + "help": "Whether or not to average tokens across devices. If enabled, will use all_reduce to " + "synchronize num_tokens_in_batch for precise loss calculation. Reference: " + "https://github.com/huggingface/transformers/issues/34242" + }, + ) + + def __post_init__(self): + # Set default output_dir if not provided + if self.output_dir is None: + self.output_dir = "trainer_output" + logger.info( + "No output directory specified, defaulting to 'trainer_output'. " + "To change this behavior, specify --output_dir when creating TrainingArguments." + ) + + # Parse in args that could be `dict` sent in from the CLI as a string + for field in self._VALID_DICT_FIELDS: + passed_value = getattr(self, field) + # We only want to do this if the str starts with a bracket to indicate a `dict` + # else its likely a filename if supported + if isinstance(passed_value, str) and passed_value.startswith("{"): + loaded_dict = json.loads(passed_value) + # Convert str values to types if applicable + loaded_dict = _convert_str_dict(loaded_dict) + setattr(self, field, loaded_dict) + + # expand paths, if not os.makedirs("~/bar") will make directory + # in the current directory instead of the actual home + # see https://github.com/huggingface/transformers/issues/10628 + if self.output_dir is not None: + self.output_dir = os.path.expanduser(self.output_dir) + if self.logging_dir is None and self.output_dir is not None: + self.logging_dir = os.path.join(self.output_dir, default_logdir()) + if self.logging_dir is not None: + self.logging_dir = os.path.expanduser(self.logging_dir) + + if self.disable_tqdm is None: + self.disable_tqdm = logger.getEffectiveLevel() > logging.WARN + + if isinstance(self.eval_strategy, EvaluationStrategy): + warnings.warn( + "using `EvaluationStrategy` for `eval_strategy` is deprecated and will be removed in version 5" + " of 🤗 Transformers. Use `IntervalStrategy` instead", + FutureWarning, + ) + # Go back to the underlying string or we won't be able to instantiate `IntervalStrategy` on it. + self.eval_strategy = self.eval_strategy.value + if self.no_cuda: + warnings.warn( + "using `no_cuda` is deprecated and will be removed in version 5.0 of 🤗 Transformers. " + "Use `use_cpu` instead", + FutureWarning, + ) + self.use_cpu = self.no_cuda + + self.eval_strategy = IntervalStrategy(self.eval_strategy) + self.logging_strategy = IntervalStrategy(self.logging_strategy) + self.save_strategy = SaveStrategy(self.save_strategy) + self.hub_strategy = HubStrategy(self.hub_strategy) + + self.lr_scheduler_type = SchedulerType(self.lr_scheduler_type) + if self.do_eval is False and self.eval_strategy != IntervalStrategy.NO: + self.do_eval = True + + if self.torch_empty_cache_steps is not None: + if not (isinstance(self.torch_empty_cache_steps, int) and self.torch_empty_cache_steps > 0): + raise ValueError( + f"`torch_empty_cache_steps` must be an integer bigger than 0, got {self.torch_empty_cache_steps}." + ) + + # eval_steps has to be defined and non-zero, fallbacks to logging_steps if the latter is non-zero + if self.eval_strategy == IntervalStrategy.STEPS and (self.eval_steps is None or self.eval_steps == 0): + if self.logging_steps > 0: + logger.info(f"using `logging_steps` to initialize `eval_steps` to {self.logging_steps}") + self.eval_steps = self.logging_steps + else: + raise ValueError( + f"evaluation strategy {self.eval_strategy} requires either non-zero --eval_steps or" + " --logging_steps" + ) + + # logging_steps must be non-zero for logging_strategy that is other than 'no' + if self.logging_strategy == IntervalStrategy.STEPS and self.logging_steps == 0: + raise ValueError(f"logging strategy {self.logging_strategy} requires non-zero --logging_steps") + + if self.logging_strategy == IntervalStrategy.STEPS and self.logging_steps > 1: + if self.logging_steps != int(self.logging_steps): + raise ValueError(f"--logging_steps must be an integer if bigger than 1: {self.logging_steps}") + self.logging_steps = int(self.logging_steps) + if self.eval_strategy == IntervalStrategy.STEPS and self.eval_steps > 1: + if self.eval_steps != int(self.eval_steps): + raise ValueError(f"--eval_steps must be an integer if bigger than 1: {self.eval_steps}") + self.eval_steps = int(self.eval_steps) + if self.save_strategy == SaveStrategy.STEPS and self.save_steps > 1: + if self.save_steps != int(self.save_steps): + raise ValueError(f"--save_steps must be an integer if bigger than 1: {self.save_steps}") + self.save_steps = int(self.save_steps) + + # Sanity checks for load_best_model_at_end: we require save and eval strategies to be compatible. + if self.load_best_model_at_end and self.save_strategy != SaveStrategy.BEST: + if self.eval_strategy != self.save_strategy: + raise ValueError( + "--load_best_model_at_end requires the save and eval strategy to match, but found\n- Evaluation " + f"strategy: {self.eval_strategy}\n- Save strategy: {self.save_strategy}" + ) + if self.eval_strategy == IntervalStrategy.STEPS and self.save_steps % self.eval_steps != 0: + if self.eval_steps < 1 or self.save_steps < 1: + if not (self.eval_steps < 1 and self.save_steps < 1): + raise ValueError( + "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation " + "steps, which cannot get guaranteed when mixing ratio and absolute steps for save_steps " + f"{self.save_steps} and eval_steps {self.eval_steps}." + ) + # Work around floating point precision issues + LARGE_MULTIPLIER = 1_000_000 + if (self.save_steps * LARGE_MULTIPLIER) % (self.eval_steps * LARGE_MULTIPLIER) != 0: + raise ValueError( + "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation " + f"steps, but found {self.save_steps}, which is not a multiple of {self.eval_steps}." + ) + raise ValueError( + "--load_best_model_at_end requires the saving steps to be a round multiple of the evaluation " + f"steps, but found {self.save_steps}, which is not a round multiple of {self.eval_steps}." + ) + + safetensors_available = is_safetensors_available() + if self.save_safetensors and not safetensors_available: + raise ValueError(f"--save_safetensors={self.save_safetensors} requires safetensors to be installed!") + if not self.save_safetensors and safetensors_available: + logger.info( + f"Found safetensors installation, but --save_safetensors={self.save_safetensors}. " + f"Safetensors should be a preferred weights saving format due to security and performance reasons. " + f"If your model cannot be saved by safetensors please feel free to open an issue at " + f"https://github.com/huggingface/safetensors!" + ) + + if ( + self.load_best_model_at_end or self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU + ) and self.metric_for_best_model is None: + self.metric_for_best_model = "loss" + if self.greater_is_better is None and self.metric_for_best_model is not None: + self.greater_is_better = not (self.metric_for_best_model.endswith("loss")) + if self.run_name is None: + self.run_name = self.output_dir + if self.framework == "pt" and is_torch_available(): + if self.fp16_backend and self.fp16_backend != "auto": + warnings.warn( + "`fp16_backend` is deprecated and will be removed in version 5 of 🤗 Transformers. Use" + " `half_precision_backend` instead", + FutureWarning, + ) + self.half_precision_backend = self.fp16_backend + + if self.bf16 or self.bf16_full_eval: + if self.use_cpu and not is_torch_bf16_cpu_available() and not is_torch_xla_available(): + # cpu + raise ValueError("Your setup doesn't support bf16/(cpu, tpu, neuroncore). You need torch>=1.10") + elif not self.use_cpu: + if torch.cuda.is_available() and not is_torch_bf16_gpu_available(): + # gpu + raise ValueError( + "Your setup doesn't support bf16/gpu. You need torch>=1.10, using Ampere GPU with cuda>=11.0" + ) + + if self.fp16 and self.bf16: + raise ValueError("At most one of fp16 and bf16 can be True, but not both") + + if self.fp16_full_eval and self.bf16_full_eval: + raise ValueError("At most one of fp16 and bf16 can be True for full eval, but not both") + + if self.bf16: + if self.half_precision_backend == "apex": + raise ValueError(" `--half_precision_backend apex`: GPU bf16 is not supported by apex.") + + if self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU: + if self.eval_strategy == IntervalStrategy.NO: + raise ValueError("lr_scheduler_type reduce_lr_on_plateau requires an eval strategy") + if not is_torch_available(): + raise ValueError("lr_scheduler_type reduce_lr_on_plateau requires torch>=0.2.0") + + self.optim = OptimizerNames(self.optim) + if self.adafactor: + warnings.warn( + "`--adafactor` is deprecated and will be removed in version 5 of 🤗 Transformers. Use `--optim" + " adafactor` instead", + FutureWarning, + ) + self.optim = OptimizerNames.ADAFACTOR + if self.optim == OptimizerNames.ADAMW_TORCH_FUSED and is_torch_available(): + if version.parse(version.parse(torch.__version__).base_version) < version.parse("2.0.0"): + raise ValueError("--optim adamw_torch_fused requires PyTorch 2.0 or higher") + # there is a bug in fp16/AMP in pt-2.0.0 + if version.parse(version.parse(torch.__version__).base_version) == version.parse("2.0.0") and self.fp16: + raise ValueError("--optim adamw_torch_fused with --fp16 requires PyTorch>2.0") + + # We need to setup the accelerator config here *before* the first call to `self.device` + if is_accelerate_available(): + if not isinstance(self.accelerator_config, AcceleratorConfig): + if self.accelerator_config is None: + self.accelerator_config = AcceleratorConfig() + elif isinstance(self.accelerator_config, dict): + self.accelerator_config = AcceleratorConfig(**self.accelerator_config) + # Check that a user didn't pass in the class instantiator + # such as `accelerator_config = AcceleratorConfig` + elif isinstance(self.accelerator_config, type): + raise NotImplementedError( + "Tried passing in a callable to `accelerator_config`, but this is not supported. " + "Please pass in a fully constructed `AcceleratorConfig` object instead." + ) + else: + self.accelerator_config = AcceleratorConfig.from_json_file(self.accelerator_config) + + # Initialize device before we proceed + if self.framework == "pt" and is_torch_available(): + self.device + + # Disable average tokens when using single device + if self.average_tokens_across_devices: + try: + if self.world_size == 1: + logger.warning( + "average_tokens_across_devices is set to True but it is invalid when world size is" + "1. Turn it to False automatically." + ) + self.average_tokens_across_devices = False + except ImportError as e: + logger.warning(f"Can not specify world size due to {e}. Turn average_tokens_across_devices to False.") + self.average_tokens_across_devices = False + + if self.torchdynamo is not None: + warnings.warn( + "`torchdynamo` is deprecated and will be removed in version 5 of 🤗 Transformers. Use" + " `torch_compile_backend` instead", + FutureWarning, + ) + self.torch_compile_backend = self.torchdynamo + if (self.torch_compile_mode is not None or self.torch_compile_backend is not None) and not self.torch_compile: + self.torch_compile = True + if self.torch_compile and self.torch_compile_backend is None: + if not self.use_cpu and is_torch_hpu_available(): + self.torch_compile_backend = "hpu_backend" + else: + self.torch_compile_backend = "inductor" + + # accelerate integration for torch compile + if self.torch_compile: + # set env vars for accelerate + prefix = "ACCELERATE_DYNAMO_" + os.environ[prefix + "BACKEND"] = self.torch_compile_backend + if self.torch_compile_mode is not None: + os.environ[prefix + "MODE"] = self.torch_compile_mode + + if self.framework == "pt" and is_torch_available() and self.torch_compile: + if is_torch_tf32_available(): + if self.tf32 is None and not self.fp16 or self.bf16: + logger.info( + "Setting TF32 in CUDA backends to speedup torch compile, you won't see any improvement" + " otherwise." + ) + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + else: + logger.warning( + "The speedups for torchdynamo mostly come with GPU Ampere or higher and which is not detected here." + ) + if self.framework == "pt" and is_torch_available() and self.tf32 is not None: + if self.tf32: + if is_torch_tf32_available(): + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + else: + raise ValueError("--tf32 requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7") + else: + if is_torch_tf32_available(): + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + # no need to assert on else + + # if training args is specified, it will override the one specified in the accelerate config + if self.half_precision_backend != "apex": + mixed_precision_dtype = os.environ.get("ACCELERATE_MIXED_PRECISION", "no") + if self.fp16: + mixed_precision_dtype = "fp16" + elif self.bf16: + mixed_precision_dtype = "bf16" + os.environ["ACCELERATE_MIXED_PRECISION"] = mixed_precision_dtype + + if self.report_to is None: + logger.info( + "The default value for the training argument `--report_to` will change in v5 (from all installed " + "integrations to none). In v5, you will need to use `--report_to all` to get the same behavior as " + "now. You should start updating your code and make this info disappear :-)." + ) + self.report_to = "all" + if self.report_to == "all" or self.report_to == ["all"]: + # Import at runtime to avoid a circular import. + from .integrations import get_available_reporting_integrations + + self.report_to = get_available_reporting_integrations() + + if "codecarbon" in self.report_to and torch.version.hip: + logger.warning( + "When using the Trainer, CodeCarbonCallback requires the `codecarbon` package, which is not compatible with AMD ROCm (https://github.com/mlco2/codecarbon/pull/490). Automatically disabling the codecarbon callback. Reference: https://huggingface.co/docs/transformers/v4.39.3/en/main_classes/trainer#transformers.TrainingArguments.report_to." + ) + self.report_to.remove("codecarbon") + + elif self.report_to == "none" or self.report_to == ["none"]: + self.report_to = [] + elif not isinstance(self.report_to, list): + self.report_to = [self.report_to] + + if self.warmup_ratio < 0 or self.warmup_ratio > 1: + raise ValueError("warmup_ratio must lie in range [0,1]") + elif self.warmup_ratio > 0 and self.warmup_steps > 0: + logger.info( + "Both warmup_ratio and warmup_steps given, warmup_steps will override any effect of warmup_ratio" + " during training" + ) + + if not isinstance(self.warmup_steps, int) or self.warmup_steps < 0: + raise ValueError("warmup_steps must be of type int and must be 0 or a positive integer.") + + if isinstance(self.fsdp, bool): + self.fsdp = [FSDPOption.FULL_SHARD] if self.fsdp else "" + if isinstance(self.fsdp, str): + self.fsdp = [FSDPOption(s) for s in self.fsdp.split()] + if self.fsdp == [FSDPOption.OFFLOAD]: + raise ValueError( + "`--fsdp offload` can't work on its own. It needs to be added to `--fsdp full_shard` or " + '`--fsdp shard_grad_op`. For example, `--fsdp "full_shard offload"`.' + ) + elif FSDPOption.FULL_SHARD in self.fsdp and FSDPOption.SHARD_GRAD_OP in self.fsdp: + raise ValueError("`--fsdp full_shard` is not compatible with `--fsdp shard_grad_op`.") + + if self.gradient_checkpointing and ( + FSDPOption.FULL_SHARD in self.fsdp or FSDPOption.HYBRID_SHARD in self.fsdp + ): + logger.warning( + "When using FSDP full shard, instead of using `gradient_checkpointing` in TrainingArguments, please" + " use `activation_checkpointing` in `fsdp_config`. The former introduces a redundant AllGather" + " operation in backward pass. Reference: https://github.com/huggingface/transformers/issues/30404" + ) + + if self.fsdp_config is None: + self.fsdp_config = {} + + if isinstance(self.fsdp_config, str): + if len(self.fsdp) == 0: + warnings.warn("`--fsdp_config` is useful only when `--fsdp` is specified.") + with open(self.fsdp_config, encoding="utf-8") as f: + self.fsdp_config = json.load(f) + for k in list(self.fsdp_config.keys()): + if k.startswith("fsdp_"): + v = self.fsdp_config.pop(k) + self.fsdp_config[k[5:]] = v + + if self.fsdp_min_num_params > 0: + warnings.warn("using `--fsdp_min_num_params` is deprecated. Use fsdp_config instead ", FutureWarning) + + self.fsdp_config["min_num_params"] = max(self.fsdp_config.get("min_num_params", 0), self.fsdp_min_num_params) + + # if fsdp_config["transformer_layer_cls_to_wrap"] is specified as a string, convert it to a list with a single object + if isinstance(self.fsdp_config.get("transformer_layer_cls_to_wrap", None), str): + self.fsdp_config["transformer_layer_cls_to_wrap"] = [self.fsdp_config["transformer_layer_cls_to_wrap"]] + + if self.fsdp_transformer_layer_cls_to_wrap is not None: + warnings.warn( + "using `--fsdp_transformer_layer_cls_to_wrap` is deprecated. Use fsdp_config instead ", FutureWarning + ) + self.fsdp_config["transformer_layer_cls_to_wrap"] = self.fsdp_config.get( + "transformer_layer_cls_to_wrap", [] + ) + [self.fsdp_transformer_layer_cls_to_wrap] + + if len(self.fsdp) == 0 and self.fsdp_config["min_num_params"] > 0: + warnings.warn("`min_num_params` is useful only when `--fsdp` is specified.") + + if len(self.fsdp) == 0 and self.fsdp_config.get("transformer_layer_cls_to_wrap", None) is not None: + warnings.warn("`transformer_layer_cls_to_wrap` is useful only when `--fsdp` is specified.") + + if ( + len(self.fsdp) > 0 + and self.fsdp_config["min_num_params"] > 0 + and self.fsdp_config.get("transformer_layer_cls_to_wrap", None) is not None + ): + raise ValueError("`min_num_params` and `transformer_layer_cls_to_wrap` are mutually exclusive.") + self.fsdp_config["xla"] = self.fsdp_config.get("xla", False) + self.fsdp_config["xla_fsdp_v2"] = self.fsdp_config.get("xla_fsdp_v2", False) + self.fsdp_config["xla_fsdp_grad_ckpt"] = self.fsdp_config.get("xla_fsdp_grad_ckpt", False) + if self.fsdp_config["xla"]: + if len(self.fsdp) > 0: + # store XLA fsdp configuration parameters into a dictionary + # Copy the config to avoid modifying the original config (which may be used for JSON serialization) + self.xla_fsdp_config = self.fsdp_config.get("xla_fsdp_settings", {}).copy() + # apply appropriate string to torch.dtype conversions for parameters + if "compute_dtype" in self.xla_fsdp_config: + self.xla_fsdp_config["compute_dtype"] = getattr(torch, self.xla_fsdp_config["compute_dtype"]) + if "buffer_dtype" in self.xla_fsdp_config: + self.xla_fsdp_config["buffer_dtype"] = getattr(torch, self.xla_fsdp_config["buffer_dtype"]) + else: + warnings.warn("XLA FSDP can be used only when `--fsdp` is specified.") + else: + if self.fsdp_config["xla_fsdp_grad_ckpt"]: + warnings.warn("`--xla_fsdp_grad_ckpt` is useful only when `--xla` is set to true.") + + if self.tp_size > 1: + if not is_accelerate_available("1.3.1"): + raise NotImplementedError( + "TP using PyTorch requires Accelerate version `accelerate` >= 1.3.1. " + "This is not supported and we recommend you to update your version." + ) + os.environ["ACCELERATE_USE_TP"] = "true" + os.environ["TP_SIZE"] = str(self.tp_size) + # accelerate integration for FSDP + if len(self.fsdp) > 0 and not self.fsdp_config["xla"]: + os.environ["ACCELERATE_USE_FSDP"] = "true" + from accelerate.utils.constants import ( + FSDP_AUTO_WRAP_POLICY, + FSDP_SHARDING_STRATEGY, + ) + + prefix = "FSDP_" + for fsdp_option in self.fsdp: + if fsdp_option.upper() in FSDP_SHARDING_STRATEGY: + # set environment variable for FSDP sharding strategy + os.environ[f"{prefix}SHARDING_STRATEGY"] = str( + FSDP_SHARDING_STRATEGY.index(fsdp_option.upper()) + 1 + ) + elif fsdp_option == FSDPOption.OFFLOAD: + os.environ[f"{prefix}OFFLOAD_PARAMS"] = "true" + elif fsdp_option == FSDPOption.AUTO_WRAP: + os.environ[f"{prefix}AUTO_WRAP_POLICY"] = FSDP_AUTO_WRAP_POLICY[0] + if self.fsdp_config["min_num_params"] > 0: + os.environ[f"{prefix}MIN_NUM_PARAMS"] = str(self.fsdp_config["min_num_params"]) + os.environ[f"{prefix}AUTO_WRAP_POLICY"] = FSDP_AUTO_WRAP_POLICY[1] + elif self.fsdp_config.get("transformer_layer_cls_to_wrap", None) is not None: + os.environ[f"{prefix}TRANSFORMER_CLS_TO_WRAP"] = ",".join( + self.fsdp_config["transformer_layer_cls_to_wrap"] + ) + prefetch_policy = self.fsdp_config.get("backward_prefetch", "NO_PREFETCH") + os.environ[f"{prefix}BACKWARD_PREFETCH"] = prefetch_policy.upper() + os.environ[f"{prefix}FORWARD_PREFETCH"] = str(self.fsdp_config.get("forward_prefetch", "false")).lower() + + sync_module_states = str(self.fsdp_config.get("sync_module_states", "true")).lower() + cpu_ram_efficient_loading = str(self.fsdp_config.get("cpu_ram_efficient_loading", "false")).lower() + + if sync_module_states == "false" and cpu_ram_efficient_loading == "true": + # In this case, all the processes except the main process would have random weights leading + # to unexpected behaviour during training, thus throwing error here to prevent it. + raise ValueError('`sync_module_states` must be `"True"` if `cpu_ram_efficient_loading` is `"True"`') + + os.environ[f"{prefix}SYNC_MODULE_STATES"] = sync_module_states + os.environ[f"{prefix}CPU_RAM_EFFICIENT_LOADING"] = cpu_ram_efficient_loading + + os.environ[f"{prefix}USE_ORIG_PARAMS"] = str(self.fsdp_config.get("use_orig_params", "true")).lower() + + if self.tpu_metrics_debug: + warnings.warn( + "using `--tpu_metrics_debug` is deprecated and will be removed in version 5 of 🤗 Transformers. Use" + " `--debug tpu_metrics_debug` instead", + FutureWarning, + ) + if self.debug is None: + self.debug = " tpu_metrics_debug" + else: + self.debug += " tpu_metrics_debug" + self.tpu_metrics_debug = False + + if isinstance(self.debug, str): + self.debug = [DebugOption(s) for s in self.debug.split()] + elif self.debug is None: + self.debug = [] + + self.deepspeed_plugin = None + if self.deepspeed: + # - must be run very last in arg parsing, since it will use a lot of these settings. + # - must be run before the model is created. + if not is_accelerate_available(): + raise ValueError( + f"--deepspeed requires Accelerate to be installed: `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`." + ) + from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig + + # will be used later by the Trainer + # note: leave self.deepspeed unmodified in case a user relies on it not to be modified) + self.hf_deepspeed_config = HfTrainerDeepSpeedConfig(self.deepspeed) + self.hf_deepspeed_config.trainer_config_process(self) + + # Accelerate DeepSpeed Plugin + from accelerate.utils import DeepSpeedPlugin + + os.environ["ACCELERATE_USE_DEEPSPEED"] = "true" + self.deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.hf_deepspeed_config) + elif strtobool(os.environ.get("ACCELERATE_USE_DEEPSPEED", "false")): + # Accelerate DeepSpeed Plugin + from accelerate.utils import DeepSpeedPlugin + + self.deepspeed_plugin = DeepSpeedPlugin() + mixed_precision = os.environ.get("ACCELERATE_MIXED_PRECISION", "no") + self.deepspeed_plugin.set_mixed_precision(mixed_precision) + self.deepspeed_plugin.set_deepspeed_weakref() + + if self.use_cpu: + self.dataloader_pin_memory = False + + if self.dataloader_num_workers == 0 and self.dataloader_prefetch_factor is not None: + raise ValueError( + "--dataloader_prefetch_factor can only be set when data is loaded in a different process, i.e." + " when --dataloader_num_workers > 1." + ) + + if self.push_to_hub_token is not None: + warnings.warn( + "`--push_to_hub_token` is deprecated and will be removed in version 5 of 🤗 Transformers. Use " + "`--hub_token` instead.", + FutureWarning, + ) + self.hub_token = self.push_to_hub_token + + if self.push_to_hub_model_id is not None: + self.hub_model_id = get_full_repo_name( + self.push_to_hub_model_id, organization=self.push_to_hub_organization, token=self.hub_token + ) + if self.push_to_hub_organization is not None: + warnings.warn( + "`--push_to_hub_model_id` and `--push_to_hub_organization` are deprecated and will be removed in " + "version 5 of 🤗 Transformers. Use `--hub_model_id` instead and pass the full repo name to this " + f"argument (in this case {self.hub_model_id}).", + FutureWarning, + ) + else: + warnings.warn( + "`--push_to_hub_model_id` is deprecated and will be removed in version 5 of 🤗 Transformers. Use " + "`--hub_model_id` instead and pass the full repo name to this argument (in this case " + f"{self.hub_model_id}).", + FutureWarning, + ) + elif self.push_to_hub_organization is not None: + self.hub_model_id = f"{self.push_to_hub_organization}/{Path(self.output_dir).name}" + warnings.warn( + "`--push_to_hub_organization` is deprecated and will be removed in version 5 of 🤗 Transformers. Use " + "`--hub_model_id` instead and pass the full repo name to this argument (in this case " + f"{self.hub_model_id}).", + FutureWarning, + ) + + if self.eval_use_gather_object and not is_accelerate_available("0.30.0"): + raise ValueError( + "--eval_use_gather_object requires Accelerate to be version of `accelerate` > 0.30.0." + "This is not supported and we recommend you to update your version." + ) + + if self.data_seed is not None: + if not is_accelerate_available("1.1.0"): + raise NotImplementedError( + "data_seed requires Accelerate version `accelerate` >= 1.1.0. " + "This is not supported and we recommend you to update your version." + ) + + if self.include_inputs_for_metrics: + logger.warning( + "Using `include_inputs_for_metrics` is deprecated and will be removed in version 5 of 🤗 Transformers. Please use `include_for_metrics` list argument instead." + ) + self.include_for_metrics.append("inputs") + + def __str__(self): + self_as_dict = asdict(self) + + # Remove deprecated arguments. That code should be removed once + # those deprecated arguments are removed from TrainingArguments. (TODO: v5) + del self_as_dict["per_gpu_train_batch_size"] + del self_as_dict["per_gpu_eval_batch_size"] + + self_as_dict = {k: f"<{k.upper()}>" if k.endswith("_token") else v for k, v in self_as_dict.items()} + + attrs_as_str = [f"{k}={v},\n" for k, v in sorted(self_as_dict.items())] + return f"{self.__class__.__name__}(\n{''.join(attrs_as_str)})" + + __repr__ = __str__ + + @property + def train_batch_size(self) -> int: + """ + The actual batch size for training (may differ from `per_gpu_train_batch_size` in distributed training). + """ + if self.per_gpu_train_batch_size: + logger.warning( + "Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future " + "version. Using `--per_device_train_batch_size` is preferred." + ) + per_device_batch_size = self.per_gpu_train_batch_size or self.per_device_train_batch_size + train_batch_size = per_device_batch_size * max(1, self.n_gpu) + return train_batch_size + + @property + def eval_batch_size(self) -> int: + """ + The actual batch size for evaluation (may differ from `per_gpu_eval_batch_size` in distributed training). + """ + if self.per_gpu_eval_batch_size: + logger.warning( + "Using deprecated `--per_gpu_eval_batch_size` argument which will be removed in a future " + "version. Using `--per_device_eval_batch_size` is preferred." + ) + per_device_batch_size = self.per_gpu_eval_batch_size or self.per_device_eval_batch_size + eval_batch_size = per_device_batch_size * max(1, self.n_gpu) + return eval_batch_size + + @property + def ddp_timeout_delta(self) -> timedelta: + """ + The actual timeout for torch.distributed.init_process_group since it expects a timedelta variable. + """ + return timedelta(seconds=self.ddp_timeout) + + @cached_property + def _setup_devices(self) -> "torch.device": + requires_backends(self, ["torch"]) + logger.info("PyTorch: setting up devices") + if not is_sagemaker_mp_enabled(): + if not is_accelerate_available(): + raise ImportError( + f"Using the `Trainer` with `PyTorch` requires `accelerate>={ACCELERATE_MIN_VERSION}`: " + f"Please run `pip install transformers[torch]` or `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`" + ) + # We delay the init of `PartialState` to the end for clarity + accelerator_state_kwargs = {"enabled": True, "use_configured_state": False} + if isinstance(self.accelerator_config, AcceleratorConfig): + accelerator_state_kwargs["use_configured_state"] = self.accelerator_config.pop( + "use_configured_state", False + ) + if accelerator_state_kwargs["use_configured_state"]: + if PartialState._shared_state == {}: + raise ValueError( + "Passing `'use_configured_state':True` to the AcceleratorConfig requires a pre-configured " + "`AcceleratorState` or `PartialState` to be defined before calling `TrainingArguments`. " + ) + # We rely on `PartialState` to yell if there's issues here (which it will) + self.distributed_state = PartialState(cpu=self.use_cpu) + if self.deepspeed and self.distributed_state.distributed_type != DistributedType.DEEPSPEED: + raise RuntimeError( + "Tried to use an already configured `Accelerator` or `PartialState` that was not initialized for DeepSpeed, " + "but also passed in a `deepspeed` configuration to the `TrainingArguments`. Please set " + "`use_configured_state:False` instead or setup your `Accelerator` or `PartialState` properly." + ) + else: + AcceleratorState._reset_state(reset_partial_state=True) + self.distributed_state = None + if not self.use_ipex and "ACCELERATE_USE_IPEX" not in os.environ: + os.environ["ACCELERATE_USE_IPEX"] = "false" + + self._n_gpu = 1 + if self.use_cpu or strtobool(os.environ.get("ACCELERATE_USE_CPU", "False")): + accelerator_state_kwargs["cpu"] = True + accelerator_state_kwargs["backend"] = self.ddp_backend + self._n_gpu = 0 + elif is_sagemaker_mp_enabled(): + accelerator_state_kwargs["enabled"] = False + local_rank = smp.local_rank() + device = torch.device("cuda", local_rank) + torch.cuda.set_device(device) + elif is_sagemaker_dp_enabled(): + accelerator_state_kwargs["_use_sagemaker_dp"] = True + elif self.deepspeed: + accelerator_state_kwargs["use_deepspeed"] = True + accelerator_state_kwargs["timeout"] = timedelta(seconds=self.ddp_timeout) + else: + accelerator_state_kwargs["backend"] = self.ddp_backend + accelerator_state_kwargs["timeout"] = timedelta(seconds=self.ddp_timeout) + + # Now we pop everything + if accelerator_state_kwargs.pop("enabled", False) and not accelerator_state_kwargs.pop( + "use_configured_state", False + ): + # We need to patch this env var when enabling to detect deepspeed + use_deepspeed = accelerator_state_kwargs.pop("use_deepspeed", False) + if use_deepspeed: + os.environ["ACCELERATE_USE_DEEPSPEED"] = "true" + self.distributed_state = PartialState(**accelerator_state_kwargs) + if use_deepspeed: + del os.environ["ACCELERATE_USE_DEEPSPEED"] + if not is_sagemaker_mp_enabled(): + device = self.distributed_state.device + self.local_rank = self.distributed_state.local_process_index + if dist.is_available() and dist.is_initialized() and self.parallel_mode != ParallelMode.DISTRIBUTED: + logger.warning( + "torch.distributed process group is initialized, but parallel_mode != ParallelMode.DISTRIBUTED. " + "In order to use Torch DDP, launch your script with `python -m torch.distributed.launch" + ) + if is_torch_xla_available(): + device = self.distributed_state.device + self._n_gpu = 0 + elif is_sagemaker_dp_enabled() or is_sagemaker_mp_enabled(): + # Already set _n_gpu + pass + elif self.distributed_state.distributed_type == DistributedType.NO: + if self.use_mps_device: + warnings.warn( + "`use_mps_device` is deprecated and will be removed in version 5.0 of 🤗 Transformers. " + "`mps` device will be used by default if available similar to the way `cuda` device is used." + "Therefore, no action from user is required. " + ) + if device.type != "mps": + raise ValueError( + "Either you do not have an MPS-enabled device on this machine or MacOS version is not 12.3+ " + "or current PyTorch install was not built with MPS enabled." + ) + if self.use_cpu: + device = torch.device("cpu") + elif is_torch_mps_available(): + device = torch.device("mps") + elif is_torch_xpu_available(): + if not is_ipex_available() and not is_accelerate_available("0.32.0.dev"): + raise ImportError("Using the XPU PyTorch backend requires `accelerate>=0.32.0.dev`") + device = torch.device("xpu:0") + torch.xpu.set_device(device) + elif is_torch_mlu_available(): + device = torch.device("mlu:0") + torch.mlu.set_device(device) + elif is_torch_musa_available(): + device = torch.device("musa:0") + torch.musa.set_device(device) + elif is_torch_npu_available(): + device = torch.device("npu:0") + torch.npu.set_device(device) + elif is_torch_hpu_available(): + device = torch.device("hpu:0") + torch.hpu.set_device(device) + else: + # if n_gpu is > 1 we'll use nn.DataParallel. + # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0` + # Explicitly set CUDA to the first (index 0) CUDA device, otherwise `set_device` will + # trigger an error that a device index is missing. Index 0 takes into account the + # GPUs available in the environment, so `CUDA_VISIBLE_DEVICES=1,2` with `cuda:0` + # will use the first GPU in that env, i.e. GPU#1 + device = torch.device( + "cuda:0" if torch.cuda.is_available() else os.environ.get("ACCELERATE_TORCH_DEVICE", "cpu") + ) + # Sometimes the line in the postinit has not been run before we end up here, so just checking we're not at + # the default value. + self._n_gpu = torch.cuda.device_count() + if device.type == "cuda": + torch.cuda.set_device(device) + return device + + @property + def device(self) -> "torch.device": + """ + The device used by this process. + """ + requires_backends(self, ["torch"]) + return self._setup_devices + + @property + def n_gpu(self): + """ + The number of GPUs used by this process. + + Note: + This will only be greater than one when you have multiple GPUs available but are not using distributed + training. For distributed training, it will always be 1. + """ + requires_backends(self, ["torch"]) + # Make sure `self._n_gpu` is properly setup. + if not hasattr(self, "_n_gpu"): + _ = self._setup_devices + return self._n_gpu + + @property + def parallel_mode(self): + """ + The current mode used for parallelism if multiple GPUs/TPU cores are available. One of: + + - `ParallelMode.NOT_PARALLEL`: no parallelism (CPU or one GPU). + - `ParallelMode.NOT_DISTRIBUTED`: several GPUs in one single process (uses `torch.nn.DataParallel`). + - `ParallelMode.DISTRIBUTED`: several GPUs, each having its own process (uses + `torch.nn.DistributedDataParallel`). + - `ParallelMode.TPU`: several TPU cores. + """ + requires_backends(self, ["torch"]) + if is_torch_xla_available(): + return ParallelMode.TPU + elif is_sagemaker_mp_enabled(): + return ParallelMode.SAGEMAKER_MODEL_PARALLEL + elif is_sagemaker_dp_enabled(): + return ParallelMode.SAGEMAKER_DATA_PARALLEL + elif ( + self.distributed_state is not None and self.distributed_state.distributed_type != DistributedType.NO + ) or (self.distributed_state is None and self.local_rank != -1): + return ParallelMode.DISTRIBUTED + elif self.n_gpu > 1: + return ParallelMode.NOT_DISTRIBUTED + else: + return ParallelMode.NOT_PARALLEL + + @property + def world_size(self): + """ + The number of processes used in parallel. + """ + requires_backends(self, ["torch"]) + if self.distributed_state is not None: + return self.distributed_state.num_processes + elif is_sagemaker_mp_enabled(): + return smp.dp_size() if not smp.state.cfg.prescaled_batch else smp.rdp_size() + return 1 + + @property + def process_index(self): + """ + The index of the current process used. + """ + requires_backends(self, ["torch"]) + if self.distributed_state is not None: + return self.distributed_state.process_index + elif is_sagemaker_mp_enabled(): + return smp.dp_rank() if not smp.state.cfg.prescaled_batch else smp.rdp_rank() + return 0 + + @property + def local_process_index(self): + """ + The index of the local process used. + """ + requires_backends(self, ["torch"]) + + if self.distributed_state is not None: + return self.distributed_state.local_process_index + elif is_sagemaker_mp_enabled(): + return smp.local_rank() + return 0 + + @property + def should_log(self): + """ + Whether or not the current process should produce log. + """ + if self.log_on_each_node: + return self.local_process_index == 0 + else: + if is_sagemaker_mp_enabled(): + return smp.rank() == 0 + else: + return self.process_index == 0 + + @property + def should_save(self): + """ + Whether or not the current process should write to disk, e.g., to save models and checkpoints. + """ + if self.save_on_each_node: + return self.local_process_index == 0 + else: + if is_sagemaker_mp_enabled(): + return smp.rank() == 0 + else: + return self.process_index == 0 + + def get_process_log_level(self): + """ + Returns the log level to be used depending on whether this process is the main process of node 0, main process + of node non-0, or a non-main process. + + For the main process the log level defaults to the logging level set (`logging.WARNING` if you didn't do + anything) unless overridden by `log_level` argument. + + For the replica processes the log level defaults to `logging.WARNING` unless overridden by `log_level_replica` + argument. + + The choice between the main and replica process settings is made according to the return value of `should_log`. + """ + + # convert to int + log_level = trainer_log_levels[self.log_level] + log_level_replica = trainer_log_levels[self.log_level_replica] + + log_level_main_node = logging.get_verbosity() if log_level == -1 else log_level + log_level_replica_node = logging.get_verbosity() if log_level_replica == -1 else log_level_replica + return log_level_main_node if self.should_log else log_level_replica_node + + @property + def place_model_on_device(self): + """ + Can be subclassed and overridden for some specific integrations. + """ + return not is_sagemaker_mp_enabled() + + @property + def _no_sync_in_gradient_accumulation(self): + """ + Whether or not to use no_sync for the gradients when doing gradient accumulation. + """ + return not ( + self.deepspeed or is_sagemaker_dp_enabled() or is_sagemaker_mp_enabled() or is_torch_neuroncore_available() + ) + + @contextlib.contextmanager + def main_process_first(self, local=True, desc="work"): + """ + A context manager for torch distributed environment where on needs to do something on the main process, while + blocking replicas, and when it's finished releasing the replicas. + + One such use is for `datasets`'s `map` feature which to be efficient should be run once on the main process, + which upon completion saves a cached version of results and which then automatically gets loaded by the + replicas. + + Args: + local (`bool`, *optional*, defaults to `True`): + if `True` first means process of rank 0 of each node if `False` first means process of rank 0 of node + rank 0 In multi-node environment with a shared filesystem you most likely will want to use + `local=False` so that only the main process of the first node will do the processing. If however, the + filesystem is not shared, then the main process of each node will need to do the processing, which is + the default behavior. + desc (`str`, *optional*, defaults to `"work"`): + a work description to be used in debug logs + + """ + if is_torch_available() and self.world_size > 1: + main_process_desc = "main local process" if local else "main process" + if self.distributed_state is not None: + is_main_process = ( + self.distributed_state.is_local_main_process if local else self.distributed_state.is_main_process + ) + elif is_sagemaker_mp_enabled(): + is_main_process = smp.rank() == 0 + + try: + if not is_main_process: + # tell all replicas to wait + logger.debug(f"{self.process_index}: waiting for the {main_process_desc} to perform {desc}") + + if is_torch_xla_available(): + xm.rendezvous(desc) + else: + dist.barrier() + yield + finally: + if is_main_process: + # the wait is over + logger.debug(f"{self.process_index}: {main_process_desc} completed {desc}, releasing all replicas") + if is_torch_xla_available(): + xm.rendezvous(desc) + else: + dist.barrier() + else: + yield + + def get_warmup_steps(self, num_training_steps: int): + """ + Get number of steps used for a linear warmup. + """ + warmup_steps = ( + self.warmup_steps if self.warmup_steps > 0 else math.ceil(num_training_steps * self.warmup_ratio) + ) + return warmup_steps + + def _dict_torch_dtype_to_str(self, d: dict[str, Any]) -> None: + """ + Checks whether the passed dictionary and its nested dicts have a *torch_dtype* key and if it's not None, + converts torch.dtype to a string of just the type. For example, `torch.float32` get converted into *"float32"* + string, which can then be stored in the json format. + """ + if d.get("torch_dtype", None) is not None and not isinstance(d["torch_dtype"], str): + d["torch_dtype"] = str(d["torch_dtype"]).split(".")[1] + for value in d.values(): + if isinstance(value, dict): + self._dict_torch_dtype_to_str(value) + + def to_dict(self): + """ + Serializes this instance while replace `Enum` by their values (for JSON serialization support). It obfuscates + the token values by removing their value. + """ + # filter out fields that are defined as field(init=False) + d = {field.name: getattr(self, field.name) for field in fields(self) if field.init} + + for k, v in d.items(): + if isinstance(v, Enum): + d[k] = v.value + if isinstance(v, list) and len(v) > 0 and isinstance(v[0], Enum): + d[k] = [x.value for x in v] + if k.endswith("_token"): + d[k] = f"<{k.upper()}>" + # Handle the accelerator_config if passed + if is_accelerate_available() and isinstance(v, AcceleratorConfig): + d[k] = v.to_dict() + self._dict_torch_dtype_to_str(d) + + return d + + def to_json_string(self): + """ + Serializes this instance to a JSON string. + """ + return json.dumps(self.to_dict(), indent=2) + + def to_sanitized_dict(self) -> dict[str, Any]: + """ + Sanitized serialization to use with TensorBoard’s hparams + """ + d = self.to_dict() + d = {**d, **{"train_batch_size": self.train_batch_size, "eval_batch_size": self.eval_batch_size}} + + valid_types = [bool, int, float, str] + if is_torch_available(): + valid_types.append(torch.Tensor) + + return {k: v if type(v) in valid_types else str(v) for k, v in d.items()} + + # The following methods are there to simplify the instantiation of `TrainingArguments` + def set_training( + self, + learning_rate: float = 5e-5, + batch_size: int = 8, + weight_decay: float = 0, + num_epochs: float = 3, + max_steps: int = -1, + gradient_accumulation_steps: int = 1, + seed: int = 42, + gradient_checkpointing: bool = False, + ): + """ + A method that regroups all basic arguments linked to the training. + + + + Calling this method will automatically set `self.do_train` to `True`. + + + + Args: + learning_rate (`float`, *optional*, defaults to 5e-5): + The initial learning rate for the optimizer. + batch_size (`int` *optional*, defaults to 8): + The batch size per device (GPU/TPU core/CPU...) used for training. + weight_decay (`float`, *optional*, defaults to 0): + The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights in the + optimizer. + num_train_epochs(`float`, *optional*, defaults to 3.0): + Total number of training epochs to perform (if not an integer, will perform the decimal part percents + of the last epoch before stopping training). + max_steps (`int`, *optional*, defaults to -1): + If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`. + For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until + `max_steps` is reached. + gradient_accumulation_steps (`int`, *optional*, defaults to 1): + Number of updates steps to accumulate the gradients for, before performing a backward/update pass. + + + + When using gradient accumulation, one step is counted as one step with backward pass. Therefore, + logging, evaluation, save will be conducted every `gradient_accumulation_steps * xxx_step` training + examples. + + + + seed (`int`, *optional*, defaults to 42): + Random seed that will be set at the beginning of training. To ensure reproducibility across runs, use + the [`~Trainer.model_init`] function to instantiate the model if it has some randomly initialized + parameters. + gradient_checkpointing (`bool`, *optional*, defaults to `False`): + If True, use gradient checkpointing to save memory at the expense of slower backward pass. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_training(learning_rate=1e-4, batch_size=32) + >>> args.learning_rate + 1e-4 + ``` + """ + self.do_train = True + self.learning_rate = learning_rate + self.per_device_train_batch_size = batch_size + self.weight_decay = weight_decay + self.num_train_epochs = num_epochs + self.max_steps = max_steps + self.gradient_accumulation_steps = gradient_accumulation_steps + self.seed = seed + self.gradient_checkpointing = gradient_checkpointing + return self + + def set_evaluate( + self, + strategy: Union[str, IntervalStrategy] = "no", + steps: int = 500, + batch_size: int = 8, + accumulation_steps: Optional[int] = None, + delay: Optional[float] = None, + loss_only: bool = False, + jit_mode: bool = False, + ): + """ + A method that regroups all arguments linked to evaluation. + + Args: + strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"no"`): + The evaluation strategy to adopt during training. Possible values are: + + - `"no"`: No evaluation is done during training. + - `"steps"`: Evaluation is done (and logged) every `steps`. + - `"epoch"`: Evaluation is done at the end of each epoch. + + Setting a `strategy` different from `"no"` will set `self.do_eval` to `True`. + steps (`int`, *optional*, defaults to 500): + Number of update steps between two evaluations if `strategy="steps"`. + batch_size (`int` *optional*, defaults to 8): + The batch size per device (GPU/TPU core/CPU...) used for evaluation. + accumulation_steps (`int`, *optional*): + Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU. + If left unset, the whole predictions are accumulated on GPU/TPU before being moved to the CPU (faster + but requires more memory). + delay (`float`, *optional*): + Number of epochs or steps to wait for before the first evaluation can be performed, depending on the + eval_strategy. + loss_only (`bool`, *optional*, defaults to `False`): + Ignores all outputs except the loss. + jit_mode (`bool`, *optional*): + Whether or not to use PyTorch jit trace for inference. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_evaluate(strategy="steps", steps=100) + >>> args.eval_steps + 100 + ``` + """ + self.eval_strategy = IntervalStrategy(strategy) + if self.eval_strategy == IntervalStrategy.STEPS and steps == 0: + raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.") + self.do_eval = self.eval_strategy != IntervalStrategy.NO + self.eval_steps = steps + self.per_device_eval_batch_size = batch_size + self.eval_accumulation_steps = accumulation_steps + self.eval_delay = delay + self.prediction_loss_only = loss_only + self.jit_mode_eval = jit_mode + return self + + def set_testing( + self, + batch_size: int = 8, + loss_only: bool = False, + jit_mode: bool = False, + ): + """ + A method that regroups all basic arguments linked to testing on a held-out dataset. + + + + Calling this method will automatically set `self.do_predict` to `True`. + + + + Args: + batch_size (`int` *optional*, defaults to 8): + The batch size per device (GPU/TPU core/CPU...) used for testing. + loss_only (`bool`, *optional*, defaults to `False`): + Ignores all outputs except the loss. + jit_mode (`bool`, *optional*): + Whether or not to use PyTorch jit trace for inference. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_testing(batch_size=32) + >>> args.per_device_eval_batch_size + 32 + ``` + """ + self.do_predict = True + self.per_device_eval_batch_size = batch_size + self.prediction_loss_only = loss_only + self.jit_mode_eval = jit_mode + return self + + def set_save( + self, + strategy: Union[str, IntervalStrategy] = "steps", + steps: int = 500, + total_limit: Optional[int] = None, + on_each_node: bool = False, + ): + """ + A method that regroups all arguments linked to checkpoint saving. + + Args: + strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`): + The checkpoint save strategy to adopt during training. Possible values are: + + - `"no"`: No save is done during training. + - `"epoch"`: Save is done at the end of each epoch. + - `"steps"`: Save is done every `save_steps`. + + steps (`int`, *optional*, defaults to 500): + Number of updates steps before two checkpoint saves if `strategy="steps"`. + total_limit (`int`, *optional*): + If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in + `output_dir`. + on_each_node (`bool`, *optional*, defaults to `False`): + When doing multi-node distributed training, whether to save models and checkpoints on each node, or + only on the main one. + + This should not be activated when the different nodes use the same storage as the files will be saved + with the same names for each node. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_save(strategy="steps", steps=100) + >>> args.save_steps + 100 + ``` + """ + self.save_strategy = SaveStrategy(strategy) + if self.save_strategy == SaveStrategy.STEPS and steps == 0: + raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.") + self.save_steps = steps + self.save_total_limit = total_limit + self.save_on_each_node = on_each_node + return self + + def set_logging( + self, + strategy: Union[str, IntervalStrategy] = "steps", + steps: int = 500, + report_to: Union[str, list[str]] = "none", + level: str = "passive", + first_step: bool = False, + nan_inf_filter: bool = False, + on_each_node: bool = False, + replica_level: str = "passive", + ): + """ + A method that regroups all arguments linked to logging. + + Args: + strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`): + The logging strategy to adopt during training. Possible values are: + + - `"no"`: No logging is done during training. + - `"epoch"`: Logging is done at the end of each epoch. + - `"steps"`: Logging is done every `logging_steps`. + + steps (`int`, *optional*, defaults to 500): + Number of update steps between two logs if `strategy="steps"`. + level (`str`, *optional*, defaults to `"passive"`): + Logger log level to use on the main process. Possible choices are the log levels as strings: `"debug"`, + `"info"`, `"warning"`, `"error"` and `"critical"`, plus a `"passive"` level which doesn't set anything + and lets the application set the level. + report_to (`str` or `List[str]`, *optional*, defaults to `"all"`): + The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`, + `"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`, + `"neptune"`, `"swanlab"`, `"tensorboard"`, and `"wandb"`. Use `"all"` to report to all integrations + installed, `"none"` for no integrations. + first_step (`bool`, *optional*, defaults to `False`): + Whether to log and evaluate the first `global_step` or not. + nan_inf_filter (`bool`, *optional*, defaults to `True`): + Whether to filter `nan` and `inf` losses for logging. If set to `True` the loss of every step that is + `nan` or `inf` is filtered and the average loss of the current logging window is taken instead. + + + + `nan_inf_filter` only influences the logging of loss values, it does not change the behavior the + gradient is computed or applied to the model. + + + + on_each_node (`bool`, *optional*, defaults to `True`): + In multinode distributed training, whether to log using `log_level` once per node, or only on the main + node. + replica_level (`str`, *optional*, defaults to `"passive"`): + Logger log level to use on replicas. Same choices as `log_level` + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_logging(strategy="steps", steps=100) + >>> args.logging_steps + 100 + ``` + """ + self.logging_strategy = IntervalStrategy(strategy) + if self.logging_strategy == IntervalStrategy.STEPS and steps == 0: + raise ValueError("Setting `strategy` as 'steps' requires a positive value for `steps`.") + self.logging_steps = steps + self.report_to = report_to + self.log_level = level + self.logging_first_step = first_step + self.logging_nan_inf_filter = nan_inf_filter + self.log_on_each_node = on_each_node + self.log_level_replica = replica_level + return self + + def set_push_to_hub( + self, + model_id: str, + strategy: Union[str, HubStrategy] = "every_save", + token: Optional[str] = None, + private_repo: Optional[bool] = None, + always_push: bool = False, + ): + """ + A method that regroups all arguments linked to synchronizing checkpoints with the Hub. + + + + Calling this method will set `self.push_to_hub` to `True`, which means the `output_dir` will begin a git + directory synced with the repo (determined by `model_id`) and the content will be pushed each time a save is + triggered (depending on your `self.save_strategy`). Calling [`~Trainer.save_model`] will also trigger a push. + + + + Args: + model_id (`str`): + The name of the repository to keep in sync with the local *output_dir*. It can be a simple model ID in + which case the model will be pushed in your namespace. Otherwise it should be the whole repository + name, for instance `"user_name/model"`, which allows you to push to an organization you are a member of + with `"organization_name/model"`. + strategy (`str` or [`~trainer_utils.HubStrategy`], *optional*, defaults to `"every_save"`): + Defines the scope of what is pushed to the Hub and when. Possible values are: + + - `"end"`: push the model, its configuration, the processing_class e.g. tokenizer (if passed along to the [`Trainer`]) and a + draft of a model card when the [`~Trainer.save_model`] method is called. + - `"every_save"`: push the model, its configuration, the processing_class e.g. tokenizer (if passed along to the [`Trainer`]) + and + a draft of a model card each time there is a model save. The pushes are asynchronous to not block + training, and in case the save are very frequent, a new push is only attempted if the previous one is + finished. A last push is made with the final model at the end of training. + - `"checkpoint"`: like `"every_save"` but the latest checkpoint is also pushed in a subfolder named + last-checkpoint, allowing you to resume training easily with + `trainer.train(resume_from_checkpoint="last-checkpoint")`. + - `"all_checkpoints"`: like `"checkpoint"` but all checkpoints are pushed like they appear in the + output + folder (so you will get one checkpoint folder per folder in your final repository) + + token (`str`, *optional*): + The token to use to push the model to the Hub. Will default to the token in the cache folder obtained + with `huggingface-cli login`. + private_repo (`bool`, *optional*, defaults to `False`): + Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. + always_push (`bool`, *optional*, defaults to `False`): + Unless this is `True`, the `Trainer` will skip pushing a checkpoint when the previous push is not + finished. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_push_to_hub("me/awesome-model") + >>> args.hub_model_id + 'me/awesome-model' + ``` + """ + self.push_to_hub = True + self.hub_model_id = model_id + self.hub_strategy = HubStrategy(strategy) + self.hub_token = token + self.hub_private_repo = private_repo + self.hub_always_push = always_push + return self + + def set_optimizer( + self, + name: Union[str, OptimizerNames] = "adamw_torch", + learning_rate: float = 5e-5, + weight_decay: float = 0, + beta1: float = 0.9, + beta2: float = 0.999, + epsilon: float = 1e-8, + args: Optional[str] = None, + ): + """ + A method that regroups all arguments linked to the optimizer and its hyperparameters. + + Args: + name (`str` or [`training_args.OptimizerNames`], *optional*, defaults to `"adamw_torch"`): + The optimizer to use: `"adamw_torch"`, `"adamw_torch_fused"`, `"adamw_apex_fused"`, + `"adamw_anyprecision"` or `"adafactor"`. + learning_rate (`float`, *optional*, defaults to 5e-5): + The initial learning rate. + weight_decay (`float`, *optional*, defaults to 0): + The weight decay to apply (if not zero) to all layers except all bias and LayerNorm weights. + beta1 (`float`, *optional*, defaults to 0.9): + The beta1 hyperparameter for the adam optimizer or its variants. + beta2 (`float`, *optional*, defaults to 0.999): + The beta2 hyperparameter for the adam optimizer or its variants. + epsilon (`float`, *optional*, defaults to 1e-8): + The epsilon hyperparameter for the adam optimizer or its variants. + args (`str`, *optional*): + Optional arguments that are supplied to AnyPrecisionAdamW (only useful when + `optim="adamw_anyprecision"`). + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_optimizer(name="adamw_torch", beta1=0.8) + >>> args.optim + 'adamw_torch' + ``` + """ + self.optim = OptimizerNames(name) + self.learning_rate = learning_rate + self.weight_decay = weight_decay + self.adam_beta1 = beta1 + self.adam_beta2 = beta2 + self.adam_epsilon = epsilon + self.optim_args = args + return self + + def set_lr_scheduler( + self, + name: Union[str, SchedulerType] = "linear", + num_epochs: float = 3.0, + max_steps: int = -1, + warmup_ratio: float = 0, + warmup_steps: int = 0, + ): + """ + A method that regroups all arguments linked to the learning rate scheduler and its hyperparameters. + + Args: + name (`str` or [`SchedulerType`], *optional*, defaults to `"linear"`): + The scheduler type to use. See the documentation of [`SchedulerType`] for all possible values. + num_epochs(`float`, *optional*, defaults to 3.0): + Total number of training epochs to perform (if not an integer, will perform the decimal part percents + of the last epoch before stopping training). + max_steps (`int`, *optional*, defaults to -1): + If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`. + For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until + `max_steps` is reached. + warmup_ratio (`float`, *optional*, defaults to 0.0): + Ratio of total training steps used for a linear warmup from 0 to `learning_rate`. + warmup_steps (`int`, *optional*, defaults to 0): + Number of steps used for a linear warmup from 0 to `learning_rate`. Overrides any effect of + `warmup_ratio`. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_lr_scheduler(name="cosine", warmup_ratio=0.05) + >>> args.warmup_ratio + 0.05 + ``` + """ + self.lr_scheduler_type = SchedulerType(name) + self.num_train_epochs = num_epochs + self.max_steps = max_steps + self.warmup_ratio = warmup_ratio + self.warmup_steps = warmup_steps + return self + + def set_dataloader( + self, + train_batch_size: int = 8, + eval_batch_size: int = 8, + drop_last: bool = False, + num_workers: int = 0, + pin_memory: bool = True, + persistent_workers: bool = False, + prefetch_factor: Optional[int] = None, + auto_find_batch_size: bool = False, + ignore_data_skip: bool = False, + sampler_seed: Optional[int] = None, + ): + """ + A method that regroups all arguments linked to the dataloaders creation. + + Args: + drop_last (`bool`, *optional*, defaults to `False`): + Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch + size) or not. + num_workers (`int`, *optional*, defaults to 0): + Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in + the main process. + pin_memory (`bool`, *optional*, defaults to `True`): + Whether you want to pin memory in data loaders or not. Will default to `True`. + persistent_workers (`bool`, *optional*, defaults to `False`): + If True, the data loader will not shut down the worker processes after a dataset has been consumed + once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training, + but will increase RAM usage. Will default to `False`. + prefetch_factor (`int`, *optional*): + Number of batches loaded in advance by each worker. + 2 means there will be a total of 2 * num_workers batches prefetched across all workers. + auto_find_batch_size (`bool`, *optional*, defaults to `False`) + Whether to find a batch size that will fit into memory automatically through exponential decay, + avoiding CUDA Out-of-Memory errors. Requires accelerate to be installed (`pip install accelerate`) + ignore_data_skip (`bool`, *optional*, defaults to `False`): + When resuming training, whether or not to skip the epochs and batches to get the data loading at the + same stage as in the previous training. If set to `True`, the training will begin faster (as that + skipping step can take a long time) but will not yield the same results as the interrupted training + would have. + sampler_seed (`int`, *optional*): + Random seed to be used with data samplers. If not set, random generators for data sampling will use the + same seed as `self.seed`. This can be used to ensure reproducibility of data sampling, independent of + the model seed. + + Example: + + ```py + >>> from transformers import TrainingArguments + + >>> args = TrainingArguments("working_dir") + >>> args = args.set_dataloader(train_batch_size=16, eval_batch_size=64) + >>> args.per_device_train_batch_size + 16 + ``` + """ + self.per_device_train_batch_size = train_batch_size + self.per_device_eval_batch_size = eval_batch_size + self.dataloader_drop_last = drop_last + self.dataloader_num_workers = num_workers + self.dataloader_pin_memory = pin_memory + self.dataloader_persistent_workers = persistent_workers + self.dataloader_prefetch_factor = prefetch_factor + self.auto_find_batch_size = auto_find_batch_size + self.ignore_data_skip = ignore_data_skip + self.data_seed = sampler_seed + return self + + +class ParallelMode(Enum): + NOT_PARALLEL = "not_parallel" + NOT_DISTRIBUTED = "not_distributed" + DISTRIBUTED = "distributed" + SAGEMAKER_MODEL_PARALLEL = "sagemaker_model_parallel" + SAGEMAKER_DATA_PARALLEL = "sagemaker_data_parallel" + TPU = "tpu" diff --git a/lib/python3.12/site-packages/transformers/training_args_seq2seq.py b/lib/python3.12/site-packages/transformers/training_args_seq2seq.py new file mode 100644 index 0000000000000000000000000000000000000000..5342b7add3932c542e35247e52920d8fc91ed325 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/training_args_seq2seq.py @@ -0,0 +1,90 @@ +# Copyright 2020 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. + +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional, Union + +from .generation.configuration_utils import GenerationConfig +from .training_args import TrainingArguments +from .utils import add_start_docstrings + + +logger = logging.getLogger(__name__) + + +@dataclass +@add_start_docstrings(TrainingArguments.__doc__) +class Seq2SeqTrainingArguments(TrainingArguments): + """ + Args: + predict_with_generate (`bool`, *optional*, defaults to `False`): + Whether to use generate to calculate generative metrics (ROUGE, BLEU). + generation_max_length (`int`, *optional*): + The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default to the + `max_length` value of the model configuration. + generation_num_beams (`int`, *optional*): + The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default to the + `num_beams` value of the model configuration. + generation_config (`str` or `Path` or [`~generation.GenerationConfig`], *optional*): + Allows to load a [`~generation.GenerationConfig`] from the `from_pretrained` method. This can be either: + + - a string, the *model id* of a pretrained model configuration hosted inside a model repo on + huggingface.co. + - a path to a *directory* containing a configuration file saved using the + [`~GenerationConfig.save_pretrained`] method, e.g., `./my_model_directory/`. + - a [`~generation.GenerationConfig`] object. + """ + + sortish_sampler: bool = field(default=False, metadata={"help": "Whether to use SortishSampler or not."}) + predict_with_generate: bool = field( + default=False, metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."} + ) + generation_max_length: Optional[int] = field( + default=None, + metadata={ + "help": ( + "The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default " + "to the `max_length` value of the model configuration." + ) + }, + ) + generation_num_beams: Optional[int] = field( + default=None, + metadata={ + "help": ( + "The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default " + "to the `num_beams` value of the model configuration." + ) + }, + ) + generation_config: Optional[Union[str, Path, GenerationConfig]] = field( + default=None, + metadata={ + "help": "Model id, file path or url pointing to a GenerationConfig json file, to use during prediction." + }, + ) + + def to_dict(self): + """ + Serializes this instance while replace `Enum` by their values and `GenerationConfig` by dictionaries (for JSON + serialization support). It obfuscates the token values by removing their value. + """ + # filter out fields that are defined as field(init=False) + d = super().to_dict() + for k, v in d.items(): + if isinstance(v, GenerationConfig): + d[k] = v.to_dict() + return d diff --git a/lib/python3.12/site-packages/transformers/training_args_tf.py b/lib/python3.12/site-packages/transformers/training_args_tf.py new file mode 100644 index 0000000000000000000000000000000000000000..6bbd4b89a7249d84db9a579fe75928a3bd5feafb --- /dev/null +++ b/lib/python3.12/site-packages/transformers/training_args_tf.py @@ -0,0 +1,299 @@ +# Copyright 2020 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. + +import warnings +from dataclasses import dataclass, field +from typing import Optional + +from .training_args import TrainingArguments +from .utils import cached_property, is_tf_available, logging, requires_backends + + +logger = logging.get_logger(__name__) + +if is_tf_available(): + import tensorflow as tf + + from .modeling_tf_utils import keras + + +@dataclass +class TFTrainingArguments(TrainingArguments): + """ + TrainingArguments is the subset of the arguments we use in our example scripts **which relate to the training loop + itself**. + + Using [`HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + output_dir (`str`): + The output directory where the model predictions and checkpoints will be written. + overwrite_output_dir (`bool`, *optional*, defaults to `False`): + If `True`, overwrite the content of the output directory. Use this to continue training if `output_dir` + points to a checkpoint directory. + do_train (`bool`, *optional*, defaults to `False`): + Whether to run training or not. This argument is not directly used by [`Trainer`], it's intended to be used + by your training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + do_eval (`bool`, *optional*): + Whether to run evaluation on the validation set or not. Will be set to `True` if `eval_strategy` is + different from `"no"`. This argument is not directly used by [`Trainer`], it's intended to be used by your + training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + do_predict (`bool`, *optional*, defaults to `False`): + Whether to run predictions on the test set or not. This argument is not directly used by [`Trainer`], it's + intended to be used by your training/evaluation scripts instead. See the [example + scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details. + eval_strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"no"`): + The evaluation strategy to adopt during training. Possible values are: + + - `"no"`: No evaluation is done during training. + - `"steps"`: Evaluation is done (and logged) every `eval_steps`. + - `"epoch"`: Evaluation is done at the end of each epoch. + + per_device_train_batch_size (`int`, *optional*, defaults to 8): + The batch size per GPU/TPU core/CPU for training. + per_device_eval_batch_size (`int`, *optional*, defaults to 8): + The batch size per GPU/TPU core/CPU for evaluation. + gradient_accumulation_steps (`int`, *optional*, defaults to 1): + Number of updates steps to accumulate the gradients for, before performing a backward/update pass. + + + + When using gradient accumulation, one step is counted as one step with backward pass. Therefore, logging, + evaluation, save will be conducted every `gradient_accumulation_steps * xxx_step` training examples. + + + + learning_rate (`float`, *optional*, defaults to 5e-5): + The initial learning rate for Adam. + weight_decay (`float`, *optional*, defaults to 0): + The weight decay to apply (if not zero). + adam_beta1 (`float`, *optional*, defaults to 0.9): + The beta1 hyperparameter for the Adam optimizer. + adam_beta2 (`float`, *optional*, defaults to 0.999): + The beta2 hyperparameter for the Adam optimizer. + adam_epsilon (`float`, *optional*, defaults to 1e-8): + The epsilon hyperparameter for the Adam optimizer. + max_grad_norm (`float`, *optional*, defaults to 1.0): + Maximum gradient norm (for gradient clipping). + num_train_epochs(`float`, *optional*, defaults to 3.0): + Total number of training epochs to perform. + max_steps (`int`, *optional*, defaults to -1): + If set to a positive number, the total number of training steps to perform. Overrides `num_train_epochs`. + For a finite dataset, training is reiterated through the dataset (if all data is exhausted) until + `max_steps` is reached. + warmup_ratio (`float`, *optional*, defaults to 0.0): + Ratio of total training steps used for a linear warmup from 0 to `learning_rate`. + warmup_steps (`int`, *optional*, defaults to 0): + Number of steps used for a linear warmup from 0 to `learning_rate`. Overrides any effect of `warmup_ratio`. + logging_dir (`str`, *optional*): + [TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to + *runs/**CURRENT_DATETIME_HOSTNAME***. + logging_strategy (`str` or [`~trainer_utils.IntervalStrategy`], *optional*, defaults to `"steps"`): + The logging strategy to adopt during training. Possible values are: + + - `"no"`: No logging is done during training. + - `"epoch"`: Logging is done at the end of each epoch. + - `"steps"`: Logging is done every `logging_steps`. + + logging_first_step (`bool`, *optional*, defaults to `False`): + Whether to log and evaluate the first `global_step` or not. + logging_steps (`int`, *optional*, defaults to 500): + Number of update steps between two logs if `logging_strategy="steps"`. + save_strategy (`str` or [`~trainer_utils.SaveStrategy`], *optional*, defaults to `"steps"`): + The checkpoint save strategy to adopt during training. Possible values are: + + - `"no"`: No save is done during training. + - `"epoch"`: Save is done at the end of each epoch. + - `"steps"`: Save is done every `save_steps`. + + save_steps (`int`, *optional*, defaults to 500): + Number of updates steps before two checkpoint saves if `save_strategy="steps"`. + save_total_limit (`int`, *optional*): + If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in + `output_dir`. + no_cuda (`bool`, *optional*, defaults to `False`): + Whether to not use CUDA even when it is available or not. + seed (`int`, *optional*, defaults to 42): + Random seed that will be set at the beginning of training. + fp16 (`bool`, *optional*, defaults to `False`): + Whether to use 16-bit (mixed) precision training (through NVIDIA Apex) instead of 32-bit training. + fp16_opt_level (`str`, *optional*, defaults to 'O1'): + For `fp16` training, Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']. See details on + the [Apex documentation](https://nvidia.github.io/apex/amp). + local_rank (`int`, *optional*, defaults to -1): + During distributed training, the rank of the process. + tpu_num_cores (`int`, *optional*): + When training on TPU, the number of TPU cores (automatically passed by launcher script). + debug (`bool`, *optional*, defaults to `False`): + Whether to activate the trace to record computation graphs and profiling information or not. + dataloader_drop_last (`bool`, *optional*, defaults to `False`): + Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size) + or not. + eval_steps (`int`, *optional*, defaults to 1000): + Number of update steps before two evaluations. + past_index (`int`, *optional*, defaults to -1): + Some models like [TransformerXL](../model_doc/transformerxl) or :doc*XLNet <../model_doc/xlnet>* can make + use of the past hidden states for their predictions. If this argument is set to a positive int, the + `Trainer` will use the corresponding output (usually index 2) as the past state and feed it to the model at + the next training step under the keyword argument `mems`. + tpu_name (`str`, *optional*): + The name of the TPU the process is running on. + tpu_zone (`str`, *optional*): + The zone of the TPU the process is running on. If not specified, we will attempt to automatically detect + from metadata. + gcp_project (`str`, *optional*): + Google Cloud Project name for the Cloud TPU-enabled project. If not specified, we will attempt to + automatically detect from metadata. + run_name (`str`, *optional*): + A descriptor for the run. Notably used for wandb, mlflow, comet and swanlab logging. + xla (`bool`, *optional*): + Whether to activate the XLA compilation or not. + """ + + framework = "tf" + tpu_name: Optional[str] = field( + default=None, + metadata={"help": "Name of TPU"}, + ) + + tpu_zone: Optional[str] = field( + default=None, + metadata={"help": "Zone of TPU"}, + ) + + gcp_project: Optional[str] = field( + default=None, + metadata={"help": "Name of Cloud TPU-enabled project"}, + ) + + poly_power: float = field( + default=1.0, + metadata={"help": "Power for the Polynomial decay LR scheduler."}, + ) + + xla: bool = field(default=False, metadata={"help": "Whether to activate the XLA compilation or not"}) + + @cached_property + def _setup_strategy(self) -> tuple["tf.distribute.Strategy", int]: + requires_backends(self, ["tf"]) + logger.info("Tensorflow: setting up strategy") + + gpus = tf.config.list_physical_devices("GPU") + + # Set to float16 at first + if self.fp16: + keras.mixed_precision.set_global_policy("mixed_float16") + + if self.no_cuda: + strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0") + else: + try: + if self.tpu_name: + tpu = tf.distribute.cluster_resolver.TPUClusterResolver( + self.tpu_name, zone=self.tpu_zone, project=self.gcp_project + ) + else: + tpu = tf.distribute.cluster_resolver.TPUClusterResolver() + except ValueError: + if self.tpu_name: + raise RuntimeError(f"Couldn't connect to TPU {self.tpu_name}!") + else: + tpu = None + + if tpu: + # Set to bfloat16 in case of TPU + if self.fp16: + keras.mixed_precision.set_global_policy("mixed_bfloat16") + + tf.config.experimental_connect_to_cluster(tpu) + tf.tpu.experimental.initialize_tpu_system(tpu) + + strategy = tf.distribute.TPUStrategy(tpu) + + elif len(gpus) == 0: + strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0") + elif len(gpus) == 1: + strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0") + elif len(gpus) > 1: + # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0` + strategy = tf.distribute.MirroredStrategy() + else: + raise ValueError("Cannot find the proper strategy, please check your environment properties.") + + return strategy + + @property + def strategy(self) -> "tf.distribute.Strategy": + """ + The strategy used for distributed training. + """ + requires_backends(self, ["tf"]) + return self._setup_strategy + + @property + def n_replicas(self) -> int: + """ + The number of replicas (CPUs, GPUs or TPU cores) used in this training. + """ + requires_backends(self, ["tf"]) + return self._setup_strategy.num_replicas_in_sync + + @property + def should_log(self): + """ + Whether or not the current process should produce log. + """ + return False # TF Logging is handled by Keras not the Trainer + + @property + def train_batch_size(self) -> int: + """ + The actual batch size for training (may differ from `per_gpu_train_batch_size` in distributed training). + """ + if self.per_gpu_train_batch_size: + logger.warning( + "Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future " + "version. Using `--per_device_train_batch_size` is preferred." + ) + per_device_batch_size = self.per_gpu_train_batch_size or self.per_device_train_batch_size + return per_device_batch_size * self.n_replicas + + @property + def eval_batch_size(self) -> int: + """ + The actual batch size for evaluation (may differ from `per_gpu_eval_batch_size` in distributed training). + """ + if self.per_gpu_eval_batch_size: + logger.warning( + "Using deprecated `--per_gpu_eval_batch_size` argument which will be removed in a future " + "version. Using `--per_device_eval_batch_size` is preferred." + ) + per_device_batch_size = self.per_gpu_eval_batch_size or self.per_device_eval_batch_size + return per_device_batch_size * self.n_replicas + + @property + def n_gpu(self) -> int: + """ + The number of replicas (CPUs, GPUs or TPU cores) used in this training. + """ + requires_backends(self, ["tf"]) + warnings.warn( + "The n_gpu argument is deprecated and will be removed in a future version, use n_replicas instead.", + FutureWarning, + ) + return self._setup_strategy.num_replicas_in_sync diff --git a/lib/python3.12/site-packages/transformers/utils/__init__.py b/lib/python3.12/site-packages/transformers/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3a1c2e59985dcf99eb88ae9c42d4255242d9672d --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/__init__.py @@ -0,0 +1,329 @@ +#!/usr/bin/env python + +# Copyright 2021 The HuggingFace Inc. 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. + +from functools import lru_cache + +from huggingface_hub import get_full_repo_name # for backward compatibility +from huggingface_hub.constants import HF_HUB_DISABLE_TELEMETRY as DISABLE_TELEMETRY # for backward compatibility +from packaging import version + +from .. import __version__ +from .backbone_utils import BackboneConfigMixin, BackboneMixin +from .chat_template_utils import DocstringParsingException, TypeHintParsingException, get_json_schema +from .constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD +from .doc import ( + add_code_sample_docstrings, + add_end_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + copy_func, + replace_return_docstrings, +) +from .generic import ( + ContextManagers, + ExplicitEnum, + LossKwargs, + ModelOutput, + PaddingStrategy, + TensorType, + add_model_info_to_auto_map, + add_model_info_to_custom_pipelines, + cached_property, + can_return_loss, + can_return_tuple, + expand_dims, + filter_out_non_signature_kwargs, + find_labels, + flatten_dict, + infer_framework, + is_jax_tensor, + is_numpy_array, + is_tensor, + is_tf_symbolic_tensor, + is_tf_tensor, + is_timm_config_dict, + is_timm_local_checkpoint, + is_torch_device, + is_torch_dtype, + is_torch_tensor, + reshape, + squeeze, + strtobool, + tensor_size, + to_numpy, + to_py_obj, + torch_float, + torch_int, + transpose, + working_or_temp_dir, +) +from .hub import ( + CLOUDFRONT_DISTRIB_PREFIX, + HF_MODULES_CACHE, + HUGGINGFACE_CO_PREFIX, + HUGGINGFACE_CO_RESOLVE_ENDPOINT, + PYTORCH_PRETRAINED_BERT_CACHE, + PYTORCH_TRANSFORMERS_CACHE, + S3_BUCKET_PREFIX, + TRANSFORMERS_CACHE, + TRANSFORMERS_DYNAMIC_MODULE_NAME, + EntryNotFoundError, + PushInProgress, + PushToHubMixin, + RepositoryNotFoundError, + RevisionNotFoundError, + cached_file, + default_cache_path, + define_sagemaker_information, + download_url, + extract_commit_hash, + has_file, + http_user_agent, + is_offline_mode, + is_remote_url, + send_example_telemetry, + try_to_load_from_cache, +) +from .import_utils import ( + ACCELERATE_MIN_VERSION, + ENV_VARS_TRUE_AND_AUTO_VALUES, + ENV_VARS_TRUE_VALUES, + GGUF_MIN_VERSION, + TORCH_FX_REQUIRED_VERSION, + USE_JAX, + USE_TF, + USE_TORCH, + XLA_FSDPV2_MIN_VERSION, + DummyObject, + OptionalDependencyNotAvailable, + _LazyModule, + ccl_version, + direct_transformers_import, + get_torch_version, + is_accelerate_available, + is_apex_available, + is_apollo_torch_available, + is_aqlm_available, + is_auto_awq_available, + is_auto_gptq_available, + is_av_available, + is_bitsandbytes_available, + is_bitsandbytes_multi_backend_available, + is_bs4_available, + is_coloredlogs_available, + is_compressed_tensors_available, + is_cv2_available, + is_cython_available, + is_datasets_available, + is_decord_available, + is_detectron2_available, + is_eetq_available, + is_essentia_available, + is_faiss_available, + is_fbgemm_gpu_available, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal, + is_flash_attn_greater_or_equal_2_10, + is_flax_available, + is_flute_available, + is_fsdp_available, + is_ftfy_available, + is_g2p_en_available, + is_galore_torch_available, + is_gguf_available, + is_gptqmodel_available, + is_grokadamw_available, + is_habana_gaudi1, + is_hadamard_available, + is_hqq_available, + is_in_notebook, + is_ipex_available, + is_jieba_available, + is_jinja_available, + is_jumanpp_available, + is_kenlm_available, + is_keras_nlp_available, + is_levenshtein_available, + is_librosa_available, + is_liger_kernel_available, + is_lomo_available, + is_mlx_available, + is_natten_available, + is_ninja_available, + is_nltk_available, + is_num2words_available, + is_onnx_available, + is_openai_available, + is_optimum_available, + is_optimum_quanto_available, + is_pandas_available, + is_peft_available, + is_phonemizer_available, + is_pretty_midi_available, + is_protobuf_available, + is_psutil_available, + is_py3nvml_available, + is_pyctcdecode_available, + is_pytesseract_available, + is_pytest_available, + is_pytorch_quantization_available, + is_quark_available, + is_rich_available, + is_rjieba_available, + is_sacremoses_available, + is_safetensors_available, + is_sagemaker_dp_enabled, + is_sagemaker_mp_enabled, + is_schedulefree_available, + is_scipy_available, + is_sentencepiece_available, + is_seqio_available, + is_sklearn_available, + is_soundfile_available, + is_spacy_available, + is_speech_available, + is_spqr_available, + is_sudachi_available, + is_sudachi_projection_available, + is_tensorflow_probability_available, + is_tensorflow_text_available, + is_tf2onnx_available, + is_tf_available, + is_tiktoken_available, + is_timm_available, + is_tokenizers_available, + is_torch_available, + is_torch_bf16_available, + is_torch_bf16_available_on_device, + is_torch_bf16_cpu_available, + is_torch_bf16_gpu_available, + is_torch_compile_available, + is_torch_cuda_available, + is_torch_deterministic, + is_torch_flex_attn_available, + is_torch_fp16_available_on_device, + is_torch_fx_available, + is_torch_fx_proxy, + is_torch_greater_or_equal, + is_torch_hpu_available, + is_torch_mlu_available, + is_torch_mps_available, + is_torch_musa_available, + is_torch_neuroncore_available, + is_torch_npu_available, + is_torch_sdpa_available, + is_torch_tensorrt_fx_available, + is_torch_tf32_available, + is_torch_xla_available, + is_torch_xpu_available, + is_torchao_available, + is_torchaudio_available, + is_torchdistx_available, + is_torchdynamo_available, + is_torchdynamo_compiling, + is_torchdynamo_exporting, + is_torchvision_available, + is_torchvision_v2_available, + is_training_run_on_sagemaker, + is_uroman_available, + is_vision_available, + is_vptq_available, + is_yt_dlp_available, + requires_backends, + torch_only_method, +) +from .peft_utils import ( + ADAPTER_CONFIG_NAME, + ADAPTER_SAFE_WEIGHTS_NAME, + ADAPTER_WEIGHTS_NAME, + check_peft_version, + find_adapter_config_file, +) + + +WEIGHTS_NAME = "pytorch_model.bin" +WEIGHTS_INDEX_NAME = "pytorch_model.bin.index.json" +TF2_WEIGHTS_NAME = "tf_model.h5" +TF2_WEIGHTS_INDEX_NAME = "tf_model.h5.index.json" +TF_WEIGHTS_NAME = "model.ckpt" +FLAX_WEIGHTS_NAME = "flax_model.msgpack" +FLAX_WEIGHTS_INDEX_NAME = "flax_model.msgpack.index.json" +SAFE_WEIGHTS_NAME = "model.safetensors" +SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json" +CONFIG_NAME = "config.json" +FEATURE_EXTRACTOR_NAME = "preprocessor_config.json" +IMAGE_PROCESSOR_NAME = FEATURE_EXTRACTOR_NAME +PROCESSOR_NAME = "processor_config.json" +CHAT_TEMPLATE_NAME = "chat_template.json" +GENERATION_CONFIG_NAME = "generation_config.json" +MODEL_CARD_NAME = "modelcard.json" + +SENTENCEPIECE_UNDERLINE = "▁" +SPIECE_UNDERLINE = SENTENCEPIECE_UNDERLINE # Kept for backward compatibility + +MULTIPLE_CHOICE_DUMMY_INPUTS = [ + [[0, 1, 0, 1], [1, 0, 0, 1]] +] * 2 # Needs to have 0s and 1s only since XLM uses it for langs too. +DUMMY_INPUTS = [[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]] +DUMMY_MASK = [[1, 1, 1, 1, 1], [1, 1, 1, 0, 0], [0, 0, 0, 1, 1]] + + +def check_min_version(min_version): + if version.parse(__version__) < version.parse(min_version): + if "dev" in min_version: + error_message = ( + "This example requires a source install from HuggingFace Transformers (see " + "`https://huggingface.co/docs/transformers/installation#install-from-source`)," + ) + else: + error_message = f"This example requires a minimum version of {min_version}," + error_message += f" but the version found is {__version__}.\n" + raise ImportError( + error_message + + "Check out https://github.com/huggingface/transformers/tree/main/examples#important-note for the examples corresponding to other " + "versions of HuggingFace Transformers." + ) + + +@lru_cache +def get_available_devices() -> frozenset[str]: + """ + Returns a frozenset of devices available for the current PyTorch installation. + """ + devices = {"cpu"} # `cpu` is always supported as a device in PyTorch + + if is_torch_cuda_available(): + devices.add("cuda") + + if is_torch_mps_available(): + devices.add("mps") + + if is_torch_xpu_available(): + devices.add("xpu") + + if is_torch_npu_available(): + devices.add("npu") + + if is_torch_hpu_available(): + devices.add("hpu") + + if is_torch_mlu_available(): + devices.add("mlu") + + if is_torch_musa_available(): + devices.add("musa") + + return frozenset(devices) diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/chat_template_utils.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/chat_template_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..db1fe83f3a251ca6eae7158738d83b6fa90ac981 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/chat_template_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/constants.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/constants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c1b7bcfbe38a37e7659cbef7f0bd497f8aaf18be Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/constants.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_essentia_and_librosa_and_pretty_midi_and_scipy_and_torch_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_essentia_and_librosa_and_pretty_midi_and_scipy_and_torch_objects.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..28ddd641ac9bcb1e168e2fec1def1a62bd23326a Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_essentia_and_librosa_and_pretty_midi_and_scipy_and_torch_objects.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_speech_objects.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_speech_objects.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52995fe92a76571dad1c2b07bbd896acee91f4a7 Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/dummy_speech_objects.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/utils/__pycache__/hp_naming.cpython-312.pyc b/lib/python3.12/site-packages/transformers/utils/__pycache__/hp_naming.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..427e729e3215e9197c719383ab10929fbe7bd23b Binary files /dev/null and b/lib/python3.12/site-packages/transformers/utils/__pycache__/hp_naming.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/transformers/utils/attention_visualizer.py b/lib/python3.12/site-packages/transformers/utils/attention_visualizer.py new file mode 100644 index 0000000000000000000000000000000000000000..bad6a7471e54f12a47f854aca51138577bb99e6d --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/attention_visualizer.py @@ -0,0 +1,228 @@ +# Copyright 2025 The HuggingFace Inc. team. +# +# 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. + + +import requests +from PIL import Image + +from ..models.auto.auto_factory import _get_model_class +from ..models.auto.configuration_auto import AutoConfig +from ..models.auto.modeling_auto import MODEL_FOR_PRETRAINING_MAPPING, MODEL_MAPPING +from ..models.auto.processing_auto import PROCESSOR_MAPPING_NAMES, AutoProcessor +from ..models.auto.tokenization_auto import TOKENIZER_MAPPING_NAMES, AutoTokenizer +from .import_utils import is_torch_available + + +if is_torch_available(): + import torch + import torch.nn as nn + +# Print the matrix with words as row labels +GREEN = "\033[92m" +YELLOW = "\033[93m" +RESET = "\033[0m" +BLACK_SQUARE = "■" +WHITE_SQUARE = "⬚" + + +def generate_attention_matrix_from_mask(words, mask, img_token="", sliding_window=None, token_type_ids=None): + """ + Generates an attention matrix from a given attention mask. + + Optionally applies a sliding window mask (e.g., for Gemma2/3) and + marks regions where image tokens occur based on the specified `img_token`. + """ + mask = mask.int() + if mask.ndim == 3: + mask = mask[0, :, :] + if mask.ndim == 4: + mask = mask[0, 0, :, :] + + n = len(words) + max_word_length = max(len(repr(word)) for word in words) + first_img_idx = 0 + output = [] + + for i, k in enumerate(words): + if k == img_token and not first_img_idx: + first_img_idx = i + mask[i, i] = 2 # Mark yellow regions + if first_img_idx > 0 and (k != img_token or i == n - 1): + if i == n - 1: + i += 1 + mask[first_img_idx:i, first_img_idx:i] = 2 # Mark yellow regions + first_img_idx = 0 + + # Generate sliding window mask (size = 4), excluding img_token + sliding_window_mask = None + if sliding_window is not None: + sliding_window_mask = [[1 if (0 <= i - j < sliding_window) else 0 for j in range(n)] for i in range(n)] + + row_dummy = " ".join( + f"{YELLOW}{BLACK_SQUARE}{RESET}" + if mask[0, j] + else f"{GREEN}{BLACK_SQUARE}{RESET}" + if 0 == j + else BLACK_SQUARE + if mask[0, j] + else WHITE_SQUARE + for j in range(n) + ) + + # Print headers + legend = f"{GREEN}{BLACK_SQUARE}{RESET}: i == j (diagonal) {YELLOW}{BLACK_SQUARE}{RESET}: token_type_ids" + output.append(" " + legend) + f_string = " " * (max_word_length + 5) + "Attention Matrix".ljust(len(row_dummy) // 2) + if sliding_window is not None: + f_string += "Sliding Window Mask" + output.append(f_string) + + vertical_header = [] + for idx, word in enumerate(words): + if mask[idx, idx] == 2: + vertical_header.append([f"{YELLOW}{k}{RESET}" for k in list(str(idx).rjust(len(str(n))))]) + else: + vertical_header.append(list(str(idx).rjust(len(str(n))))) + + vertical_header = list(map(list, zip(*vertical_header))) # Transpose + + for row in vertical_header: + output.append( + (max_word_length + 5) * " " + " ".join(row) + " | " + " ".join(row) + if sliding_window is not None + else "" + ) + + for i, word in enumerate(words): + word_repr = repr(word).ljust(max_word_length) + colored_word = f"{YELLOW}{word_repr}{RESET}" if img_token in word else word_repr + row_display = " ".join( + f"{YELLOW}{BLACK_SQUARE}{RESET}" + if img_token in words[j] and mask[i, j] and img_token in words[i] + else f"{GREEN}{BLACK_SQUARE}{RESET}" + if i == j + else BLACK_SQUARE + if mask[i, j] + else WHITE_SQUARE + for j in range(n) + ) + sliding_window_row = "" + if sliding_window is not None: + sliding_window_row = " ".join( + f"{YELLOW}{BLACK_SQUARE}{RESET}" + if img_token in words[j] and img_token in words[i] + else f"{GREEN}{BLACK_SQUARE}{RESET}" + if i == j + else BLACK_SQUARE + if sliding_window_mask[i][j] + else WHITE_SQUARE + for j in range(n) + ) + + output.append(f"{colored_word}: {str(i).rjust(2)} {row_display} | {sliding_window_row}") + + return "\n".join(output) + + +class AttentionMaskVisualizer: + def __init__(self, model_name: str): + config = AutoConfig.from_pretrained(model_name) + self.image_token = "" + if hasattr(config.get_text_config(), "sliding_window"): + config.sliding_window = 5 + try: + mapped_cls = _get_model_class(config, MODEL_MAPPING) + except Exception: + mapped_cls = _get_model_class(config, MODEL_FOR_PRETRAINING_MAPPING) + + if mapped_cls is None: + raise ValueError(f"Model name {model_name} is not supported for attention visualization") + self.mapped_cls = mapped_cls + + class _ModelWrapper(mapped_cls, nn.Module): + def __init__(self, config, model_name): + nn.Module.__init__(self) + self.dummy_module = nn.Linear(1, 1) + self.config = config + + self.model = _ModelWrapper(config, model_name) + self.model.to(config.torch_dtype) + self.repo_id = model_name + self.config = config + + def __call__(self, input_sentence: str, suffix=""): + self.visualize_attention_mask(input_sentence, suffix=suffix) + + def visualize_attention_mask(self, input_sentence: str, suffix=""): + model = self.model + kwargs = {} + if self.config.model_type in PROCESSOR_MAPPING_NAMES: + img = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg?download=true" + img = Image.open(requests.get(img, stream=True).raw) + processor = AutoProcessor.from_pretrained(self.repo_id, image_seq_length=5) + if hasattr(processor, "image_token"): + image_token = processor.image_token + else: + image_token = processor.tokenizer.convert_ids_to_tokens([processor.image_token_id])[0] + + if image_token: + input_sentence = input_sentence.replace("", image_token) + + inputs = processor(img, input_sentence, suffix=suffix, return_tensors="pt") + + self.image_token = processor.tokenizer.convert_ids_to_tokens([processor.image_token_id])[0] + + attention_mask = inputs["attention_mask"] + if "token_type_ids" in inputs: # TODO inspect signature of update causal mask + kwargs["token_type_ids"] = inputs["token_type_ids"] + tokens = processor.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]) + elif self.config.model_type in TOKENIZER_MAPPING_NAMES: + tokenizer = AutoTokenizer.from_pretrained(self.repo_id) + tokens = tokenizer.tokenize(input_sentence) + attention_mask = tokenizer(input_sentence, return_tensors="pt")["attention_mask"] + else: + raise ValueError(f"Model type {model.config.model_type} does not support attention visualization") + + model.config._attn_implementation = "eager" + model.train() + attention_mask = ~model._update_causal_mask( + attention_mask=attention_mask, + input_tensor=attention_mask.to(self.model.dtype), + cache_position=torch.arange(attention_mask.shape[1]), + past_key_values=None, + **kwargs, + ).bool() + top_bottom_border = "##" * ( + len(f"Attention visualization for {self.config.model_type} | {self.mapped_cls}") + 4 + ) # Box width adjusted to text length + side_border = "##" + print(f"\n{top_bottom_border}") + print( + "##" + + f" Attention visualization for \033[1m{self.config.model_type}:{self.repo_id}\033[0m {self.mapped_cls.__name__}".center( + len(top_bottom_border) + ) + + " " + + side_border + ) + print(f"{top_bottom_border}") + f_string = generate_attention_matrix_from_mask( + tokens, + attention_mask, + img_token=self.image_token, + sliding_window=getattr(self.config, "sliding_window", None), + token_type_ids=kwargs.get("token_type_ids", None), + ) + print(f_string) + print(f"{top_bottom_border}") diff --git a/lib/python3.12/site-packages/transformers/utils/backbone_utils.py b/lib/python3.12/site-packages/transformers/utils/backbone_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a6ab0046bc5c77d9a18450bd7518df16c31faabd --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/backbone_utils.py @@ -0,0 +1,377 @@ +# Copyright 2023 The HuggingFace Inc. team. +# +# 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. + +"""Collection of utils to be used by backbones and their components.""" + +import enum +import inspect +from collections.abc import Iterable +from typing import TYPE_CHECKING, Optional, Union + + +if TYPE_CHECKING: + from ..configuration_utils import PretrainedConfig + + +class BackboneType(enum.Enum): + TIMM = "timm" + TRANSFORMERS = "transformers" + + +def verify_out_features_out_indices( + out_features: Optional[Iterable[str]], out_indices: Optional[Iterable[int]], stage_names: Optional[Iterable[str]] +): + """ + Verify that out_indices and out_features are valid for the given stage_names. + """ + if stage_names is None: + raise ValueError("Stage_names must be set for transformers backbones") + + if out_features is not None: + if not isinstance(out_features, (list,)): + raise ValueError(f"out_features must be a list got {type(out_features)}") + if any(feat not in stage_names for feat in out_features): + raise ValueError(f"out_features must be a subset of stage_names: {stage_names} got {out_features}") + if len(out_features) != len(set(out_features)): + raise ValueError(f"out_features must not contain any duplicates, got {out_features}") + if out_features != (sorted_feats := [feat for feat in stage_names if feat in out_features]): + raise ValueError( + f"out_features must be in the same order as stage_names, expected {sorted_feats} got {out_features}" + ) + + if out_indices is not None: + if not isinstance(out_indices, list): + raise ValueError(f"out_indices must be a list, got {type(out_indices)}") + # Convert negative indices to their positive equivalent: [-1,] -> [len(stage_names) - 1,] + positive_indices = tuple(idx % len(stage_names) if idx < 0 else idx for idx in out_indices) + if any(idx for idx in positive_indices if idx not in range(len(stage_names))): + raise ValueError(f"out_indices must be valid indices for stage_names {stage_names}, got {out_indices}") + if len(positive_indices) != len(set(positive_indices)): + msg = f"out_indices must not contain any duplicates, got {out_indices}" + msg += f"(equivalent to {positive_indices}))" if positive_indices != out_indices else "" + raise ValueError(msg) + if positive_indices != tuple(sorted(positive_indices)): + sorted_negative = [idx for _, idx in sorted(zip(positive_indices, out_indices), key=lambda x: x[0])] + raise ValueError( + f"out_indices must be in the same order as stage_names, expected {sorted_negative} got {out_indices}" + ) + + if out_features is not None and out_indices is not None: + if len(out_features) != len(out_indices): + raise ValueError("out_features and out_indices should have the same length if both are set") + if out_features != [stage_names[idx] for idx in out_indices]: + raise ValueError("out_features and out_indices should correspond to the same stages if both are set") + + +def _align_output_features_output_indices( + out_features: Optional[list[str]], + out_indices: Optional[Union[list[int], tuple[int]]], + stage_names: list[str], +): + """ + Finds the corresponding `out_features` and `out_indices` for the given `stage_names`. + + The logic is as follows: + - `out_features` not set, `out_indices` set: `out_features` is set to the `out_features` corresponding to the + `out_indices`. + - `out_indices` not set, `out_features` set: `out_indices` is set to the `out_indices` corresponding to the + `out_features`. + - `out_indices` and `out_features` not set: `out_indices` and `out_features` are set to the last stage. + - `out_indices` and `out_features` set: input `out_indices` and `out_features` are returned. + + Args: + out_features (`List[str]`): The names of the features for the backbone to output. + out_indices (`List[int]` or `Tuple[int]`): The indices of the features for the backbone to output. + stage_names (`List[str]`): The names of the stages of the backbone. + """ + if out_indices is None and out_features is None: + out_indices = [len(stage_names) - 1] + out_features = [stage_names[-1]] + elif out_indices is None and out_features is not None: + out_indices = [stage_names.index(layer) for layer in out_features] + elif out_features is None and out_indices is not None: + out_features = [stage_names[idx] for idx in out_indices] + return out_features, out_indices + + +def get_aligned_output_features_output_indices( + out_features: Optional[list[str]], + out_indices: Optional[Union[list[int], tuple[int]]], + stage_names: list[str], +) -> tuple[list[str], list[int]]: + """ + Get the `out_features` and `out_indices` so that they are aligned. + + The logic is as follows: + - `out_features` not set, `out_indices` set: `out_features` is set to the `out_features` corresponding to the + `out_indices`. + - `out_indices` not set, `out_features` set: `out_indices` is set to the `out_indices` corresponding to the + `out_features`. + - `out_indices` and `out_features` not set: `out_indices` and `out_features` are set to the last stage. + - `out_indices` and `out_features` set: they are verified to be aligned. + + Args: + out_features (`List[str]`): The names of the features for the backbone to output. + out_indices (`List[int]` or `Tuple[int]`): The indices of the features for the backbone to output. + stage_names (`List[str]`): The names of the stages of the backbone. + """ + out_indices = list(out_indices) if out_indices is not None else None + # First verify that the out_features and out_indices are valid + verify_out_features_out_indices(out_features=out_features, out_indices=out_indices, stage_names=stage_names) + output_features, output_indices = _align_output_features_output_indices( + out_features=out_features, out_indices=out_indices, stage_names=stage_names + ) + # Verify that the aligned out_features and out_indices are valid + verify_out_features_out_indices(out_features=output_features, out_indices=output_indices, stage_names=stage_names) + return output_features, output_indices + + +class BackboneMixin: + backbone_type: Optional[BackboneType] = None + + def _init_timm_backbone(self, config) -> None: + """ + Initialize the backbone model from timm The backbone must already be loaded to self._backbone + """ + if getattr(self, "_backbone", None) is None: + raise ValueError("self._backbone must be set before calling _init_timm_backbone") + + # These will diagree with the defaults for the transformers models e.g. for resnet50 + # the transformer model has out_features = ['stem', 'stage1', 'stage2', 'stage3', 'stage4'] + # the timm model has out_features = ['act', 'layer1', 'layer2', 'layer3', 'layer4'] + self.stage_names = [stage["module"] for stage in self._backbone.feature_info.info] + self.num_features = [stage["num_chs"] for stage in self._backbone.feature_info.info] + + # In some timm versions, out_indices reflects the input type of out_indices on the `create_model` call, + # in later versions >= 1, it is always a tuple + out_indices = list(self._backbone.feature_info.out_indices) + out_features = self._backbone.feature_info.module_name() + + # We verify the out indices and out features are valid + verify_out_features_out_indices( + out_features=out_features, out_indices=out_indices, stage_names=self.stage_names + ) + self._out_features, self._out_indices = out_features, out_indices + + def _init_transformers_backbone(self, config) -> None: + stage_names = getattr(config, "stage_names") + out_features = getattr(config, "out_features", None) + out_indices = getattr(config, "out_indices", None) + + self.stage_names = stage_names + self._out_features, self._out_indices = get_aligned_output_features_output_indices( + out_features=out_features, out_indices=out_indices, stage_names=stage_names + ) + # Number of channels for each stage. This is set in the transformer backbone model init + self.num_features = None + + def _init_backbone(self, config) -> None: + """ + Method to initialize the backbone. This method is called by the constructor of the base class after the + pretrained model weights have been loaded. + """ + self.config = config + + self.use_timm_backbone = getattr(config, "use_timm_backbone", False) + self.backbone_type = BackboneType.TIMM if self.use_timm_backbone else BackboneType.TRANSFORMERS + + if self.backbone_type == BackboneType.TIMM: + self._init_timm_backbone(config) + elif self.backbone_type == BackboneType.TRANSFORMERS: + self._init_transformers_backbone(config) + else: + raise ValueError(f"backbone_type {self.backbone_type} not supported.") + + @property + def out_features(self): + return self._out_features + + @out_features.setter + def out_features(self, out_features: list[str]): + """ + Set the out_features attribute. This will also update the out_indices attribute to match the new out_features. + """ + self._out_features, self._out_indices = get_aligned_output_features_output_indices( + out_features=out_features, out_indices=None, stage_names=self.stage_names + ) + + @property + def out_indices(self): + return self._out_indices + + @out_indices.setter + def out_indices(self, out_indices: Union[tuple[int], list[int]]): + """ + Set the out_indices attribute. This will also update the out_features attribute to match the new out_indices. + """ + self._out_features, self._out_indices = get_aligned_output_features_output_indices( + out_features=None, out_indices=out_indices, stage_names=self.stage_names + ) + + @property + def out_feature_channels(self): + # the current backbones will output the number of channels for each stage + # even if that stage is not in the out_features list. + return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names)} + + @property + def channels(self): + return [self.out_feature_channels[name] for name in self.out_features] + + def forward_with_filtered_kwargs(self, *args, **kwargs): + signature = dict(inspect.signature(self.forward).parameters) + filtered_kwargs = {k: v for k, v in kwargs.items() if k in signature} + return self(*args, **filtered_kwargs) + + def forward( + self, + pixel_values, + output_hidden_states: Optional[bool] = None, + output_attentions: Optional[bool] = None, + return_dict: Optional[bool] = None, + ): + raise NotImplementedError("This method should be implemented by the derived class.") + + def to_dict(self): + """ + Serializes this instance to a Python dictionary. Override the default `to_dict()` from `PretrainedConfig` to + include the `out_features` and `out_indices` attributes. + """ + output = super().to_dict() + output["out_features"] = output.pop("_out_features") + output["out_indices"] = output.pop("_out_indices") + return output + + +class BackboneConfigMixin: + """ + A Mixin to support handling the `out_features` and `out_indices` attributes for the backbone configurations. + """ + + @property + def out_features(self): + return self._out_features + + @out_features.setter + def out_features(self, out_features: list[str]): + """ + Set the out_features attribute. This will also update the out_indices attribute to match the new out_features. + """ + self._out_features, self._out_indices = get_aligned_output_features_output_indices( + out_features=out_features, out_indices=None, stage_names=self.stage_names + ) + + @property + def out_indices(self): + return self._out_indices + + @out_indices.setter + def out_indices(self, out_indices: Union[tuple[int], list[int]]): + """ + Set the out_indices attribute. This will also update the out_features attribute to match the new out_indices. + """ + self._out_features, self._out_indices = get_aligned_output_features_output_indices( + out_features=None, out_indices=out_indices, stage_names=self.stage_names + ) + + def to_dict(self): + """ + Serializes this instance to a Python dictionary. Override the default `to_dict()` from `PretrainedConfig` to + include the `out_features` and `out_indices` attributes. + """ + output = super().to_dict() + output["out_features"] = output.pop("_out_features") + output["out_indices"] = output.pop("_out_indices") + return output + + +def load_backbone(config): + """ + Loads the backbone model from a config object. + + If the config is from the backbone model itself, then we return a backbone model with randomly initialized + weights. + + If the config is from the parent model of the backbone model itself, then we load the pretrained backbone weights + if specified. + """ + from transformers import AutoBackbone, AutoConfig + + backbone_config = getattr(config, "backbone_config", None) + use_timm_backbone = getattr(config, "use_timm_backbone", None) + use_pretrained_backbone = getattr(config, "use_pretrained_backbone", None) + backbone_checkpoint = getattr(config, "backbone", None) + backbone_kwargs = getattr(config, "backbone_kwargs", None) + backbone_kwargs = {} if backbone_kwargs is None else backbone_kwargs + + if backbone_kwargs and backbone_config is not None: + raise ValueError("You can't specify both `backbone_kwargs` and `backbone_config`.") + + # If there is a backbone_config and a backbone checkpoint, and use_pretrained_backbone=False then the desired + # behaviour is ill-defined: do you want to load from the checkpoint's config or the backbone_config? + if backbone_config is not None and backbone_checkpoint is not None and use_pretrained_backbone is not None: + raise ValueError("Cannot specify both config.backbone_config and config.backbone") + + # If any of thhe following are set, then the config passed in is from a model which contains a backbone. + if ( + backbone_config is None + and use_timm_backbone is None + and backbone_checkpoint is None + and backbone_checkpoint is None + ): + return AutoBackbone.from_config(config=config, **backbone_kwargs) + + # config from the parent model that has a backbone + if use_timm_backbone: + if backbone_checkpoint is None: + raise ValueError("config.backbone must be set if use_timm_backbone is True") + # Because of how timm backbones were originally added to models, we need to pass in use_pretrained_backbone + # to determine whether to load the pretrained weights. + backbone = AutoBackbone.from_pretrained( + backbone_checkpoint, + use_timm_backbone=use_timm_backbone, + use_pretrained_backbone=use_pretrained_backbone, + **backbone_kwargs, + ) + elif use_pretrained_backbone: + if backbone_checkpoint is None: + raise ValueError("config.backbone must be set if use_pretrained_backbone is True") + backbone = AutoBackbone.from_pretrained(backbone_checkpoint, **backbone_kwargs) + else: + if backbone_config is None and backbone_checkpoint is None: + raise ValueError("Either config.backbone_config or config.backbone must be set") + if backbone_config is None: + backbone_config = AutoConfig.from_pretrained(backbone_checkpoint, **backbone_kwargs) + backbone = AutoBackbone.from_config(config=backbone_config) + return backbone + + +def verify_backbone_config_arguments( + use_timm_backbone: bool, + use_pretrained_backbone: bool, + backbone: Optional[str], + backbone_config: Optional[Union[dict, "PretrainedConfig"]], + backbone_kwargs: Optional[dict], +): + """ + Verify that the config arguments to be passed to load_backbone are valid + """ + if backbone_config is not None and backbone is not None: + raise ValueError("You can't specify both `backbone` and `backbone_config`.") + + if backbone_config is not None and use_timm_backbone: + raise ValueError("You can't specify both `backbone_config` and `use_timm_backbone`.") + + if backbone_kwargs is not None and backbone_kwargs and backbone_config is not None: + raise ValueError("You can't specify both `backbone_kwargs` and `backbone_config`.") diff --git a/lib/python3.12/site-packages/transformers/utils/bitsandbytes.py b/lib/python3.12/site-packages/transformers/utils/bitsandbytes.py new file mode 100644 index 0000000000000000000000000000000000000000..71707cf5659909f7e28f939e91df6c48e64aba43 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/bitsandbytes.py @@ -0,0 +1,28 @@ +# Copyright 2023 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. +import warnings + + +warnings.warn( + "transformers.utils.bitsandbytes module is deprecated and will be removed in a future version. Please import bitsandbytes modules directly from transformers.integrations", + FutureWarning, +) + +from ..integrations import ( # noqa + get_keys_to_not_convert, + replace_8bit_linear, + replace_with_bnb_linear, + set_module_8bit_tensor_to_device, + set_module_quantized_tensor_to_device, +) diff --git a/lib/python3.12/site-packages/transformers/utils/chat_template_utils.py b/lib/python3.12/site-packages/transformers/utils/chat_template_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f96b5d8ddecab445a6706d9562de3e1c6af43d0c --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/chat_template_utils.py @@ -0,0 +1,435 @@ +# 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. + +import inspect +import json +import re +import types +from contextlib import contextmanager +from datetime import datetime +from functools import lru_cache +from typing import Any, Callable, Optional, Union, get_args, get_origin, get_type_hints + +from packaging import version + +from .import_utils import is_jinja_available, is_torch_available, is_vision_available + + +if is_jinja_available(): + import jinja2 + from jinja2.ext import Extension + from jinja2.sandbox import ImmutableSandboxedEnvironment +else: + jinja2 = None + +if is_vision_available(): + from PIL.Image import Image + +if is_torch_available(): + from torch import Tensor + + +BASIC_TYPES = (int, float, str, bool, Any, type(None), ...) +# Extracts the initial segment of the docstring, containing the function description +description_re = re.compile(r"^(.*?)[\n\s]*(Args:|Returns:|Raises:|\Z)", re.DOTALL) +# Extracts the Args: block from the docstring +args_re = re.compile(r"\n\s*Args:\n\s*(.*?)[\n\s]*(Returns:|Raises:|\Z)", re.DOTALL) +# Splits the Args: block into individual arguments +args_split_re = re.compile( + r""" +(?:^|\n) # Match the start of the args block, or a newline +\s*(\w+):\s* # Capture the argument name and strip spacing +(.*?)\s* # Capture the argument description, which can span multiple lines, and strip trailing spacing +(?=\n\s*\w+:|\Z) # Stop when you hit the next argument or the end of the block +""", + re.DOTALL | re.VERBOSE, +) +# Extracts the Returns: block from the docstring, if present. Note that most chat templates ignore the return type/doc! +returns_re = re.compile(r"\n\s*Returns:\n\s*(.*?)[\n\s]*(Raises:|\Z)", re.DOTALL) + + +class TypeHintParsingException(Exception): + """Exception raised for errors in parsing type hints to generate JSON schemas""" + + pass + + +class DocstringParsingException(Exception): + """Exception raised for errors in parsing docstrings to generate JSON schemas""" + + pass + + +def _get_json_schema_type(param_type: str) -> dict[str, str]: + type_mapping = { + int: {"type": "integer"}, + float: {"type": "number"}, + str: {"type": "string"}, + bool: {"type": "boolean"}, + type(None): {"type": "null"}, + Any: {}, + } + if is_vision_available(): + type_mapping[Image] = {"type": "image"} + if is_torch_available(): + type_mapping[Tensor] = {"type": "audio"} + return type_mapping.get(param_type, {"type": "object"}) + + +def _parse_type_hint(hint: str) -> dict: + origin = get_origin(hint) + args = get_args(hint) + + if origin is None: + try: + return _get_json_schema_type(hint) + except KeyError: + raise TypeHintParsingException( + "Couldn't parse this type hint, likely due to a custom class or object: ", hint + ) + + elif origin is Union or (hasattr(types, "UnionType") and origin is types.UnionType): + # Recurse into each of the subtypes in the Union, except None, which is handled separately at the end + subtypes = [_parse_type_hint(t) for t in args if t is not type(None)] + if len(subtypes) == 1: + # A single non-null type can be expressed directly + return_dict = subtypes[0] + elif all(isinstance(subtype["type"], str) for subtype in subtypes): + # A union of basic types can be expressed as a list in the schema + return_dict = {"type": sorted([subtype["type"] for subtype in subtypes])} + else: + # A union of more complex types requires "anyOf" + return_dict = {"anyOf": subtypes} + if type(None) in args: + return_dict["nullable"] = True + return return_dict + + elif origin is list: + if not args: + return {"type": "array"} + else: + # Lists can only have a single type argument, so recurse into it + return {"type": "array", "items": _parse_type_hint(args[0])} + + elif origin is tuple: + if not args: + return {"type": "array"} + if len(args) == 1: + raise TypeHintParsingException( + f"The type hint {str(hint).replace('typing.', '')} is a Tuple with a single element, which " + "we do not automatically convert to JSON schema as it is rarely necessary. If this input can contain " + "more than one element, we recommend " + "using a List[] type instead, or if it really is a single element, remove the Tuple[] wrapper and just " + "pass the element directly." + ) + if ... in args: + raise TypeHintParsingException( + "Conversion of '...' is not supported in Tuple type hints. " + "Use List[] types for variable-length" + " inputs instead." + ) + return {"type": "array", "prefixItems": [_parse_type_hint(t) for t in args]} + + elif origin is dict: + # The JSON equivalent to a dict is 'object', which mandates that all keys are strings + # However, we can specify the type of the dict values with "additionalProperties" + out = {"type": "object"} + if len(args) == 2: + out["additionalProperties"] = _parse_type_hint(args[1]) + return out + + raise TypeHintParsingException("Couldn't parse this type hint, likely due to a custom class or object: ", hint) + + +def _convert_type_hints_to_json_schema(func: Callable) -> dict: + type_hints = get_type_hints(func) + signature = inspect.signature(func) + required = [] + for param_name, param in signature.parameters.items(): + if param.annotation == inspect.Parameter.empty: + raise TypeHintParsingException(f"Argument {param.name} is missing a type hint in function {func.__name__}") + if param.default == inspect.Parameter.empty: + required.append(param_name) + + properties = {} + for param_name, param_type in type_hints.items(): + properties[param_name] = _parse_type_hint(param_type) + + schema = {"type": "object", "properties": properties} + if required: + schema["required"] = required + + return schema + + +def parse_google_format_docstring(docstring: str) -> tuple[Optional[str], Optional[dict], Optional[str]]: + """ + Parses a Google-style docstring to extract the function description, + argument descriptions, and return description. + + Args: + docstring (str): The docstring to parse. + + Returns: + The function description, arguments, and return description. + """ + + # Extract the sections + description_match = description_re.search(docstring) + args_match = args_re.search(docstring) + returns_match = returns_re.search(docstring) + + # Clean and store the sections + description = description_match.group(1).strip() if description_match else None + docstring_args = args_match.group(1).strip() if args_match else None + returns = returns_match.group(1).strip() if returns_match else None + + # Parsing the arguments into a dictionary + if docstring_args is not None: + docstring_args = "\n".join([line for line in docstring_args.split("\n") if line.strip()]) # Remove blank lines + matches = args_split_re.findall(docstring_args) + args_dict = {match[0]: re.sub(r"\s*\n+\s*", " ", match[1].strip()) for match in matches} + else: + args_dict = {} + + return description, args_dict, returns + + +def get_json_schema(func: Callable) -> dict: + """ + This function generates a JSON schema for a given function, based on its docstring and type hints. This is + mostly used for passing lists of tools to a chat template. The JSON schema contains the name and description of + the function, as well as the names, types and descriptions for each of its arguments. `get_json_schema()` requires + that the function has a docstring, and that each argument has a description in the docstring, in the standard + Google docstring format shown below. It also requires that all the function arguments have a valid Python type hint. + + Although it is not required, a `Returns` block can also be added, which will be included in the schema. This is + optional because most chat templates ignore the return value of the function. + + Args: + func: The function to generate a JSON schema for. + + Returns: + A dictionary containing the JSON schema for the function. + + Examples: + ```python + >>> def multiply(x: float, y: float): + >>> ''' + >>> A function that multiplies two numbers + >>> + >>> Args: + >>> x: The first number to multiply + >>> y: The second number to multiply + >>> ''' + >>> return x * y + >>> + >>> print(get_json_schema(multiply)) + { + "name": "multiply", + "description": "A function that multiplies two numbers", + "parameters": { + "type": "object", + "properties": { + "x": {"type": "number", "description": "The first number to multiply"}, + "y": {"type": "number", "description": "The second number to multiply"} + }, + "required": ["x", "y"] + } + } + ``` + + The general use for these schemas is that they are used to generate tool descriptions for chat templates that + support them, like so: + + ```python + >>> from transformers import AutoTokenizer + >>> from transformers.utils import get_json_schema + >>> + >>> def multiply(x: float, y: float): + >>> ''' + >>> A function that multiplies two numbers + >>> + >>> Args: + >>> x: The first number to multiply + >>> y: The second number to multiply + >>> return x * y + >>> ''' + >>> + >>> multiply_schema = get_json_schema(multiply) + >>> tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-v01") + >>> messages = [{"role": "user", "content": "What is 179 x 4571?"}] + >>> formatted_chat = tokenizer.apply_chat_template( + >>> messages, + >>> tools=[multiply_schema], + >>> chat_template="tool_use", + >>> return_dict=True, + >>> return_tensors="pt", + >>> add_generation_prompt=True + >>> ) + >>> # The formatted chat can now be passed to model.generate() + ``` + + Each argument description can also have an optional `(choices: ...)` block at the end, such as + `(choices: ["tea", "coffee"])`, which will be parsed into an `enum` field in the schema. Note that this will + only be parsed correctly if it is at the end of the line: + + ```python + >>> def drink_beverage(beverage: str): + >>> ''' + >>> A function that drinks a beverage + >>> + >>> Args: + >>> beverage: The beverage to drink (choices: ["tea", "coffee"]) + >>> ''' + >>> pass + >>> + >>> print(get_json_schema(drink_beverage)) + ``` + { + 'name': 'drink_beverage', + 'description': 'A function that drinks a beverage', + 'parameters': { + 'type': 'object', + 'properties': { + 'beverage': { + 'type': 'string', + 'enum': ['tea', 'coffee'], + 'description': 'The beverage to drink' + } + }, + 'required': ['beverage'] + } + } + """ + doc = inspect.getdoc(func) + if not doc: + raise DocstringParsingException( + f"Cannot generate JSON schema for {func.__name__} because it has no docstring!" + ) + doc = doc.strip() + main_doc, param_descriptions, return_doc = parse_google_format_docstring(doc) + + json_schema = _convert_type_hints_to_json_schema(func) + if (return_dict := json_schema["properties"].pop("return", None)) is not None: + if return_doc is not None: # We allow a missing return docstring since most templates ignore it + return_dict["description"] = return_doc + for arg, schema in json_schema["properties"].items(): + if arg not in param_descriptions: + raise DocstringParsingException( + f"Cannot generate JSON schema for {func.__name__} because the docstring has no description for the argument '{arg}'" + ) + desc = param_descriptions[arg] + enum_choices = re.search(r"\(choices:\s*(.*?)\)\s*$", desc, flags=re.IGNORECASE) + if enum_choices: + schema["enum"] = [c.strip() for c in json.loads(enum_choices.group(1))] + desc = enum_choices.string[: enum_choices.start()].strip() + schema["description"] = desc + + output = {"name": func.__name__, "description": main_doc, "parameters": json_schema} + if return_dict is not None: + output["return"] = return_dict + return {"type": "function", "function": output} + + +def _render_with_assistant_indices( + compiled_template, messages, tools, documents, add_generation_prompt, **template_kwargs +): + rendered_blocks = [] + generation_indices = [] + with compiled_template.environment.activate_tracker(rendered_blocks, generation_indices): + for block in compiled_template.generate( + messages=messages, + tools=tools, + documents=documents, + add_generation_prompt=add_generation_prompt, + **template_kwargs, + ): + rendered_blocks.append(block) + rendered_chat = "".join(rendered_blocks) + return rendered_chat, generation_indices + + +@lru_cache +def _compile_jinja_template(chat_template): + if not is_jinja_available(): + raise ImportError( + "apply_chat_template requires jinja2 to be installed. Please install it using `pip install jinja2`." + ) + + class AssistantTracker(Extension): + # This extension is used to track the indices of assistant-generated tokens in the rendered chat + tags = {"generation"} + + def __init__(self, environment: ImmutableSandboxedEnvironment): + # The class is only initiated by jinja. + super().__init__(environment) + environment.extend(activate_tracker=self.activate_tracker) + self._rendered_blocks = None + self._generation_indices = None + + def parse(self, parser: jinja2.parser.Parser) -> jinja2.nodes.CallBlock: + lineno = next(parser.stream).lineno + body = parser.parse_statements(["name:endgeneration"], drop_needle=True) + return jinja2.nodes.CallBlock(self.call_method("_generation_support"), [], [], body).set_lineno(lineno) + + @jinja2.pass_eval_context + def _generation_support(self, context: jinja2.nodes.EvalContext, caller: jinja2.runtime.Macro) -> str: + rv = caller() + if self.is_active(): + # Only track generation indices if the tracker is active + start_index = len("".join(self._rendered_blocks)) + end_index = start_index + len(rv) + self._generation_indices.append((start_index, end_index)) + return rv + + def is_active(self) -> bool: + return self._rendered_blocks or self._generation_indices + + @contextmanager + def activate_tracker(self, rendered_blocks: list[int], generation_indices: list[int]): + try: + if self.is_active(): + raise ValueError("AssistantTracker should not be reused before closed") + self._rendered_blocks = rendered_blocks + self._generation_indices = generation_indices + + yield + finally: + self._rendered_blocks = None + self._generation_indices = None + + if version.parse(jinja2.__version__) < version.parse("3.1.0"): + raise ImportError( + f"apply_chat_template requires jinja2>=3.1.0 to be installed. Your version is {jinja2.__version__}." + ) + + def raise_exception(message): + raise jinja2.exceptions.TemplateError(message) + + def tojson(x, ensure_ascii=False, indent=None, separators=None, sort_keys=False): + # We override the built-in tojson filter because Jinja's default filter escapes HTML characters + # We also expose some options like custom indents and separators + return json.dumps(x, ensure_ascii=ensure_ascii, indent=indent, separators=separators, sort_keys=sort_keys) + + def strftime_now(format): + return datetime.now().strftime(format) + + jinja_env = ImmutableSandboxedEnvironment( + trim_blocks=True, lstrip_blocks=True, extensions=[AssistantTracker, jinja2.ext.loopcontrols] + ) + jinja_env.filters["tojson"] = tojson + jinja_env.globals["raise_exception"] = raise_exception + jinja_env.globals["strftime_now"] = strftime_now + return jinja_env.from_string(chat_template) diff --git a/lib/python3.12/site-packages/transformers/utils/constants.py b/lib/python3.12/site-packages/transformers/utils/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..fefd1b4601da04e073ff2880099ccaf87d0b1666 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/constants.py @@ -0,0 +1,6 @@ +IMAGENET_DEFAULT_MEAN = [0.485, 0.456, 0.406] +IMAGENET_DEFAULT_STD = [0.229, 0.224, 0.225] +IMAGENET_STANDARD_MEAN = [0.5, 0.5, 0.5] +IMAGENET_STANDARD_STD = [0.5, 0.5, 0.5] +OPENAI_CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] +OPENAI_CLIP_STD = [0.26862954, 0.26130258, 0.27577711] diff --git a/lib/python3.12/site-packages/transformers/utils/deprecation.py b/lib/python3.12/site-packages/transformers/utils/deprecation.py new file mode 100644 index 0000000000000000000000000000000000000000..064decb14dba8f7a0a9eb860b450823c425ca43a --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/deprecation.py @@ -0,0 +1,176 @@ +# 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. +import inspect +import warnings +from functools import wraps +from typing import Optional + +import packaging.version + +from .. import __version__ +from . import ExplicitEnum, is_torch_available, is_torchdynamo_compiling + + +# This is needed in case we deprecate a kwarg of a function/method being compiled +if is_torch_available(): + import torch # noqa: F401 + + +class Action(ExplicitEnum): + NONE = "none" + NOTIFY = "notify" + NOTIFY_ALWAYS = "notify_always" + RAISE = "raise" + + +def deprecate_kwarg( + old_name: str, + version: str, + new_name: Optional[str] = None, + warn_if_greater_or_equal_version: bool = False, + raise_if_greater_or_equal_version: bool = False, + raise_if_both_names: bool = False, + additional_message: Optional[str] = None, +): + """ + Function or method decorator to notify users about deprecated keyword arguments, replacing them with a new name if specified. + Note that is decorator is `torch.compile`-safe, i.e. it will not cause graph breaks (but no warning will be displayed if compiling). + + This decorator allows you to: + - Notify users when a keyword argument is deprecated. + - Automatically replace deprecated keyword arguments with new ones. + - Raise an error if deprecated arguments are used, depending on the specified conditions. + + By default, the decorator notifies the user about the deprecated argument while the `transformers.__version__` < specified `version` + in the decorator. To keep notifications with any version `warn_if_greater_or_equal_version=True` can be set. + + Parameters: + old_name (`str`): + Name of the deprecated keyword argument. + version (`str`): + The version in which the keyword argument was (or will be) deprecated. + new_name (`Optional[str]`, *optional*): + The new name for the deprecated keyword argument. If specified, the deprecated keyword argument will be replaced with this new name. + warn_if_greater_or_equal_version (`bool`, *optional*, defaults to `False`): + Whether to show warning if current `transformers` version is greater or equal to the deprecated version. + raise_if_greater_or_equal_version (`bool`, *optional*, defaults to `False`): + Whether to raise `ValueError` if current `transformers` version is greater or equal to the deprecated version. + raise_if_both_names (`bool`, *optional*, defaults to `False`): + Whether to raise `ValueError` if both deprecated and new keyword arguments are set. + additional_message (`Optional[str]`, *optional*): + An additional message to append to the default deprecation message. + + Raises: + ValueError: + If raise_if_greater_or_equal_version is True and the current version is greater than or equal to the deprecated version, or if raise_if_both_names is True and both old and new keyword arguments are provided. + + Returns: + Callable: + A wrapped function that handles the deprecated keyword arguments according to the specified parameters. + + Example usage with renaming argument: + + ```python + @deprecate_kwarg("reduce_labels", new_name="do_reduce_labels", version="6.0.0") + def my_function(do_reduce_labels): + print(do_reduce_labels) + + my_function(reduce_labels=True) # Will show a deprecation warning and use do_reduce_labels=True + ``` + + Example usage without renaming argument: + + ```python + @deprecate_kwarg("max_size", version="6.0.0") + def my_function(max_size): + print(max_size) + + my_function(max_size=1333) # Will show a deprecation warning + ``` + + """ + + deprecated_version = packaging.version.parse(version) + current_version = packaging.version.parse(__version__) + is_greater_or_equal_version = current_version >= deprecated_version + + if is_greater_or_equal_version: + version_message = f"and removed starting from version {version}" + else: + version_message = f"and will be removed in version {version}" + + def wrapper(func): + # Required for better warning message + sig = inspect.signature(func) + function_named_args = set(sig.parameters.keys()) + is_instance_method = "self" in function_named_args + is_class_method = "cls" in function_named_args + + @wraps(func) + def wrapped_func(*args, **kwargs): + # Get class + function name (just for better warning message) + func_name = func.__name__ + if is_instance_method: + func_name = f"{args[0].__class__.__name__}.{func_name}" + elif is_class_method: + func_name = f"{args[0].__name__}.{func_name}" + + minimum_action = Action.NONE + message = None + + # deprecated kwarg and its new version are set for function call -> replace it with new name + if old_name in kwargs and new_name in kwargs: + minimum_action = Action.RAISE if raise_if_both_names else Action.NOTIFY_ALWAYS + message = f"Both `{old_name}` and `{new_name}` are set for `{func_name}`. Using `{new_name}={kwargs[new_name]}` and ignoring deprecated `{old_name}={kwargs[old_name]}`." + kwargs.pop(old_name) + + # only deprecated kwarg is set for function call -> replace it with new name + elif old_name in kwargs and new_name is not None and new_name not in kwargs: + minimum_action = Action.NOTIFY + message = f"`{old_name}` is deprecated {version_message} for `{func_name}`. Use `{new_name}` instead." + kwargs[new_name] = kwargs.pop(old_name) + + # deprecated kwarg is not set for function call and new name is not specified -> just notify + elif old_name in kwargs: + minimum_action = Action.NOTIFY + message = f"`{old_name}` is deprecated {version_message} for `{func_name}`." + + if message is not None and additional_message is not None: + message = f"{message} {additional_message}" + + # update minimum_action if argument is ALREADY deprecated (current version >= deprecated version) + if is_greater_or_equal_version: + # change to (NOTIFY, NOTIFY_ALWAYS) -> RAISE if specified + # in case we want to raise error for already deprecated arguments + if raise_if_greater_or_equal_version and minimum_action != Action.NONE: + minimum_action = Action.RAISE + + # change to NOTIFY -> NONE if specified (NOTIFY_ALWAYS can't be changed to NONE) + # in case we want to ignore notifications for already deprecated arguments + elif not warn_if_greater_or_equal_version and minimum_action == Action.NOTIFY: + minimum_action = Action.NONE + + # raise error or notify user + if minimum_action == Action.RAISE: + raise ValueError(message) + # If we are compiling, we do not raise the warning as it would break compilation + elif minimum_action in (Action.NOTIFY, Action.NOTIFY_ALWAYS) and not is_torchdynamo_compiling(): + # DeprecationWarning is ignored by default, so we use FutureWarning instead + warnings.warn(message, FutureWarning, stacklevel=2) + + return func(*args, **kwargs) + + return wrapped_func + + return wrapper diff --git a/lib/python3.12/site-packages/transformers/utils/doc.py b/lib/python3.12/site-packages/transformers/utils/doc.py new file mode 100644 index 0000000000000000000000000000000000000000..f01ffc28442cd5a5dcda3ea55d3ea08eb249b54f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/doc.py @@ -0,0 +1,1218 @@ +# Copyright 2022 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. +""" +Doc utilities: Utilities related to documentation +""" + +import functools +import inspect +import re +import textwrap +import types + + +def get_docstring_indentation_level(func): + """Return the indentation level of the start of the docstring of a class or function (or method).""" + # We assume classes are always defined in the global scope + if inspect.isclass(func): + return 4 + source = inspect.getsource(func) + first_line = source.splitlines()[0] + function_def_level = len(first_line) - len(first_line.lstrip()) + return 4 + function_def_level + + +def add_start_docstrings(*docstr): + def docstring_decorator(fn): + fn.__doc__ = "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "") + return fn + + return docstring_decorator + + +def add_start_docstrings_to_model_forward(*docstr): + def docstring_decorator(fn): + class_name = f"[`{fn.__qualname__.split('.')[0]}`]" + intro = rf""" The {class_name} forward method, overrides the `__call__` special method. + + + + Although the recipe for forward pass needs to be defined within this function, one should call the [`Module`] + instance afterwards instead of this since the former takes care of running the pre and post processing steps while + the latter silently ignores them. + + +""" + + correct_indentation = get_docstring_indentation_level(fn) + current_doc = fn.__doc__ if fn.__doc__ is not None else "" + try: + first_non_empty = next(line for line in current_doc.splitlines() if line.strip() != "") + doc_indentation = len(first_non_empty) - len(first_non_empty.lstrip()) + except StopIteration: + doc_indentation = correct_indentation + + docs = docstr + # In this case, the correct indentation level (class method, 2 Python levels) was respected, and we should + # correctly reindent everything. Otherwise, the doc uses a single indentation level + if doc_indentation == 4 + correct_indentation: + docs = [textwrap.indent(textwrap.dedent(doc), " " * correct_indentation) for doc in docstr] + intro = textwrap.indent(textwrap.dedent(intro), " " * correct_indentation) + + docstring = "".join(docs) + current_doc + fn.__doc__ = intro + docstring + return fn + + return docstring_decorator + + +def add_end_docstrings(*docstr): + def docstring_decorator(fn): + fn.__doc__ = (fn.__doc__ if fn.__doc__ is not None else "") + "".join(docstr) + return fn + + return docstring_decorator + + +PT_RETURN_INTRODUCTION = r""" + Returns: + [`{full_output_type}`] or `tuple(torch.FloatTensor)`: A [`{full_output_type}`] or a tuple of + `torch.FloatTensor` (if `return_dict=False` is passed or when `config.return_dict=False`) comprising various + elements depending on the configuration ([`{config_class}`]) and inputs. + +""" + + +TF_RETURN_INTRODUCTION = r""" + Returns: + [`{full_output_type}`] or `tuple(tf.Tensor)`: A [`{full_output_type}`] or a tuple of `tf.Tensor` (if + `return_dict=False` is passed or when `config.return_dict=False`) comprising various elements depending on the + configuration ([`{config_class}`]) and inputs. + +""" + + +def _get_indent(t): + """Returns the indentation in the first line of t""" + search = re.search(r"^(\s*)\S", t) + return "" if search is None else search.groups()[0] + + +def _convert_output_args_doc(output_args_doc): + """Convert output_args_doc to display properly.""" + # Split output_arg_doc in blocks argument/description + indent = _get_indent(output_args_doc) + blocks = [] + current_block = "" + for line in output_args_doc.split("\n"): + # If the indent is the same as the beginning, the line is the name of new arg. + if _get_indent(line) == indent: + if len(current_block) > 0: + blocks.append(current_block[:-1]) + current_block = f"{line}\n" + else: + # Otherwise it's part of the description of the current arg. + # We need to remove 2 spaces to the indentation. + current_block += f"{line[2:]}\n" + blocks.append(current_block[:-1]) + + # Format each block for proper rendering + for i in range(len(blocks)): + blocks[i] = re.sub(r"^(\s+)(\S+)(\s+)", r"\1- **\2**\3", blocks[i]) + blocks[i] = re.sub(r":\s*\n\s*(\S)", r" -- \1", blocks[i]) + + return "\n".join(blocks) + + +def _prepare_output_docstrings(output_type, config_class, min_indent=None): + """ + Prepares the return part of the docstring using `output_type`. + """ + output_docstring = output_type.__doc__ + + # Remove the head of the docstring to keep the list of args only + lines = output_docstring.split("\n") + i = 0 + while i < len(lines) and re.search(r"^\s*(Args|Parameters):\s*$", lines[i]) is None: + i += 1 + if i < len(lines): + params_docstring = "\n".join(lines[(i + 1) :]) + params_docstring = _convert_output_args_doc(params_docstring) + else: + raise ValueError( + f"No `Args` or `Parameters` section is found in the docstring of `{output_type.__name__}`. Make sure it has " + "docstring and contain either `Args` or `Parameters`." + ) + + # Add the return introduction + full_output_type = f"{output_type.__module__}.{output_type.__name__}" + intro = TF_RETURN_INTRODUCTION if output_type.__name__.startswith("TF") else PT_RETURN_INTRODUCTION + intro = intro.format(full_output_type=full_output_type, config_class=config_class) + result = intro + params_docstring + + # Apply minimum indent if necessary + if min_indent is not None: + lines = result.split("\n") + # Find the indent of the first nonempty line + i = 0 + while len(lines[i]) == 0: + i += 1 + indent = len(_get_indent(lines[i])) + # If too small, add indentation to all nonempty lines + if indent < min_indent: + to_add = " " * (min_indent - indent) + lines = [(f"{to_add}{line}" if len(line) > 0 else line) for line in lines] + result = "\n".join(lines) + + return result + + +FAKE_MODEL_DISCLAIMER = """ + + + This example uses a random model as the real ones are all very big. To get proper results, you should use + {real_checkpoint} instead of {fake_checkpoint}. If you get out-of-memory when loading that checkpoint, you can try + adding `device_map="auto"` in the `from_pretrained` call. + + +""" + + +PT_TOKEN_CLASSIFICATION_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer( + ... "HuggingFace is a company based in Paris and New York", add_special_tokens=False, return_tensors="pt" + ... ) + + >>> with torch.no_grad(): + ... logits = model(**inputs).logits + + >>> predicted_token_class_ids = logits.argmax(-1) + + >>> # Note that tokens are classified rather then input words which means that + >>> # there might be more predicted token classes than words. + >>> # Multiple token classes might account for the same word + >>> predicted_tokens_classes = [model.config.id2label[t.item()] for t in predicted_token_class_ids[0]] + >>> predicted_tokens_classes + {expected_output} + + >>> labels = predicted_token_class_ids + >>> loss = model(**inputs, labels=labels).loss + >>> round(loss.item(), 2) + {expected_loss} + ``` +""" + +PT_QUESTION_ANSWERING_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet" + + >>> inputs = tokenizer(question, text, return_tensors="pt") + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> answer_start_index = outputs.start_logits.argmax() + >>> answer_end_index = outputs.end_logits.argmax() + + >>> predict_answer_tokens = inputs.input_ids[0, answer_start_index : answer_end_index + 1] + >>> tokenizer.decode(predict_answer_tokens, skip_special_tokens=True) + {expected_output} + + >>> # target is "nice puppet" + >>> target_start_index = torch.tensor([{qa_target_start_index}]) + >>> target_end_index = torch.tensor([{qa_target_end_index}]) + + >>> outputs = model(**inputs, start_positions=target_start_index, end_positions=target_end_index) + >>> loss = outputs.loss + >>> round(loss.item(), 2) + {expected_loss} + ``` +""" + +PT_SEQUENCE_CLASSIFICATION_SAMPLE = r""" + Example of single-label classification: + + ```python + >>> import torch + >>> from transformers import AutoTokenizer, {model_class} + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") + + >>> with torch.no_grad(): + ... logits = model(**inputs).logits + + >>> predicted_class_id = logits.argmax().item() + >>> model.config.id2label[predicted_class_id] + {expected_output} + + >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)` + >>> num_labels = len(model.config.id2label) + >>> model = {model_class}.from_pretrained("{checkpoint}", num_labels=num_labels) + + >>> labels = torch.tensor([1]) + >>> loss = model(**inputs, labels=labels).loss + >>> round(loss.item(), 2) + {expected_loss} + ``` + + Example of multi-label classification: + + ```python + >>> import torch + >>> from transformers import AutoTokenizer, {model_class} + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}", problem_type="multi_label_classification") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") + + >>> with torch.no_grad(): + ... logits = model(**inputs).logits + + >>> predicted_class_ids = torch.arange(0, logits.shape[-1])[torch.sigmoid(logits).squeeze(dim=0) > 0.5] + + >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)` + >>> num_labels = len(model.config.id2label) + >>> model = {model_class}.from_pretrained( + ... "{checkpoint}", num_labels=num_labels, problem_type="multi_label_classification" + ... ) + + >>> labels = torch.sum( + ... torch.nn.functional.one_hot(predicted_class_ids[None, :].clone(), num_classes=num_labels), dim=1 + ... ).to(torch.float) + >>> loss = model(**inputs, labels=labels).loss + ``` +""" + +PT_MASKED_LM_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("The capital of France is {mask}.", return_tensors="pt") + + >>> with torch.no_grad(): + ... logits = model(**inputs).logits + + >>> # retrieve index of {mask} + >>> mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0] + + >>> predicted_token_id = logits[0, mask_token_index].argmax(axis=-1) + >>> tokenizer.decode(predicted_token_id) + {expected_output} + + >>> labels = tokenizer("The capital of France is Paris.", return_tensors="pt")["input_ids"] + >>> # mask labels of non-{mask} tokens + >>> labels = torch.where(inputs.input_ids == tokenizer.mask_token_id, labels, -100) + + >>> outputs = model(**inputs, labels=labels) + >>> round(outputs.loss.item(), 2) + {expected_loss} + ``` +""" + +PT_BASE_MODEL_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") + >>> outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + ``` +""" + +PT_MULTIPLE_CHOICE_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." + >>> choice0 = "It is eaten with a fork and a knife." + >>> choice1 = "It is eaten while held in the hand." + >>> labels = torch.tensor(0).unsqueeze(0) # choice0 is correct (according to Wikipedia ;)), batch size 1 + + >>> encoding = tokenizer([prompt, prompt], [choice0, choice1], return_tensors="pt", padding=True) + >>> outputs = model(**{{k: v.unsqueeze(0) for k, v in encoding.items()}}, labels=labels) # batch size is 1 + + >>> # the linear classifier still needs to be trained + >>> loss = outputs.loss + >>> logits = outputs.logits + ``` +""" + +PT_CAUSAL_LM_SAMPLE = r""" + Example: + + ```python + >>> import torch + >>> from transformers import AutoTokenizer, {model_class} + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt") + >>> outputs = model(**inputs, labels=inputs["input_ids"]) + >>> loss = outputs.loss + >>> logits = outputs.logits + ``` +""" + +PT_SPEECH_BASE_MODEL_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoProcessor, {model_class} + >>> import torch + >>> from datasets import load_dataset + + >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation", trust_remote_code=True) + >>> dataset = dataset.sort("id") + >>> sampling_rate = dataset.features["audio"].sampling_rate + + >>> processor = AutoProcessor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> # audio file is decoded on the fly + >>> inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt") + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + >>> list(last_hidden_states.shape) + {expected_output} + ``` +""" + +PT_SPEECH_CTC_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoProcessor, {model_class} + >>> from datasets import load_dataset + >>> import torch + + >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation", trust_remote_code=True) + >>> dataset = dataset.sort("id") + >>> sampling_rate = dataset.features["audio"].sampling_rate + + >>> processor = AutoProcessor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> # audio file is decoded on the fly + >>> inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt") + >>> with torch.no_grad(): + ... logits = model(**inputs).logits + >>> predicted_ids = torch.argmax(logits, dim=-1) + + >>> # transcribe speech + >>> transcription = processor.batch_decode(predicted_ids) + >>> transcription[0] + {expected_output} + + >>> inputs["labels"] = processor(text=dataset[0]["text"], return_tensors="pt").input_ids + + >>> # compute loss + >>> loss = model(**inputs).loss + >>> round(loss.item(), 2) + {expected_loss} + ``` +""" + +PT_SPEECH_SEQ_CLASS_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoFeatureExtractor, {model_class} + >>> from datasets import load_dataset + >>> import torch + + >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation", trust_remote_code=True) + >>> dataset = dataset.sort("id") + >>> sampling_rate = dataset.features["audio"].sampling_rate + + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> # audio file is decoded on the fly + >>> inputs = feature_extractor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt") + + >>> with torch.no_grad(): + ... logits = model(**inputs).logits + + >>> predicted_class_ids = torch.argmax(logits, dim=-1).item() + >>> predicted_label = model.config.id2label[predicted_class_ids] + >>> predicted_label + {expected_output} + + >>> # compute loss - target_label is e.g. "down" + >>> target_label = model.config.id2label[0] + >>> inputs["labels"] = torch.tensor([model.config.label2id[target_label]]) + >>> loss = model(**inputs).loss + >>> round(loss.item(), 2) + {expected_loss} + ``` +""" + + +PT_SPEECH_FRAME_CLASS_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoFeatureExtractor, {model_class} + >>> from datasets import load_dataset + >>> import torch + + >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation", trust_remote_code=True) + >>> dataset = dataset.sort("id") + >>> sampling_rate = dataset.features["audio"].sampling_rate + + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> # audio file is decoded on the fly + >>> inputs = feature_extractor(dataset[0]["audio"]["array"], return_tensors="pt", sampling_rate=sampling_rate) + >>> with torch.no_grad(): + ... logits = model(**inputs).logits + + >>> probabilities = torch.sigmoid(logits[0]) + >>> # labels is a one-hot array of shape (num_frames, num_speakers) + >>> labels = (probabilities > 0.5).long() + >>> labels[0].tolist() + {expected_output} + ``` +""" + + +PT_SPEECH_XVECTOR_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoFeatureExtractor, {model_class} + >>> from datasets import load_dataset + >>> import torch + + >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation", trust_remote_code=True) + >>> dataset = dataset.sort("id") + >>> sampling_rate = dataset.features["audio"].sampling_rate + + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> # audio file is decoded on the fly + >>> inputs = feature_extractor( + ... [d["array"] for d in dataset[:2]["audio"]], sampling_rate=sampling_rate, return_tensors="pt", padding=True + ... ) + >>> with torch.no_grad(): + ... embeddings = model(**inputs).embeddings + + >>> embeddings = torch.nn.functional.normalize(embeddings, dim=-1).cpu() + + >>> # the resulting embeddings can be used for cosine similarity-based retrieval + >>> cosine_sim = torch.nn.CosineSimilarity(dim=-1) + >>> similarity = cosine_sim(embeddings[0], embeddings[1]) + >>> threshold = 0.7 # the optimal threshold is dataset-dependent + >>> if similarity < threshold: + ... print("Speakers are not the same!") + >>> round(similarity.item(), 2) + {expected_output} + ``` +""" + +PT_VISION_BASE_MODEL_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoImageProcessor, {model_class} + >>> import torch + >>> from datasets import load_dataset + + >>> dataset = load_dataset("huggingface/cats-image", trust_remote_code=True) + >>> image = dataset["test"]["image"][0] + + >>> image_processor = AutoImageProcessor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = image_processor(image, return_tensors="pt") + + >>> with torch.no_grad(): + ... outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + >>> list(last_hidden_states.shape) + {expected_output} + ``` +""" + +PT_VISION_SEQ_CLASS_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoImageProcessor, {model_class} + >>> import torch + >>> from datasets import load_dataset + + >>> dataset = load_dataset("huggingface/cats-image", trust_remote_code=True) + >>> image = dataset["test"]["image"][0] + + >>> image_processor = AutoImageProcessor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = image_processor(image, return_tensors="pt") + + >>> with torch.no_grad(): + ... logits = model(**inputs).logits + + >>> # model predicts one of the 1000 ImageNet classes + >>> predicted_label = logits.argmax(-1).item() + >>> print(model.config.id2label[predicted_label]) + {expected_output} + ``` +""" + + +PT_SAMPLE_DOCSTRINGS = { + "SequenceClassification": PT_SEQUENCE_CLASSIFICATION_SAMPLE, + "QuestionAnswering": PT_QUESTION_ANSWERING_SAMPLE, + "TokenClassification": PT_TOKEN_CLASSIFICATION_SAMPLE, + "MultipleChoice": PT_MULTIPLE_CHOICE_SAMPLE, + "MaskedLM": PT_MASKED_LM_SAMPLE, + "LMHead": PT_CAUSAL_LM_SAMPLE, + "BaseModel": PT_BASE_MODEL_SAMPLE, + "SpeechBaseModel": PT_SPEECH_BASE_MODEL_SAMPLE, + "CTC": PT_SPEECH_CTC_SAMPLE, + "AudioClassification": PT_SPEECH_SEQ_CLASS_SAMPLE, + "AudioFrameClassification": PT_SPEECH_FRAME_CLASS_SAMPLE, + "AudioXVector": PT_SPEECH_XVECTOR_SAMPLE, + "VisionBaseModel": PT_VISION_BASE_MODEL_SAMPLE, + "ImageClassification": PT_VISION_SEQ_CLASS_SAMPLE, +} + + +TF_TOKEN_CLASSIFICATION_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import tensorflow as tf + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer( + ... "HuggingFace is a company based in Paris and New York", add_special_tokens=False, return_tensors="tf" + ... ) + + >>> logits = model(**inputs).logits + >>> predicted_token_class_ids = tf.math.argmax(logits, axis=-1) + + >>> # Note that tokens are classified rather then input words which means that + >>> # there might be more predicted token classes than words. + >>> # Multiple token classes might account for the same word + >>> predicted_tokens_classes = [model.config.id2label[t] for t in predicted_token_class_ids[0].numpy().tolist()] + >>> predicted_tokens_classes + {expected_output} + ``` + + ```python + >>> labels = predicted_token_class_ids + >>> loss = tf.math.reduce_mean(model(**inputs, labels=labels).loss) + >>> round(float(loss), 2) + {expected_loss} + ``` +""" + +TF_QUESTION_ANSWERING_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import tensorflow as tf + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet" + + >>> inputs = tokenizer(question, text, return_tensors="tf") + >>> outputs = model(**inputs) + + >>> answer_start_index = int(tf.math.argmax(outputs.start_logits, axis=-1)[0]) + >>> answer_end_index = int(tf.math.argmax(outputs.end_logits, axis=-1)[0]) + + >>> predict_answer_tokens = inputs.input_ids[0, answer_start_index : answer_end_index + 1] + >>> tokenizer.decode(predict_answer_tokens) + {expected_output} + ``` + + ```python + >>> # target is "nice puppet" + >>> target_start_index = tf.constant([{qa_target_start_index}]) + >>> target_end_index = tf.constant([{qa_target_end_index}]) + + >>> outputs = model(**inputs, start_positions=target_start_index, end_positions=target_end_index) + >>> loss = tf.math.reduce_mean(outputs.loss) + >>> round(float(loss), 2) + {expected_loss} + ``` +""" + +TF_SEQUENCE_CLASSIFICATION_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import tensorflow as tf + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="tf") + + >>> logits = model(**inputs).logits + + >>> predicted_class_id = int(tf.math.argmax(logits, axis=-1)[0]) + >>> model.config.id2label[predicted_class_id] + {expected_output} + ``` + + ```python + >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)` + >>> num_labels = len(model.config.id2label) + >>> model = {model_class}.from_pretrained("{checkpoint}", num_labels=num_labels) + + >>> labels = tf.constant(1) + >>> loss = model(**inputs, labels=labels).loss + >>> round(float(loss), 2) + {expected_loss} + ``` +""" + +TF_MASKED_LM_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import tensorflow as tf + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("The capital of France is {mask}.", return_tensors="tf") + >>> logits = model(**inputs).logits + + >>> # retrieve index of {mask} + >>> mask_token_index = tf.where((inputs.input_ids == tokenizer.mask_token_id)[0]) + >>> selected_logits = tf.gather_nd(logits[0], indices=mask_token_index) + + >>> predicted_token_id = tf.math.argmax(selected_logits, axis=-1) + >>> tokenizer.decode(predicted_token_id) + {expected_output} + ``` + + ```python + >>> labels = tokenizer("The capital of France is Paris.", return_tensors="tf")["input_ids"] + >>> # mask labels of non-{mask} tokens + >>> labels = tf.where(inputs.input_ids == tokenizer.mask_token_id, labels, -100) + + >>> outputs = model(**inputs, labels=labels) + >>> round(float(outputs.loss), 2) + {expected_loss} + ``` +""" + +TF_BASE_MODEL_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import tensorflow as tf + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="tf") + >>> outputs = model(inputs) + + >>> last_hidden_states = outputs.last_hidden_state + ``` +""" + +TF_MULTIPLE_CHOICE_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import tensorflow as tf + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." + >>> choice0 = "It is eaten with a fork and a knife." + >>> choice1 = "It is eaten while held in the hand." + + >>> encoding = tokenizer([prompt, prompt], [choice0, choice1], return_tensors="tf", padding=True) + >>> inputs = {{k: tf.expand_dims(v, 0) for k, v in encoding.items()}} + >>> outputs = model(inputs) # batch size is 1 + + >>> # the linear classifier still needs to be trained + >>> logits = outputs.logits + ``` +""" + +TF_CAUSAL_LM_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + >>> import tensorflow as tf + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="tf") + >>> outputs = model(inputs) + >>> logits = outputs.logits + ``` +""" + +TF_SPEECH_BASE_MODEL_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoProcessor, {model_class} + >>> from datasets import load_dataset + + >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation", trust_remote_code=True) + >>> dataset = dataset.sort("id") + >>> sampling_rate = dataset.features["audio"].sampling_rate + + >>> processor = AutoProcessor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> # audio file is decoded on the fly + >>> inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="tf") + >>> outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + >>> list(last_hidden_states.shape) + {expected_output} + ``` +""" + +TF_SPEECH_CTC_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoProcessor, {model_class} + >>> from datasets import load_dataset + >>> import tensorflow as tf + + >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation", trust_remote_code=True) + >>> dataset = dataset.sort("id") + >>> sampling_rate = dataset.features["audio"].sampling_rate + + >>> processor = AutoProcessor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> # audio file is decoded on the fly + >>> inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="tf") + >>> logits = model(**inputs).logits + >>> predicted_ids = tf.math.argmax(logits, axis=-1) + + >>> # transcribe speech + >>> transcription = processor.batch_decode(predicted_ids) + >>> transcription[0] + {expected_output} + ``` + + ```python + >>> inputs["labels"] = processor(text=dataset[0]["text"], return_tensors="tf").input_ids + + >>> # compute loss + >>> loss = model(**inputs).loss + >>> round(float(loss), 2) + {expected_loss} + ``` +""" + +TF_VISION_BASE_MODEL_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoImageProcessor, {model_class} + >>> from datasets import load_dataset + + >>> dataset = load_dataset("huggingface/cats-image", trust_remote_code=True) + >>> image = dataset["test"]["image"][0] + + >>> image_processor = AutoImageProcessor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = image_processor(image, return_tensors="tf") + >>> outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + >>> list(last_hidden_states.shape) + {expected_output} + ``` +""" + +TF_VISION_SEQ_CLASS_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoImageProcessor, {model_class} + >>> import tensorflow as tf + >>> from datasets import load_dataset + + >>> dataset = load_dataset("huggingface/cats-image", trust_remote_code=True) + >>> image = dataset["test"]["image"][0] + + >>> image_processor = AutoImageProcessor.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = image_processor(image, return_tensors="tf") + >>> logits = model(**inputs).logits + + >>> # model predicts one of the 1000 ImageNet classes + >>> predicted_label = int(tf.math.argmax(logits, axis=-1)) + >>> print(model.config.id2label[predicted_label]) + {expected_output} + ``` +""" + +TF_SAMPLE_DOCSTRINGS = { + "SequenceClassification": TF_SEQUENCE_CLASSIFICATION_SAMPLE, + "QuestionAnswering": TF_QUESTION_ANSWERING_SAMPLE, + "TokenClassification": TF_TOKEN_CLASSIFICATION_SAMPLE, + "MultipleChoice": TF_MULTIPLE_CHOICE_SAMPLE, + "MaskedLM": TF_MASKED_LM_SAMPLE, + "LMHead": TF_CAUSAL_LM_SAMPLE, + "BaseModel": TF_BASE_MODEL_SAMPLE, + "SpeechBaseModel": TF_SPEECH_BASE_MODEL_SAMPLE, + "CTC": TF_SPEECH_CTC_SAMPLE, + "VisionBaseModel": TF_VISION_BASE_MODEL_SAMPLE, + "ImageClassification": TF_VISION_SEQ_CLASS_SAMPLE, +} + + +FLAX_TOKEN_CLASSIFICATION_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax") + + >>> outputs = model(**inputs) + >>> logits = outputs.logits + ``` +""" + +FLAX_QUESTION_ANSWERING_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet" + >>> inputs = tokenizer(question, text, return_tensors="jax") + + >>> outputs = model(**inputs) + >>> start_scores = outputs.start_logits + >>> end_scores = outputs.end_logits + ``` +""" + +FLAX_SEQUENCE_CLASSIFICATION_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax") + + >>> outputs = model(**inputs) + >>> logits = outputs.logits + ``` +""" + +FLAX_MASKED_LM_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("The capital of France is {mask}.", return_tensors="jax") + + >>> outputs = model(**inputs) + >>> logits = outputs.logits + ``` +""" + +FLAX_BASE_MODEL_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="jax") + >>> outputs = model(**inputs) + + >>> last_hidden_states = outputs.last_hidden_state + ``` +""" + +FLAX_MULTIPLE_CHOICE_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." + >>> choice0 = "It is eaten with a fork and a knife." + >>> choice1 = "It is eaten while held in the hand." + + >>> encoding = tokenizer([prompt, prompt], [choice0, choice1], return_tensors="jax", padding=True) + >>> outputs = model(**{{k: v[None, :] for k, v in encoding.items()}}) + + >>> logits = outputs.logits + ``` +""" + +FLAX_CAUSAL_LM_SAMPLE = r""" + Example: + + ```python + >>> from transformers import AutoTokenizer, {model_class} + + >>> tokenizer = AutoTokenizer.from_pretrained("{checkpoint}") + >>> model = {model_class}.from_pretrained("{checkpoint}") + + >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="np") + >>> outputs = model(**inputs) + + >>> # retrieve logts for next token + >>> next_token_logits = outputs.logits[:, -1] + ``` +""" + +FLAX_SAMPLE_DOCSTRINGS = { + "SequenceClassification": FLAX_SEQUENCE_CLASSIFICATION_SAMPLE, + "QuestionAnswering": FLAX_QUESTION_ANSWERING_SAMPLE, + "TokenClassification": FLAX_TOKEN_CLASSIFICATION_SAMPLE, + "MultipleChoice": FLAX_MULTIPLE_CHOICE_SAMPLE, + "MaskedLM": FLAX_MASKED_LM_SAMPLE, + "BaseModel": FLAX_BASE_MODEL_SAMPLE, + "LMHead": FLAX_CAUSAL_LM_SAMPLE, +} + + +def filter_outputs_from_example(docstring, **kwargs): + """ + Removes the lines testing an output with the doctest syntax in a code sample when it's set to `None`. + """ + for key, value in kwargs.items(): + if value is not None: + continue + + doc_key = "{" + key + "}" + docstring = re.sub(rf"\n([^\n]+)\n\s+{doc_key}\n", "\n", docstring) + + return docstring + + +def add_code_sample_docstrings( + *docstr, + processor_class=None, + checkpoint=None, + output_type=None, + config_class=None, + mask="[MASK]", + qa_target_start_index=14, + qa_target_end_index=15, + model_cls=None, + modality=None, + expected_output=None, + expected_loss=None, + real_checkpoint=None, + revision=None, +): + def docstring_decorator(fn): + # model_class defaults to function's class if not specified otherwise + model_class = fn.__qualname__.split(".")[0] if model_cls is None else model_cls + + if model_class[:2] == "TF": + sample_docstrings = TF_SAMPLE_DOCSTRINGS + elif model_class[:4] == "Flax": + sample_docstrings = FLAX_SAMPLE_DOCSTRINGS + else: + sample_docstrings = PT_SAMPLE_DOCSTRINGS + + # putting all kwargs for docstrings in a dict to be used + # with the `.format(**doc_kwargs)`. Note that string might + # be formatted with non-existing keys, which is fine. + doc_kwargs = { + "model_class": model_class, + "processor_class": processor_class, + "checkpoint": checkpoint, + "mask": mask, + "qa_target_start_index": qa_target_start_index, + "qa_target_end_index": qa_target_end_index, + "expected_output": expected_output, + "expected_loss": expected_loss, + "real_checkpoint": real_checkpoint, + "fake_checkpoint": checkpoint, + "true": "{true}", # For syntax that conflicts with formatting. + } + + if ("SequenceClassification" in model_class or "AudioClassification" in model_class) and modality == "audio": + code_sample = sample_docstrings["AudioClassification"] + elif "SequenceClassification" in model_class: + code_sample = sample_docstrings["SequenceClassification"] + elif "QuestionAnswering" in model_class: + code_sample = sample_docstrings["QuestionAnswering"] + elif "TokenClassification" in model_class: + code_sample = sample_docstrings["TokenClassification"] + elif "MultipleChoice" in model_class: + code_sample = sample_docstrings["MultipleChoice"] + elif "MaskedLM" in model_class or model_class in ["FlaubertWithLMHeadModel", "XLMWithLMHeadModel"]: + code_sample = sample_docstrings["MaskedLM"] + elif "LMHead" in model_class or "CausalLM" in model_class: + code_sample = sample_docstrings["LMHead"] + elif "CTC" in model_class: + code_sample = sample_docstrings["CTC"] + elif "AudioFrameClassification" in model_class: + code_sample = sample_docstrings["AudioFrameClassification"] + elif "XVector" in model_class and modality == "audio": + code_sample = sample_docstrings["AudioXVector"] + elif "Model" in model_class and modality == "audio": + code_sample = sample_docstrings["SpeechBaseModel"] + elif "Model" in model_class and modality == "vision": + code_sample = sample_docstrings["VisionBaseModel"] + elif "Model" in model_class or "Encoder" in model_class: + code_sample = sample_docstrings["BaseModel"] + elif "ImageClassification" in model_class: + code_sample = sample_docstrings["ImageClassification"] + else: + raise ValueError(f"Docstring can't be built for model {model_class}") + + code_sample = filter_outputs_from_example( + code_sample, expected_output=expected_output, expected_loss=expected_loss + ) + if real_checkpoint is not None: + code_sample = FAKE_MODEL_DISCLAIMER + code_sample + func_doc = (fn.__doc__ or "") + "".join(docstr) + output_doc = "" if output_type is None else _prepare_output_docstrings(output_type, config_class) + built_doc = code_sample.format(**doc_kwargs) + if revision is not None: + if re.match(r"^refs/pr/\\d+", revision): + raise ValueError( + f"The provided revision '{revision}' is incorrect. It should point to" + " a pull request reference on the hub like 'refs/pr/6'" + ) + built_doc = built_doc.replace( + f'from_pretrained("{checkpoint}")', f'from_pretrained("{checkpoint}", revision="{revision}")' + ) + + fn.__doc__ = func_doc + output_doc + built_doc + return fn + + return docstring_decorator + + +def replace_return_docstrings(output_type=None, config_class=None): + def docstring_decorator(fn): + func_doc = fn.__doc__ + lines = func_doc.split("\n") + i = 0 + while i < len(lines) and re.search(r"^\s*Returns?:\s*$", lines[i]) is None: + i += 1 + if i < len(lines): + indent = len(_get_indent(lines[i])) + lines[i] = _prepare_output_docstrings(output_type, config_class, min_indent=indent) + func_doc = "\n".join(lines) + else: + raise ValueError( + f"The function {fn} should have an empty 'Return:' or 'Returns:' in its docstring as placeholder, " + f"current docstring is:\n{func_doc}" + ) + fn.__doc__ = func_doc + return fn + + return docstring_decorator + + +def copy_func(f): + """Returns a copy of a function f.""" + # Based on http://stackoverflow.com/a/6528148/190597 (Glenn Maynard) + g = types.FunctionType(f.__code__, f.__globals__, name=f.__name__, argdefs=f.__defaults__, closure=f.__closure__) + g = functools.update_wrapper(g, f) + g.__kwdefaults__ = f.__kwdefaults__ + return g diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_detectron2_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_detectron2_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..22ec32fe30a1b9abbc5d810510d97d5c81082561 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_detectron2_objects.py @@ -0,0 +1,11 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import requires_backends + + +class LayoutLMv2Model: + def __init__(self, *args, **kwargs): + requires_backends(self, ["detectron2"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["detectron2"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_essentia_and_librosa_and_pretty_midi_and_scipy_and_torch_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_essentia_and_librosa_and_pretty_midi_and_scipy_and_torch_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..e6d75a6ec22e90427c972a753a24afd1a780758f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_essentia_and_librosa_and_pretty_midi_and_scipy_and_torch_objects.py @@ -0,0 +1,23 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class Pop2PianoFeatureExtractor(metaclass=DummyObject): + _backends = ["essentia", "librosa", "pretty_midi", "scipy", "torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["essentia", "librosa", "pretty_midi", "scipy", "torch"]) + + +class Pop2PianoTokenizer(metaclass=DummyObject): + _backends = ["essentia", "librosa", "pretty_midi", "scipy", "torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["essentia", "librosa", "pretty_midi", "scipy", "torch"]) + + +class Pop2PianoProcessor(metaclass=DummyObject): + _backends = ["essentia", "librosa", "pretty_midi", "scipy", "torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["essentia", "librosa", "pretty_midi", "scipy", "torch"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_flax_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_flax_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..0f2390fb694b7ac3bb2a64c60ef8b383389ac616 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_flax_objects.py @@ -0,0 +1,1430 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class FlaxForcedBOSTokenLogitsProcessor(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxForcedEOSTokenLogitsProcessor(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxForceTokensLogitsProcessor(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGenerationMixin(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxLogitsProcessor(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxLogitsProcessorList(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxLogitsWarper(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMinLengthLogitsProcessor(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxSuppressTokensAtBeginLogitsProcessor(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxSuppressTokensLogitsProcessor(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxTemperatureLogitsWarper(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxTopKLogitsWarper(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxTopPLogitsWarper(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxWhisperTimeStampLogitsProcessor(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAlbertForMaskedLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAlbertForMultipleChoice(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAlbertForPreTraining(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAlbertForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAlbertForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAlbertForTokenClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAlbertModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAlbertPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +FLAX_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING = None + + +FLAX_MODEL_FOR_CAUSAL_LM_MAPPING = None + + +FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING = None + + +FLAX_MODEL_FOR_MASKED_LM_MAPPING = None + + +FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING = None + + +FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING = None + + +FLAX_MODEL_FOR_PRETRAINING_MAPPING = None + + +FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING = None + + +FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = None + + +FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = None + + +FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING = None + + +FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = None + + +FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING = None + + +FLAX_MODEL_MAPPING = None + + +class FlaxAutoModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForImageClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForMaskedLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForMultipleChoice(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForNextSentencePrediction(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForPreTraining(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForSeq2SeqLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForSpeechSeq2Seq(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForTokenClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxAutoModelForVision2Seq(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBartDecoderPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBartForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBartForConditionalGeneration(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBartForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBartForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBartModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBartPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBeitForImageClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBeitForMaskedImageModeling(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBeitModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBeitPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBertForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBertForMaskedLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBertForMultipleChoice(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBertForNextSentencePrediction(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBertForPreTraining(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBertForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBertForTokenClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBertModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBertPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBigBirdForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBigBirdForMaskedLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBigBirdForMultipleChoice(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBigBirdForPreTraining(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBigBirdForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBigBirdForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBigBirdForTokenClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBigBirdModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBigBirdPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBlenderbotForConditionalGeneration(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBlenderbotModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBlenderbotPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBlenderbotSmallForConditionalGeneration(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBlenderbotSmallModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBlenderbotSmallPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBloomForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBloomModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxBloomPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxCLIPModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxCLIPPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxCLIPTextModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxCLIPTextModelWithProjection(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxCLIPTextPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxCLIPVisionModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxCLIPVisionPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxDinov2ForImageClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxDinov2Model(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxDinov2PreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxDistilBertForMaskedLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxDistilBertForMultipleChoice(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxDistilBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxDistilBertForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxDistilBertForTokenClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxDistilBertModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxDistilBertPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxElectraForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxElectraForMaskedLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxElectraForMultipleChoice(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxElectraForPreTraining(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxElectraForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxElectraForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxElectraForTokenClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxElectraModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxElectraPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxEncoderDecoderModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGemmaForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGemmaModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGemmaPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGPT2LMHeadModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGPT2Model(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGPT2PreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGPTNeoForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGPTNeoModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGPTNeoPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGPTJForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGPTJModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxGPTJPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxLlamaForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxLlamaModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxLlamaPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxLongT5ForConditionalGeneration(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxLongT5Model(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxLongT5PreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMarianModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMarianMTModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMarianPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMBartForConditionalGeneration(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMBartForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMBartForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMBartModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMBartPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMistralForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMistralModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMistralPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMT5EncoderModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMT5ForConditionalGeneration(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxMT5Model(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxOPTForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxOPTModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxOPTPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxPegasusForConditionalGeneration(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxPegasusModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxPegasusPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRegNetForImageClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRegNetModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRegNetPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxResNetForImageClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxResNetModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxResNetPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaForMaskedLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaForMultipleChoice(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaForTokenClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaPreLayerNormForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaPreLayerNormForMaskedLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaPreLayerNormForMultipleChoice(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaPreLayerNormForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaPreLayerNormForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaPreLayerNormForTokenClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaPreLayerNormModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRobertaPreLayerNormPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRoFormerForMaskedLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRoFormerForMultipleChoice(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRoFormerForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRoFormerForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRoFormerForTokenClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRoFormerModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxRoFormerPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxSpeechEncoderDecoderModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxT5EncoderModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxT5ForConditionalGeneration(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxT5Model(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxT5PreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxVisionEncoderDecoderModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxVisionTextDualEncoderModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxViTForImageClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxViTModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxViTPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxWav2Vec2ForCTC(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxWav2Vec2ForPreTraining(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxWav2Vec2Model(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxWav2Vec2PreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxWhisperForAudioClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxWhisperForConditionalGeneration(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxWhisperModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxWhisperPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXGLMForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXGLMModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXGLMPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXLMRobertaForCausalLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXLMRobertaForMaskedLM(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXLMRobertaForMultipleChoice(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXLMRobertaForQuestionAnswering(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXLMRobertaForSequenceClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXLMRobertaForTokenClassification(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXLMRobertaModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) + + +class FlaxXLMRobertaPreTrainedModel(metaclass=DummyObject): + _backends = ["flax"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["flax"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_keras_nlp_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_keras_nlp_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..c6bb86a6d9b49e78f8936f3c1eb3cfc8b8db7951 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_keras_nlp_objects.py @@ -0,0 +1,9 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class TFGPT2Tokenizer(metaclass=DummyObject): + _backends = ["keras_nlp"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["keras_nlp"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_music_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_music_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..89052be47c1d32bac5cbd6fceab183fc1d75d3bf --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_music_objects.py @@ -0,0 +1,16 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class Pop2PianoFeatureExtractor(metaclass=DummyObject): + _backends = ["music"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["music"]) + + +class Pop2PianoTokenizer(metaclass=DummyObject): + _backends = ["music"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["music"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_pt_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_pt_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..b03455c89e6ec30a8a6efc2fa883436559196f8f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_pt_objects.py @@ -0,0 +1,11168 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class Cache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CacheConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DynamicCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EncoderDecoderCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HQQQuantizedCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HybridCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MambaCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OffloadedCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OffloadedStaticCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QuantizedCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QuantizedCacheConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QuantoQuantizedCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SinkCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SlidingWindowCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class StaticCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GlueDataset(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GlueDataTrainingArguments(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LineByLineTextDataset(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LineByLineWithRefDataset(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LineByLineWithSOPTextDataset(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SquadDataset(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SquadDataTrainingArguments(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TextDataset(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TextDatasetForNextSentencePrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlternatingCodebooksLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BayesianDetectorConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BayesianDetectorModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BeamScorer(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BeamSearchScorer(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClassifierFreeGuidanceLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConstrainedBeamSearchScorer(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Constraint(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConstraintListState(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DisjunctiveConstraint(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EncoderNoRepeatNGramLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EncoderRepetitionPenaltyLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EosTokenCriteria(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EpsilonLogitsWarper(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EtaLogitsWarper(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ExponentialDecayLengthPenalty(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ForcedBOSTokenLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ForcedEOSTokenLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GenerationMixin(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HammingDiversityLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InfNanRemoveLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LogitNormalization(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LogitsProcessorList(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MaxLengthCriteria(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MaxTimeCriteria(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MinLengthLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MinNewTokensLengthLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MinPLogitsWarper(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NoBadWordsLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NoRepeatNGramLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhrasalConstraint(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PrefixConstrainedLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RepetitionPenaltyLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SequenceBiasLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class StoppingCriteria(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class StoppingCriteriaList(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class StopStringCriteria(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SuppressTokensAtBeginLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SuppressTokensLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SynthIDTextWatermarkDetector(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SynthIDTextWatermarkingConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SynthIDTextWatermarkLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TemperatureLogitsWarper(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TopKLogitsWarper(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TopPLogitsWarper(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TypicalLogitsWarper(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UnbatchedClassifierFreeGuidanceLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WatermarkDetector(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WatermarkLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WhisperTimeStampLogitsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TorchExportableModuleWithStaticCache(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def convert_and_export_with_cache(*args, **kwargs): + requires_backends(convert_and_export_with_cache, ["torch"]) + + +def model_addition_debugger(*args, **kwargs): + requires_backends(model_addition_debugger, ["torch"]) + + +def model_addition_debugger_context(*args, **kwargs): + requires_backends(model_addition_debugger_context, ["torch"]) + + +ROPE_INIT_FUNCTIONS = None + + +def dynamic_rope_update(*args, **kwargs): + requires_backends(dynamic_rope_update, ["torch"]) + + +class AttentionInterface(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlbertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlbertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlbertForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlbertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlbertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlbertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlbertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlbertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_albert(*args, **kwargs): + requires_backends(load_tf_weights_in_albert, ["torch"]) + + +class AlignModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlignPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlignTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AlignVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AltCLIPModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AltCLIPPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AltCLIPTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AltCLIPVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AriaForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AriaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AriaTextForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AriaTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AriaTextPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ASTForAudioClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ASTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ASTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING = None + + +MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING = None + + +MODEL_FOR_AUDIO_XVECTOR_MAPPING = None + + +MODEL_FOR_BACKBONE_MAPPING = None + + +MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING = None + + +MODEL_FOR_CAUSAL_LM_MAPPING = None + + +MODEL_FOR_CTC_MAPPING = None + + +MODEL_FOR_DEPTH_ESTIMATION_MAPPING = None + + +MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING = None + + +MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING = None + + +MODEL_FOR_IMAGE_MAPPING = None + + +MODEL_FOR_IMAGE_SEGMENTATION_MAPPING = None + + +MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING = None + + +MODEL_FOR_IMAGE_TO_IMAGE_MAPPING = None + + +MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING = None + + +MODEL_FOR_KEYPOINT_DETECTION_MAPPING = None + + +MODEL_FOR_MASK_GENERATION_MAPPING = None + + +MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING = None + + +MODEL_FOR_MASKED_LM_MAPPING = None + + +MODEL_FOR_MULTIPLE_CHOICE_MAPPING = None + + +MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING = None + + +MODEL_FOR_OBJECT_DETECTION_MAPPING = None + + +MODEL_FOR_PRETRAINING_MAPPING = None + + +MODEL_FOR_QUESTION_ANSWERING_MAPPING = None + + +MODEL_FOR_RETRIEVAL_MAPPING = None + + +MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING = None + + +MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = None + + +MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = None + + +MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING = None + + +MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING = None + + +MODEL_FOR_TEXT_ENCODING_MAPPING = None + + +MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING = None + + +MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING = None + + +MODEL_FOR_TIME_SERIES_CLASSIFICATION_MAPPING = None + + +MODEL_FOR_TIME_SERIES_REGRESSION_MAPPING = None + + +MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = None + + +MODEL_FOR_UNIVERSAL_SEGMENTATION_MAPPING = None + + +MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING = None + + +MODEL_FOR_VISION_2_SEQ_MAPPING = None + + +MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING = None + + +MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING = None + + +MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING = None + + +MODEL_MAPPING = None + + +MODEL_WITH_LM_HEAD_MAPPING = None + + +class AutoBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForAudioClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForAudioFrameClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForAudioXVector(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForDepthEstimation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForDocumentQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForImageSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForImageTextToText(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForImageToImage(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForInstanceSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForKeypointDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForMaskedImageModeling(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForMaskGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForNextSentencePrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForSemanticSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForSeq2SeqLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForSpeechSeq2Seq(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForTableQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForTextEncoding(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForTextToSpectrogram(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForTextToWaveform(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForUniversalSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForVideoClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForVision2Seq(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForVisualQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForZeroShotImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelForZeroShotObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoModelWithLMHead(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoformerForPrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AutoformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AyaVisionForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AyaVisionPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BambaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BambaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BambaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BarkCausalModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BarkCoarseModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BarkFineModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BarkModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BarkPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BarkSemanticModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BartForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BartForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BartForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BartForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BartModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BartPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BartPretrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PretrainedBartModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BeitBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BeitForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BeitForMaskedImageModeling(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BeitForSemanticSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BeitModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BeitPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertForNextSentencePrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertLMHeadModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_bert(*args, **kwargs): + requires_backends(load_tf_weights_in_bert, ["torch"]) + + +class BertGenerationDecoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertGenerationEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BertGenerationPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_bert_generation(*args, **kwargs): + requires_backends(load_tf_weights_in_bert_generation, ["torch"]) + + +class BigBirdForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_big_bird(*args, **kwargs): + requires_backends(load_tf_weights_in_big_bird, ["torch"]) + + +class BigBirdPegasusForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdPegasusForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdPegasusForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdPegasusForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdPegasusModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BigBirdPegasusPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BioGptForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BioGptForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BioGptForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BioGptModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BioGptPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BitBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BitForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BitModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BitPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlenderbotForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlenderbotForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlenderbotModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlenderbotPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlenderbotSmallForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlenderbotSmallForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlenderbotSmallModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlenderbotSmallPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlipForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlipForImageTextRetrieval(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlipForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlipModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlipPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlipTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BlipVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Blip2ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Blip2ForImageTextRetrieval(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Blip2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Blip2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Blip2QFormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Blip2TextModelWithProjection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Blip2VisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Blip2VisionModelWithProjection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BloomForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BloomForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BloomForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BloomForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BloomModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BloomPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BridgeTowerForContrastiveLearning(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BridgeTowerForImageAndTextRetrieval(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BridgeTowerForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BridgeTowerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BridgeTowerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BrosForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BrosModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BrosPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BrosProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BrosSpadeEEForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class BrosSpadeELForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CamembertForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CamembertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CamembertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CamembertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CamembertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CamembertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CamembertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CamembertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CanineForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CanineForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CanineForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CanineForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CanineModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CaninePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_canine(*args, **kwargs): + requires_backends(load_tf_weights_in_canine, ["torch"]) + + +class ChameleonForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ChameleonModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ChameleonPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ChameleonProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ChameleonVQVAE(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ChineseCLIPModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ChineseCLIPPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ChineseCLIPTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ChineseCLIPVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClapAudioModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClapAudioModelWithProjection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClapFeatureExtractor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClapModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClapPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClapTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClapTextModelWithProjection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPTextModelWithProjection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPVisionModelWithProjection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPSegForImageSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPSegModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPSegPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPSegTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CLIPSegVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClvpDecoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClvpEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClvpForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClvpModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClvpModelForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ClvpPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CodeGenForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CodeGenModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CodeGenPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CohereForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CohereModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CoherePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Cohere2ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Cohere2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Cohere2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ColPaliForRetrieval(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ColPaliPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConditionalDetrForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConditionalDetrForSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConditionalDetrModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConditionalDetrPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvBertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvBertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvBertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvBertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_convbert(*args, **kwargs): + requires_backends(load_tf_weights_in_convbert, ["torch"]) + + +class ConvNextBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvNextForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvNextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvNextPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvNextV2Backbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvNextV2ForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvNextV2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ConvNextV2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CpmAntForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CpmAntModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CpmAntPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CTRLForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CTRLLMHeadModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CTRLModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CTRLPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CvtForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CvtModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class CvtPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DabDetrForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DabDetrModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DabDetrPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DacModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DacPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecAudioForAudioFrameClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecAudioForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecAudioForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecAudioForXVector(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecAudioModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecAudioPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecTextForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecTextForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecTextForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecTextForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecTextForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecTextForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecTextPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecVisionForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecVisionForSemanticSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Data2VecVisionPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DbrxForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DbrxModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DbrxPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaV2ForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaV2ForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaV2ForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaV2ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaV2ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaV2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DebertaV2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DecisionTransformerGPT2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DecisionTransformerGPT2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DecisionTransformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DecisionTransformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeepseekV3ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeepseekV3Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeepseekV3PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeformableDetrForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeformableDetrModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeformableDetrPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeiTForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeiTForImageClassificationWithTeacher(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeiTForMaskedImageModeling(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeiTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DeiTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DetaForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DetaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DetaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EfficientFormerForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EfficientFormerForImageClassificationWithTeacher(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EfficientFormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EfficientFormerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieMForInformationExtraction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieMForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieMForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieMForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieMForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieMModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieMPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTSanJapaneseForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTSanJapaneseModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTSanJapanesePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraphormerForGraphClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraphormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraphormerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JukeboxModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JukeboxPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JukeboxPrior(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JukeboxVQVAE(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MCTCTForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MCTCTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MCTCTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegaForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegaForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegaForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegaForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MMBTForClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MMBTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ModalEmbeddings(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NatBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NatForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NatModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NatPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NezhaForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NezhaForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NezhaForNextSentencePrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NezhaForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NezhaForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NezhaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NezhaForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NezhaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NezhaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OpenLlamaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OpenLlamaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OpenLlamaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OpenLlamaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QDQBertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QDQBertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QDQBertForNextSentencePrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QDQBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QDQBertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QDQBertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QDQBertLMHeadModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QDQBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class QDQBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_qdqbert(*args, **kwargs): + requires_backends(load_tf_weights_in_qdqbert, ["torch"]) + + +class RealmEmbedder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RealmForOpenQA(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RealmKnowledgeAugEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RealmPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RealmReader(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RealmRetriever(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RealmScorer(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_realm(*args, **kwargs): + requires_backends(load_tf_weights_in_realm, ["torch"]) + + +class RetriBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RetriBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Speech2Text2ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Speech2Text2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TrajectoryTransformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TrajectoryTransformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class AdaptiveEmbedding(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TransfoXLForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TransfoXLLMHeadModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TransfoXLModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TransfoXLPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_transfo_xl(*args, **kwargs): + requires_backends(load_tf_weights_in_transfo_xl, ["torch"]) + + +class TvltForAudioVisualClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TvltForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TvltModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TvltPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VanForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VanModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VanPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTHybridForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTHybridModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTHybridPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMProphetNetDecoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMProphetNetEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMProphetNetForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMProphetNetForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMProphetNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMProphetNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DepthAnythingForDepthEstimation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DepthAnythingPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DepthProForDepthEstimation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DepthProModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DepthProPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DetrForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DetrForSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DetrModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DetrPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DiffLlamaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DiffLlamaForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DiffLlamaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DiffLlamaForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DiffLlamaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DiffLlamaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DinatBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DinatForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DinatModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DinatPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Dinov2Backbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Dinov2ForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Dinov2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Dinov2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Dinov2WithRegistersBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Dinov2WithRegistersForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Dinov2WithRegistersModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Dinov2WithRegistersPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DistilBertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DistilBertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DistilBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DistilBertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DistilBertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DistilBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DistilBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DonutSwinModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DonutSwinPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPRContextEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPRPretrainedContextEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPRPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPRPretrainedQuestionEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPRPretrainedReader(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPRQuestionEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPRReader(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPTForDepthEstimation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPTForSemanticSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class DPTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EfficientNetForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EfficientNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EfficientNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ElectraForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ElectraForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ElectraForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ElectraForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ElectraForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ElectraForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ElectraForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ElectraModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ElectraPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_electra(*args, **kwargs): + requires_backends(load_tf_weights_in_electra, ["torch"]) + + +class Emu3ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Emu3ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Emu3PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Emu3TextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Emu3VQVAE(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EncodecModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EncodecPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EncoderDecoderModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieForNextSentencePrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErnieModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ErniePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EsmFoldPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EsmForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EsmForProteinFolding(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EsmForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EsmForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EsmModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class EsmPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FalconForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FalconForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FalconForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FalconForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FalconModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FalconPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FalconMambaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FalconMambaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FalconMambaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FastSpeech2ConformerHifiGan(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FastSpeech2ConformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FastSpeech2ConformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FastSpeech2ConformerWithHifiGan(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlaubertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlaubertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlaubertForQuestionAnsweringSimple(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlaubertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlaubertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlaubertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlaubertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlaubertWithLMHeadModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlavaForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlavaImageCodebook(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlavaImageModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlavaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlavaMultimodalModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlavaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FlavaTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FNetForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FNetForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FNetForNextSentencePrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FNetForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FNetForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FNetForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FNetForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FocalNetBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FocalNetForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FocalNetForMaskedImageModeling(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FocalNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FocalNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FSMTForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FSMTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PretrainedFSMTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FunnelBaseModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FunnelForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FunnelForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FunnelForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FunnelForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FunnelForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FunnelForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FunnelModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FunnelPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_funnel(*args, **kwargs): + requires_backends(load_tf_weights_in_funnel, ["torch"]) + + +class FuyuForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class FuyuPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GemmaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GemmaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GemmaForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GemmaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GemmaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Gemma2ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Gemma2ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Gemma2ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Gemma2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Gemma2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Gemma3ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Gemma3ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Gemma3PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Gemma3TextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GitForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GitModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GitPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GitVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GlmForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GlmForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GlmForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GlmModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GlmPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GLPNForDepthEstimation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GLPNModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GLPNPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GotOcr2ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GotOcr2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPT2DoubleHeadsModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPT2ForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPT2ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPT2ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPT2LMHeadModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPT2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPT2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_gpt2(*args, **kwargs): + requires_backends(load_tf_weights_in_gpt2, ["torch"]) + + +class GPTBigCodeForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTBigCodeForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTBigCodeForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTBigCodeModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTBigCodePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_gpt_neo(*args, **kwargs): + requires_backends(load_tf_weights_in_gpt_neo, ["torch"]) + + +class GPTNeoXForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoXForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoXForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoXForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoXModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoXPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoXJapaneseForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoXJapaneseModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTNeoXJapanesePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTJForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTJForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTJForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTJModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GPTJPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraniteForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraniteModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GranitePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraniteMoeForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraniteMoeModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraniteMoePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraniteMoeSharedForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraniteMoeSharedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GraniteMoeSharedPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GroundingDinoForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GroundingDinoModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GroundingDinoPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GroupViTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GroupViTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GroupViTTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class GroupViTVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HeliumForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HeliumForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HeliumForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HeliumModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HeliumPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HieraBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HieraForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HieraForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HieraModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HieraPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HubertForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HubertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HubertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class HubertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IBertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IBertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IBertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IBertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IdeficsForVisionText2Text(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IdeficsModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IdeficsPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IdeficsProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Idefics2ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Idefics2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Idefics2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Idefics2Processor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Idefics3ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Idefics3Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Idefics3PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Idefics3Processor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Idefics3VisionConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Idefics3VisionTransformer(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IJepaForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IJepaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class IJepaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ImageGPTForCausalImageModeling(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ImageGPTForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ImageGPTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ImageGPTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_imagegpt(*args, **kwargs): + requires_backends(load_tf_weights_in_imagegpt, ["torch"]) + + +class InformerForPrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InstructBlipForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InstructBlipPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InstructBlipQFormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InstructBlipVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InstructBlipVideoForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InstructBlipVideoPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InstructBlipVideoQFormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class InstructBlipVideoVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JambaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JambaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JambaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JambaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JetMoeForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JetMoeForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JetMoeModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class JetMoePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Kosmos2ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Kosmos2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Kosmos2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMv2ForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMv2ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMv2ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMv2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMv2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMv3ForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMv3ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMv3ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMv3Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LayoutLMv3PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LEDForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LEDForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LEDForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LEDModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LEDPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LevitForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LevitForImageClassificationWithTeacher(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LevitModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LevitPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LiltForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LiltForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LiltForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LiltModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LiltPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlamaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlamaForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlamaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlamaForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlamaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlamaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Llama4ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Llama4ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Llama4PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Llama4TextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Llama4VisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlavaForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlavaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlavaNextForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlavaNextPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlavaNextVideoForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlavaNextVideoPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlavaOnevisionForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LlavaOnevisionPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongformerForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongformerForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongformerForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongformerForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongformerForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongT5EncoderModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongT5ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongT5Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LongT5PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LukeForEntityClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LukeForEntityPairClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LukeForEntitySpanClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LukeForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LukeForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LukeForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LukeForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LukeForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LukeModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LukePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LxmertEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LxmertForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LxmertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LxmertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LxmertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class LxmertVisualFeatureEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class M2M100ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class M2M100Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class M2M100PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MambaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MambaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MambaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Mamba2ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Mamba2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Mamba2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MarianForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MarianModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MarianMTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MarianPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MarkupLMForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MarkupLMForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MarkupLMForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MarkupLMModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MarkupLMPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Mask2FormerForUniversalSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Mask2FormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Mask2FormerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MaskFormerForInstanceSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MaskFormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MaskFormerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MaskFormerSwinBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MBartForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MBartForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MBartForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MBartForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MBartModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MBartPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegatronBertForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegatronBertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegatronBertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegatronBertForNextSentencePrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegatronBertForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegatronBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegatronBertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegatronBertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegatronBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MegatronBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MgpstrForSceneTextRecognition(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MgpstrModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MgpstrPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MimiModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MimiPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MistralForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MistralForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MistralForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MistralForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MistralModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MistralPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Mistral3ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Mistral3PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MixtralForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MixtralForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MixtralForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MixtralForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MixtralModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MixtralPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MllamaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MllamaForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MllamaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MllamaProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MllamaTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MllamaVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileBertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileBertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileBertForNextSentencePrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileBertForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileBertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileBertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_mobilebert(*args, **kwargs): + requires_backends(load_tf_weights_in_mobilebert, ["torch"]) + + +class MobileNetV1ForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileNetV1Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileNetV1PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_mobilenet_v1(*args, **kwargs): + requires_backends(load_tf_weights_in_mobilenet_v1, ["torch"]) + + +class MobileNetV2ForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileNetV2ForSemanticSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileNetV2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileNetV2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_mobilenet_v2(*args, **kwargs): + requires_backends(load_tf_weights_in_mobilenet_v2, ["torch"]) + + +class MobileViTForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileViTForSemanticSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileViTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileViTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileViTV2ForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileViTV2ForSemanticSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileViTV2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MobileViTV2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ModernBertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ModernBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ModernBertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ModernBertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ModernBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ModernBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MoonshineForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MoonshineModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MoonshinePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MoshiForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MoshiForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MoshiModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MoshiPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MPNetForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MPNetForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MPNetForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MPNetForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MPNetForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MPNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MPNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MptForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MptForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MptForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MptForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MptModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MptPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MraForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MraForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MraForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MraForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MraForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MraModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MraPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MT5EncoderModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MT5ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MT5ForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MT5ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MT5ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MT5Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MT5PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MusicgenForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MusicgenForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MusicgenModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MusicgenPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MusicgenProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MusicgenMelodyForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MusicgenMelodyForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MusicgenMelodyModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MusicgenMelodyPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MvpForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MvpForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MvpForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MvpForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MvpModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class MvpPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NemotronForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NemotronForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NemotronForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NemotronForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NemotronModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NemotronPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NllbMoeForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NllbMoeModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NllbMoePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NllbMoeSparseMLP(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NllbMoeTop2Router(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NystromformerForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NystromformerForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NystromformerForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NystromformerForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NystromformerForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NystromformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class NystromformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OlmoForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OlmoModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OlmoPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Olmo2ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Olmo2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Olmo2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OlmoeForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OlmoeModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OlmoePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OmDetTurboForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OmDetTurboPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OneFormerForUniversalSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OneFormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OneFormerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OpenAIGPTDoubleHeadsModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OpenAIGPTForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OpenAIGPTLMHeadModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OpenAIGPTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OpenAIGPTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_openai_gpt(*args, **kwargs): + requires_backends(load_tf_weights_in_openai_gpt, ["torch"]) + + +class OPTForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OPTForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OPTForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OPTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OPTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Owlv2ForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Owlv2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Owlv2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Owlv2TextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Owlv2VisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OwlViTForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OwlViTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OwlViTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OwlViTTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class OwlViTVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PaliGemmaForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PaliGemmaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PaliGemmaProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSMixerForPrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSMixerForPretraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSMixerForRegression(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSMixerForTimeSeriesClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSMixerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSMixerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSTForClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSTForPrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSTForPretraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSTForRegression(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PatchTSTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PegasusForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PegasusForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PegasusModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PegasusPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PegasusXForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PegasusXModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PegasusXPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PerceiverForImageClassificationConvProcessing(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PerceiverForImageClassificationFourier(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PerceiverForImageClassificationLearned(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PerceiverForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PerceiverForMultimodalAutoencoding(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PerceiverForOpticalFlow(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PerceiverForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PerceiverModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PerceiverPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PersimmonForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PersimmonForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PersimmonForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PersimmonModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PersimmonPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhiForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhiForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhiForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhiModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhiPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi3ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi3ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi3ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi3Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi3PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi4MultimodalAudioModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi4MultimodalAudioPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi4MultimodalForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi4MultimodalModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi4MultimodalPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi4MultimodalVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Phi4MultimodalVisionPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhimoeForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhimoeForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhimoeModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PhimoePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Pix2StructForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Pix2StructPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Pix2StructTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Pix2StructVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PixtralPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PixtralVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PLBartForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PLBartForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PLBartForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PLBartModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PLBartPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PoolFormerForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PoolFormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PoolFormerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Pop2PianoForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Pop2PianoPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PromptDepthAnythingForDepthEstimation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PromptDepthAnythingPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ProphetNetDecoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ProphetNetEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ProphetNetForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ProphetNetForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ProphetNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ProphetNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PvtForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PvtModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PvtPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PvtV2Backbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PvtV2ForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PvtV2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class PvtV2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2ForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2_5_VLForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2_5_VLModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2_5_VLPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2AudioEncoder(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2AudioForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2AudioPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2MoeForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2MoeForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2MoeForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2MoeForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2MoeModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2MoePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2VLForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2VLModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen2VLPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3ForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3MoeForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3MoeForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3MoeForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3MoeForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3MoeModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Qwen3MoePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RagModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RagPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RagSequenceForGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RagTokenForGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RecurrentGemmaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RecurrentGemmaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RecurrentGemmaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ReformerForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ReformerForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ReformerForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ReformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ReformerModelWithLMHead(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ReformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RegNetForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RegNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RegNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RemBertForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RemBertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RemBertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RemBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RemBertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RemBertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RemBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RemBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_rembert(*args, **kwargs): + requires_backends(load_tf_weights_in_rembert, ["torch"]) + + +class ResNetBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ResNetForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ResNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ResNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaPreLayerNormForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaPreLayerNormForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaPreLayerNormForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaPreLayerNormForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaPreLayerNormForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaPreLayerNormForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaPreLayerNormModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RobertaPreLayerNormPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoCBertForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoCBertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoCBertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoCBertForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoCBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoCBertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoCBertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoCBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoCBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_roc_bert(*args, **kwargs): + requires_backends(load_tf_weights_in_roc_bert, ["torch"]) + + +class RoFormerForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoFormerForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoFormerForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoFormerForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoFormerForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoFormerForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoFormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RoFormerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_roformer(*args, **kwargs): + requires_backends(load_tf_weights_in_roformer, ["torch"]) + + +class RTDetrForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RTDetrModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RTDetrPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RTDetrResNetBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RTDetrResNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RTDetrV2ForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RTDetrV2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RTDetrV2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RwkvForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RwkvModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class RwkvPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SamModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SamPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SamVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4TCodeHifiGan(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4TForSpeechToSpeech(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4TForSpeechToText(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4TForTextToSpeech(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4TForTextToText(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4THifiGan(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4TModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4TPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4TTextToUnitForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4TTextToUnitModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4Tv2ForSpeechToSpeech(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4Tv2ForSpeechToText(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4Tv2ForTextToSpeech(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4Tv2ForTextToText(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4Tv2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SeamlessM4Tv2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SegformerDecodeHead(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SegformerForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SegformerForSemanticSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SegformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SegformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SegGptForImageSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SegGptModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SegGptPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SEWForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SEWForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SEWModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SEWPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SEWDForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SEWDForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SEWDModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SEWDPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ShieldGemma2ForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SiglipForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SiglipModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SiglipPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SiglipTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SiglipVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Siglip2ForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Siglip2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Siglip2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Siglip2TextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Siglip2VisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SmolVLMForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SmolVLMModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SmolVLMPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SmolVLMProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SmolVLMVisionConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SmolVLMVisionTransformer(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SpeechEncoderDecoderModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Speech2TextForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Speech2TextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Speech2TextPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SpeechT5ForSpeechToSpeech(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SpeechT5ForSpeechToText(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SpeechT5ForTextToSpeech(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SpeechT5HifiGan(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SpeechT5Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SpeechT5PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SplinterForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SplinterForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SplinterModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SplinterPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SqueezeBertForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SqueezeBertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SqueezeBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SqueezeBertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SqueezeBertForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SqueezeBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SqueezeBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class StableLmForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class StableLmForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class StableLmForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class StableLmModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class StableLmPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Starcoder2ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Starcoder2ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Starcoder2ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Starcoder2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Starcoder2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SuperGlueForKeypointMatching(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SuperGluePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SuperPointForKeypointDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SuperPointPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwiftFormerForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwiftFormerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwiftFormerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwinBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwinForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwinForMaskedImageModeling(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwinModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwinPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Swin2SRForImageSuperResolution(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Swin2SRModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Swin2SRPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Swinv2Backbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Swinv2ForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Swinv2ForMaskedImageModeling(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Swinv2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Swinv2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwitchTransformersEncoderModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwitchTransformersForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwitchTransformersModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwitchTransformersPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwitchTransformersSparseMLP(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class SwitchTransformersTop1Router(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class T5EncoderModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class T5ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class T5ForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class T5ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class T5ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class T5Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class T5PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_t5(*args, **kwargs): + requires_backends(load_tf_weights_in_t5, ["torch"]) + + +class TableTransformerForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TableTransformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TableTransformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TapasForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TapasForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TapasForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TapasModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TapasPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_tapas(*args, **kwargs): + requires_backends(load_tf_weights_in_tapas, ["torch"]) + + +class TextNetBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TextNetForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TextNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TextNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TimeSeriesTransformerForPrediction(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TimeSeriesTransformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TimeSeriesTransformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TimesformerForVideoClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TimesformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TimesformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TimmBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TimmWrapperForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TimmWrapperModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TimmWrapperPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TrOCRForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TrOCRPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TvpForVideoGrounding(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TvpModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class TvpPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UdopEncoderModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UdopForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UdopModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UdopPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UMT5EncoderModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UMT5ForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UMT5ForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UMT5ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UMT5ForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UMT5Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UMT5PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechSatForAudioFrameClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechSatForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechSatForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechSatForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechSatForXVector(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechSatModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UniSpeechSatPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UnivNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UperNetForSemanticSegmentation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class UperNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VideoLlavaForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VideoLlavaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VideoLlavaProcessor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VideoMAEForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VideoMAEForVideoClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VideoMAEModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VideoMAEPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViltForImageAndTextRetrieval(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViltForImagesAndTextClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViltForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViltForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViltForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViltModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViltPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VipLlavaForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VipLlavaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VisionEncoderDecoderModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VisionTextDualEncoderModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VisualBertForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VisualBertForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VisualBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VisualBertForRegionToPhraseAlignment(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VisualBertForVisualReasoning(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VisualBertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VisualBertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTForMaskedImageModeling(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTMAEForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTMAEModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTMAEPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTMSNForImageClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTMSNModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ViTMSNPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitDetBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitDetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitDetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitMatteForImageMatting(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitMattePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitPoseForPoseEstimation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitPosePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitPoseBackbone(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitPoseBackbonePreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitsModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VitsPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VivitForVideoClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VivitModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class VivitPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ForAudioFrameClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ForXVector(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2BertForAudioFrameClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2BertForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2BertForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2BertForXVector(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2BertModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2BertPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ConformerForAudioFrameClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ConformerForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ConformerForPreTraining(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ConformerForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ConformerForXVector(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ConformerModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Wav2Vec2ConformerPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WavLMForAudioFrameClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WavLMForCTC(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WavLMForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WavLMForXVector(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WavLMModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WavLMPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WhisperForAudioClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WhisperForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WhisperForConditionalGeneration(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WhisperModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class WhisperPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XCLIPModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XCLIPPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XCLIPTextModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XCLIPVisionModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XGLMForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XGLMModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XGLMPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMForQuestionAnsweringSimple(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMWithLMHeadModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaXLForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaXLForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaXLForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaXLForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaXLForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaXLForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaXLModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLMRobertaXLPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLNetForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLNetForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLNetForQuestionAnsweringSimple(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLNetForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLNetForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLNetLMHeadModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLNetModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XLNetPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def load_tf_weights_in_xlnet(*args, **kwargs): + requires_backends(load_tf_weights_in_xlnet, ["torch"]) + + +class XmodForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XmodForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XmodForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XmodForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XmodForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XmodForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XmodModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class XmodPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class YolosForObjectDetection(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class YolosModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class YolosPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class YosoForMaskedLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class YosoForMultipleChoice(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class YosoForQuestionAnswering(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class YosoForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class YosoForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class YosoModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class YosoPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ZambaForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ZambaForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ZambaModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ZambaPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Zamba2ForCausalLM(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Zamba2ForSequenceClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Zamba2Model(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Zamba2PreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ZoeDepthForDepthEstimation(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class ZoeDepthPreTrainedModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +class Adafactor(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def get_constant_schedule(*args, **kwargs): + requires_backends(get_constant_schedule, ["torch"]) + + +def get_constant_schedule_with_warmup(*args, **kwargs): + requires_backends(get_constant_schedule_with_warmup, ["torch"]) + + +def get_cosine_schedule_with_warmup(*args, **kwargs): + requires_backends(get_cosine_schedule_with_warmup, ["torch"]) + + +def get_cosine_with_hard_restarts_schedule_with_warmup(*args, **kwargs): + requires_backends(get_cosine_with_hard_restarts_schedule_with_warmup, ["torch"]) + + +def get_inverse_sqrt_schedule(*args, **kwargs): + requires_backends(get_inverse_sqrt_schedule, ["torch"]) + + +def get_linear_schedule_with_warmup(*args, **kwargs): + requires_backends(get_linear_schedule_with_warmup, ["torch"]) + + +def get_polynomial_decay_schedule_with_warmup(*args, **kwargs): + requires_backends(get_polynomial_decay_schedule_with_warmup, ["torch"]) + + +def get_scheduler(*args, **kwargs): + requires_backends(get_scheduler, ["torch"]) + + +def get_wsd_schedule(*args, **kwargs): + requires_backends(get_wsd_schedule, ["torch"]) + + +class Conv1D(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def apply_chunking_to_forward(*args, **kwargs): + requires_backends(apply_chunking_to_forward, ["torch"]) + + +def prune_layer(*args, **kwargs): + requires_backends(prune_layer, ["torch"]) + + +class Trainer(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + +def torch_distributed_zero_first(*args, **kwargs): + requires_backends(torch_distributed_zero_first, ["torch"]) + + +class Seq2SeqTrainer(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_sentencepiece_and_tokenizers_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_sentencepiece_and_tokenizers_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..38775330a81d91030f000e58c0e6035bba1c0f31 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_sentencepiece_and_tokenizers_objects.py @@ -0,0 +1,9 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +SLOW_TO_FAST_CONVERTERS = None + + +def convert_slow_tokenizer(*args, **kwargs): + requires_backends(convert_slow_tokenizer, ["sentencepiece", "tokenizers"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_sentencepiece_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_sentencepiece_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..7931e0fe6584bbe85ed966911a14bc57c0d43c48 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_sentencepiece_objects.py @@ -0,0 +1,254 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class AlbertTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class BarthezTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class BartphoTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class BertGenerationTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class BigBirdTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class CamembertTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class CodeLlamaTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class CpmTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class DebertaV2Tokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class ErnieMTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class XLMProphetNetTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class FNetTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class GemmaTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class GPTSw3Tokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class LayoutXLMTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class LlamaTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class M2M100Tokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class MarianTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class MBartTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class MBart50Tokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class MLukeTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class MT5Tokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class NllbTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class PegasusTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class PLBartTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class ReformerTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class RemBertTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class SeamlessM4TTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class SiglipTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class Speech2TextTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class SpeechT5Tokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class T5Tokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class UdopTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class XGLMTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class XLMRobertaTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) + + +class XLNetTokenizer(metaclass=DummyObject): + _backends = ["sentencepiece"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["sentencepiece"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_speech_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_speech_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..0bf08ebea42b4595ae1f8bbc2afcddf0630dcf4b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_speech_objects.py @@ -0,0 +1,16 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class ASTFeatureExtractor(metaclass=DummyObject): + _backends = ["speech"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["speech"]) + + +class Speech2TextFeatureExtractor(metaclass=DummyObject): + _backends = ["speech"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["speech"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_tensorflow_text_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_tensorflow_text_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..70c7ad5cbf4077609e36592566e461c1a1ded28a --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_tensorflow_text_objects.py @@ -0,0 +1,9 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class TFBertTokenizer(metaclass=DummyObject): + _backends = ["tensorflow_text"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tensorflow_text"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_tf_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_tf_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..985445fbba567b87bea76a81e1c980cdaf4b2daf --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_tf_objects.py @@ -0,0 +1,2869 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class TFForcedBOSTokenLogitsProcessor(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFForcedEOSTokenLogitsProcessor(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFForceTokensLogitsProcessor(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGenerationMixin(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLogitsProcessor(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLogitsProcessorList(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLogitsWarper(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMinLengthLogitsProcessor(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFNoBadWordsLogitsProcessor(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFNoRepeatNGramLogitsProcessor(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRepetitionPenaltyLogitsProcessor(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSuppressTokensAtBeginLogitsProcessor(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSuppressTokensLogitsProcessor(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTemperatureLogitsWarper(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTopKLogitsWarper(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTopPLogitsWarper(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class KerasMetricCallback(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class PushToHubCallback(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSequenceSummary(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSharedEmbeddings(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +def shape_list(*args, **kwargs): + requires_backends(shape_list, ["tf"]) + + +class TFAlbertForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAlbertForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAlbertForPreTraining(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAlbertForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAlbertForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAlbertForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAlbertMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAlbertModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAlbertPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING = None + + +TF_MODEL_FOR_CAUSAL_LM_MAPPING = None + + +TF_MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING = None + + +TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING = None + + +TF_MODEL_FOR_MASK_GENERATION_MAPPING = None + + +TF_MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING = None + + +TF_MODEL_FOR_MASKED_LM_MAPPING = None + + +TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING = None + + +TF_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING = None + + +TF_MODEL_FOR_PRETRAINING_MAPPING = None + + +TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING = None + + +TF_MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING = None + + +TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = None + + +TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = None + + +TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING = None + + +TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING = None + + +TF_MODEL_FOR_TEXT_ENCODING_MAPPING = None + + +TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = None + + +TF_MODEL_FOR_VISION_2_SEQ_MAPPING = None + + +TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING = None + + +TF_MODEL_MAPPING = None + + +TF_MODEL_WITH_LM_HEAD_MAPPING = None + + +class TFAutoModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForAudioClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForDocumentQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForMaskedImageModeling(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForMaskGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForNextSentencePrediction(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForPreTraining(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForSemanticSegmentation(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForSeq2SeqLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForSpeechSeq2Seq(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForTableQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForTextEncoding(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForVision2Seq(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelForZeroShotImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAutoModelWithLMHead(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBartForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBartForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBartModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBartPretrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertForNextSentencePrediction(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertForPreTraining(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertLMHeadModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBertPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlenderbotForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlenderbotModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlenderbotPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlenderbotSmallForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlenderbotSmallModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlenderbotSmallPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlipForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlipForImageTextRetrieval(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlipForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlipModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlipPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlipTextModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFBlipVisionModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCamembertForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCamembertForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCamembertForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCamembertForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCamembertForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCamembertForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCamembertModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCamembertPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCLIPModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCLIPPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCLIPTextModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCLIPVisionModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvBertForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvBertForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvBertForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvBertForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvBertModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvBertPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvNextForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvNextModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvNextPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvNextV2ForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvNextV2Model(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFConvNextV2PreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCTRLForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCTRLLMHeadModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCTRLModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCTRLPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCvtForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCvtModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFCvtPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFData2VecVisionForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFData2VecVisionForSemanticSegmentation(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFData2VecVisionModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFData2VecVisionPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaV2ForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaV2ForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaV2ForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaV2ForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaV2ForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaV2Model(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDebertaV2PreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDeiTForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDeiTForImageClassificationWithTeacher(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDeiTForMaskedImageModeling(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDeiTModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDeiTPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFEfficientFormerForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFEfficientFormerForImageClassificationWithTeacher(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFEfficientFormerModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFEfficientFormerPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFAdaptiveEmbedding(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTransfoXLForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTransfoXLLMHeadModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTransfoXLMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTransfoXLModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTransfoXLPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDistilBertForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDistilBertForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDistilBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDistilBertForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDistilBertForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDistilBertMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDistilBertModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDistilBertPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDPRContextEncoder(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDPRPretrainedContextEncoder(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDPRPretrainedQuestionEncoder(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDPRPretrainedReader(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDPRQuestionEncoder(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFDPRReader(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFElectraForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFElectraForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFElectraForPreTraining(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFElectraForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFElectraForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFElectraForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFElectraModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFElectraPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFEncoderDecoderModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFEsmForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFEsmForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFEsmForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFEsmModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFEsmPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFlaubertForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFlaubertForQuestionAnsweringSimple(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFlaubertForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFlaubertForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFlaubertModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFlaubertPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFlaubertWithLMHeadModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFunnelBaseModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFunnelForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFunnelForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFunnelForPreTraining(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFunnelForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFunnelForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFunnelForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFunnelModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFFunnelPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPT2DoubleHeadsModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPT2ForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPT2LMHeadModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPT2MainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPT2Model(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPT2PreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPTJForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPTJForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPTJForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPTJModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGPTJPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGroupViTModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGroupViTPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGroupViTTextModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFGroupViTVisionModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFHubertForCTC(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFHubertModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFHubertPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFIdeficsForVisionText2Text(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFIdeficsModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFIdeficsPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMv3ForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMv3ForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMv3ForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMv3Model(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLayoutLMv3PreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLEDForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLEDModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLEDPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLongformerForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLongformerForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLongformerForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLongformerForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLongformerForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLongformerModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLongformerPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLxmertForPreTraining(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLxmertMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLxmertModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLxmertPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFLxmertVisualFeatureEncoder(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMarianModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMarianMTModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMarianPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMBartForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMBartModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMBartPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMistralForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMistralForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMistralModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMistralPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileBertForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileBertForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileBertForNextSentencePrediction(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileBertForPreTraining(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileBertForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileBertForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileBertMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileBertModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileBertPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileViTForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileViTForSemanticSegmentation(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileViTModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMobileViTPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMPNetForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMPNetForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMPNetForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMPNetForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMPNetForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMPNetMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMPNetModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMPNetPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMT5EncoderModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMT5ForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFMT5Model(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFOpenAIGPTDoubleHeadsModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFOpenAIGPTForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFOpenAIGPTLMHeadModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFOpenAIGPTMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFOpenAIGPTModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFOpenAIGPTPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFOPTForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFOPTModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFOPTPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFPegasusForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFPegasusModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFPegasusPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRagModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRagPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRagSequenceForGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRagTokenForGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRegNetForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRegNetModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRegNetPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRemBertForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRemBertForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRemBertForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRemBertForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRemBertForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRemBertForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRemBertModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRemBertPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFResNetForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFResNetModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFResNetPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaPreLayerNormForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaPreLayerNormForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaPreLayerNormForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaPreLayerNormForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaPreLayerNormForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaPreLayerNormForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaPreLayerNormMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaPreLayerNormModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRobertaPreLayerNormPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRoFormerForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRoFormerForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRoFormerForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRoFormerForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRoFormerForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRoFormerForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRoFormerModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFRoFormerPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSamModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSamPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSamVisionModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSegformerDecodeHead(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSegformerForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSegformerForSemanticSegmentation(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSegformerModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSegformerPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSpeech2TextForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSpeech2TextModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSpeech2TextPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSwiftFormerForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSwiftFormerModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSwiftFormerPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSwinForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSwinForMaskedImageModeling(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSwinModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFSwinPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFT5EncoderModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFT5ForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFT5Model(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFT5PreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTapasForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTapasForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTapasForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTapasModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFTapasPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFVisionEncoderDecoderModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFVisionTextDualEncoderModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFViTForImageClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFViTModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFViTPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFViTMAEForPreTraining(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFViTMAEModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFViTMAEPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFWav2Vec2ForCTC(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFWav2Vec2ForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFWav2Vec2Model(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFWav2Vec2PreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFWhisperForConditionalGeneration(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFWhisperModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFWhisperPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXGLMForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXGLMModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXGLMPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMForQuestionAnsweringSimple(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMWithLMHeadModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMRobertaForCausalLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMRobertaForMaskedLM(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMRobertaForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMRobertaForQuestionAnswering(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMRobertaForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMRobertaForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMRobertaModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLMRobertaPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLNetForMultipleChoice(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLNetForQuestionAnsweringSimple(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLNetForSequenceClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLNetForTokenClassification(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLNetLMHeadModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLNetMainLayer(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLNetModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class TFXLNetPreTrainedModel(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class AdamWeightDecay(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class GradientAccumulator(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +class WarmUp(metaclass=DummyObject): + _backends = ["tf"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tf"]) + + +def create_optimizer(*args, **kwargs): + requires_backends(create_optimizer, ["tf"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_timm_and_torchvision_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_timm_and_torchvision_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..8b67b5dac58db13278706e0838aebbe495a9d2fc --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_timm_and_torchvision_objects.py @@ -0,0 +1,9 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class TimmWrapperImageProcessor(metaclass=DummyObject): + _backends = ["timm", "torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["timm", "torchvision"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_tokenizers_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_tokenizers_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..df83e6fa6478e6c192ff5f80ccd0d0fe8916180b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_tokenizers_objects.py @@ -0,0 +1,457 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class AlbertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class BartTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class BarthezTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class BertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class BigBirdTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class BlenderbotTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class BlenderbotSmallTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class BloomTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class CamembertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class CLIPTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class CodeLlamaTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class CodeGenTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class CohereTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class ConvBertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class CpmTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class DebertaTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class DebertaV2TokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class RealmTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class RetriBertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class DistilBertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class DPRContextEncoderTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class DPRQuestionEncoderTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class DPRReaderTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class ElectraTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class FNetTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class FunnelTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class GemmaTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class GPT2TokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class GPTNeoXTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class GPTNeoXJapaneseTokenizer(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class HerbertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class LayoutLMTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class LayoutLMv2TokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class LayoutLMv3TokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class LayoutXLMTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class LEDTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class LlamaTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class LongformerTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class LxmertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class MarkupLMTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class MBartTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class MBart50TokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class MobileBertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class MPNetTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class MT5TokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class MvpTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class NllbTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class NougatTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class OpenAIGPTTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class PegasusTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class Qwen2TokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class ReformerTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class RemBertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class RobertaTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class RoFormerTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class SeamlessM4TTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class SplinterTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class SqueezeBertTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class T5TokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class UdopTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class WhisperTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class XGLMTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class XLMRobertaTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class XLNetTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) + + +class PreTrainedTokenizerFast(metaclass=DummyObject): + _backends = ["tokenizers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["tokenizers"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_torchaudio_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_torchaudio_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..58b01f06a8ab5e2ea0711586a8ef5726ba5cc256 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_torchaudio_objects.py @@ -0,0 +1,16 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class MusicgenMelodyFeatureExtractor(metaclass=DummyObject): + _backends = ["torchaudio"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchaudio"]) + + +class MusicgenMelodyProcessor(metaclass=DummyObject): + _backends = ["torchaudio"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchaudio"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_torchvision_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_torchvision_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..c59c9b4bdd55f77c8215b788f6626ea5b5073563 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_torchvision_objects.py @@ -0,0 +1,149 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class BaseImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class BlipImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class CLIPImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class ConvNextImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class DeformableDetrImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class DeiTImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class DepthProImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class DetrImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class Gemma3ImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class GotOcr2ImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class Llama4ImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class LlavaImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class LlavaNextImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class LlavaOnevisionImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class Phi4MultimodalImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class PixtralImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class Qwen2VLImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class RTDetrImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class SiglipImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class Siglip2ImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) + + +class ViTImageProcessorFast(metaclass=DummyObject): + _backends = ["torchvision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torchvision"]) diff --git a/lib/python3.12/site-packages/transformers/utils/dummy_vision_objects.py b/lib/python3.12/site-packages/transformers/utils/dummy_vision_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..4e7293d50ae99e8343bee3dbc2e4ddc301198b57 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/dummy_vision_objects.py @@ -0,0 +1,807 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class ImageProcessingMixin(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class BaseImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ImageFeatureExtractionMixin(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class AriaImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class BeitFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class BeitImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class BitImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class BlipImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class BridgeTowerImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ChameleonImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ChineseCLIPFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ChineseCLIPImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class CLIPFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class CLIPImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ConditionalDetrFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ConditionalDetrImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ConvNextFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ConvNextImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DeformableDetrFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DeformableDetrImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DeiTFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DeiTImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DetaImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class EfficientFormerImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class TvltImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ViTHybridImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DepthProImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DepthProImageProcessorFast(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DetrFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DetrImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DonutFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DonutImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DPTFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class DPTImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class EfficientNetImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Emu3ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class FlavaFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class FlavaImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class FlavaProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class FuyuImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class FuyuProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Gemma3ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class GLPNFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class GLPNImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class GotOcr2ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class GroundingDinoImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class IdeficsImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Idefics2ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Idefics3ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ImageGPTFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ImageGPTImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class InstructBlipVideoImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LayoutLMv2FeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LayoutLMv2ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LayoutLMv3FeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LayoutLMv3ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LevitFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LevitImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Llama4ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LlavaImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LlavaNextImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LlavaNextVideoImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LlavaOnevisionImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class LlavaOnevisionVideoProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Mask2FormerImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class MaskFormerFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class MaskFormerImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class MllamaImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class MobileNetV1FeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class MobileNetV1ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class MobileNetV2FeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class MobileNetV2ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class MobileViTFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class MobileViTImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class NougatImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class OneFormerImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Owlv2ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class OwlViTFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class OwlViTImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class PerceiverFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class PerceiverImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Pix2StructImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class PixtralImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class PoolFormerFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class PoolFormerImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class PromptDepthAnythingImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class PvtImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Qwen2VLImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class RTDetrImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class SamImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class SegformerFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class SegformerImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class SegGptImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class SiglipImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Siglip2ImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class SmolVLMImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class SuperGlueImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class SuperPointImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class Swin2SRImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class TextNetImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class TvpImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class VideoLlavaImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class VideoMAEFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class VideoMAEImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ViltFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ViltImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ViltProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ViTFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ViTImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class VitMatteImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class VitPoseImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class VivitImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class YolosFeatureExtractor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class YolosImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) + + +class ZoeDepthImageProcessor(metaclass=DummyObject): + _backends = ["vision"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["vision"]) diff --git a/lib/python3.12/site-packages/transformers/utils/fx.py b/lib/python3.12/site-packages/transformers/utils/fx.py new file mode 100644 index 0000000000000000000000000000000000000000..6ab0c45d996ab7374bc9489e2a63cf2a43d3ada6 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/fx.py @@ -0,0 +1,1512 @@ +# Copyright 2021 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. + +import builtins +import collections +import contextlib +import functools +import inspect +import math +import operator +import os +import random +import sys +import warnings +from typing import Any, Callable, Literal, Optional, Union + +import torch +import torch.utils._pytree as pytree +from torch import nn +from torch.fx import Graph, GraphModule, Node, Proxy, Tracer +from torch.fx._compatibility import compatibility +from torch.fx._symbolic_trace import is_fx_tracing +from torch.fx.proxy import ParameterProxy + +from .. import logging +from ..cache_utils import Cache, DynamicCache, SinkCache, StaticCache +from ..modeling_utils import PretrainedConfig, PreTrainedModel +from ..models.auto import get_values +from ..models.auto.modeling_auto import ( + MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES, + MODEL_FOR_BACKBONE_MAPPING_NAMES, + MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, + MODEL_FOR_CTC_MAPPING_NAMES, + MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES, + MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES, + MODEL_FOR_IMAGE_MAPPING_NAMES, + MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING_NAMES, + MODEL_FOR_MASKED_LM_MAPPING_NAMES, + MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES, + MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES, + MODEL_FOR_PRETRAINING_MAPPING_NAMES, + MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES, + MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES, + MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES, + MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES, + MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES, + MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES, + MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES, + MODEL_MAPPING_NAMES, +) +from .import_utils import ( + ENV_VARS_TRUE_VALUES, + TORCH_FX_REQUIRED_VERSION, + get_torch_version, + is_peft_available, + is_torch_fx_available, +) + + +if is_peft_available(): + from peft import PeftModel + + +logger = logging.get_logger(__name__) +_IS_IN_DEBUG_MODE = os.environ.get("FX_DEBUG_MODE", "").upper() in ENV_VARS_TRUE_VALUES + + +def _generate_supported_model_class_names( + model_name: type[PretrainedConfig], + supported_tasks: Optional[Union[str, list[str]]] = None, +) -> list[str]: + task_mapping = { + "default": MODEL_MAPPING_NAMES, + "pretraining": MODEL_FOR_PRETRAINING_MAPPING_NAMES, + "next-sentence-prediction": MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES, + "masked-lm": MODEL_FOR_MASKED_LM_MAPPING_NAMES, + "causal-lm": MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, + "seq2seq-lm": MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES, + "speech-seq2seq": MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES, + "multiple-choice": MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES, + "document-question-answering": MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES, + "question-answering": MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES, + "sequence-classification": MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES, + "token-classification": MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES, + "masked-image-modeling": MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING_NAMES, + "image-classification": MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES, + "zero-shot-image-classification": MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES, + "ctc": MODEL_FOR_CTC_MAPPING_NAMES, + "audio-classification": MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES, + "semantic-segmentation": MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES, + "backbone": MODEL_FOR_BACKBONE_MAPPING_NAMES, + "image-feature-extraction": MODEL_FOR_IMAGE_MAPPING_NAMES, + } + + if supported_tasks is None: + supported_tasks = task_mapping.keys() + if isinstance(supported_tasks, str): + supported_tasks = [supported_tasks] + + model_class_names = [] + for task in supported_tasks: + class_name = task_mapping[task].get(model_name, None) + if class_name: + model_class_names.append(class_name) + + return model_class_names + + +_REGULAR_SUPPORTED_MODEL_NAMES_AND_TASKS = [ + "altclip", + "albert", + "bart", + "bert", + "blenderbot", + "blenderbot-small", + "bloom", + "clip", + "convnext", + "deberta", + "deberta-v2", + "dinov2", + "distilbert", + "donut-swin", + "electra", + "gpt2", + "gpt_neo", + "gptj", + "hiera", + "hubert", + "ijepa", + "layoutlm", + "llama", + "cohere", + "lxmert", + "m2m_100", + "marian", + "mbart", + "megatron-bert", + "mistral", + "mixtral", + "mobilebert", + "mt5", + "nezha", + "opt", + "pegasus", + "plbart", + "qwen2", + "qwen2_moe", + "qwen3", + "qwen3_moe", + "resnet", + "roberta", + "segformer", + "speech_to_text", + "speech_to_text_2", + "swin", + "t5", + "trocr", + "vit", + "xglm", + "wav2vec2", + # "xlnet", +] + +_FX_SUPPORTED_MODELS_WITH_KV_CACHE = ["llama", "opt"] + +_REGULAR_SUPPORTED_MODELS = [] +for item in _REGULAR_SUPPORTED_MODEL_NAMES_AND_TASKS: + if isinstance(item, dict): + _REGULAR_SUPPORTED_MODELS.extend(_generate_supported_model_class_names(**item)) + else: + _REGULAR_SUPPORTED_MODELS.extend(_generate_supported_model_class_names(item)) + +_SPECIAL_SUPPORTED_MODELS = [ + "CLIPTextModel", + "CLIPTextModelWithProjection", + "CLIPVisionModel", + "CLIPVisionModelWithProjection", + "AltCLIPTextModel", + "AltCLIPVisionModel", + "GitVisionModel", + "GPT2DoubleHeadsModel", + "Speech2Text2Decoder", + "TrOCRDecoder", + "PeftModelForCausalLM", + "PeftModelForSeq2SeqLM", + # TODO: add support for them as it should be quite easy to do so (small blocking issues). + # XLNetForQuestionAnswering, +] +_SUPPORTED_MODELS = tuple(sorted(set(_REGULAR_SUPPORTED_MODELS + _SPECIAL_SUPPORTED_MODELS))) + +_CURRENT_TRACER = None + + +def torch_nn_embedding(self, input): + return torch.empty(*input.shape, self.weight.shape[-1], device="meta", dtype=self.weight.dtype) + + +def torch_nn_functional_embedding( + input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False +): + return torch.empty(*input.shape, weight.shape[-1], device="meta", dtype=weight.dtype) + + +def torch_nn_layernorm(self, input): + return input + + +def torch_nn_groupnorm(self, input): + return input + + +def torch_nn_linear(self, input): + return torch.empty(input.shape[:-1] + (self.out_features,), device="meta") + + +def torch_relu(x): + return x + + +def torch_nn_relu(self, x): + return x + + +def torch_nn_functional_relu(x, inplace=False): + if not inplace: + raise ValueError("Don't support in-place functional.relu for MetaTensor analysis") + return x + + +def torch_where(condition, x, y): + # torch.where returns the broadcasted tensor of condition, x, and y, + # so hack it by using addition + return condition.to(device="meta") + x.to(device="meta") + y.to(device="meta") + + +def torch_abs(input, *, out=None): + if out is not None: + raise ValueError("Don't support in-place abs for MetaTensor analysis") + return input + + +def torch_arange(*args, **kwargs): + n = len(args) + step = 1 + if n == 1: + start = 0 + end = args[0] + elif n == 2: + start, end = args + else: + start, end, step = args + if isinstance(start, float): + start = int(start) + if isinstance(end, float): + start = int(end) + if isinstance(step, float): + step = int(step) + step = kwargs.get("step", step) + dtype = kwargs.get("dtype") + return torch.empty((end - start) // step, dtype=dtype, device="meta") + + +def torch_full(*args, **kwargs): + args = list(args) + # We set the fill value to 1 as its value is not important as long as it's not a tensor on the `meta` device. + if len(args) > 1: + args[1] = 1 + else: + kwargs["fill_value"] = 1 + kwargs_without_device = dict(kwargs) + kwargs_without_device.pop("device", None) + return torch.full(*args, **kwargs_without_device, device="meta") + + +def torch_cat(tensors, dim=None, axis=None, *, out=None): + if dim is None and axis is None: + dim = 0 + if dim is None and axis is not None: + dim = axis + if dim < 0: + dim = tensors[0].dim() + dim + shapes = [t.shape for t in tensors] + shape = list(shapes[0]) + concatenated_dim = sum(shape[dim] for shape in shapes) + final_shape = shape[:dim] + [concatenated_dim] + shape[dim + 1 :] + return torch.empty(final_shape, device="meta") + + +def torch_stack(tensors, dim=None, axis=None, *, out=None): + if dim is None and axis is None: + dim = 0 + if dim is None and axis is not None: + dim = axis + if dim < 0: + dim = tensors[0].dim() + 1 + dim + shape = list(tensors[0].shape) + shape.insert(dim, len(tensors)) + return torch.empty(shape, device="meta") + + +def torch_add(input, other, *, alpha=1, out=None): + if not isinstance(input, torch.Tensor): + return torch.empty_like(other, device="meta") + if not isinstance(other, torch.Tensor): + return torch.empty_like(input, device="meta") + max_length = max(input.dim(), other.dim()) + input_shape = list(input.shape) + [1] * (max_length - input.dim()) + other_shape = list(other.shape) + [1] * (max_length - other.dim()) + shape = [] + for i in range(max_length): + shape.append(max(input_shape[i], other_shape[i])) + return torch.empty(shape, device="meta") + + +def torch_mul(input, other, *, out=None): + return torch_add(input, other, out=out) + + +def torch_tensor_mul(self, other): + return torch_mul(self, other) + + +def torch_matmul(input, other, *, out=None): + d1 = input.dim() + d2 = other.dim() + shape = None + if d1 == 1 and d2 == 1: + shape = None + elif d1 == 2 and d2 == 2: + shape = (input.size(0), other.size(1)) + elif d1 == 1 and d2 == 2: + shape = (other.size(1),) + elif d1 == 2 and d1 == 1: + shape = (input.size(0),) + else: + max_length = max(input.dim(), other.dim()) + shape1 = list(input.shape) + shape2 = list(other.shape) + if d1 == 1: + shape1 = [1] + shape1 + if d2 == 1: + shape2.append(1) + shape1 = [-1] * (max_length - d1) + list(input.shape) + shape2 = [-1] * (max_length - d2) + list(other.shape) + shape = [] + for i in range(max_length): + shape.append(max(shape1[i], shape2[i])) + shape[-2] = shape1[-2] + shape[-1] = shape2[-1] + if d1 == 1: + shape.pop(-2) + if d2 == 1: + shape.pop(-1) + if shape is None: + return torch.tensor(0.0, device="meta") + return torch.empty(*shape, device="meta") + + +def torch_bmm(input, mat2, *, out=None): + if out is not None: + raise ValueError("Don't support in-place bmm for MetaTensor analysis") + batch_size, n, m = input.shape + _, _, p = mat2.shape + return torch.empty(batch_size, n, p, device="meta") + + +def torch_baddbmm(input, batch1, batch2, *, beta=1, alpha=1, out=None): + if out is not None: + raise ValueError("Don't support in-place baddbmm for MetaTensor analysis") + return torch_bmm(batch1, batch2) + + +def torch_tensor_baddbmm(self, batch1, batch2, *, beta=1, alpha=1, out=None): + return torch_baddbmm(self, batch1, batch2, beta=beta, alpha=alpha, out=out) + + +def torch_einsum(equation, *operands): + # TODO: infer shape without performing the computation, this might be quite hard. + concrete_operands = (torch.empty_like(operand, device="cpu") for operand in operands) + return torch.einsum(equation, *concrete_operands).to("meta") + + +def torch_tensor_repeat(self, *sizes): + shape = list(self.shape) + for i, x in enumerate(sizes): + shape[i] *= x + return torch.empty(shape, device="meta") + + +def torch_repeat_interleave(*args, dim=None, output_size=None): + num_args = len(args) + if num_args == 1: + shape = [output_size if output_size is not None else args[0].sum()] + else: + shape = list(args[0].shape) + if dim is None: + if num_args > 2: + dim = args[2] + else: + shape = [sum(shape)] + dim = 0 + repeats = args[1] + if isinstance(repeats, int) or torch.numel(repeats) == 1: + shape[dim] *= int(repeats) + else: + shape[dim] = output_size if output_size is not None else repeats.sum() + return torch.empty(*shape, device="meta") + + +def torch_index_select(input, dim, index, *, out=None): + shape = list(input.shape) + shape[dim] = len(index) + return torch.empty(*shape, device="meta") + + +def torch_tensor_index_select(self, dim, index): + return torch_index_select(self, dim, index) + + +def torch_gather(input, dim, index, *, sparse_grad=False, out=None): + shape = list(input.shape) + shape[dim] = index.shape[dim] + return torch.empty(*shape, device="meta") + + +def torch_tensor_gather(self, dim, index): + return torch_gather(self, dim, index) + + +def torch_roll(input, shifts, dims=None): + return input + + +def torch_flip(input, dims): + return input + + +def torch_tensor_flip(self, dims): + return self + + +def torch_nn_conv1d(self, input): + l_in = input.shape[-1] + shape = None + padding = self.padding + if padding == "valid": + padding = (0, 0) + if padding == "same": + shape = list(input.shape) + if shape is None: + shape = list(input.shape) + l_out = math.floor( + (l_in + 2 * padding[0] - self.dilation[0] * (self.kernel_size[0] - 1) - 1) / self.stride[0] + 1 + ) + shape[-1] = l_out + shape[-2] = self.out_channels + return torch.empty(shape, device="meta") + + +def torch_nn_conv2d(self, input): + h_in, w_in = input.shape[-2:] + shape = None + padding = self.padding + if padding == "valid": + padding = (0, 0) + if padding == "same": + shape = list(input.shape) + if shape is None: + shape = list(input.shape) + h_out = math.floor( + (h_in + 2 * padding[0] - self.dilation[0] * (self.kernel_size[0] - 1) - 1) / self.stride[0] + 1 + ) + w_out = math.floor( + (w_in + 2 * padding[1] - self.dilation[1] * (self.kernel_size[1] - 1) - 1) / self.stride[1] + 1 + ) + shape[-2:] = [h_out, w_out] + shape[-3] = self.out_channels + return torch.empty(shape, device="meta") + + +def torch_squeeze(input, dim=None): + shape = list(input.shape) + if dim is not None: + if dim < 0: + dim = input.dim() + dim + if shape[dim] == 1: + shape.pop(dim) + else: + new_shape = [] + for dim_value in shape: + if dim_value == 1: + continue + new_shape.append(dim_value) + shape = new_shape + return torch.empty(shape, device="meta") + + +def torch_tensor_squeeze(self, dim=None): + return torch_squeeze(self, dim) + + +def torch_unsqueeze(input, dim): + shape = list(input.shape) + if dim < 0: + dim = input.dim() + 1 + dim + shape.insert(dim, 1) + return torch.empty(shape, device="meta") + + +def torch_tensor_unsqueeze(self, dim): + return torch_unsqueeze(self, dim) + + +def torch_unique_consecutive(input, **kwargs): + output = torch.unique_consecutive(torch.zeros_like(input, device="cpu"), **kwargs) + if isinstance(output, torch.Tensor): + return output.to("meta") + else: + return tuple(map(output, lambda x: x.to("meta"))) + + +def torch_nn_functional_one_hot(tensor, num_classes=-1): + if num_classes < 0: + raise ValueError("Don't support automatic num_classes inference for MetaTensor analysis") + shape = list(tensor.shape) + [num_classes] + return torch.empty(shape, device="meta") + + +def torch_nn_functional_scaled_dot_product_attention( + query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None +): + target_length = query.shape[-2] + head_dim = value.shape[-1] + return torch.empty((*query.shape[:-2], target_length, head_dim), device="meta") + + +def torch_nn_mseloss(self, input, target): + if self.reduction == "none": + shape = target.shape + else: + shape = (1,) + return torch.empty(shape, device="meta") + + +def torch_nn_crossentropyloss(self, input, target): + if self.reduction == "none": + shape = target.shape + else: + shape = (1,) + return torch.empty(shape, device="meta") + + +def torch_nn_bcewithlogitsloss(self, input, target): + if self.reduction == "none": + shape = target.shape + else: + shape = (1,) + return torch.empty(shape, device="meta") + + +def operator_getitem(a, b): + def to_concrete(t): + if isinstance(t, torch.Tensor): + concrete = torch.ones_like(t, device="cpu") + if concrete.dtype in [torch.float16, torch.float32, torch.float64, torch.int32]: + concrete = concrete.to(torch.int64) + return concrete + return t + + if isinstance(a, torch.Tensor): + # TODO: infer shape without performing the computation. + if isinstance(b, tuple): + b = tuple(map(to_concrete, b)) + else: + b = to_concrete(b) + return operator.getitem(torch.empty_like(a, device="cpu"), b).to("meta") + return operator.getitem(a, b) + + +_MANUAL_META_OVERRIDES: dict[Callable, Callable] = { + torch.nn.Embedding: torch_nn_embedding, + torch.nn.functional.embedding: torch_nn_functional_embedding, + torch.nn.LayerNorm: torch_nn_layernorm, + torch.nn.GroupNorm: torch_nn_groupnorm, + torch.nn.Linear: torch_nn_linear, + torch.relu: torch_relu, + torch.nn.functional.relu: torch_nn_functional_relu, + torch.nn.ReLU: torch_nn_relu, + torch.where: torch_where, + torch.abs: torch_abs, + torch.arange: torch_arange, + torch.full: torch_full, + torch.cat: torch_cat, + torch.stack: torch_stack, + torch.add: torch_add, + torch.mul: torch_mul, + torch.Tensor.mul: torch_tensor_mul, + torch.matmul: torch_matmul, + torch.bmm: torch_bmm, + torch.baddbmm: torch_baddbmm, + torch.Tensor.baddbmm: torch_tensor_baddbmm, + torch.einsum: torch_einsum, + torch.Tensor.repeat: torch_tensor_repeat, + torch.repeat_interleave: torch_repeat_interleave, + torch.roll: torch_roll, + torch.flip: torch_flip, + torch.Tensor.flip: torch_tensor_flip, + torch.index_select: torch_index_select, + torch.Tensor.index_select: torch_tensor_index_select, + torch.gather: torch_gather, + torch.Tensor.gather: torch_tensor_gather, + torch.nn.Conv1d: torch_nn_conv1d, + torch.nn.Conv2d: torch_nn_conv2d, + torch.squeeze: torch_squeeze, + torch.Tensor.squeeze: torch_tensor_squeeze, + torch.unsqueeze: torch_unsqueeze, + torch.Tensor.unsqueeze: torch_tensor_unsqueeze, + torch.unique_consecutive: torch_unique_consecutive, + torch.nn.functional.one_hot: torch_nn_functional_one_hot, + torch.nn.MSELoss: torch_nn_mseloss, + torch.nn.CrossEntropyLoss: torch_nn_crossentropyloss, + torch.nn.BCEWithLogitsLoss: torch_nn_bcewithlogitsloss, + operator.getitem: operator_getitem, +} + +_MANUAL_META_OVERRIDES[torch.nn.functional.scaled_dot_product_attention] = ( + torch_nn_functional_scaled_dot_product_attention +) + + +class HFProxy(Proxy): + """ + Proxy that uses metadata to handle data-dependent control-flow. + """ + + def install_metadata(self, metadata): + self._metadata = metadata + + @property + def shape(self): + return self.tracer.create_proxy("call_method", "size", (self,), {}) + + @property + def device(self): + # Hack so we can track when devices are used. During meta-tensor propagation, + # replace these values with a constant 'meta' + return MetaDeviceAttribute(self, "device") + + def __len__(self): + if hasattr(self, "_metadata") and self._metadata is not None: + return len(self._metadata) + return super().__len__() + + def __bool__(self): + if hasattr(self, "_metadata") and self._metadata is not None: + return self._metadata + return super().__bool__() + + def __getattr__(self, k): + if k == "_metadata": + return self.__getattribute__(k) + # note: not added to the graph yet, if this is a method call + # we peephole optimize to the method invocation + return HFAttribute(self, k) + + def __setitem__(self, indices, values): + return self.tracer.create_proxy("call_function", operator.setitem, (self, indices, values), {}) + + def __contains__(self, key): + if hasattr(self, "_metadata") and self._metadata is not None: + return key in self._metadata + return super().__contains__(key) + + +class HFAttribute(HFProxy): + def __init__(self, root, attr: str): + self.root = root + self.attr = attr + self.tracer = root.tracer + self._node = None + + if hasattr(self.root, "_metadata"): + self.install_metadata(getattr(self.root._metadata, attr)) + + @property + def node(self): + # the node for attributes is added lazily, since most will just be method calls + # which do not rely on the getitem call + if self._node is None: + self._node = self.tracer.create_proxy("call_function", builtins.getattr, (self.root, self.attr), {}).node + return self._node + + def __call__(self, *args, **kwargs): + return self.tracer.create_proxy("call_method", self.attr, (self.root,) + args, kwargs) + + +class MetaDeviceAttribute(HFAttribute): + pass + + +class HFCacheProxy(HFProxy): + """ + Proxy that represents an instance of `transformers.cache_utils.Cache`. + """ + + def install_orig_cache_cls(self, orig_cache_cls: type[Cache]): + self._orig_cache_cls = orig_cache_cls + + @property + def __class__(self): + if not hasattr(self, "_orig_cache_cls"): + raise RuntimeError("The original Cache class must be installed to the HFCacheProxy.") + return self.tracer._CLASSES_TO_PATCH[self._orig_cache_cls] + + +def create_wrapper( + function: Callable, + op_type: Union[Literal["call_function"], Literal["call_method"], Literal["get_attr"]], + proxy_factory_fn: Optional[Callable[[Node], Proxy]] = None, +) -> Callable: + @functools.wraps(function) + def wrapper(*args, **kwargs): + if not is_fx_tracing(): + return function(*args, **kwargs) + + found_proxies = [] + + def check_proxy(a): + if isinstance(a, Proxy): + found_proxies.append(a) + + torch.fx.node.map_aggregate(args, check_proxy) + torch.fx.node.map_aggregate(kwargs, check_proxy) + + if len(found_proxies) > 0: + tracer = found_proxies[0].tracer + if op_type == "call_function": + target = function + elif op_type == "call_method": + target = function.__name__ + elif op_type == "get_attr": + target = function.__name__ + else: + raise ValueError(f"op_type {op_type} not supported.") + return tracer.create_proxy(op_type, target, args, kwargs, proxy_factory_fn=proxy_factory_fn) + else: + return function(*args, **kwargs) + + return wrapper + + +class HFProxyableClassMeta(type): + """ + Metaclass that creates a class with its main methods wrapped to be proxyable. + """ + + def __new__( + cls, + name: str, + bases: tuple[type, ...], + attrs: dict[str, Any], + proxy_factory_fn: Optional[Callable[[Node], Proxy]] = None, + ): + cls = super().__new__(cls, name, bases, attrs) + for attr_name in dir(cls): + attr = getattr(cls, attr_name, None) + if attr is None: + continue + if attr_name == "__init__": + op_type = "call_function" + elif attr_name.startswith("__"): + op_type = None + elif inspect.ismethod(attr): + op_type = "call_function" + elif inspect.isfunction(attr): + op_type = "call_method" + else: + op_type = None + if op_type is not None: + setattr(cls, attr_name, create_wrapper(attr, op_type, proxy_factory_fn=proxy_factory_fn)) + return cls + + +def gen_constructor_wrapper(target: Callable) -> tuple[Callable, Callable]: + """ + Wraps `target` to be proxyable. Used for tensor creators like `torch.ones`, `torch.arange` and so on. + """ + wrapper = create_wrapper(target, "call_function") + return wrapper, target + + +def _proxies_to_metas(v): + """Returns the underlying metadata for HFProxies, and behaves like the identity for the others.""" + if isinstance(v, MetaDeviceAttribute): + return "meta" + if isinstance(v, torch.fx.Proxy): + if not (isinstance(v, HFProxy) and hasattr(v, "_metadata")): + raise RuntimeError(f"No metadata was found for {v}") + return v._metadata + return v + + +def create_cache_proxy_factory_fn(orig_cache_cls: type[Cache]) -> Callable[[Node], HFCacheProxy]: + def cache_proxy_factory_fn(n: Node) -> HFCacheProxy: + global _CURRENT_TRACER + if not isinstance(_CURRENT_TRACER, HFTracer): + raise RuntimeError("Cannot create HFCacheProxy because there is no HFTracer currently tracing.") + cache_proxy = HFCacheProxy(n, _CURRENT_TRACER) + cache_proxy.install_orig_cache_cls(orig_cache_cls) + return cache_proxy + + return cache_proxy_factory_fn + + +# Proxyable equivalent of the cache classes defined in `transformers.cache_utils`. +ProxyableCache = HFProxyableClassMeta( + "ProxyableCache", (Cache,), {}, proxy_factory_fn=create_cache_proxy_factory_fn(Cache) +) +ProxyableDynamicCache = HFProxyableClassMeta( + "ProxyableDynamicCache", + (DynamicCache,), + {}, + proxy_factory_fn=create_cache_proxy_factory_fn(DynamicCache), +) +ProxyableSinkCache = HFProxyableClassMeta( + "ProxyableSinkCache", + (SinkCache,), + {}, + proxy_factory_fn=create_cache_proxy_factory_fn(SinkCache), +) +ProxyableStaticCache = HFProxyableClassMeta( + "ProxyableStaticCache", + (StaticCache,), + {}, + proxy_factory_fn=create_cache_proxy_factory_fn(StaticCache), +) + + +def _generate_random_int(low: int = 10, high: int = 20, forbidden_values: Optional[list[int]] = None): + if forbidden_values is None: + forbidden_values = [] + value = random.randint(low, high) + while value in forbidden_values: + value = random.randint(low, high) + return value + + +class HFTracer(Tracer): + """ + Tracer that is able to symbolically trace models from the library. To do that, it uses the HFProxy instead of the + regular PyTorch torch.fx.Proxy. + """ + + # Feature flag for proxying accesses to buffer values + proxy_buffer_attributes: bool = True + allow_insert_stateless_mods: bool = True + _TORCH_METHODS_TO_PATCH = [ + "arange", + "zeros", + "ones", + "full", + "full_like", + "eye", + "empty", + "tensor", + "clamp", + "finfo", + "tril", + ] + _CLASSES_TO_PATCH = { + Cache: ProxyableCache, + DynamicCache: ProxyableDynamicCache, + SinkCache: ProxyableSinkCache, + StaticCache: ProxyableStaticCache, + } + + supported_archs = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel) + + def __init__(self, autowrap_modules=(math,), autowrap_functions=()): + super().__init__(autowrap_modules=autowrap_modules, autowrap_functions=autowrap_functions) + + if not is_torch_fx_available(): + raise ImportError( + f"Found an incompatible version of torch. Found version {get_torch_version()}, but only version " + f"{TORCH_FX_REQUIRED_VERSION} is supported." + ) + + def _generate_dummy_input( + self, model: "PreTrainedModel", input_name: str, shape: list[int], input_names: list[str] + ) -> dict[str, torch.Tensor]: + """Generates dummy input for model inference recording.""" + # Retrieving the model class, either from the "class_for_deserialization" attribute if the model was restored + # from pickle, or from the "__class__" attribute in the general case. + model_class_name = getattr(model, "class_for_deserialization", model.__class__).__name__ + device = model.device + inputs_dict = {} + + # when tracing a model with KV cache, we simply need to unsure that the KV cache length is larger than one to + # rightfully pass certain controlflows (Example: https://github.com/huggingface/transformers/blob/5c8d941d66734811d2ef6f57f15b44f7fb7a98c4/src/transformers/modeling_attn_mask_utils.py#L162). + # After tracing, the model can then still be used with arbitrary lengths different than the one used during tracing. + kv_cache_length = 5 + + if input_name in ["labels", "start_positions", "end_positions"]: + batch_size = shape[0] + if model_class_name in [ + *get_values(MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES), + *get_values(MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES), + *get_values(MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES), + *get_values(MODEL_FOR_BACKBONE_MAPPING_NAMES), + *get_values(MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES), + ]: + inputs_dict["labels"] = torch.zeros(batch_size, dtype=torch.long, device=device) + elif model_class_name in [ + *get_values(MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES), + *get_values(MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES), + "XLNetForQuestionAnswering", + ]: + inputs_dict["start_positions"] = torch.zeros(batch_size, dtype=torch.long, device=device) + inputs_dict["end_positions"] = torch.zeros(batch_size, dtype=torch.long, device=device) + elif model_class_name in get_values(MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES): + if not hasattr(model.config, "problem_type") or model.config.problem_type is None: + raise ValueError( + "Could not retrieve the problem type for the sequence classification task, please set " + 'model.config.problem_type to one of the following values: "regression", ' + '"single_label_classification", or "multi_label_classification".' + ) + + if model.config.problem_type == "regression": + labels_shape = (batch_size, model.config.num_labels) + labels_dtype = torch.float32 + elif model.config.problem_type == "single_label_classification": + labels_shape = (batch_size,) + labels_dtype = torch.long + elif model.config.problem_type == "multi_label_classification": + labels_shape = (batch_size, model.config.num_labels) + labels_dtype = torch.float32 + else: + raise ValueError( + 'Expected model.config.problem_type to be either: "regression", "single_label_classification"' + f', or "multi_label_classification", but "{model.config.problem_type}" was provided.' + ) + inputs_dict["labels"] = torch.zeros(*labels_shape, dtype=labels_dtype, device=device) + + elif model_class_name in [ + *get_values(MODEL_FOR_PRETRAINING_MAPPING_NAMES), + *get_values(MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES), + *get_values(MODEL_FOR_CAUSAL_LM_MAPPING_NAMES), + *get_values(MODEL_FOR_MASKED_LM_MAPPING_NAMES), + *get_values(MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES), + *get_values(MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES), + "GPT2DoubleHeadsModel", + "PeftModelForCausalLM", + "PeftModelForSeq2SeqLM", + ]: + inputs_dict["labels"] = torch.zeros(shape, dtype=torch.long, device=device) + elif model_class_name in [*get_values(MODEL_FOR_CTC_MAPPING_NAMES)]: + inputs_dict["labels"] = torch.zeros(shape, dtype=torch.float32, device=device) + else: + raise NotImplementedError( + f"Generating the dummy input named {input_name} for {model_class_name} is not supported yet." + ) + elif "pixel_values" in input_name: + batch_size = shape[0] + image_size = getattr(model.config, "image_size", None) + if image_size is None: + if hasattr(model.config, "vision_config"): + image_size = model.config.vision_config.image_size + elif hasattr(model.config, "encoder"): + image_size = model.config.encoder.image_size + else: + image_size = (_generate_random_int(), _generate_random_int()) + + # If no num_channels is in the config, use some arbitrary value. + num_channels = getattr(model.config, "num_channels", 3) + if not isinstance(image_size, collections.abc.Iterable): + image_size = (image_size, image_size) + height, width = image_size + inputs_dict[input_name] = torch.zeros( + batch_size, num_channels, height, width, dtype=torch.float32, device=device + ) + elif "bbox" in input_name: + inputs_dict[input_name] = torch.zeros(*shape, 4, dtype=torch.float, device=device) + elif "input_features" in input_name: + inputs_dict[input_name] = torch.zeros( + *shape, model.config.input_feat_per_channel, dtype=torch.float, device=device + ) + elif "inputs_embeds" in input_name: + batch_size = shape[0] + + if ( + getattr(model.config, "embedding_size", None) is not None + and model.config.model_type != "megatron-bert" + ): + embedding_size = model.config.embedding_size + else: + embedding_size = model.config.hidden_size + + if len(shape) == 3: + # (batch_size, num_choices, sequence_length, embedding_size) + embedding_shape = (batch_size, shape[1], shape[2], embedding_size) + else: + # (batch_size, sequence_length, embedding_size) + embedding_shape = (batch_size, shape[1], embedding_size) + + inputs_dict[input_name] = torch.zeros(embedding_shape, dtype=torch.float, device=device) + elif "visual_feats" in input_name: + inputs_dict[input_name] = torch.zeros( + shape + + [ + model.config.visual_feat_dim, + ], + dtype=torch.float, + device=device, + ) + elif "visual_pos" in input_name: + inputs_dict[input_name] = torch.zeros( + shape + + [ + model.config.visual_pos_dim, + ], + dtype=torch.float, + device=device, + ) + elif "inputs" in input_name: + inputs_dict[input_name] = torch.zeros(*shape, dtype=torch.float, device=device) + elif "input_values" in input_name: + batch_size, _ = shape + # Generating big sequence length for audio inputs. + seq_length = _generate_random_int(low=10000, high=20000) + inputs_dict[input_name] = torch.zeros(batch_size, seq_length, dtype=torch.float, device=device) + elif "mask" in input_name: + if "past_key_values" in input_names: + mask_shape = [shape[0], shape[1] + kv_cache_length] + else: + mask_shape = shape + + inputs_dict[input_name] = torch.zeros(mask_shape, dtype=torch.long, device=device) + elif "ids" in input_name: + inputs_dict[input_name] = torch.zeros(shape, dtype=torch.long, device=device) + elif "past_key_values" in input_name: + if model.config.model_type not in _FX_SUPPORTED_MODELS_WITH_KV_CACHE: + raise NotImplementedError( + f"Symbolic trace with past_key_values input is not supported yet for the model {model.config.model_type}. Please open an issue or a PR in Transformers repository if you would like to see the support added." + ) + num_heads = model.config.num_attention_heads + head_dim = model.config.hidden_size // model.config.num_attention_heads + + cache_shape = (shape[0], num_heads, kv_cache_length, head_dim) + pkv = tuple( + ( + torch.rand(cache_shape, dtype=torch.float, device=device), + torch.rand(cache_shape, dtype=torch.float, device=device), + ) + for i in range(model.config.num_hidden_layers) + ) + inputs_dict[input_name] = pkv + else: + shape_with_hidden_size = shape + [model.config.hidden_size] + inputs_dict[input_name] = torch.zeros(shape_with_hidden_size, dtype=torch.float, device=device) + + return inputs_dict + + def create_proxy(self, kind, target, args, kwargs, name=None, type_expr=None, proxy_factory_fn=None): + rv = super().create_proxy(kind, target, args, kwargs, name, type_expr, proxy_factory_fn) + + if kind == "placeholder" and target in self.meta_args: + rv.install_metadata(self.meta_args[target]) + return rv + + if target in self.orig_fns: + # NOTE: tensor constructors in PyTorch define the `device` argument as + # *kwargs-only*. That is why this works. If you add methods to + # _TORCH_METHODS_TO_PATCH that do not define `device` as kwarg-only, + # this will break and you will likely see issues where we cannot infer + # the size of the output. + if "device" in kwargs: + kwargs["device"] = "meta" + + try: + args_metas = torch.fx.node.map_aggregate(args, _proxies_to_metas) + kwargs_metas = torch.fx.node.map_aggregate(kwargs, _proxies_to_metas) + + should_install_metadata = True + + self._disable_module_getattr = True + self._disable_call_module = True + + if kind == "call_function": + meta_target = _MANUAL_META_OVERRIDES.get(target, target) + meta_out = meta_target(*args_metas, **kwargs_metas) + if isinstance(meta_out, torch.Tensor): + meta_out = meta_out.to(device="meta") + elif kind == "call_method": + method = getattr(args_metas[0].__class__, target) + meta_target = _MANUAL_META_OVERRIDES.get(method, method) + meta_out = meta_target(*args_metas, **kwargs_metas) + elif kind == "call_module": + if not hasattr(self, "orig_forward"): + raise AttributeError(f"{self} does not have an attribute called orig_forward") + mod = self.root.get_submodule(target) + mod_type = type(mod) + if mod_type in _MANUAL_META_OVERRIDES: + meta_out = _MANUAL_META_OVERRIDES[mod_type](mod, *args_metas, **kwargs_metas) + else: + meta_out = self.orig_forward(*args_metas, **kwargs_metas) + elif kind == "get_attr": + attr_itr = self.root + atoms = target.split(".") + for atom in atoms: + attr_itr = getattr(attr_itr, atom) + if isinstance(attr_itr, torch.Tensor): + meta_out = attr_itr.to(device="meta") + else: + meta_out = attr_itr + else: + should_install_metadata = False + + if should_install_metadata: + if not isinstance(rv, Proxy): + raise ValueError("Don't support composite output yet") + rv.install_metadata(meta_out) + + except Exception as e: + if _IS_IN_DEBUG_MODE: + warnings.warn(f"Could not compute metadata for {kind} target {target}: {e}") + + self._disable_module_getattr = False + self._disable_call_module = False + + return rv + + # Replaced by .getattr from PyTorch 1.13 + def _module_getattr(self, attr, attr_val, parameter_proxy_cache): + if getattr(self, "_disable_module_getattr", False): + return attr_val + else: + + def maybe_get_proxy_for_attr(attr_val, collection_to_search, parameter_proxy_cache): + for n, p in collection_to_search: + if attr_val is p: + if n not in parameter_proxy_cache: + kwargs = {} + if "proxy_factory_fn" in inspect.signature(self.create_proxy).parameters: + kwargs["proxy_factory_fn"] = ( + None + if not self.param_shapes_constant + else lambda node: ParameterProxy(self, node, n, attr_val) + ) + val_proxy = self.create_proxy("get_attr", n, (), {}, **kwargs) # type: ignore[arg-type] + parameter_proxy_cache[n] = val_proxy + return parameter_proxy_cache[n] + return None + + if isinstance(attr_val, torch.nn.Parameter): + maybe_parameter_proxy = maybe_get_proxy_for_attr( + attr_val, self.root.named_parameters(), parameter_proxy_cache + ) + if maybe_parameter_proxy is not None: + return maybe_parameter_proxy + + if self.proxy_buffer_attributes and isinstance(attr_val, torch.Tensor): + maybe_buffer_proxy = maybe_get_proxy_for_attr( + attr_val, self.root.named_buffers(), parameter_proxy_cache + ) + if maybe_buffer_proxy is not None: + return maybe_buffer_proxy + + return attr_val + + # Needed for PyTorch 1.13+ + def getattr(self, attr: str, attr_val: Any, parameter_proxy_cache: dict[str, Any]): + return self._module_getattr(attr, attr_val, parameter_proxy_cache) + + def call_module(self, m, forward, args, kwargs): + if getattr(self, "_disable_call_module", False): + return forward(*args, **kwargs) + self.orig_forward = forward + return super().call_module(m, forward, args, kwargs) + + def proxy(self, node): + return HFProxy(node, self) + + @contextlib.contextmanager + def patch_for_tracing(self, root: Union[torch.nn.Module, Callable[..., Any]]): + # Patching torch functions + self.patched_torch_methods = { + target: gen_constructor_wrapper(getattr(torch, target)) for target in self._TORCH_METHODS_TO_PATCH + } + self.orig_fns = set() + + for name, (wrapper, orig) in self.patched_torch_methods.items(): + setattr(torch, name, wrapper) + self.orig_fns.add(orig) + + # Patching classes + patched = [] + module_of_model = inspect.getmodule(root) + for name, mod in sys.modules.items(): + if module_of_model is not None and mod is not module_of_model: + continue + if not name.startswith("transformers"): + continue + for orig_cls, patched_cls in self._CLASSES_TO_PATCH.items(): + for attr_name, attr in mod.__dict__.items(): + if attr is orig_cls: + patched.append((mod, attr_name, orig_cls)) + setattr(mod, attr_name, patched_cls) + + yield + + # Restoring patched functions and classes. + for name, (_, orig) in self.patched_torch_methods.items(): + setattr(torch, name, orig) + self.patched_torch_methods = {} + self.orig_fns = set() + + for mod, attr_name, orig_cls in patched: + setattr(mod, attr_name, orig_cls) + + def trace( + self, + root: Union[torch.nn.Module, Callable[..., Any]], + concrete_args: Optional[dict[str, Any]] = None, + dummy_inputs: Optional[dict[str, Any]] = None, + complete_concrete_args_with_inputs_not_in_dummy_inputs: bool = True, + ) -> Graph: + """ + Traces `root` and returns the corresponding FX `torch.fx.Graph` representation. `root` can either be a + `torch.nn.Module` instance or a Python callable. Note that after this call, `self.root` may be different from + the `root` passed in here. For example, when a free function is passed to `trace()`, we will create a + `torch.nn.Module` instance to use as the root and add embedded constants to. + + Args: + root (`torch.nn.Module` or `Callable`): + Either a `torch.nn.Module`` or a function to be traced through. If root is not a + [`~transformers.PreTrainedModel`], then `dummy_inputs` must be passed, otherwise tracing will fail. + concrete_args (`Dict[str, Any], *optional*): + Concrete arguments that should not be treated as Proxies + dummy_inputs (`Dict[str, Any]`, *optional*): + The dummy inputs needed to handle data-dependent control-flow if `root` is not a + [`~transformers.PreTrainedModel`]. It can also be used when `root` is a + [`~transformers.PreTrainedModel`] to specify custom dummy inputs for a subset or all the model inputs. + complete_concrete_args_with_inputs_not_in_dummy_inputs (`bool`, *optional*, defaults to `True`): + If `True`, and `dummy_inputs` is specified, every argument that `root` can take that is not in + `dummy_inputs` and not in `concrete_args` will be added to `concrete_args`, otherwise does nothing. + + Returns: + `torch.fx.Graph`: + A FX `torch.fx.Graph` representing the semantics of the passed-in `root`. + + """ + sig = inspect.signature(root.forward if isinstance(root, torch.nn.Module) else root) + + if concrete_args is None: + concrete_args = {} + + if dummy_inputs is not None and complete_concrete_args_with_inputs_not_in_dummy_inputs: + for param in sig.parameters.values(): + if param.name in dummy_inputs: + continue + if param.default is inspect.Parameter.empty: + raise ValueError(f"You need to specify a default value for the parameter {param.name}.") + concrete_args.update( + { + p.name: p.default + for p in sig.parameters.values() + if (p.name not in dummy_inputs and p.name not in concrete_args) + } + ) + + input_names = sig.parameters.keys() - concrete_args.keys() + + # Creating a random input shape to generate dummy inputs. + batch_size = _generate_random_int() + sequence_length = _generate_random_int() + shape = [batch_size, sequence_length] + + if root.__class__.__name__ in get_values(MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES): + num_choices = _generate_random_int(low=2, high=5) + shape.insert(1, num_choices) + + inputs = dict(dummy_inputs) if dummy_inputs is not None else {} + for input_name in input_names: + if input_name in inputs: + continue + # We enforce that root must either be a PreTrainedModel or deserialized from a serialized traced model to + # be able to use HFTracer._generate_dummy_input. + if isinstance(root, self.supported_archs) or type(root).__qualname__.startswith( + ("_deserialize_graph_module", "_CodeOnlyModule") + ): + inputs.update(self._generate_dummy_input(root, input_name, shape, input_names=input_names)) + else: + raise RuntimeError( + f"Could not generate input named {input_name} for because root is not a" + " transformers.PreTrainedModel." + ) + + def to_meta(value): + if isinstance(value, torch.Tensor): + return value.to("meta") + return value + + concrete_metas = pytree.tree_map(to_meta, inputs) + + for param in sig.parameters.values(): + if param.kind == inspect.Parameter.VAR_KEYWORD and param.name not in input_names: + concrete_metas[f"**{param.name}"] = {} + self.meta_args = concrete_metas + + global _CURRENT_TRACER + _CURRENT_TRACER = self + with self.patch_for_tracing(root): + try: + self.graph = super().trace(root, concrete_args=concrete_args) + finally: + _CURRENT_TRACER = None + + # This is necessary because concrete args are added as input to the traced module since + # https://github.com/pytorch/pytorch/pull/55888. + for node in self.graph.nodes: + if node.op == "placeholder": + # Removing default values for inputs as the forward pass will fail with them. + if node.target in input_names: + node.args = () + # Without this, torch.jit.script fails because the inputs type is Optional[torch.Tensor]. + # It cannot infer on the attributes and methods the input should have, and fails. + node.type = torch.Tensor + # It is a concrete arg so it is not used and should be removed. + else: + to_visit = [node] + to_delete = collections.OrderedDict() + while to_visit: + n = to_visit.pop(0) + to_delete[n] = None + to_visit += list(n.users.keys()) + + for user in reversed(to_delete.keys()): + self.graph.erase_node(user) + + # TODO: solves GraphModule creation. + # Without this, return type annotation "Tuple" is causing code execution failure. + if node.op == "output": + node.type = None + + return self.graph + + def _stateless_mod_instanciation_depends_on_proxies(self, mod: nn.Module) -> bool: + """ + Whether the module was instantiated with Proxies. If that is the case, such module cannot be a leaf module + because its attributes are input-dependent. + """ + return any(isinstance(attr, Proxy) for attr in mod.__dict__.values()) + + def _insert_module_as_submodule(self, mod: nn.Module) -> str: + """ + Helper method which tries to insert a module that was not declared as submodule. + """ + # If one of the module attributes is a Proxy, it means that its instantiation is input-dependent. + # It is not possible to insert such modules, those should be traced through. + if self._stateless_mod_instanciation_depends_on_proxies(mod): + return "" + idx = 0 + mod_name = mod.__class__.__name__.lower() + path = f"{mod_name}_{idx}" + already_inserted = False + while hasattr(self.root, path): + if getattr(self.root, path) is mod: + already_inserted = True + break + path = f"{mod_name}_{idx}" + idx += 1 + + # No need to add multiple instances of the same module. + if not already_inserted: + self.root.add_module(path, mod) + return path + + def path_of_module(self, mod: nn.Module) -> str: + """ + Helper method to find the qualified name of `mod` in the Module hierarchy of `root`. For example, if `root` has + a submodule named `foo`, which has a submodule named `bar`, passing `bar` into this function will return the + string "foo.bar". + + Args: + mod (str): The `Module` to retrieve the qualified name for. + """ + try: + return super().path_of_module(mod) + except NameError as e: + if self.allow_insert_stateless_mods and len(list(mod.parameters())) == 0 and len(list(mod.buffers())) == 0: + path = self._insert_module_as_submodule(mod) + return path + raise e + + def is_leaf_module(self, m: torch.nn.Module, module_qualified_name: str) -> bool: + return (not self._stateless_mod_instanciation_depends_on_proxies(m)) and super().is_leaf_module( + m, module_qualified_name + ) + + @compatibility(is_backward_compatible=True) + def keys(self, obj: "Proxy") -> Any: + """Called when a proxy object is has the keys() method called. + This is what happens when ** is called on a proxy. This should return an iterator if ** is supposed to work in + your custom tracer. + """ + attribute = HFAttribute(obj, "keys")() + if obj.node.target.startswith("**"): + return attribute._metadata + return attribute + + +def get_concrete_args(model: nn.Module, input_names: list[str]): + sig = inspect.signature(model.forward) + + if not (set(input_names) <= set(sig.parameters.keys())): + formatted_input_names = input_names[0] if len(input_names) == 1 else ", ".join(input_names) + formatted_allowed_input_names = ", ".join(sig.parameters.keys()) + raise ValueError( + f"The model does not have input(s) named: {formatted_input_names}, expected a subset of the following:" + f" {formatted_allowed_input_names}" + ) + + return {p.name: p.default for p in sig.parameters.values() if p.name not in input_names} + + +def is_model_supported(model: "PreTrainedModel"): + return model.__class__.__name__ in _SUPPORTED_MODELS + + +def check_if_model_is_supported(model: "PreTrainedModel"): + if not is_model_supported(model): + supported_model_names = ", ".join(_SUPPORTED_MODELS) + raise NotImplementedError( + f"Model {model.__class__.__name__} is not supported yet, supported models: {supported_model_names}" + ) + + +def symbolic_trace( + model: "PreTrainedModel", + input_names: Optional[list[str]] = None, + disable_check: bool = False, + tracer_cls: type[HFTracer] = HFTracer, +) -> GraphModule: + """ + Performs symbolic tracing on the model. + + Args: + model ([`PretrainedModel`]): + The model to trace. + input_names (`List[str]`, *optional*): + The names of the inputs of the traced model. If unset, model.dummy_inputs.keys() are used instead. + disable_check (`bool`, *optional*, defaults to `False`): + If `True`, no check is done before trying to trace the model, this is mostly usesul for debugging purposes. + tracer_cls (`Type[HFTracer]`, *optional*, defaults to `HFTracer`): + The tracer class to use for instantiating the tracer. If unset, `HFTracer` is used instead. + + Returns: + `torch.fx.GraphModule`: A GraphModule constructed by recording operations seen while tracing the model. + + Example: + + ```python + from transformers.utils.fx import symbolic_trace + + traced_model = symbolic_trace(model, input_names=["input_ids", "attention_mask", "token_type_ids"]) + ``` + """ + if input_names is None: + input_names = model.dummy_inputs.keys() + + input_names = list(input_names) + concrete_args = get_concrete_args(model, input_names) + + if not disable_check: + check_if_model_is_supported(model) + + if "past_key_values" in input_names and not getattr(model.config, "use_cache", False): + logger.warning( + "`past_key_values` were specified as input names, but model.config.use_cache = False, this might lead to " + "unexpected behavior." + ) + if "past_key_values" not in input_names and getattr(model.config, "use_cache", False): + logger.warning( + "`past_key_values` were not specified as input names, but model.config.use_cache = True. Setting " + "model.config.use_cache = False." + ) + model.config.use_cache = False + + # Tracing. + tracer = tracer_cls() + traced_graph = tracer.trace(model, concrete_args=concrete_args) + traced = torch.fx.GraphModule(model, traced_graph) + + traced.config = model.config + # The model class must be stored as an attribute to allow model deserialization, which uses trace, and thus + # _generate_dummy_input, where the model class is needed. + traced.class_for_deserialization = model.__class__ + traced.device = model.device + + return traced diff --git a/lib/python3.12/site-packages/transformers/utils/generic.py b/lib/python3.12/site-packages/transformers/utils/generic.py new file mode 100644 index 0000000000000000000000000000000000000000..65a3efaed5ab5cfaf0935638e93c20c23aecef8d --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/generic.py @@ -0,0 +1,975 @@ +# Copyright 2022 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. +""" +Generic utilities +""" + +import inspect +import json +import os +import tempfile +import warnings +from collections import OrderedDict, UserDict +from collections.abc import Iterable, MutableMapping +from contextlib import ExitStack, contextmanager +from dataclasses import fields, is_dataclass +from enum import Enum +from functools import partial, wraps +from typing import Any, ContextManager, Optional, TypedDict + +import numpy as np +from packaging import version + +from .import_utils import ( + get_torch_version, + is_flax_available, + is_mlx_available, + is_tf_available, + is_torch_available, + is_torch_fx_proxy, +) + + +if is_torch_available(): + # required for @can_return_tuple decorator to work with torchdynamo + import torch # noqa: F401 + + +class cached_property(property): + """ + Descriptor that mimics @property but caches output in member variable. + + From tensorflow_datasets + + Built-in in functools from Python 3.8. + """ + + def __get__(self, obj, objtype=None): + # See docs.python.org/3/howto/descriptor.html#properties + if obj is None: + return self + if self.fget is None: + raise AttributeError("unreadable attribute") + attr = "__cached_" + self.fget.__name__ + cached = getattr(obj, attr, None) + if cached is None: + cached = self.fget(obj) + setattr(obj, attr, cached) + return cached + + +# vendored from distutils.util +def strtobool(val): + """Convert a string representation of truth to true (1) or false (0). + + True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values are 'n', 'no', 'f', 'false', 'off', and '0'. + Raises ValueError if 'val' is anything else. + """ + val = val.lower() + if val in {"y", "yes", "t", "true", "on", "1"}: + return 1 + if val in {"n", "no", "f", "false", "off", "0"}: + return 0 + raise ValueError(f"invalid truth value {val!r}") + + +def infer_framework_from_repr(x): + """ + Tries to guess the framework of an object `x` from its repr (brittle but will help in `is_tensor` to try the + frameworks in a smart order, without the need to import the frameworks). + """ + representation = str(type(x)) + if representation.startswith(" + + You can't unpack a `ModelOutput` directly. Use the [`~utils.ModelOutput.to_tuple`] method to convert it to a tuple + before. + + + """ + + def __init_subclass__(cls) -> None: + """Register subclasses as pytree nodes. + + This is necessary to synchronize gradients when using `torch.nn.parallel.DistributedDataParallel` with + `static_graph=True` with modules that output `ModelOutput` subclasses. + """ + if is_torch_available(): + if version.parse(get_torch_version()) >= version.parse("2.2"): + _torch_pytree.register_pytree_node( + cls, + _model_output_flatten, + partial(_model_output_unflatten, output_type=cls), + serialized_type_name=f"{cls.__module__}.{cls.__name__}", + ) + else: + _torch_pytree._register_pytree_node( + cls, + _model_output_flatten, + partial(_model_output_unflatten, output_type=cls), + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # Subclasses of ModelOutput must use the @dataclass decorator + # This check is done in __init__ because the @dataclass decorator operates after __init_subclass__ + # issubclass() would return True for issubclass(ModelOutput, ModelOutput) when False is needed + # Just need to check that the current class is not ModelOutput + is_modeloutput_subclass = self.__class__ != ModelOutput + + if is_modeloutput_subclass and not is_dataclass(self): + raise TypeError( + f"{self.__module__}.{self.__class__.__name__} is not a dataclass." + " This is a subclass of ModelOutput and so must use the @dataclass decorator." + ) + + def __post_init__(self): + """Check the ModelOutput dataclass. + + Only occurs if @dataclass decorator has been used. + """ + class_fields = fields(self) + + # Safety and consistency checks + if not len(class_fields): + raise ValueError(f"{self.__class__.__name__} has no fields.") + if not all(field.default is None for field in class_fields[1:]): + raise ValueError(f"{self.__class__.__name__} should not have more than one required field.") + + first_field = getattr(self, class_fields[0].name) + other_fields_are_none = all(getattr(self, field.name) is None for field in class_fields[1:]) + + if other_fields_are_none and not is_tensor(first_field): + if isinstance(first_field, dict): + iterator = first_field.items() + first_field_iterator = True + else: + try: + iterator = iter(first_field) + first_field_iterator = True + except TypeError: + first_field_iterator = False + + # if we provided an iterator as first field and the iterator is a (key, value) iterator + # set the associated fields + if first_field_iterator: + for idx, element in enumerate(iterator): + if ( + not isinstance(element, (list, tuple)) + or not len(element) == 2 + or not isinstance(element[0], str) + ): + if idx == 0: + # If we do not have an iterator of key/values, set it as attribute + self[class_fields[0].name] = first_field + else: + # If we have a mixed iterator, raise an error + raise ValueError( + f"Cannot set key/value for {element}. It needs to be a tuple (key, value)." + ) + break + setattr(self, element[0], element[1]) + if element[1] is not None: + self[element[0]] = element[1] + elif first_field is not None: + self[class_fields[0].name] = first_field + else: + for field in class_fields: + v = getattr(self, field.name) + if v is not None: + self[field.name] = v + + def __delitem__(self, *args, **kwargs): + raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.") + + def setdefault(self, *args, **kwargs): + raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.") + + def pop(self, *args, **kwargs): + raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.") + + def update(self, *args, **kwargs): + raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.") + + def __getitem__(self, k): + if isinstance(k, str): + inner_dict = dict(self.items()) + return inner_dict[k] + else: + return self.to_tuple()[k] + + def __setattr__(self, name, value): + if name in self.keys() and value is not None: + # Don't call self.__setitem__ to avoid recursion errors + super().__setitem__(name, value) + super().__setattr__(name, value) + + def __setitem__(self, key, value): + # Will raise a KeyException if needed + super().__setitem__(key, value) + # Don't call self.__setattr__ to avoid recursion errors + super().__setattr__(key, value) + + def __reduce__(self): + if not is_dataclass(self): + return super().__reduce__() + callable, _args, *remaining = super().__reduce__() + args = tuple(getattr(self, field.name) for field in fields(self)) + return callable, args, *remaining + + def to_tuple(self) -> tuple[Any]: + """ + Convert self to a tuple containing all the attributes/keys that are not `None`. + """ + return tuple(self[k] for k in self.keys()) + + +if is_torch_available(): + import torch.utils._pytree as _torch_pytree + + def _model_output_flatten(output: ModelOutput) -> tuple[list[Any], "_torch_pytree.Context"]: + return list(output.values()), list(output.keys()) + + def _model_output_unflatten( + values: Iterable[Any], + context: "_torch_pytree.Context", + output_type=None, + ) -> ModelOutput: + return output_type(**dict(zip(context, values))) + + if version.parse(get_torch_version()) >= version.parse("2.2"): + _torch_pytree.register_pytree_node( + ModelOutput, + _model_output_flatten, + partial(_model_output_unflatten, output_type=ModelOutput), + serialized_type_name=f"{ModelOutput.__module__}.{ModelOutput.__name__}", + ) + else: + _torch_pytree._register_pytree_node( + ModelOutput, + _model_output_flatten, + partial(_model_output_unflatten, output_type=ModelOutput), + ) + + +class ExplicitEnum(str, Enum): + """ + Enum with more explicit error message for missing values. + """ + + @classmethod + def _missing_(cls, value): + raise ValueError( + f"{value} is not a valid {cls.__name__}, please select one of {list(cls._value2member_map_.keys())}" + ) + + +class PaddingStrategy(ExplicitEnum): + """ + Possible values for the `padding` argument in [`PreTrainedTokenizerBase.__call__`]. Useful for tab-completion in an + IDE. + """ + + LONGEST = "longest" + MAX_LENGTH = "max_length" + DO_NOT_PAD = "do_not_pad" + + +class TensorType(ExplicitEnum): + """ + Possible values for the `return_tensors` argument in [`PreTrainedTokenizerBase.__call__`]. Useful for + tab-completion in an IDE. + """ + + PYTORCH = "pt" + TENSORFLOW = "tf" + NUMPY = "np" + JAX = "jax" + MLX = "mlx" + + +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 can_return_loss(model_class): + """ + Check if a given model can return loss. + + Args: + model_class (`type`): The class of the model. + """ + framework = infer_framework(model_class) + if framework == "tf": + signature = inspect.signature(model_class.call) # TensorFlow models + elif framework == "pt": + signature = inspect.signature(model_class.forward) # PyTorch models + else: + signature = inspect.signature(model_class.__call__) # Flax models + + for p in signature.parameters: + if p == "return_loss" and signature.parameters[p].default is True: + return True + + return False + + +def find_labels(model_class): + """ + Find the labels used by a given model. + + Args: + model_class (`type`): The class of the model. + """ + model_name = model_class.__name__ + framework = infer_framework(model_class) + if framework == "tf": + signature = inspect.signature(model_class.call) # TensorFlow models + elif framework == "pt": + signature = inspect.signature(model_class.forward) # PyTorch models + else: + signature = inspect.signature(model_class.__call__) # Flax models + + if "QuestionAnswering" in model_name: + return [p for p in signature.parameters if "label" in p or p in ("start_positions", "end_positions")] + else: + return [p for p in signature.parameters if "label" in p] + + +def flatten_dict(d: MutableMapping, parent_key: str = "", delimiter: str = "."): + """Flatten a nested dict into a single level dict.""" + + def _flatten_dict(d, parent_key="", delimiter="."): + for k, v in d.items(): + key = str(parent_key) + delimiter + str(k) if parent_key else k + if v and isinstance(v, MutableMapping): + yield from flatten_dict(v, key, delimiter=delimiter).items() + else: + yield key, v + + return dict(_flatten_dict(d, parent_key, delimiter)) + + +@contextmanager +def working_or_temp_dir(working_dir, use_temp_dir: bool = False): + if use_temp_dir: + with tempfile.TemporaryDirectory() as tmp_dir: + yield tmp_dir + else: + yield working_dir + + +def transpose(array, axes=None): + """ + Framework-agnostic version of `numpy.transpose` that will work on torch/TensorFlow/Jax tensors as well as NumPy + arrays. + """ + if is_numpy_array(array): + return np.transpose(array, axes=axes) + elif is_torch_tensor(array): + return array.T if axes is None else array.permute(*axes) + elif is_tf_tensor(array): + import tensorflow as tf + + return tf.transpose(array, perm=axes) + elif is_jax_tensor(array): + import jax.numpy as jnp + + return jnp.transpose(array, axes=axes) + else: + raise ValueError(f"Type not supported for transpose: {type(array)}.") + + +def reshape(array, newshape): + """ + Framework-agnostic version of `numpy.reshape` that will work on torch/TensorFlow/Jax tensors as well as NumPy + arrays. + """ + if is_numpy_array(array): + return np.reshape(array, newshape) + elif is_torch_tensor(array): + return array.reshape(*newshape) + elif is_tf_tensor(array): + import tensorflow as tf + + return tf.reshape(array, newshape) + elif is_jax_tensor(array): + import jax.numpy as jnp + + return jnp.reshape(array, newshape) + else: + raise ValueError(f"Type not supported for reshape: {type(array)}.") + + +def squeeze(array, axis=None): + """ + Framework-agnostic version of `numpy.squeeze` that will work on torch/TensorFlow/Jax tensors as well as NumPy + arrays. + """ + if is_numpy_array(array): + return np.squeeze(array, axis=axis) + elif is_torch_tensor(array): + return array.squeeze() if axis is None else array.squeeze(dim=axis) + elif is_tf_tensor(array): + import tensorflow as tf + + return tf.squeeze(array, axis=axis) + elif is_jax_tensor(array): + import jax.numpy as jnp + + return jnp.squeeze(array, axis=axis) + else: + raise ValueError(f"Type not supported for squeeze: {type(array)}.") + + +def expand_dims(array, axis): + """ + Framework-agnostic version of `numpy.expand_dims` that will work on torch/TensorFlow/Jax tensors as well as NumPy + arrays. + """ + if is_numpy_array(array): + return np.expand_dims(array, axis) + elif is_torch_tensor(array): + return array.unsqueeze(dim=axis) + elif is_tf_tensor(array): + import tensorflow as tf + + return tf.expand_dims(array, axis=axis) + elif is_jax_tensor(array): + import jax.numpy as jnp + + return jnp.expand_dims(array, axis=axis) + else: + raise ValueError(f"Type not supported for expand_dims: {type(array)}.") + + +def tensor_size(array): + """ + Framework-agnostic version of `numpy.size` that will work on torch/TensorFlow/Jax tensors as well as NumPy arrays. + """ + if is_numpy_array(array): + return np.size(array) + elif is_torch_tensor(array): + return array.numel() + elif is_tf_tensor(array): + import tensorflow as tf + + return tf.size(array) + elif is_jax_tensor(array): + return array.size + else: + raise ValueError(f"Type not supported for tensor_size: {type(array)}.") + + +def add_model_info_to_auto_map(auto_map, repo_id): + """ + Adds the information of the repo_id to a given auto map. + """ + for key, value in auto_map.items(): + if isinstance(value, (tuple, list)): + auto_map[key] = [f"{repo_id}--{v}" if (v is not None and "--" not in v) else v for v in value] + elif value is not None and "--" not in value: + auto_map[key] = f"{repo_id}--{value}" + + return auto_map + + +def add_model_info_to_custom_pipelines(custom_pipeline, repo_id): + """ + Adds the information of the repo_id to a given custom pipeline. + """ + # {custom_pipelines : {task: {"impl": "path.to.task"},...} } + for task in custom_pipeline.keys(): + if "impl" in custom_pipeline[task]: + module = custom_pipeline[task]["impl"] + if "--" not in module: + custom_pipeline[task]["impl"] = f"{repo_id}--{module}" + return custom_pipeline + + +def infer_framework(model_class): + """ + Infers the framework of a given model without using isinstance(), because we cannot guarantee that the relevant + classes are imported or available. + """ + for base_class in inspect.getmro(model_class): + module = base_class.__module__ + name = base_class.__name__ + if module.startswith("tensorflow") or module.startswith("keras") or name == "TFPreTrainedModel": + return "tf" + elif module.startswith("torch") or name == "PreTrainedModel": + return "pt" + elif module.startswith("flax") or module.startswith("jax") or name == "FlaxPreTrainedModel": + return "flax" + else: + raise TypeError(f"Could not infer framework from class {model_class}.") + + +def torch_int(x): + """ + Casts an input to a torch int64 tensor if we are in a tracing context, otherwise to a Python int. + """ + if not is_torch_available(): + return int(x) + + import torch + + return x.to(torch.int64) if torch.jit.is_tracing() and isinstance(x, torch.Tensor) else int(x) + + +def torch_float(x): + """ + Casts an input to a torch float32 tensor if we are in a tracing context, otherwise to a Python float. + """ + if not is_torch_available(): + return int(x) + + import torch + + return x.to(torch.float32) if torch.jit.is_tracing() and isinstance(x, torch.Tensor) else int(x) + + +def filter_out_non_signature_kwargs(extra: Optional[list] = None): + """ + Decorator to filter out named arguments that are not in the function signature. + + This decorator ensures that only the keyword arguments that match the function's signature, or are specified in the + `extra` list, are passed to the function. Any additional keyword arguments are filtered out and a warning is issued. + + Parameters: + extra (`Optional[list]`, *optional*): + A list of extra keyword argument names that are allowed even if they are not in the function's signature. + + Returns: + Callable: + A decorator that wraps the function and filters out invalid keyword arguments. + + Example usage: + + ```python + @filter_out_non_signature_kwargs(extra=["allowed_extra_arg"]) + def my_function(arg1, arg2, **kwargs): + print(arg1, arg2, kwargs) + + my_function(arg1=1, arg2=2, allowed_extra_arg=3, invalid_arg=4) + # This will print: 1 2 {"allowed_extra_arg": 3} + # And issue a warning: "The following named arguments are not valid for `my_function` and were ignored: 'invalid_arg'" + ``` + """ + extra = extra or [] + extra_params_to_pass = set(extra) + + def decorator(func): + sig = inspect.signature(func) + function_named_args = set(sig.parameters.keys()) + valid_kwargs_to_pass = function_named_args.union(extra_params_to_pass) + + # Required for better warning message + is_instance_method = "self" in function_named_args + is_class_method = "cls" in function_named_args + + # Mark function as decorated + func._filter_out_non_signature_kwargs = True + + @wraps(func) + def wrapper(*args, **kwargs): + valid_kwargs = {} + invalid_kwargs = {} + + for k, v in kwargs.items(): + if k in valid_kwargs_to_pass: + valid_kwargs[k] = v + else: + invalid_kwargs[k] = v + + if invalid_kwargs: + invalid_kwargs_names = [f"'{k}'" for k in invalid_kwargs.keys()] + invalid_kwargs_names = ", ".join(invalid_kwargs_names) + + # Get the class name for better warning message + if is_instance_method: + cls_prefix = args[0].__class__.__name__ + "." + elif is_class_method: + cls_prefix = args[0].__name__ + "." + else: + cls_prefix = "" + + warnings.warn( + f"The following named arguments are not valid for `{cls_prefix}{func.__name__}`" + f" and were ignored: {invalid_kwargs_names}", + UserWarning, + stacklevel=2, + ) + + return func(*args, **valid_kwargs) + + return wrapper + + return decorator + + +class LossKwargs(TypedDict, total=False): + """ + Keyword arguments to be passed to the loss function + + Attributes: + num_items_in_batch (`int`, *optional*): + Number of items in the batch. It is recommended to pass it when + you are doing gradient accumulation. + """ + + num_items_in_batch: Optional[int] + + +def is_timm_config_dict(config_dict: dict[str, Any]) -> bool: + """Checks whether a config dict is a timm config dict.""" + return "pretrained_cfg" in config_dict + + +def is_timm_local_checkpoint(pretrained_model_path: str) -> bool: + """ + Checks whether a checkpoint is a timm model checkpoint. + """ + if pretrained_model_path is None: + return False + + # in case it's Path, not str + pretrained_model_path = str(pretrained_model_path) + + is_file = os.path.isfile(pretrained_model_path) + is_dir = os.path.isdir(pretrained_model_path) + + # pretrained_model_path is a file + if is_file and pretrained_model_path.endswith(".json"): + with open(pretrained_model_path) as f: + config_dict = json.load(f) + return is_timm_config_dict(config_dict) + + # pretrained_model_path is a directory with a config.json + if is_dir and os.path.exists(os.path.join(pretrained_model_path, "config.json")): + with open(os.path.join(pretrained_model_path, "config.json")) as f: + config_dict = json.load(f) + return is_timm_config_dict(config_dict) + + return False + + +def set_attribute_for_modules(module: "torch.nn.Module", key: str, value: Any): + """ + Set a value to a module and all submodules. + """ + setattr(module, key, value) + for submodule in module.children(): + set_attribute_for_modules(submodule, key, value) + + +def del_attribute_from_modules(module: "torch.nn.Module", key: str): + """ + Delete a value from a module and all submodules. + """ + # because we might remove it previously in case it's a shared module, e.g. activation function + if hasattr(module, key): + delattr(module, key) + + for submodule in module.children(): + del_attribute_from_modules(submodule, key) + + +def can_return_tuple(func): + """ + Decorator to wrap model method, to call output.to_tuple() if return_dict=False passed as a kwarg or + use_return_dict=False is set in the config. + + Note: + output.to_tuple() convert output to tuple skipping all `None` values. + """ + + @wraps(func) + def wrapper(self, *args, **kwargs): + is_requested_to_return_tuple = kwargs.pop("return_dict", True) is False + is_configured_to_return_tuple = self.config.use_return_dict is False if hasattr(self, "config") else False + + # The following allows to convert output to tuple ONLY on top level forward call, + # while internal modules of the model will return Output objects + # to be able to use name-based attribute access in modeling code. + + # We will check if we are on top level module, if so, turn off to tuple conversion for all + # underling calls. + is_top_level_module = getattr(self, "_is_top_level_module", True) + if is_configured_to_return_tuple and is_top_level_module: + set_attribute_for_modules(self, "_is_top_level_module", False) + + try: + output = func(self, *args, **kwargs) + if is_requested_to_return_tuple or (is_configured_to_return_tuple and is_top_level_module): + output = output.to_tuple() + finally: + # Remove the flag after the model forward call is finished. + if is_configured_to_return_tuple and is_top_level_module: + del_attribute_from_modules(self, "_is_top_level_module") + + return output + + return wrapper diff --git a/lib/python3.12/site-packages/transformers/utils/hp_naming.py b/lib/python3.12/site-packages/transformers/utils/hp_naming.py new file mode 100644 index 0000000000000000000000000000000000000000..f7c5cb5259f8452b09cc910aee1fec7f1ba438c8 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/hp_naming.py @@ -0,0 +1,162 @@ +# Copyright 2020 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. + +import copy +import re + + +class TrialShortNamer: + PREFIX = "hp" + DEFAULTS = {} + NAMING_INFO = None + + @classmethod + def set_defaults(cls, prefix, defaults): + cls.PREFIX = prefix + cls.DEFAULTS = defaults + cls.build_naming_info() + + @staticmethod + def shortname_for_word(info, word): + if len(word) == 0: + return "" + short_word = None + if any(char.isdigit() for char in word): + raise Exception(f"Parameters should not contain numbers: '{word}' contains a number") + if word in info["short_word"]: + return info["short_word"][word] + for prefix_len in range(1, len(word) + 1): + prefix = word[:prefix_len] + if prefix in info["reverse_short_word"]: + continue + else: + short_word = prefix + break + + if short_word is None: + # Paranoid fallback + def int_to_alphabetic(integer): + s = "" + while integer != 0: + s = chr(ord("A") + integer % 10) + s + integer //= 10 + return s + + i = 0 + while True: + sword = word + "#" + int_to_alphabetic(i) + if sword in info["reverse_short_word"]: + continue + else: + short_word = sword + break + + info["short_word"][word] = short_word + info["reverse_short_word"][short_word] = word + return short_word + + @staticmethod + def shortname_for_key(info, param_name): + words = param_name.split("_") + + shortname_parts = [TrialShortNamer.shortname_for_word(info, word) for word in words] + + # We try to create a separatorless short name, but if there is a collision we have to fallback + # to a separated short name + separators = ["", "_"] + + for separator in separators: + shortname = separator.join(shortname_parts) + if shortname not in info["reverse_short_param"]: + info["short_param"][param_name] = shortname + info["reverse_short_param"][shortname] = param_name + return shortname + + return param_name + + @staticmethod + def add_new_param_name(info, param_name): + short_name = TrialShortNamer.shortname_for_key(info, param_name) + info["short_param"][param_name] = short_name + info["reverse_short_param"][short_name] = param_name + + @classmethod + def build_naming_info(cls): + if cls.NAMING_INFO is not None: + return + + info = { + "short_word": {}, + "reverse_short_word": {}, + "short_param": {}, + "reverse_short_param": {}, + } + + field_keys = list(cls.DEFAULTS.keys()) + + for k in field_keys: + cls.add_new_param_name(info, k) + + cls.NAMING_INFO = info + + @classmethod + def shortname(cls, params): + cls.build_naming_info() + assert cls.PREFIX is not None + name = [copy.copy(cls.PREFIX)] + + for k, v in params.items(): + if k not in cls.DEFAULTS: + raise Exception(f"You should provide a default value for the param name {k} with value {v}") + if v == cls.DEFAULTS[k]: + # The default value is not added to the name + continue + + key = cls.NAMING_INFO["short_param"][k] + + if isinstance(v, bool): + v = 1 if v else 0 + + sep = "" if isinstance(v, (int, float)) else "-" + e = f"{key}{sep}{v}" + name.append(e) + + return "_".join(name) + + @classmethod + def parse_repr(cls, repr): + repr = repr[len(cls.PREFIX) + 1 :] + if repr == "": + values = [] + else: + values = repr.split("_") + + parameters = {} + + for value in values: + if "-" in value: + p_k, p_v = value.split("-") + else: + p_k = re.sub("[0-9.]", "", value) + p_v = float(re.sub("[^0-9.]", "", value)) + + key = cls.NAMING_INFO["reverse_short_param"][p_k] + + parameters[key] = p_v + + for k in cls.DEFAULTS: + if k not in parameters: + parameters[k] = cls.DEFAULTS[k] + + return parameters diff --git a/lib/python3.12/site-packages/transformers/utils/hub.py b/lib/python3.12/site-packages/transformers/utils/hub.py new file mode 100644 index 0000000000000000000000000000000000000000..57a267b72889d35074a71352d0630dfcefc9f747 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/hub.py @@ -0,0 +1,1193 @@ +# Copyright 2022 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. +""" +Hub utilities: utilities related to download and cache models +""" + +import json +import os +import re +import sys +import tempfile +import warnings +from concurrent import futures +from pathlib import Path +from typing import Optional, Union +from urllib.parse import urlparse +from uuid import uuid4 + +import huggingface_hub +import requests +from huggingface_hub import ( + _CACHED_NO_EXIST, + CommitOperationAdd, + ModelCard, + ModelCardData, + constants, + create_branch, + create_commit, + create_repo, + hf_hub_download, + hf_hub_url, + snapshot_download, + try_to_load_from_cache, +) +from huggingface_hub.file_download import REGEX_COMMIT_HASH, http_get +from huggingface_hub.utils import ( + EntryNotFoundError, + GatedRepoError, + HfHubHTTPError, + LocalEntryNotFoundError, + OfflineModeIsEnabled, + RepositoryNotFoundError, + RevisionNotFoundError, + build_hf_headers, + get_session, + hf_raise_for_status, + send_telemetry, +) +from requests.exceptions import HTTPError + +from . import __version__, logging +from .generic import working_or_temp_dir +from .import_utils import ( + ENV_VARS_TRUE_VALUES, + _tf_version, + _torch_version, + is_tf_available, + is_torch_available, + is_training_run_on_sagemaker, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +_is_offline_mode = huggingface_hub.constants.HF_HUB_OFFLINE + + +def is_offline_mode(): + return _is_offline_mode + + +torch_cache_home = os.getenv("TORCH_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "torch")) +default_cache_path = constants.default_cache_path + +# Determine default cache directory. Lots of legacy environment variables to ensure backward compatibility. +# The best way to set the cache path is with the environment variable HF_HOME. For more details, checkout this +# documentation page: https://huggingface.co/docs/huggingface_hub/package_reference/environment_variables. +# +# In code, use `HF_HUB_CACHE` as the default cache path. This variable is set by the library and is guaranteed +# to be set to the right value. +# +# TODO: clean this for v5? +PYTORCH_PRETRAINED_BERT_CACHE = os.getenv("PYTORCH_PRETRAINED_BERT_CACHE", constants.HF_HUB_CACHE) +PYTORCH_TRANSFORMERS_CACHE = os.getenv("PYTORCH_TRANSFORMERS_CACHE", PYTORCH_PRETRAINED_BERT_CACHE) +TRANSFORMERS_CACHE = os.getenv("TRANSFORMERS_CACHE", PYTORCH_TRANSFORMERS_CACHE) + +HF_MODULES_CACHE = os.getenv("HF_MODULES_CACHE", os.path.join(constants.HF_HOME, "modules")) +TRANSFORMERS_DYNAMIC_MODULE_NAME = "transformers_modules" +SESSION_ID = uuid4().hex + +# Add deprecation warning for old environment variables. +for key in ("PYTORCH_PRETRAINED_BERT_CACHE", "PYTORCH_TRANSFORMERS_CACHE", "TRANSFORMERS_CACHE"): + if os.getenv(key) is not None: + warnings.warn( + f"Using `{key}` is deprecated and will be removed in v5 of Transformers. Use `HF_HOME` instead.", + FutureWarning, + ) + + +S3_BUCKET_PREFIX = "https://s3.amazonaws.com/models.huggingface.co/bert" +CLOUDFRONT_DISTRIB_PREFIX = "https://cdn.huggingface.co" + +_staging_mode = os.environ.get("HUGGINGFACE_CO_STAGING", "NO").upper() in ENV_VARS_TRUE_VALUES +_default_endpoint = "https://hub-ci.huggingface.co" if _staging_mode else "https://huggingface.co" + +HUGGINGFACE_CO_RESOLVE_ENDPOINT = _default_endpoint +if os.environ.get("HUGGINGFACE_CO_RESOLVE_ENDPOINT", None) is not None: + warnings.warn( + "Using the environment variable `HUGGINGFACE_CO_RESOLVE_ENDPOINT` is deprecated and will be removed in " + "Transformers v5. Use `HF_ENDPOINT` instead.", + FutureWarning, + ) + HUGGINGFACE_CO_RESOLVE_ENDPOINT = os.environ.get("HUGGINGFACE_CO_RESOLVE_ENDPOINT", None) +HUGGINGFACE_CO_RESOLVE_ENDPOINT = os.environ.get("HF_ENDPOINT", HUGGINGFACE_CO_RESOLVE_ENDPOINT) +HUGGINGFACE_CO_PREFIX = HUGGINGFACE_CO_RESOLVE_ENDPOINT + "/{model_id}/resolve/{revision}/{filename}" +HUGGINGFACE_CO_EXAMPLES_TELEMETRY = HUGGINGFACE_CO_RESOLVE_ENDPOINT + "/api/telemetry/examples" + + +def _get_cache_file_to_return( + path_or_repo_id: str, full_filename: str, cache_dir: Union[str, Path, None] = None, revision: Optional[str] = None +): + # 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, revision=revision) + if resolved_file is not None and resolved_file != _CACHED_NO_EXIST: + return resolved_file + return None + + +def is_remote_url(url_or_filename): + parsed = urlparse(url_or_filename) + return parsed.scheme in ("http", "https") + + +def define_sagemaker_information(): + try: + instance_data = requests.get(os.environ["ECS_CONTAINER_METADATA_URI"]).json() + dlc_container_used = instance_data["Image"] + dlc_tag = instance_data["Image"].split(":")[1] + except Exception: + dlc_container_used = None + dlc_tag = None + + sagemaker_params = json.loads(os.getenv("SM_FRAMEWORK_PARAMS", "{}")) + runs_distributed_training = True if "sagemaker_distributed_dataparallel_enabled" in sagemaker_params else False + account_id = os.getenv("TRAINING_JOB_ARN").split(":")[4] if "TRAINING_JOB_ARN" in os.environ else None + + sagemaker_object = { + "sm_framework": os.getenv("SM_FRAMEWORK_MODULE", None), + "sm_region": os.getenv("AWS_REGION", None), + "sm_number_gpu": os.getenv("SM_NUM_GPUS", 0), + "sm_number_cpu": os.getenv("SM_NUM_CPUS", 0), + "sm_distributed_training": runs_distributed_training, + "sm_deep_learning_container": dlc_container_used, + "sm_deep_learning_container_tag": dlc_tag, + "sm_account_id": account_id, + } + return sagemaker_object + + +def http_user_agent(user_agent: Union[dict, str, None] = None) -> str: + """ + Formats a user-agent string with basic info about a request. + """ + ua = f"transformers/{__version__}; python/{sys.version.split()[0]}; session_id/{SESSION_ID}" + if is_torch_available(): + ua += f"; torch/{_torch_version}" + if is_tf_available(): + ua += f"; tensorflow/{_tf_version}" + if constants.HF_HUB_DISABLE_TELEMETRY: + return ua + "; telemetry/off" + if is_training_run_on_sagemaker(): + ua += "; " + "; ".join(f"{k}/{v}" for k, v in define_sagemaker_information().items()) + # CI will set this value to True + if os.environ.get("TRANSFORMERS_IS_CI", "").upper() in ENV_VARS_TRUE_VALUES: + ua += "; is_ci/true" + if isinstance(user_agent, dict): + ua += "; " + "; ".join(f"{k}/{v}" for k, v in user_agent.items()) + elif isinstance(user_agent, str): + ua += "; " + user_agent + return ua + + +def extract_commit_hash(resolved_file: Optional[str], commit_hash: Optional[str]) -> Optional[str]: + """ + Extracts the commit hash from a resolved filename toward a cache file. + """ + if resolved_file is None or commit_hash is not None: + return commit_hash + resolved_file = str(Path(resolved_file).as_posix()) + search = re.search(r"snapshots/([^/]+)/", resolved_file) + if search is None: + return None + commit_hash = search.groups()[0] + return commit_hash if REGEX_COMMIT_HASH.match(commit_hash) else None + + +def cached_file( + path_or_repo_id: Union[str, os.PathLike], + filename: str, + **kwargs, +) -> Optional[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. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the configuration files and override the cached versions if they + exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, will only try to load the tokenizer configuration from local files. + 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. + repo_type (`str`, *optional*): + Specify the repo type (useful when downloading from a space for instance). + + + + Passing `token=True` is required when you want to use a private model. + + + + 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("google-bert/bert-base-uncased", "pytorch_model.bin") + ``` + """ + file = cached_files(path_or_repo_id=path_or_repo_id, filenames=[filename], **kwargs) + file = file[0] if file is not None else file + return file + + +def cached_files( + path_or_repo_id: Union[str, os.PathLike], + filenames: list[str], + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + resume_download: Optional[bool] = None, + proxies: Optional[dict[str, str]] = None, + token: Optional[Union[bool, str]] = None, + revision: Optional[str] = None, + local_files_only: bool = False, + subfolder: str = "", + repo_type: Optional[str] = None, + user_agent: Optional[Union[str, dict[str, str]]] = None, + _raise_exceptions_for_gated_repo: bool = True, + _raise_exceptions_for_missing_entries: bool = True, + _raise_exceptions_for_connection_errors: bool = True, + _commit_hash: Optional[str] = None, + **deprecated_kwargs, +) -> Optional[str]: + """ + Tries to locate several files in a local folder and repo, downloads and cache them 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. + filenames (`List[str]`): + The name of all the files 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. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the configuration files and override the cached versions if they + exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, will only try to load the tokenizer configuration from local files. + 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. + repo_type (`str`, *optional*): + Specify the repo type (useful when downloading from a space for instance). + + Private args: + _raise_exceptions_for_gated_repo (`bool`): + if False, do not raise an exception for gated repo error but return None. + _raise_exceptions_for_missing_entries (`bool`): + if False, do not raise an exception for missing entries but return None. + _raise_exceptions_for_connection_errors (`bool`): + if False, do not raise an exception for connection errors but return None. + _commit_hash (`str`, *optional*): + passed when we are chaining several calls to various files (e.g. when loading a tokenizer or + a pipeline). If files are cached for this commit hash, avoid calls to head and get from the cache. + + + + Passing `token=True` is required when you want to use a private model. + + + + 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("google-bert/bert-base-uncased", "pytorch_model.bin") + ``` + """ + use_auth_token = deprecated_kwargs.pop("use_auth_token", None) + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.") + token = use_auth_token + + if is_offline_mode() and not local_files_only: + logger.info("Offline mode: forcing local_files_only=True") + local_files_only = True + if subfolder is None: + subfolder = "" + + # Add folder to filenames + full_filenames = [os.path.join(subfolder, file) for file in filenames] + + path_or_repo_id = str(path_or_repo_id) + existing_files = [] + for filename in full_filenames: + if os.path.isdir(path_or_repo_id): + resolved_file = os.path.join(path_or_repo_id, filename) + if not os.path.isfile(resolved_file): + if _raise_exceptions_for_missing_entries and filename != os.path.join(subfolder, "config.json"): + revision_ = "main" if revision is None else revision + raise OSError( + f"{path_or_repo_id} does not appear to have a file named {filename}. Checkout " + f"'https://huggingface.co/{path_or_repo_id}/tree/{revision_}' for available files." + ) + else: + return None + existing_files.append(resolved_file) + + # All files exist + if len(existing_files) == len(full_filenames): + return existing_files + + if cache_dir is None: + cache_dir = TRANSFORMERS_CACHE + if isinstance(cache_dir, Path): + cache_dir = str(cache_dir) + + existing_files = [] + file_counter = 0 + if _commit_hash is not None and not force_download: + for filename in full_filenames: + # If the file is cached under that commit hash, we return it directly. + resolved_file = try_to_load_from_cache( + path_or_repo_id, filename, cache_dir=cache_dir, revision=_commit_hash, repo_type=repo_type + ) + if resolved_file is not None: + if resolved_file is not _CACHED_NO_EXIST: + file_counter += 1 + existing_files.append(resolved_file) + elif not _raise_exceptions_for_missing_entries: + file_counter += 1 + else: + raise OSError(f"Could not locate {filename} inside {path_or_repo_id}.") + + # Either all the files were found, or some were _CACHED_NO_EXIST but we do not raise for missing entries + if file_counter == len(full_filenames): + return existing_files if len(existing_files) > 0 else None + + user_agent = http_user_agent(user_agent) + # download the files if needed + try: + if len(full_filenames) == 1: + # This is slightly better for only 1 file + hf_hub_download( + path_or_repo_id, + filenames[0], + subfolder=None if len(subfolder) == 0 else subfolder, + repo_type=repo_type, + revision=revision, + cache_dir=cache_dir, + user_agent=user_agent, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + else: + snapshot_download( + path_or_repo_id, + allow_patterns=full_filenames, + repo_type=repo_type, + revision=revision, + cache_dir=cache_dir, + user_agent=user_agent, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + + except Exception as e: + # We cannot recover from them + if isinstance(e, RepositoryNotFoundError) and not isinstance(e, GatedRepoError): + raise OSError( + f"{path_or_repo_id} is not a local folder and is not a valid model identifier " + "listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a token " + "having permission to this repo either by logging in with `huggingface-cli login` or by passing " + "`token=`" + ) from e + elif isinstance(e, RevisionNotFoundError): + raise OSError( + f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists " + "for this model name. Check the model page at " + f"'https://huggingface.co/{path_or_repo_id}' for available revisions." + ) from e + + # Now we try to recover if we can find all files correctly in the cache + resolved_files = [ + _get_cache_file_to_return(path_or_repo_id, filename, cache_dir, revision) for filename in full_filenames + ] + if all(file is not None for file in resolved_files): + return resolved_files + + # Raise based on the flags. Note that we will raise for missing entries at the very end, even when + # not entering this Except block, as it may also happen when `snapshot_download` does not raise + if isinstance(e, GatedRepoError): + if not _raise_exceptions_for_gated_repo: + return None + raise OSError( + "You are trying to access a gated repo.\nMake sure to have access to it at " + f"https://huggingface.co/{path_or_repo_id}.\n{str(e)}" + ) from e + elif isinstance(e, LocalEntryNotFoundError): + if not _raise_exceptions_for_connection_errors: + return None + # Here we only raise if both flags for missing entry and connection errors are True (because it can be raised + # even when `local_files_only` is True, in which case raising for connections errors only would not make sense) + elif _raise_exceptions_for_missing_entries: + raise OSError( + f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load the files, and couldn't find them in the" + f" cached files.\nCheckout your internet connection or see how to run the library in offline mode at" + " 'https://huggingface.co/docs/transformers/installation#offline-mode'." + ) from e + # snapshot_download will not raise EntryNotFoundError, but hf_hub_download can. If this is the case, it will be treated + # later on anyway and re-raised if needed + elif isinstance(e, HTTPError) and not isinstance(e, EntryNotFoundError): + if not _raise_exceptions_for_connection_errors: + return None + raise OSError(f"There was a specific connection error when trying to load {path_or_repo_id}:\n{e}") + + resolved_files = [ + _get_cache_file_to_return(path_or_repo_id, filename, cache_dir, revision) for filename in full_filenames + ] + # If there are any missing file and the flag is active, raise + if any(file is None for file in resolved_files) and _raise_exceptions_for_missing_entries: + missing_entries = [original for original, resolved in zip(full_filenames, resolved_files) if resolved is None] + # Last escape + if len(resolved_files) == 1 and missing_entries[0] == os.path.join(subfolder, "config.json"): + return None + # Now we raise for missing entries + revision_ = "main" if revision is None else revision + msg = ( + f"a file named {missing_entries[0]}" if len(missing_entries) == 1 else f"files named {(*missing_entries,)}" + ) + raise EnvironmentError( + f"{path_or_repo_id} does not appear to have {msg}. Checkout 'https://huggingface.co/{path_or_repo_id}/tree/{revision_}'" + "for available files." + ) + + # Remove potential missing entries (we can silently remove them at this point based on the flags) + resolved_files = [file for file in resolved_files if file is not None] + # Return `None` if the list is empty, coherent with other Exception when the flag is not active + resolved_files = None if len(resolved_files) == 0 else resolved_files + + return resolved_files + + +# TODO cyril: Deprecated and should be removed in 4.51 +def get_file_from_repo( + *args, + **kwargs, +): + """ + Tries to locate a file in a local folder and repo, downloads and cache it if necessary. + + Args: + path_or_repo (`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. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the configuration files and override the cached versions if they + exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, will only try to load the tokenizer configuration from local files. + 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. + + + + Passing `token=True` is required when you want to use a private model. + + + + Returns: + `Optional[str]`: Returns the resolved file (to the cache folder if downloaded from a repo) or `None` if the + file does not exist. + + Examples: + + ```python + # Download a tokenizer configuration from huggingface.co and cache. + tokenizer_config = get_file_from_repo("google-bert/bert-base-uncased", "tokenizer_config.json") + # This model does not have a tokenizer config so the result will be None. + tokenizer_config = get_file_from_repo("FacebookAI/xlm-roberta-base", "tokenizer_config.json") + ``` + """ + logger.warning( + "`get_file_from_repo` is deprecated and will be removed in version 4.51. Use `cached_file` instead." + ) + return cached_file( + *args, + _raise_exceptions_for_gated_repo=False, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + **kwargs, + ) + + +def download_url(url, proxies=None): + """ + Downloads a given url in a temporary file. This function is not safe to use in multiple processes. Its only use is + for deprecated behavior allowing to download config/models with a single url instead of using the Hub. + + Args: + url (`str`): The url of the file to download. + proxies (`Dict[str, str]`, *optional*): + 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. + + Returns: + `str`: The location of the temporary file where the url was downloaded. + """ + warnings.warn( + f"Using `from_pretrained` with the url of a file (here {url}) is deprecated and won't be possible anymore in" + " v5 of Transformers. You should host your file on the Hub (hf.co) instead and use the repository ID. Note" + " that this is not compatible with the caching system (your file will be downloaded at each execution) or" + " multiple processes (each process will download the file in a different temporary file).", + FutureWarning, + ) + tmp_fd, tmp_file = tempfile.mkstemp() + with os.fdopen(tmp_fd, "wb") as f: + http_get(url, f, proxies=proxies) + return tmp_file + + +def has_file( + path_or_repo: Union[str, os.PathLike], + filename: str, + revision: Optional[str] = None, + proxies: Optional[dict[str, str]] = None, + token: Optional[Union[bool, str]] = None, + *, + local_files_only: bool = False, + cache_dir: Union[str, Path, None] = None, + repo_type: Optional[str] = None, + **deprecated_kwargs, +): + """ + Checks if a repo contains a given file without downloading it. Works for remote repos and local folders. + + If offline mode is enabled, checks if the file exists in the cache. + + + + This function will raise an error if the repository `path_or_repo` is not valid or if `revision` does not exist for + this repo, but will return False for regular connection errors. + + + """ + use_auth_token = deprecated_kwargs.pop("use_auth_token", None) + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.") + token = use_auth_token + + # If path to local directory, check if the file exists + if os.path.isdir(path_or_repo): + return os.path.isfile(os.path.join(path_or_repo, filename)) + + # Else it's a repo => let's check if the file exists in local cache or on the Hub + + # Check if file exists in cache + # This information might be outdated so it's best to also make a HEAD call (if allowed). + cached_path = try_to_load_from_cache( + repo_id=path_or_repo, + filename=filename, + revision=revision, + repo_type=repo_type, + cache_dir=cache_dir, + ) + has_file_in_cache = isinstance(cached_path, str) + + # If local_files_only, don't try the HEAD call + if local_files_only: + return has_file_in_cache + + # Check if the file exists + try: + response = get_session().head( + hf_hub_url(path_or_repo, filename=filename, revision=revision, repo_type=repo_type), + headers=build_hf_headers(token=token, user_agent=http_user_agent()), + allow_redirects=False, + proxies=proxies, + timeout=10, + ) + except (requests.exceptions.SSLError, requests.exceptions.ProxyError): + # Actually raise for those subclasses of ConnectionError + raise + except ( + requests.exceptions.ConnectionError, + requests.exceptions.Timeout, + OfflineModeIsEnabled, + ): + return has_file_in_cache + + try: + hf_raise_for_status(response) + return True + except GatedRepoError as e: + logger.error(e) + raise OSError( + f"{path_or_repo} is a gated repository. Make sure to request access at " + f"https://huggingface.co/{path_or_repo} and pass a token having permission to this repo either by " + "logging in with `huggingface-cli login` or by passing `token=`." + ) from e + except RepositoryNotFoundError as e: + logger.error(e) + raise OSError(f"{path_or_repo} is not a local folder or a valid repository name on 'https://hf.co'.") from e + except RevisionNotFoundError as e: + logger.error(e) + raise OSError( + f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for this " + f"model name. Check the model page at 'https://huggingface.co/{path_or_repo}' for available revisions." + ) from e + except EntryNotFoundError: + return False # File does not exist + except requests.HTTPError: + # Any authentication/authorization error will be caught here => default to cache + return has_file_in_cache + + +class PushToHubMixin: + """ + A Mixin containing the functionality to push a model or tokenizer to the hub. + """ + + def _create_repo( + self, + repo_id: str, + private: Optional[bool] = None, + token: Optional[Union[bool, str]] = None, + repo_url: Optional[str] = None, + organization: Optional[str] = None, + ) -> str: + """ + Create the repo if needed, cleans up repo_id with deprecated kwargs `repo_url` and `organization`, retrieves + the token. + """ + if repo_url is not None: + warnings.warn( + "The `repo_url` argument is deprecated and will be removed in v5 of Transformers. Use `repo_id` " + "instead." + ) + if repo_id is not None: + raise ValueError( + "`repo_id` and `repo_url` are both specified. Please set only the argument `repo_id`." + ) + repo_id = repo_url.replace(f"{HUGGINGFACE_CO_RESOLVE_ENDPOINT}/", "") + if organization is not None: + warnings.warn( + "The `organization` argument is deprecated and will be removed in v5 of Transformers. Set your " + "organization directly in the `repo_id` passed instead (`repo_id={organization}/{model_id}`)." + ) + if not repo_id.startswith(organization): + if "/" in repo_id: + repo_id = repo_id.split("/")[-1] + repo_id = f"{organization}/{repo_id}" + + url = create_repo(repo_id=repo_id, token=token, private=private, exist_ok=True) + return url.repo_id + + def _get_files_timestamps(self, working_dir: Union[str, os.PathLike]): + """ + Returns the list of files with their last modification timestamp. + """ + return {f: os.path.getmtime(os.path.join(working_dir, f)) for f in os.listdir(working_dir)} + + def _upload_modified_files( + self, + working_dir: Union[str, os.PathLike], + repo_id: str, + files_timestamps: dict[str, float], + commit_message: Optional[str] = None, + token: Optional[Union[bool, str]] = None, + create_pr: bool = False, + revision: Optional[str] = None, + commit_description: Optional[str] = None, + ): + """ + Uploads all modified files in `working_dir` to `repo_id`, based on `files_timestamps`. + """ + if commit_message is None: + if "Model" in self.__class__.__name__: + commit_message = "Upload model" + elif "Config" in self.__class__.__name__: + commit_message = "Upload config" + elif "Tokenizer" in self.__class__.__name__: + commit_message = "Upload tokenizer" + elif "FeatureExtractor" in self.__class__.__name__: + commit_message = "Upload feature extractor" + elif "Processor" in self.__class__.__name__: + commit_message = "Upload processor" + else: + commit_message = f"Upload {self.__class__.__name__}" + modified_files = [ + f + for f in os.listdir(working_dir) + if f not in files_timestamps or os.path.getmtime(os.path.join(working_dir, f)) > files_timestamps[f] + ] + + # filter for actual files + folders at the root level + modified_files = [ + f + for f in modified_files + if os.path.isfile(os.path.join(working_dir, f)) or os.path.isdir(os.path.join(working_dir, f)) + ] + + operations = [] + # upload standalone files + for file in modified_files: + if os.path.isdir(os.path.join(working_dir, file)): + # go over individual files of folder + for f in os.listdir(os.path.join(working_dir, file)): + operations.append( + CommitOperationAdd( + path_or_fileobj=os.path.join(working_dir, file, f), path_in_repo=os.path.join(file, f) + ) + ) + else: + operations.append( + CommitOperationAdd(path_or_fileobj=os.path.join(working_dir, file), path_in_repo=file) + ) + + if revision is not None and not revision.startswith("refs/pr"): + try: + create_branch(repo_id=repo_id, branch=revision, token=token, exist_ok=True) + except HfHubHTTPError as e: + if e.response.status_code == 403 and create_pr: + # If we are creating a PR on a repo we don't have access to, we can't create the branch. + # so let's assume the branch already exists. If it's not the case, an error will be raised when + # calling `create_commit` below. + pass + else: + raise + + logger.info(f"Uploading the following files to {repo_id}: {','.join(modified_files)}") + return create_commit( + repo_id=repo_id, + operations=operations, + commit_message=commit_message, + commit_description=commit_description, + token=token, + create_pr=create_pr, + revision=revision, + ) + + def push_to_hub( + self, + repo_id: str, + use_temp_dir: Optional[bool] = None, + commit_message: Optional[str] = None, + private: Optional[bool] = None, + token: Optional[Union[bool, str]] = None, + max_shard_size: Optional[Union[int, str]] = "5GB", + create_pr: bool = False, + safe_serialization: bool = True, + revision: Optional[str] = None, + commit_description: Optional[str] = None, + tags: Optional[list[str]] = None, + **deprecated_kwargs, + ) -> str: + """ + Upload the {object_files} to the 🤗 Model Hub. + + Parameters: + repo_id (`str`): + The name of the repository you want to push your {object} to. It should contain your organization name + when pushing to a given organization. + use_temp_dir (`bool`, *optional*): + Whether or not to use a temporary directory to store the files saved before they are pushed to the Hub. + Will default to `True` if there is no directory named like `repo_id`, `False` otherwise. + commit_message (`str`, *optional*): + Message to commit while pushing. Will default to `"Upload {object}"`. + private (`bool`, *optional*): + Whether to make the repo private. If `None` (default), the repo will be public unless the organization's default is private. This value is ignored if the repo already exists. + token (`bool` or `str`, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). Will default to `True` if `repo_url` + is not specified. + max_shard_size (`int` or `str`, *optional*, defaults to `"5GB"`): + Only applicable for models. The maximum size for a checkpoint before being sharded. Checkpoints shard + will then be each of size lower than this size. If expressed as a string, needs to be digits followed + by a unit (like `"5MB"`). We default it to `"5GB"` so that users can easily load models on free-tier + Google Colab instances without any CPU OOM issues. + create_pr (`bool`, *optional*, defaults to `False`): + Whether or not to create a PR with the uploaded files or directly commit. + safe_serialization (`bool`, *optional*, defaults to `True`): + Whether or not to convert the model weights in safetensors format for safer serialization. + revision (`str`, *optional*): + Branch to push the uploaded files to. + commit_description (`str`, *optional*): + The description of the commit that will be created + tags (`List[str]`, *optional*): + List of tags to push on the Hub. + + Examples: + + ```python + from transformers import {object_class} + + {object} = {object_class}.from_pretrained("google-bert/bert-base-cased") + + # Push the {object} to your namespace with the name "my-finetuned-bert". + {object}.push_to_hub("my-finetuned-bert") + + # Push the {object} to an organization with the name "my-finetuned-bert". + {object}.push_to_hub("huggingface/my-finetuned-bert") + ``` + """ + use_auth_token = deprecated_kwargs.pop("use_auth_token", None) + ignore_metadata_errors = deprecated_kwargs.pop("ignore_metadata_errors", False) + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError( + "`token` and `use_auth_token` are both specified. Please set only the argument `token`." + ) + token = use_auth_token + + repo_path_or_name = deprecated_kwargs.pop("repo_path_or_name", None) + if repo_path_or_name is not None: + # Should use `repo_id` instead of `repo_path_or_name`. When using `repo_path_or_name`, we try to infer + # repo_id from the folder path, if it exists. + warnings.warn( + "The `repo_path_or_name` argument is deprecated and will be removed in v5 of Transformers. Use " + "`repo_id` instead.", + FutureWarning, + ) + if repo_id is not None: + raise ValueError( + "`repo_id` and `repo_path_or_name` are both specified. Please set only the argument `repo_id`." + ) + if os.path.isdir(repo_path_or_name): + # repo_path: infer repo_id from the path + repo_id = repo_id.split(os.path.sep)[-1] + working_dir = repo_id + else: + # repo_name: use it as repo_id + repo_id = repo_path_or_name + working_dir = repo_id.split("/")[-1] + else: + # Repo_id is passed correctly: infer working_dir from it + working_dir = repo_id.split("/")[-1] + + # Deprecation warning will be sent after for repo_url and organization + repo_url = deprecated_kwargs.pop("repo_url", None) + organization = deprecated_kwargs.pop("organization", None) + + repo_id = self._create_repo( + repo_id, private=private, token=token, repo_url=repo_url, organization=organization + ) + + # Create a new empty model card and eventually tag it + model_card = create_and_tag_model_card( + repo_id, tags, token=token, ignore_metadata_errors=ignore_metadata_errors + ) + + if use_temp_dir is None: + use_temp_dir = not os.path.isdir(working_dir) + + with working_or_temp_dir(working_dir=working_dir, use_temp_dir=use_temp_dir) as work_dir: + files_timestamps = self._get_files_timestamps(work_dir) + + # Save all files. + self.save_pretrained(work_dir, max_shard_size=max_shard_size, safe_serialization=safe_serialization) + + # Update model card if needed: + model_card.save(os.path.join(work_dir, "README.md")) + + return self._upload_modified_files( + work_dir, + repo_id, + files_timestamps, + commit_message=commit_message, + token=token, + create_pr=create_pr, + revision=revision, + commit_description=commit_description, + ) + + +def send_example_telemetry(example_name, *example_args, framework="pytorch"): + """ + Sends telemetry that helps tracking the examples use. + + Args: + example_name (`str`): The name of the example. + *example_args (dataclasses or `argparse.ArgumentParser`): The arguments to the script. This function will only + try to extract the model and dataset name from those. Nothing else is tracked. + framework (`str`, *optional*, defaults to `"pytorch"`): The framework for the example. + """ + if is_offline_mode(): + return + + data = {"example": example_name, "framework": framework} + for args in example_args: + args_as_dict = {k: v for k, v in args.__dict__.items() if not k.startswith("_") and v is not None} + if "model_name_or_path" in args_as_dict: + model_name = args_as_dict["model_name_or_path"] + # Filter out local paths + if not os.path.isdir(model_name): + data["model_name"] = args_as_dict["model_name_or_path"] + if "dataset_name" in args_as_dict: + data["dataset_name"] = args_as_dict["dataset_name"] + elif "task_name" in args_as_dict: + # Extract script name from the example_name + script_name = example_name.replace("tf_", "").replace("flax_", "").replace("run_", "") + script_name = script_name.replace("_no_trainer", "") + data["dataset_name"] = f"{script_name}-{args_as_dict['task_name']}" + + # Send telemetry in the background + send_telemetry( + topic="examples", library_name="transformers", library_version=__version__, user_agent=http_user_agent(data) + ) + + +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 get_checkpoint_shard_files( + pretrained_model_name_or_path, + index_filename, + cache_dir=None, + force_download=False, + proxies=None, + resume_download=None, + local_files_only=False, + token=None, + user_agent=None, + revision=None, + subfolder="", + _commit_hash=None, + **deprecated_kwargs, +): + """ + 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 + + use_auth_token = deprecated_kwargs.pop("use_auth_token", None) + if use_auth_token is not None: + warnings.warn( + "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.", + FutureWarning, + ) + if token is not None: + raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.") + token = use_auth_token + + 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) 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() + + # 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. Try to get everything from cache, + # or download the files + cached_filenames = cached_files( + pretrained_model_name_or_path, + shard_filenames, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + local_files_only=local_files_only, + token=token, + user_agent=user_agent, + revision=revision, + subfolder=subfolder, + _commit_hash=_commit_hash, + ) + + return cached_filenames, sharded_metadata + + +def create_and_tag_model_card( + repo_id: str, + tags: Optional[list[str]] = None, + token: Optional[str] = None, + ignore_metadata_errors: bool = False, +): + """ + Creates or loads an existing model card and tags it. + + Args: + repo_id (`str`): + The repo_id where to look for the model card. + tags (`List[str]`, *optional*): + The list of tags to add in the model card + token (`str`, *optional*): + Authentication token, obtained with `huggingface_hub.HfApi.login` method. Will default to the stored token. + ignore_metadata_errors (`bool`, *optional*, defaults to `False`): + If True, errors while parsing the metadata section will be ignored. Some information might be lost during + the process. Use it at your own risk. + """ + try: + # Check if the model card is present on the remote repo + model_card = ModelCard.load(repo_id, token=token, ignore_metadata_errors=ignore_metadata_errors) + except EntryNotFoundError: + # Otherwise create a simple model card from template + model_description = "This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated." + card_data = ModelCardData(tags=[] if tags is None else tags, library_name="transformers") + model_card = ModelCard.from_template(card_data, model_description=model_description) + + if tags is not None: + # Ensure model_card.data.tags is a list and not None + if model_card.data.tags is None: + model_card.data.tags = [] + for model_tag in tags: + if model_tag not in model_card.data.tags: + model_card.data.tags.append(model_tag) + + return model_card + + +class PushInProgress: + """ + Internal class to keep track of a push in progress (which might contain multiple `Future` jobs). + """ + + def __init__(self, jobs: Optional[futures.Future] = None) -> None: + self.jobs = [] if jobs is None else jobs + + def is_done(self): + return all(job.done() for job in self.jobs) + + def wait_until_done(self): + futures.wait(self.jobs) + + def cancel(self) -> None: + self.jobs = [ + job + for job in self.jobs + # Cancel the job if it wasn't started yet and remove cancelled/done jobs from the list + if not (job.cancel() or job.done()) + ] diff --git a/lib/python3.12/site-packages/transformers/utils/import_utils.py b/lib/python3.12/site-packages/transformers/utils/import_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..61fb91e5a844a04fcb7f2eab9b259754964ffde6 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/import_utils.py @@ -0,0 +1,2417 @@ +# Copyright 2022 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. +""" +Import utilities: Utilities related to imports and our lazy inits. +""" + +import importlib.machinery +import importlib.metadata +import importlib.util +import json +import os +import shutil +import subprocess +import sys +import warnings +from collections import OrderedDict +from functools import lru_cache +from itertools import chain +from types import ModuleType +from typing import Any, Dict, FrozenSet, Optional, Set, Tuple, Union + +from packaging import version + +from . import logging + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# TODO: This doesn't work for all packages (`bs4`, `faiss`, etc.) Talk to Sylvain to see how to do with it better. +def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[Tuple[bool, str], bool]: + # Check if the package spec exists and grab its version to avoid importing a local directory + package_exists = importlib.util.find_spec(pkg_name) is not None + package_version = "N/A" + if package_exists: + try: + # TODO: Once python 3.9 support is dropped, `importlib.metadata.packages_distributions()` + # should be used here to map from package name to distribution names + # e.g. PIL -> Pillow, Pillow-SIMD; quark -> amd-quark; onnxruntime -> onnxruntime-gpu. + # `importlib.metadata.packages_distributions()` is not available in Python 3.9. + + # Primary method to get the package version + package_version = importlib.metadata.version(pkg_name) + except importlib.metadata.PackageNotFoundError: + # Fallback method: Only for "torch" and versions containing "dev" + if pkg_name == "torch": + try: + package = importlib.import_module(pkg_name) + temp_version = getattr(package, "__version__", "N/A") + # Check if the version contains "dev" + if "dev" in temp_version: + package_version = temp_version + package_exists = True + else: + package_exists = False + except ImportError: + # If the package can't be imported, it's not available + package_exists = False + elif pkg_name == "quark": + # TODO: remove once `importlib.metadata.packages_distributions()` is supported. + try: + package_version = importlib.metadata.version("amd-quark") + except Exception: + package_exists = False + else: + # For packages other than "torch", don't attempt the fallback and set as not available + package_exists = False + logger.debug(f"Detected {pkg_name} version: {package_version}") + if return_version: + return package_exists, package_version + else: + return package_exists + + +ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"} +ENV_VARS_TRUE_AND_AUTO_VALUES = ENV_VARS_TRUE_VALUES.union({"AUTO"}) + +USE_TF = os.environ.get("USE_TF", "AUTO").upper() +USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper() +USE_JAX = os.environ.get("USE_FLAX", "AUTO").upper() + +# Try to run a native pytorch job in an environment with TorchXLA installed by setting this value to 0. +USE_TORCH_XLA = os.environ.get("USE_TORCH_XLA", "1").upper() + +FORCE_TF_AVAILABLE = os.environ.get("FORCE_TF_AVAILABLE", "AUTO").upper() + +# `transformers` requires `torch>=1.11` but this variable is exposed publicly, and we can't simply remove it. +# This is the version of torch required to run torch.fx features and torch.onnx with dictionary inputs. +TORCH_FX_REQUIRED_VERSION = version.parse("1.10") + +ACCELERATE_MIN_VERSION = "0.26.0" +SCHEDULEFREE_MIN_VERSION = "1.2.6" +FSDP_MIN_VERSION = "1.12.0" +GGUF_MIN_VERSION = "0.10.0" +XLA_FSDPV2_MIN_VERSION = "2.2.0" +HQQ_MIN_VERSION = "0.2.1" +VPTQ_MIN_VERSION = "0.0.4" +TORCHAO_MIN_VERSION = "0.4.0" + + +_accelerate_available, _accelerate_version = _is_package_available("accelerate", return_version=True) +_apex_available = _is_package_available("apex") +_apollo_torch_available = _is_package_available("apollo_torch") +_aqlm_available = _is_package_available("aqlm") +_vptq_available, _vptq_version = _is_package_available("vptq", return_version=True) +_av_available = importlib.util.find_spec("av") is not None +_decord_available = importlib.util.find_spec("decord") is not None +_bitsandbytes_available = _is_package_available("bitsandbytes") +_eetq_available = _is_package_available("eetq") +_fbgemm_gpu_available = _is_package_available("fbgemm_gpu") +_galore_torch_available = _is_package_available("galore_torch") +_lomo_available = _is_package_available("lomo_optim") +_grokadamw_available = _is_package_available("grokadamw") +_schedulefree_available, _schedulefree_version = _is_package_available("schedulefree", return_version=True) +# `importlib.metadata.version` doesn't work with `bs4` but `beautifulsoup4`. For `importlib.util.find_spec`, reversed. +_bs4_available = importlib.util.find_spec("bs4") is not None +_coloredlogs_available = _is_package_available("coloredlogs") +# `importlib.metadata.util` doesn't work with `opencv-python-headless`. +_cv2_available = importlib.util.find_spec("cv2") is not None +_yt_dlp_available = importlib.util.find_spec("yt_dlp") is not None +_datasets_available = _is_package_available("datasets") +_detectron2_available = _is_package_available("detectron2") +# We need to check both `faiss` and `faiss-cpu`. +_faiss_available = importlib.util.find_spec("faiss") is not None +try: + _faiss_version = importlib.metadata.version("faiss") + logger.debug(f"Successfully imported faiss version {_faiss_version}") +except importlib.metadata.PackageNotFoundError: + try: + _faiss_version = importlib.metadata.version("faiss-cpu") + logger.debug(f"Successfully imported faiss version {_faiss_version}") + except importlib.metadata.PackageNotFoundError: + _faiss_available = False +_ftfy_available = _is_package_available("ftfy") +_g2p_en_available = _is_package_available("g2p_en") +_hadamard_available = _is_package_available("fast_hadamard_transform") +_ipex_available, _ipex_version = _is_package_available("intel_extension_for_pytorch", return_version=True) +_jieba_available = _is_package_available("jieba") +_jinja_available = _is_package_available("jinja2") +_kenlm_available = _is_package_available("kenlm") +_keras_nlp_available = _is_package_available("keras_nlp") +_levenshtein_available = _is_package_available("Levenshtein") +_librosa_available = _is_package_available("librosa") +_natten_available = _is_package_available("natten") +_nltk_available = _is_package_available("nltk") +_onnx_available = _is_package_available("onnx") +_openai_available = _is_package_available("openai") +_optimum_available = _is_package_available("optimum") +_auto_gptq_available = _is_package_available("auto_gptq") +_gptqmodel_available = _is_package_available("gptqmodel") +# `importlib.metadata.version` doesn't work with `awq` +_auto_awq_available = importlib.util.find_spec("awq") is not None +_quark_available = _is_package_available("quark") +_is_optimum_quanto_available = False +try: + importlib.metadata.version("optimum_quanto") + _is_optimum_quanto_available = True +except importlib.metadata.PackageNotFoundError: + _is_optimum_quanto_available = False +# For compressed_tensors, only check spec to allow compressed_tensors-nightly package +_compressed_tensors_available = importlib.util.find_spec("compressed_tensors") is not None +_pandas_available = _is_package_available("pandas") +_peft_available = _is_package_available("peft") +_phonemizer_available = _is_package_available("phonemizer") +_uroman_available = _is_package_available("uroman") +_psutil_available = _is_package_available("psutil") +_py3nvml_available = _is_package_available("py3nvml") +_pyctcdecode_available = _is_package_available("pyctcdecode") +_pygments_available = _is_package_available("pygments") +_pytesseract_available = _is_package_available("pytesseract") +_pytest_available = _is_package_available("pytest") +_pytorch_quantization_available = _is_package_available("pytorch_quantization") +_rjieba_available = _is_package_available("rjieba") +_sacremoses_available = _is_package_available("sacremoses") +_safetensors_available = _is_package_available("safetensors") +_scipy_available = _is_package_available("scipy") +_sentencepiece_available = _is_package_available("sentencepiece") +_is_seqio_available = _is_package_available("seqio") +_is_gguf_available, _gguf_version = _is_package_available("gguf", return_version=True) +_sklearn_available = importlib.util.find_spec("sklearn") is not None +if _sklearn_available: + try: + importlib.metadata.version("scikit-learn") + except importlib.metadata.PackageNotFoundError: + _sklearn_available = False +_smdistributed_available = importlib.util.find_spec("smdistributed") is not None +_soundfile_available = _is_package_available("soundfile") +_spacy_available = _is_package_available("spacy") +_sudachipy_available, _sudachipy_version = _is_package_available("sudachipy", return_version=True) +_tensorflow_probability_available = _is_package_available("tensorflow_probability") +_tensorflow_text_available = _is_package_available("tensorflow_text") +_tf2onnx_available = _is_package_available("tf2onnx") +_timm_available = _is_package_available("timm") +_tokenizers_available = _is_package_available("tokenizers") +_torchaudio_available = _is_package_available("torchaudio") +_torchao_available, _torchao_version = _is_package_available("torchao", return_version=True) +_torchdistx_available = _is_package_available("torchdistx") +_torchvision_available, _torchvision_version = _is_package_available("torchvision", return_version=True) +_mlx_available = _is_package_available("mlx") +_num2words_available = _is_package_available("num2words") +_hqq_available, _hqq_version = _is_package_available("hqq", return_version=True) +_tiktoken_available = _is_package_available("tiktoken") +_blobfile_available = _is_package_available("blobfile") +_liger_kernel_available = _is_package_available("liger_kernel") +_triton_available = _is_package_available("triton") +_spqr_available = _is_package_available("spqr_quant") +_rich_available = _is_package_available("rich") + +_torch_version = "N/A" +_torch_available = False +if USE_TORCH in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TF not in ENV_VARS_TRUE_VALUES: + _torch_available, _torch_version = _is_package_available("torch", return_version=True) +else: + logger.info("Disabling PyTorch because USE_TF is set") + _torch_available = False + + +_tf_version = "N/A" +_tf_available = False +if FORCE_TF_AVAILABLE in ENV_VARS_TRUE_VALUES: + _tf_available = True +else: + if USE_TF in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TORCH not in ENV_VARS_TRUE_VALUES: + # Note: _is_package_available("tensorflow") fails for tensorflow-cpu. Please test any changes to the line below + # with tensorflow-cpu to make sure it still works! + _tf_available = importlib.util.find_spec("tensorflow") is not None + if _tf_available: + candidates = ( + "tensorflow", + "tensorflow-cpu", + "tensorflow-gpu", + "tf-nightly", + "tf-nightly-cpu", + "tf-nightly-gpu", + "tf-nightly-rocm", + "intel-tensorflow", + "intel-tensorflow-avx512", + "tensorflow-rocm", + "tensorflow-macos", + "tensorflow-aarch64", + ) + _tf_version = None + # For the metadata, we have to look for both tensorflow and tensorflow-cpu + for pkg in candidates: + try: + _tf_version = importlib.metadata.version(pkg) + break + except importlib.metadata.PackageNotFoundError: + pass + _tf_available = _tf_version is not None + if _tf_available: + if version.parse(_tf_version) < version.parse("2"): + logger.info( + f"TensorFlow found but with version {_tf_version}. Transformers requires version 2 minimum." + ) + _tf_available = False + else: + logger.info("Disabling Tensorflow because USE_TORCH is set") + + +_essentia_available = importlib.util.find_spec("essentia") is not None +try: + _essentia_version = importlib.metadata.version("essentia") + logger.debug(f"Successfully imported essentia version {_essentia_version}") +except importlib.metadata.PackageNotFoundError: + _essentia_version = False + + +_pretty_midi_available = importlib.util.find_spec("pretty_midi") is not None +try: + _pretty_midi_version = importlib.metadata.version("pretty_midi") + logger.debug(f"Successfully imported pretty_midi version {_pretty_midi_version}") +except importlib.metadata.PackageNotFoundError: + _pretty_midi_available = False + + +ccl_version = "N/A" +_is_ccl_available = ( + importlib.util.find_spec("torch_ccl") is not None + or importlib.util.find_spec("oneccl_bindings_for_pytorch") is not None +) +try: + ccl_version = importlib.metadata.version("oneccl_bind_pt") + logger.debug(f"Detected oneccl_bind_pt version {ccl_version}") +except importlib.metadata.PackageNotFoundError: + _is_ccl_available = False + + +_flax_available = False +if USE_JAX in ENV_VARS_TRUE_AND_AUTO_VALUES: + _flax_available, _flax_version = _is_package_available("flax", return_version=True) + if _flax_available: + _jax_available, _jax_version = _is_package_available("jax", return_version=True) + if _jax_available: + logger.info(f"JAX version {_jax_version}, Flax version {_flax_version} available.") + else: + _flax_available = _jax_available = False + _jax_version = _flax_version = "N/A" + + +_torch_fx_available = False +if _torch_available: + torch_version = version.parse(_torch_version) + _torch_fx_available = (torch_version.major, torch_version.minor) >= ( + TORCH_FX_REQUIRED_VERSION.major, + TORCH_FX_REQUIRED_VERSION.minor, + ) + + +_torch_xla_available = False +if USE_TORCH_XLA in ENV_VARS_TRUE_VALUES: + _torch_xla_available, _torch_xla_version = _is_package_available("torch_xla", return_version=True) + if _torch_xla_available: + logger.info(f"Torch XLA version {_torch_xla_version} available.") + + +def is_kenlm_available(): + return _kenlm_available + + +def is_cv2_available(): + return _cv2_available + + +def is_yt_dlp_available(): + return _yt_dlp_available + + +def is_torch_available(): + return _torch_available + + +def is_accelerate_available(min_version: str = ACCELERATE_MIN_VERSION): + return _accelerate_available and version.parse(_accelerate_version) >= version.parse(min_version) + + +def is_torch_deterministic(): + """ + Check whether pytorch uses deterministic algorithms by looking if torch.set_deterministic_debug_mode() is set to 1 or 2" + """ + import torch + + if torch.get_deterministic_debug_mode() == 0: + return False + else: + return True + + +def is_hadamard_available(): + return _hadamard_available + + +def is_hqq_available(min_version: str = HQQ_MIN_VERSION): + return _hqq_available and version.parse(_hqq_version) >= version.parse(min_version) + + +def is_pygments_available(): + return _pygments_available + + +def get_torch_version(): + return _torch_version + + +def is_torch_sdpa_available(): + if not is_torch_available(): + return False + elif _torch_version == "N/A": + return False + + # NOTE: We require torch>=2.1 (and not torch>=2.0) to use SDPA in Transformers for two reasons: + # - Allow the global use of the `scale` argument introduced in https://github.com/pytorch/pytorch/pull/95259 + # - Memory-efficient attention supports arbitrary attention_mask: https://github.com/pytorch/pytorch/pull/104310 + # NOTE: MLU is OK with non-contiguous inputs. + if is_torch_mlu_available(): + return version.parse(_torch_version) >= version.parse("2.1.0") + # NOTE: NPU can use SDPA in Transformers with torch>=2.1.0. + if is_torch_npu_available(): + return version.parse(_torch_version) >= version.parse("2.1.0") + # NOTE: We require torch>=2.1.1 to avoid a numerical issue in SDPA with non-contiguous inputs: https://github.com/pytorch/pytorch/issues/112577 + return version.parse(_torch_version) >= version.parse("2.1.1") + + +def is_torch_flex_attn_available(): + if not is_torch_available(): + return False + elif _torch_version == "N/A": + return False + + # TODO check if some bugs cause push backs on the exact version + # NOTE: We require torch>=2.5.0 as it is the first release + return version.parse(_torch_version) >= version.parse("2.5.0") + + +def is_torchvision_available(): + return _torchvision_available + + +def is_torchvision_v2_available(): + if not is_torchvision_available(): + return False + + # NOTE: We require torchvision>=0.15 as v2 transforms are available from this version: https://pytorch.org/vision/stable/transforms.html#v1-or-v2-which-one-should-i-use + return version.parse(_torchvision_version) >= version.parse("0.15") + + +def is_galore_torch_available(): + return _galore_torch_available + + +def is_apollo_torch_available(): + return _apollo_torch_available + + +def is_lomo_available(): + return _lomo_available + + +def is_grokadamw_available(): + return _grokadamw_available + + +def is_schedulefree_available(min_version: str = SCHEDULEFREE_MIN_VERSION): + return _schedulefree_available and version.parse(_schedulefree_version) >= version.parse(min_version) + + +def is_pyctcdecode_available(): + return _pyctcdecode_available + + +def is_librosa_available(): + return _librosa_available + + +def is_essentia_available(): + return _essentia_available + + +def is_pretty_midi_available(): + return _pretty_midi_available + + +def is_torch_cuda_available(): + if is_torch_available(): + import torch + + return torch.cuda.is_available() + else: + return False + + +def is_mamba_ssm_available(): + if is_torch_available(): + import torch + + if not torch.cuda.is_available(): + return False + else: + return _is_package_available("mamba_ssm") + return False + + +def is_mamba_2_ssm_available(): + if is_torch_available(): + import torch + + if not torch.cuda.is_available(): + return False + else: + if _is_package_available("mamba_ssm"): + import mamba_ssm + + if version.parse(mamba_ssm.__version__) >= version.parse("2.0.4"): + return True + return False + + +def is_causal_conv1d_available(): + if is_torch_available(): + import torch + + if not torch.cuda.is_available(): + return False + return _is_package_available("causal_conv1d") + return False + + +def is_mambapy_available(): + if is_torch_available(): + return _is_package_available("mambapy") + return False + + +def is_torch_mps_available(min_version: Optional[str] = None): + if is_torch_available(): + import torch + + if hasattr(torch.backends, "mps"): + backend_available = torch.backends.mps.is_available() and torch.backends.mps.is_built() + if min_version is not None: + flag = version.parse(_torch_version) >= version.parse(min_version) + backend_available = backend_available and flag + return backend_available + return False + + +def is_torch_bf16_gpu_available(): + if not is_torch_available(): + return False + + import torch + + return torch.cuda.is_available() and torch.cuda.is_bf16_supported() + + +def is_torch_bf16_cpu_available(): + if not is_torch_available(): + return False + + import torch + + try: + # multiple levels of AttributeError depending on the pytorch version so do them all in one check + _ = torch.cpu.amp.autocast + except AttributeError: + return False + + return True + + +def is_torch_bf16_available(): + # the original bf16 check was for gpu only, but later a cpu/bf16 combo has emerged so this util + # has become ambiguous and therefore deprecated + warnings.warn( + "The util is_torch_bf16_available is deprecated, please use is_torch_bf16_gpu_available " + "or is_torch_bf16_cpu_available instead according to whether it's used with cpu or gpu", + FutureWarning, + ) + return is_torch_bf16_gpu_available() + + +@lru_cache() +def is_torch_fp16_available_on_device(device): + if not is_torch_available(): + return False + + if is_torch_hpu_available(): + if is_habana_gaudi1(): + return False + else: + return True + + import torch + + try: + x = torch.zeros(2, 2, dtype=torch.float16, device=device) + _ = x @ x + + # At this moment, let's be strict of the check: check if `LayerNorm` is also supported on device, because many + # models use this layer. + batch, sentence_length, embedding_dim = 3, 4, 5 + embedding = torch.randn(batch, sentence_length, embedding_dim, dtype=torch.float16, device=device) + layer_norm = torch.nn.LayerNorm(embedding_dim, dtype=torch.float16, device=device) + _ = layer_norm(embedding) + + except: # noqa: E722 + # TODO: more precise exception matching, if possible. + # most backends should return `RuntimeError` however this is not guaranteed. + return False + + return True + + +@lru_cache() +def is_torch_bf16_available_on_device(device): + if not is_torch_available(): + return False + + import torch + + if device == "cuda": + return is_torch_bf16_gpu_available() + + if device == "hpu": + return True + + try: + x = torch.zeros(2, 2, dtype=torch.bfloat16, device=device) + _ = x @ x + except: # noqa: E722 + # TODO: more precise exception matching, if possible. + # most backends should return `RuntimeError` however this is not guaranteed. + return False + + return True + + +def is_torch_tf32_available(): + if not is_torch_available(): + return False + + import torch + + if not torch.cuda.is_available() or torch.version.cuda is None: + return False + if torch.cuda.get_device_properties(torch.cuda.current_device()).major < 8: + return False + if int(torch.version.cuda.split(".")[0]) < 11: + return False + if version.parse(version.parse(torch.__version__).base_version) < version.parse("1.7"): + return False + + return True + + +def is_torch_fx_available(): + return _torch_fx_available + + +def is_peft_available(): + return _peft_available + + +def is_bs4_available(): + return _bs4_available + + +def is_tf_available(): + return _tf_available + + +def is_coloredlogs_available(): + return _coloredlogs_available + + +def is_tf2onnx_available(): + return _tf2onnx_available + + +def is_onnx_available(): + return _onnx_available + + +def is_openai_available(): + return _openai_available + + +def is_flax_available(): + return _flax_available + + +def is_flute_available(): + try: + return importlib.util.find_spec("flute") is not None and importlib.metadata.version("flute-kernel") >= "0.4.1" + except importlib.metadata.PackageNotFoundError: + return False + + +def is_ftfy_available(): + return _ftfy_available + + +def is_g2p_en_available(): + return _g2p_en_available + + +@lru_cache +def is_torch_xla_available(check_is_tpu=False, check_is_gpu=False): + """ + Check if `torch_xla` is available. To train a native pytorch job in an environment with torch xla installed, set + the USE_TORCH_XLA to false. + """ + assert not (check_is_tpu and check_is_gpu), "The check_is_tpu and check_is_gpu cannot both be true." + + if not _torch_xla_available: + return False + + import torch_xla + + if check_is_gpu: + return torch_xla.runtime.device_type() in ["GPU", "CUDA"] + elif check_is_tpu: + return torch_xla.runtime.device_type() == "TPU" + + return True + + +@lru_cache() +def is_torch_neuroncore_available(check_device=True): + if importlib.util.find_spec("torch_neuronx") is not None: + return is_torch_xla_available() + return False + + +@lru_cache() +def is_torch_npu_available(check_device=False): + "Checks if `torch_npu` is installed and potentially if a NPU is in the environment" + if not _torch_available or importlib.util.find_spec("torch_npu") is None: + return False + + import torch + import torch_npu # noqa: F401 + + if check_device: + try: + # Will raise a RuntimeError if no NPU is found + _ = torch.npu.device_count() + return torch.npu.is_available() + except RuntimeError: + return False + return hasattr(torch, "npu") and torch.npu.is_available() + + +@lru_cache() +def is_torch_mlu_available(check_device=False): + """ + Checks if `mlu` is available via an `cndev-based` check which won't trigger the drivers and leave mlu + uninitialized. + """ + if not _torch_available or importlib.util.find_spec("torch_mlu") is None: + return False + + import torch + import torch_mlu # noqa: F401 + + pytorch_cndev_based_mlu_check_previous_value = os.environ.get("PYTORCH_CNDEV_BASED_MLU_CHECK") + try: + os.environ["PYTORCH_CNDEV_BASED_MLU_CHECK"] = str(1) + available = torch.mlu.is_available() + finally: + if pytorch_cndev_based_mlu_check_previous_value: + os.environ["PYTORCH_CNDEV_BASED_MLU_CHECK"] = pytorch_cndev_based_mlu_check_previous_value + else: + os.environ.pop("PYTORCH_CNDEV_BASED_MLU_CHECK", None) + + return available + + +@lru_cache() +def is_torch_musa_available(check_device=False): + "Checks if `torch_musa` is installed and potentially if a MUSA is in the environment" + if not _torch_available or importlib.util.find_spec("torch_musa") is None: + return False + + import torch + import torch_musa # noqa: F401 + + torch_musa_min_version = "0.33.0" + if _accelerate_available and version.parse(_accelerate_version) < version.parse(torch_musa_min_version): + return False + + if check_device: + try: + # Will raise a RuntimeError if no MUSA is found + _ = torch.musa.device_count() + return torch.musa.is_available() + except RuntimeError: + return False + return hasattr(torch, "musa") and torch.musa.is_available() + + +@lru_cache +def is_torch_hpu_available(): + "Checks if `torch.hpu` is available and potentially if a HPU is in the environment" + if ( + not _torch_available + or importlib.util.find_spec("habana_frameworks") is None + or importlib.util.find_spec("habana_frameworks.torch") is None + ): + return False + + torch_hpu_min_version = "1.5.0" + if _accelerate_available and version.parse(_accelerate_version) < version.parse(torch_hpu_min_version): + return False + + import torch + + if os.environ.get("PT_HPU_LAZY_MODE", "1") == "1": + # import habana_frameworks.torch in case of lazy mode to patch torch with torch.hpu + import habana_frameworks.torch # noqa: F401 + + if not hasattr(torch, "hpu") or not torch.hpu.is_available(): + return False + + import habana_frameworks.torch.utils.experimental as htexp # noqa: F401 + + # IlyasMoutawwakil: We patch masked_fill_ for int64 tensors to avoid a bug on Gaudi1 + # synNodeCreateWithId failed for node: masked_fill_fwd_i64 with synStatus 26 [Generic failure] + # This can be removed once Gaudi1 support is discontinued but for now we need it to keep using + # dl1.24xlarge Gaudi1 instances on AWS for testing. + # check if the device is Gaudi1 (vs Gaudi2, Gaudi3). + if htexp._get_device_type() == htexp.synDeviceType.synDeviceGaudi: + original_masked_fill_ = torch.Tensor.masked_fill_ + + def patched_masked_fill_(self, mask, value): + if self.dtype == torch.int64: + logger.warning_once( + "In-place tensor.masked_fill_(mask, value) is not supported for int64 tensors on Gaudi1. " + "This operation will be performed out-of-place using tensor[mask] = value." + ) + self[mask] = value + else: + original_masked_fill_(self, mask, value) + + torch.Tensor.masked_fill_ = patched_masked_fill_ + + return True + + +@lru_cache +def is_habana_gaudi1(): + if not is_torch_hpu_available(): + return False + + import habana_frameworks.torch.utils.experimental as htexp # noqa: F401 + + # Check if the device is Gaudi1 (vs Gaudi2, Gaudi3) + return htexp._get_device_type() == htexp.synDeviceType.synDeviceGaudi + + +def is_torchdynamo_available(): + if not is_torch_available(): + return False + + return version.parse(_torch_version) >= version.parse("2.0.0") + + +def is_torch_compile_available(): + if not is_torch_available(): + return False + + import torch + + # We don't do any version check here to support nighlies marked as 1.14. Ultimately needs to check version against + # 2.0 but let's do it later. + return hasattr(torch, "compile") + + +def is_torchdynamo_compiling(): + if not is_torch_available(): + return False + + # Importing torch._dynamo causes issues with PyTorch profiler (https://github.com/pytorch/pytorch/issues/130622) + # hence rather relying on `torch.compiler.is_compiling()` when possible (torch>=2.3) + try: + import torch + + return torch.compiler.is_compiling() + except Exception: + try: + import torch._dynamo as dynamo # noqa: F401 + + return dynamo.is_compiling() + except Exception: + return False + + +def is_torchdynamo_exporting(): + if not is_torch_available(): + return False + + try: + import torch + + return torch.compiler.is_exporting() + except Exception: + try: + import torch._dynamo as dynamo # noqa: F401 + + return dynamo.is_exporting() + except Exception: + return False + + +def is_torch_tensorrt_fx_available(): + if importlib.util.find_spec("torch_tensorrt") is None: + return False + return importlib.util.find_spec("torch_tensorrt.fx") is not None + + +def is_datasets_available(): + return _datasets_available + + +def is_detectron2_available(): + return _detectron2_available + + +def is_rjieba_available(): + return _rjieba_available + + +def is_psutil_available(): + return _psutil_available + + +def is_py3nvml_available(): + return _py3nvml_available + + +def is_sacremoses_available(): + return _sacremoses_available + + +def is_apex_available(): + return _apex_available + + +def is_aqlm_available(): + return _aqlm_available + + +def is_vptq_available(min_version: str = VPTQ_MIN_VERSION): + return _vptq_available and version.parse(_vptq_version) >= version.parse(min_version) + + +def is_av_available(): + return _av_available + + +def is_decord_available(): + return _decord_available + + +def is_ninja_available(): + r""" + Code comes from *torch.utils.cpp_extension.is_ninja_available()*. Returns `True` if the + [ninja](https://ninja-build.org/) build system is available on the system, `False` otherwise. + """ + try: + subprocess.check_output("ninja --version".split()) + except Exception: + return False + else: + return True + + +def is_ipex_available(min_version: str = ""): + def get_major_and_minor_from_version(full_version): + return str(version.parse(full_version).major) + "." + str(version.parse(full_version).minor) + + if not is_torch_available() or not _ipex_available: + return False + + torch_major_and_minor = get_major_and_minor_from_version(_torch_version) + ipex_major_and_minor = get_major_and_minor_from_version(_ipex_version) + if torch_major_and_minor != ipex_major_and_minor: + logger.warning( + f"Intel Extension for PyTorch {ipex_major_and_minor} needs to work with PyTorch {ipex_major_and_minor}.*," + f" but PyTorch {_torch_version} is found. Please switch to the matching version and run again." + ) + return False + if min_version: + return version.parse(_ipex_version) >= version.parse(min_version) + return True + + +@lru_cache +def is_torch_xpu_available(check_device=False): + """ + Checks if XPU acceleration is available either via native PyTorch (>=2.6), + `intel_extension_for_pytorch` or via stock PyTorch (>=2.4) and potentially + if a XPU is in the environment. + """ + if not is_torch_available(): + return False + + torch_version = version.parse(_torch_version) + if torch_version.major < 2 or (torch_version.major == 2 and torch_version.minor < 6): + if is_ipex_available(): + import intel_extension_for_pytorch # noqa: F401 + elif torch_version.major < 2 or (torch_version.major == 2 and torch_version.minor < 4): + return False + + import torch + + if check_device: + try: + # Will raise a RuntimeError if no XPU is found + _ = torch.xpu.device_count() + return torch.xpu.is_available() + except RuntimeError: + return False + return hasattr(torch, "xpu") and torch.xpu.is_available() + + +@lru_cache() +def is_bitsandbytes_available(): + if not is_torch_available() or not _bitsandbytes_available: + return False + + import torch + + # `bitsandbytes` versions older than 0.43.1 eagerly require CUDA at import time, + # so those versions of the library are practically only available when CUDA is too. + if version.parse(importlib.metadata.version("bitsandbytes")) < version.parse("0.43.1"): + return torch.cuda.is_available() + + # Newer versions of `bitsandbytes` can be imported on systems without CUDA. + return True + + +def is_bitsandbytes_multi_backend_available() -> bool: + if not is_bitsandbytes_available(): + return False + + import bitsandbytes as bnb + + return "multi_backend" in getattr(bnb, "features", set()) + + +def is_flash_attn_2_available(): + if not is_torch_available(): + return False + + if not _is_package_available("flash_attn"): + return False + + # Let's add an extra check to see if cuda is available + import torch + + if not (torch.cuda.is_available() or is_torch_mlu_available()): + return False + + if torch.version.cuda: + return version.parse(importlib.metadata.version("flash_attn")) >= version.parse("2.1.0") + elif torch.version.hip: + # TODO: Bump the requirement to 2.1.0 once released in https://github.com/ROCmSoftwarePlatform/flash-attention + return version.parse(importlib.metadata.version("flash_attn")) >= version.parse("2.0.4") + elif is_torch_mlu_available(): + return version.parse(importlib.metadata.version("flash_attn")) >= version.parse("2.3.3") + else: + return False + + +@lru_cache() +def is_flash_attn_greater_or_equal_2_10(): + if not _is_package_available("flash_attn"): + return False + + return version.parse(importlib.metadata.version("flash_attn")) >= version.parse("2.1.0") + + +@lru_cache() +def is_flash_attn_greater_or_equal(library_version: str): + if not _is_package_available("flash_attn"): + return False + + return version.parse(importlib.metadata.version("flash_attn")) >= version.parse(library_version) + + +@lru_cache() +def is_torch_greater_or_equal(library_version: str, accept_dev: bool = False): + """ + Accepts a library version and returns True if the current version of the library is greater than or equal to the + given version. If `accept_dev` is True, it will also accept development versions (e.g. 2.7.0.dev20250320 matches + 2.7.0). + """ + if not _is_package_available("torch"): + return False + + if accept_dev: + return version.parse(version.parse(importlib.metadata.version("torch")).base_version) >= version.parse( + library_version + ) + else: + return version.parse(importlib.metadata.version("torch")) >= version.parse(library_version) + + +def is_torchdistx_available(): + return _torchdistx_available + + +def is_faiss_available(): + return _faiss_available + + +def is_scipy_available(): + return _scipy_available + + +def is_sklearn_available(): + return _sklearn_available + + +def is_sentencepiece_available(): + return _sentencepiece_available + + +def is_seqio_available(): + return _is_seqio_available + + +def is_gguf_available(min_version: str = GGUF_MIN_VERSION): + return _is_gguf_available and version.parse(_gguf_version) >= version.parse(min_version) + + +def is_protobuf_available(): + if importlib.util.find_spec("google") is None: + return False + return importlib.util.find_spec("google.protobuf") is not None + + +def is_fsdp_available(min_version: str = FSDP_MIN_VERSION): + return is_torch_available() and version.parse(_torch_version) >= version.parse(min_version) + + +def is_optimum_available(): + return _optimum_available + + +def is_auto_awq_available(): + return _auto_awq_available + + +def is_optimum_quanto_available(): + # `importlib.metadata.version` doesn't work with `optimum.quanto`, need to put `optimum_quanto` + return _is_optimum_quanto_available + + +def is_quark_available(): + return _quark_available + + +def is_compressed_tensors_available(): + return _compressed_tensors_available + + +def is_auto_gptq_available(): + return _auto_gptq_available + + +def is_gptqmodel_available(): + return _gptqmodel_available + + +def is_eetq_available(): + return _eetq_available + + +def is_fbgemm_gpu_available(): + return _fbgemm_gpu_available + + +def is_levenshtein_available(): + return _levenshtein_available + + +def is_optimum_neuron_available(): + return _optimum_available and _is_package_available("optimum.neuron") + + +def is_safetensors_available(): + return _safetensors_available + + +def is_tokenizers_available(): + return _tokenizers_available + + +@lru_cache +def is_vision_available(): + _pil_available = importlib.util.find_spec("PIL") is not None + if _pil_available: + try: + package_version = importlib.metadata.version("Pillow") + except importlib.metadata.PackageNotFoundError: + try: + package_version = importlib.metadata.version("Pillow-SIMD") + except importlib.metadata.PackageNotFoundError: + return False + logger.debug(f"Detected PIL version {package_version}") + return _pil_available + + +def is_pytesseract_available(): + return _pytesseract_available + + +def is_pytest_available(): + return _pytest_available + + +def is_spacy_available(): + return _spacy_available + + +def is_tensorflow_text_available(): + return is_tf_available() and _tensorflow_text_available + + +def is_keras_nlp_available(): + return is_tensorflow_text_available() and _keras_nlp_available + + +def is_in_notebook(): + try: + # Test adapted from tqdm.autonotebook: https://github.com/tqdm/tqdm/blob/master/tqdm/autonotebook.py + get_ipython = sys.modules["IPython"].get_ipython + if "IPKernelApp" not in get_ipython().config: + raise ImportError("console") + # Removed the lines to include VSCode + if "DATABRICKS_RUNTIME_VERSION" in os.environ and os.environ["DATABRICKS_RUNTIME_VERSION"] < "11.0": + # Databricks Runtime 11.0 and above uses IPython kernel by default so it should be compatible with Jupyter notebook + # https://docs.microsoft.com/en-us/azure/databricks/notebooks/ipython-kernel + raise ImportError("databricks") + + return importlib.util.find_spec("IPython") is not None + except (AttributeError, ImportError, KeyError): + return False + + +def is_pytorch_quantization_available(): + return _pytorch_quantization_available + + +def is_tensorflow_probability_available(): + return _tensorflow_probability_available + + +def is_pandas_available(): + return _pandas_available + + +def is_sagemaker_dp_enabled(): + # Get the sagemaker specific env variable. + sagemaker_params = os.getenv("SM_FRAMEWORK_PARAMS", "{}") + try: + # Parse it and check the field "sagemaker_distributed_dataparallel_enabled". + sagemaker_params = json.loads(sagemaker_params) + if not sagemaker_params.get("sagemaker_distributed_dataparallel_enabled", False): + return False + except json.JSONDecodeError: + return False + # Lastly, check if the `smdistributed` module is present. + return _smdistributed_available + + +def is_sagemaker_mp_enabled(): + # Get the sagemaker specific mp parameters from smp_options variable. + smp_options = os.getenv("SM_HP_MP_PARAMETERS", "{}") + try: + # Parse it and check the field "partitions" is included, it is required for model parallel. + smp_options = json.loads(smp_options) + if "partitions" not in smp_options: + return False + except json.JSONDecodeError: + return False + + # Get the sagemaker specific framework parameters from mpi_options variable. + mpi_options = os.getenv("SM_FRAMEWORK_PARAMS", "{}") + try: + # Parse it and check the field "sagemaker_distributed_dataparallel_enabled". + mpi_options = json.loads(mpi_options) + if not mpi_options.get("sagemaker_mpi_enabled", False): + return False + except json.JSONDecodeError: + return False + # Lastly, check if the `smdistributed` module is present. + return _smdistributed_available + + +def is_training_run_on_sagemaker(): + return "SAGEMAKER_JOB_NAME" in os.environ + + +def is_soundfile_available(): + return _soundfile_available + + +def is_timm_available(): + return _timm_available + + +def is_natten_available(): + return _natten_available + + +def is_nltk_available(): + return _nltk_available + + +def is_torchaudio_available(): + return _torchaudio_available + + +def is_torchao_available(min_version: str = TORCHAO_MIN_VERSION): + return _torchao_available and version.parse(_torchao_version) >= version.parse(min_version) + + +def is_speech_available(): + # For now this depends on torchaudio but the exact dependency might evolve in the future. + return _torchaudio_available + + +def is_spqr_available(): + return _spqr_available + + +def is_phonemizer_available(): + return _phonemizer_available + + +def is_uroman_available(): + return _uroman_available + + +def torch_only_method(fn): + def wrapper(*args, **kwargs): + if not _torch_available: + raise ImportError( + "You need to install pytorch to use this method or class, " + "or activate it with environment variables USE_TORCH=1 and USE_TF=0." + ) + else: + return fn(*args, **kwargs) + + return wrapper + + +def is_ccl_available(): + return _is_ccl_available + + +def is_sudachi_available(): + return _sudachipy_available + + +def get_sudachi_version(): + return _sudachipy_version + + +def is_sudachi_projection_available(): + if not is_sudachi_available(): + return False + + # NOTE: We require sudachipy>=0.6.8 to use projection option in sudachi_kwargs for the constructor of BertJapaneseTokenizer. + # - `projection` option is not supported in sudachipy<0.6.8, see https://github.com/WorksApplications/sudachi.rs/issues/230 + return version.parse(_sudachipy_version) >= version.parse("0.6.8") + + +def is_jumanpp_available(): + return (importlib.util.find_spec("rhoknp") is not None) and (shutil.which("jumanpp") is not None) + + +def is_cython_available(): + return importlib.util.find_spec("pyximport") is not None + + +def is_jieba_available(): + return _jieba_available + + +def is_jinja_available(): + return _jinja_available + + +def is_mlx_available(): + return _mlx_available + + +def is_num2words_available(): + return _num2words_available + + +def is_tiktoken_available(): + return _tiktoken_available and _blobfile_available + + +def is_liger_kernel_available(): + if not _liger_kernel_available: + return False + + return version.parse(importlib.metadata.version("liger_kernel")) >= version.parse("0.3.0") + + +def is_triton_available(): + return _triton_available + + +def is_rich_available(): + return _rich_available + + +# docstyle-ignore +AV_IMPORT_ERROR = """ +{0} requires the PyAv library but it was not found in your environment. You can install it with: +``` +pip install av +``` +Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +YT_DLP_IMPORT_ERROR = """ +{0} requires the YT-DLP library but it was not found in your environment. You can install it with: +``` +pip install yt-dlp +``` +Please note that you may need to restart your runtime after installation. +""" + +DECORD_IMPORT_ERROR = """ +{0} requires the PyAv library but it was not found in your environment. You can install it with: +``` +pip install decord +``` +Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +CV2_IMPORT_ERROR = """ +{0} requires the OpenCV library but it was not found in your environment. You can install it with: +``` +pip install opencv-python +``` +Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +DATASETS_IMPORT_ERROR = """ +{0} requires the 🤗 Datasets library but it was not found in your environment. You can install it with: +``` +pip install datasets +``` +In a notebook or a colab, you can install it by executing a cell with +``` +!pip install datasets +``` +then restarting your kernel. + +Note that if you have a local folder named `datasets` or a local python file named `datasets.py` in your current +working directory, python may try to import this instead of the 🤗 Datasets library. You should rename this folder or +that python file if that's the case. Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +TOKENIZERS_IMPORT_ERROR = """ +{0} requires the 🤗 Tokenizers library but it was not found in your environment. You can install it with: +``` +pip install tokenizers +``` +In a notebook or a colab, you can install it by executing a cell with +``` +!pip install tokenizers +``` +Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +SENTENCEPIECE_IMPORT_ERROR = """ +{0} requires the SentencePiece library but it was not found in your environment. Checkout the instructions on the +installation page of its repo: https://github.com/google/sentencepiece#installation and follow the ones +that match your environment. Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +PROTOBUF_IMPORT_ERROR = """ +{0} requires the protobuf library but it was not found in your environment. Checkout the instructions on the +installation page of its repo: https://github.com/protocolbuffers/protobuf/tree/master/python#installation and follow the ones +that match your environment. Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +FAISS_IMPORT_ERROR = """ +{0} requires the faiss library but it was not found in your environment. Checkout the instructions on the +installation page of its repo: https://github.com/facebookresearch/faiss/blob/master/INSTALL.md and follow the ones +that match your environment. Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +PYTORCH_IMPORT_ERROR = """ +{0} requires the PyTorch library but it was not found in your environment. Checkout the instructions on the +installation page: https://pytorch.org/get-started/locally/ and follow the ones that match your environment. +Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +TORCHVISION_IMPORT_ERROR = """ +{0} requires the Torchvision library but it was not found in your environment. Checkout the instructions on the +installation page: https://pytorch.org/get-started/locally/ and follow the ones that match your environment. +Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +PYTORCH_IMPORT_ERROR_WITH_TF = """ +{0} requires the PyTorch library but it was not found in your environment. +However, we were able to find a TensorFlow installation. TensorFlow classes begin +with "TF", but are otherwise identically named to our PyTorch classes. This +means that the TF equivalent of the class you tried to import would be "TF{0}". +If you want to use TensorFlow, please use TF classes instead! + +If you really do want to use PyTorch please go to +https://pytorch.org/get-started/locally/ and follow the instructions that +match your environment. +""" + +# docstyle-ignore +TF_IMPORT_ERROR_WITH_PYTORCH = """ +{0} requires the TensorFlow library but it was not found in your environment. +However, we were able to find a PyTorch installation. PyTorch classes do not begin +with "TF", but are otherwise identically named to our TF classes. +If you want to use PyTorch, please use those classes instead! + +If you really do want to use TensorFlow, please follow the instructions on the +installation page https://www.tensorflow.org/install that match your environment. +""" + +# docstyle-ignore +BS4_IMPORT_ERROR = """ +{0} requires the Beautiful Soup library but it was not found in your environment. You can install it with pip: +`pip install beautifulsoup4`. Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +SKLEARN_IMPORT_ERROR = """ +{0} requires the scikit-learn library but it was not found in your environment. You can install it with: +``` +pip install -U scikit-learn +``` +In a notebook or a colab, you can install it by executing a cell with +``` +!pip install -U scikit-learn +``` +Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +TENSORFLOW_IMPORT_ERROR = """ +{0} requires the TensorFlow library but it was not found in your environment. Checkout the instructions on the +installation page: https://www.tensorflow.org/install and follow the ones that match your environment. +Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +DETECTRON2_IMPORT_ERROR = """ +{0} requires the detectron2 library but it was not found in your environment. Checkout the instructions on the +installation page: https://github.com/facebookresearch/detectron2/blob/master/INSTALL.md and follow the ones +that match your environment. Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +FLAX_IMPORT_ERROR = """ +{0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the +installation page: https://github.com/google/flax and follow the ones that match your environment. +Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +FTFY_IMPORT_ERROR = """ +{0} requires the ftfy library but it was not found in your environment. Checkout the instructions on the +installation section: https://github.com/rspeer/python-ftfy/tree/master#installing and follow the ones +that match your environment. Please note that you may need to restart your runtime after installation. +""" + +LEVENSHTEIN_IMPORT_ERROR = """ +{0} requires the python-Levenshtein library but it was not found in your environment. You can install it with pip: `pip +install python-Levenshtein`. Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +G2P_EN_IMPORT_ERROR = """ +{0} requires the g2p-en library but it was not found in your environment. You can install it with pip: +`pip install g2p-en`. Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +PYTORCH_QUANTIZATION_IMPORT_ERROR = """ +{0} requires the pytorch-quantization library but it was not found in your environment. You can install it with pip: +`pip install pytorch-quantization --extra-index-url https://pypi.ngc.nvidia.com` +Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +TENSORFLOW_PROBABILITY_IMPORT_ERROR = """ +{0} requires the tensorflow_probability library but it was not found in your environment. You can install it with pip as +explained here: https://github.com/tensorflow/probability. Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +TENSORFLOW_TEXT_IMPORT_ERROR = """ +{0} requires the tensorflow_text library but it was not found in your environment. You can install it with pip as +explained here: https://www.tensorflow.org/text/guide/tf_text_intro. +Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +TORCHAUDIO_IMPORT_ERROR = """ +{0} requires the torchaudio library but it was not found in your environment. Please install it and restart your +runtime. +""" + +# docstyle-ignore +PANDAS_IMPORT_ERROR = """ +{0} requires the pandas library but it was not found in your environment. You can install it with pip as +explained here: https://pandas.pydata.org/pandas-docs/stable/getting_started/install.html. +Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +PHONEMIZER_IMPORT_ERROR = """ +{0} requires the phonemizer library but it was not found in your environment. You can install it with pip: +`pip install phonemizer`. Please note that you may need to restart your runtime after installation. +""" +# docstyle-ignore +UROMAN_IMPORT_ERROR = """ +{0} requires the uroman library but it was not found in your environment. You can install it with pip: +`pip install uroman`. Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +SACREMOSES_IMPORT_ERROR = """ +{0} requires the sacremoses library but it was not found in your environment. You can install it with pip: +`pip install sacremoses`. Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +SCIPY_IMPORT_ERROR = """ +{0} requires the scipy library but it was not found in your environment. You can install it with pip: +`pip install scipy`. Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +SPEECH_IMPORT_ERROR = """ +{0} requires the torchaudio library but it was not found in your environment. You can install it with pip: +`pip install torchaudio`. Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +TIMM_IMPORT_ERROR = """ +{0} requires the timm library but it was not found in your environment. You can install it with pip: +`pip install timm`. Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +NATTEN_IMPORT_ERROR = """ +{0} requires the natten library but it was not found in your environment. You can install it by referring to: +shi-labs.com/natten . You can also install it with pip (may take longer to build): +`pip install natten`. Please note that you may need to restart your runtime after installation. +""" + +NUMEXPR_IMPORT_ERROR = """ +{0} requires the numexpr library but it was not found in your environment. You can install it by referring to: +https://numexpr.readthedocs.io/en/latest/index.html. +""" + + +# docstyle-ignore +NLTK_IMPORT_ERROR = """ +{0} requires the NLTK library but it was not found in your environment. You can install it by referring to: +https://www.nltk.org/install.html. Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +VISION_IMPORT_ERROR = """ +{0} requires the PIL library but it was not found in your environment. You can install it with pip: +`pip install pillow`. Please note that you may need to restart your runtime after installation. +""" + + +# docstyle-ignore +PYTESSERACT_IMPORT_ERROR = """ +{0} requires the PyTesseract library but it was not found in your environment. You can install it with pip: +`pip install pytesseract`. Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +PYCTCDECODE_IMPORT_ERROR = """ +{0} requires the pyctcdecode library but it was not found in your environment. You can install it with pip: +`pip install pyctcdecode`. Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +ACCELERATE_IMPORT_ERROR = """ +{0} requires the accelerate library >= {ACCELERATE_MIN_VERSION} it was not found in your environment. +You can install or update it with pip: `pip install --upgrade accelerate`. Please note that you may need to restart your +runtime after installation. +""" + +# docstyle-ignore +CCL_IMPORT_ERROR = """ +{0} requires the torch ccl library but it was not found in your environment. You can install it with pip: +`pip install oneccl_bind_pt -f https://developer.intel.com/ipex-whl-stable` +Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +ESSENTIA_IMPORT_ERROR = """ +{0} requires essentia library. But that was not found in your environment. You can install them with pip: +`pip install essentia==2.1b6.dev1034` +Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +LIBROSA_IMPORT_ERROR = """ +{0} requires thes librosa library. But that was not found in your environment. You can install them with pip: +`pip install librosa` +Please note that you may need to restart your runtime after installation. +""" + +# docstyle-ignore +PRETTY_MIDI_IMPORT_ERROR = """ +{0} requires thes pretty_midi library. But that was not found in your environment. You can install them with pip: +`pip install pretty_midi` +Please note that you may need to restart your runtime after installation. +""" + + +CYTHON_IMPORT_ERROR = """ +{0} requires the Cython library but it was not found in your environment. You can install it with pip: `pip install +Cython`. Please note that you may need to restart your runtime after installation. +""" + +JIEBA_IMPORT_ERROR = """ +{0} requires the jieba library but it was not found in your environment. You can install it with pip: `pip install +jieba`. Please note that you may need to restart your runtime after installation. +""" + +PEFT_IMPORT_ERROR = """ +{0} requires the peft library but it was not found in your environment. You can install it with pip: `pip install +peft`. Please note that you may need to restart your runtime after installation. +""" + +JINJA_IMPORT_ERROR = """ +{0} requires the jinja library but it was not found in your environment. You can install it with pip: `pip install +jinja2`. Please note that you may need to restart your runtime after installation. +""" + +RICH_IMPORT_ERROR = """ +{0} requires the rich library but it was not found in your environment. You can install it with pip: `pip install +rich`. Please note that you may need to restart your runtime after installation. +""" + +BACKENDS_MAPPING = OrderedDict( + [ + ("av", (is_av_available, AV_IMPORT_ERROR)), + ("bs4", (is_bs4_available, BS4_IMPORT_ERROR)), + ("cv2", (is_cv2_available, CV2_IMPORT_ERROR)), + ("datasets", (is_datasets_available, DATASETS_IMPORT_ERROR)), + ("decord", (is_decord_available, DECORD_IMPORT_ERROR)), + ("detectron2", (is_detectron2_available, DETECTRON2_IMPORT_ERROR)), + ("essentia", (is_essentia_available, ESSENTIA_IMPORT_ERROR)), + ("faiss", (is_faiss_available, FAISS_IMPORT_ERROR)), + ("flax", (is_flax_available, FLAX_IMPORT_ERROR)), + ("ftfy", (is_ftfy_available, FTFY_IMPORT_ERROR)), + ("g2p_en", (is_g2p_en_available, G2P_EN_IMPORT_ERROR)), + ("pandas", (is_pandas_available, PANDAS_IMPORT_ERROR)), + ("phonemizer", (is_phonemizer_available, PHONEMIZER_IMPORT_ERROR)), + ("uroman", (is_uroman_available, UROMAN_IMPORT_ERROR)), + ("pretty_midi", (is_pretty_midi_available, PRETTY_MIDI_IMPORT_ERROR)), + ("levenshtein", (is_levenshtein_available, LEVENSHTEIN_IMPORT_ERROR)), + ("librosa", (is_librosa_available, LIBROSA_IMPORT_ERROR)), + ("protobuf", (is_protobuf_available, PROTOBUF_IMPORT_ERROR)), + ("pyctcdecode", (is_pyctcdecode_available, PYCTCDECODE_IMPORT_ERROR)), + ("pytesseract", (is_pytesseract_available, PYTESSERACT_IMPORT_ERROR)), + ("sacremoses", (is_sacremoses_available, SACREMOSES_IMPORT_ERROR)), + ("pytorch_quantization", (is_pytorch_quantization_available, PYTORCH_QUANTIZATION_IMPORT_ERROR)), + ("sentencepiece", (is_sentencepiece_available, SENTENCEPIECE_IMPORT_ERROR)), + ("sklearn", (is_sklearn_available, SKLEARN_IMPORT_ERROR)), + ("speech", (is_speech_available, SPEECH_IMPORT_ERROR)), + ("tensorflow_probability", (is_tensorflow_probability_available, TENSORFLOW_PROBABILITY_IMPORT_ERROR)), + ("tf", (is_tf_available, TENSORFLOW_IMPORT_ERROR)), + ("tensorflow_text", (is_tensorflow_text_available, TENSORFLOW_TEXT_IMPORT_ERROR)), + ("timm", (is_timm_available, TIMM_IMPORT_ERROR)), + ("torchaudio", (is_torchaudio_available, TORCHAUDIO_IMPORT_ERROR)), + ("natten", (is_natten_available, NATTEN_IMPORT_ERROR)), + ("nltk", (is_nltk_available, NLTK_IMPORT_ERROR)), + ("tokenizers", (is_tokenizers_available, TOKENIZERS_IMPORT_ERROR)), + ("torch", (is_torch_available, PYTORCH_IMPORT_ERROR)), + ("torchvision", (is_torchvision_available, TORCHVISION_IMPORT_ERROR)), + ("vision", (is_vision_available, VISION_IMPORT_ERROR)), + ("scipy", (is_scipy_available, SCIPY_IMPORT_ERROR)), + ("accelerate", (is_accelerate_available, ACCELERATE_IMPORT_ERROR)), + ("oneccl_bind_pt", (is_ccl_available, CCL_IMPORT_ERROR)), + ("cython", (is_cython_available, CYTHON_IMPORT_ERROR)), + ("jieba", (is_jieba_available, JIEBA_IMPORT_ERROR)), + ("peft", (is_peft_available, PEFT_IMPORT_ERROR)), + ("jinja", (is_jinja_available, JINJA_IMPORT_ERROR)), + ("yt_dlp", (is_yt_dlp_available, YT_DLP_IMPORT_ERROR)), + ("rich", (is_rich_available, RICH_IMPORT_ERROR)), + ] +) + + +def requires_backends(obj, backends): + if not isinstance(backends, (list, tuple)): + backends = [backends] + + name = obj.__name__ if hasattr(obj, "__name__") else obj.__class__.__name__ + + # Raise an error for users who might not realize that classes without "TF" are torch-only + if "torch" in backends and "tf" not in backends and not is_torch_available() and is_tf_available(): + raise ImportError(PYTORCH_IMPORT_ERROR_WITH_TF.format(name)) + + # Raise the inverse error for PyTorch users trying to load TF classes + if "tf" in backends and "torch" not in backends and is_torch_available() and not is_tf_available(): + raise ImportError(TF_IMPORT_ERROR_WITH_PYTORCH.format(name)) + + checks = (BACKENDS_MAPPING[backend] for backend in backends) + failed = [msg.format(name) for available, msg in checks if not available()] + if failed: + raise ImportError("".join(failed)) + + +class DummyObject(type): + """ + Metaclass for the dummy objects. Any class inheriting from it will return the ImportError generated by + `requires_backend` each time a user tries to access any method of that class. + """ + + def __getattribute__(cls, key): + if key.startswith("_") and key != "_from_config": + return super().__getattribute__(key) + requires_backends(cls, cls._backends) + + +def is_torch_fx_proxy(x): + if is_torch_fx_available(): + import torch.fx + + return isinstance(x, torch.fx.Proxy) + return False + + +BACKENDS_T = FrozenSet[str] +IMPORT_STRUCTURE_T = Dict[BACKENDS_T, Dict[str, Set[str]]] + + +class _LazyModule(ModuleType): + """ + Module class that surfaces all objects but only performs associated imports when the objects are requested. + """ + + # Very heavily inspired by optuna.integration._IntegrationModule + # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py + def __init__( + self, + name: str, + module_file: str, + import_structure: IMPORT_STRUCTURE_T, + module_spec: Optional[importlib.machinery.ModuleSpec] = None, + extra_objects: Dict[str, object] = None, + ): + super().__init__(name) + + self._object_missing_backend = {} + if any(isinstance(key, frozenset) for key in import_structure.keys()): + self._modules = set() + self._class_to_module = {} + self.__all__ = [] + + _import_structure = {} + + for backends, module in import_structure.items(): + missing_backends = [] + for backend in backends: + if backend not in BACKENDS_MAPPING: + raise ValueError( + f"Error: the following backend: '{backend}' was specified around object {module} but isn't specified in the backends mapping." + ) + callable, error = BACKENDS_MAPPING[backend] + if not callable(): + missing_backends.append(backend) + self._modules = self._modules.union(set(module.keys())) + + for key, values in module.items(): + if len(missing_backends): + self._object_missing_backend[key] = missing_backends + + for value in values: + self._class_to_module[value] = key + if len(missing_backends): + self._object_missing_backend[value] = missing_backends + _import_structure.setdefault(key, []).extend(values) + + # Needed for autocompletion in an IDE + self.__all__.extend(list(module.keys()) + list(chain(*module.values()))) + + self.__file__ = module_file + self.__spec__ = module_spec + self.__path__ = [os.path.dirname(module_file)] + self._objects = {} if extra_objects is None else extra_objects + self._name = name + self._import_structure = _import_structure + + # This can be removed once every exportable object has a `export()` export. + else: + self._modules = set(import_structure.keys()) + self._class_to_module = {} + for key, values in import_structure.items(): + for value in values: + self._class_to_module[value] = key + # Needed for autocompletion in an IDE + self.__all__ = list(import_structure.keys()) + list(chain(*import_structure.values())) + self.__file__ = module_file + self.__spec__ = module_spec + self.__path__ = [os.path.dirname(module_file)] + self._objects = {} if extra_objects is None else extra_objects + self._name = name + self._import_structure = import_structure + + # Needed for autocompletion in an IDE + def __dir__(self): + result = super().__dir__() + # The elements of self.__all__ that are submodules may or may not be in the dir already, depending on whether + # they have been accessed or not. So we only add the elements of self.__all__ that are not already in the dir. + for attr in self.__all__: + if attr not in result: + result.append(attr) + return result + + def __getattr__(self, name: str) -> Any: + if name in self._objects: + return self._objects[name] + if name in self._object_missing_backend.keys(): + missing_backends = self._object_missing_backend[name] + + class Placeholder(metaclass=DummyObject): + _backends = missing_backends + + def __init__(self, *args, **kwargs): + requires_backends(self, missing_backends) + + Placeholder.__name__ = name + Placeholder.__module__ = self.__spec__ + + value = Placeholder + elif name in self._class_to_module.keys(): + module = self._get_module(self._class_to_module[name]) + value = getattr(module, name) + elif name in self._modules: + value = self._get_module(name) + else: + raise AttributeError(f"module {self.__name__} has no attribute {name}") + + setattr(self, name, value) + return value + + def _get_module(self, module_name: str): + try: + return importlib.import_module("." + module_name, self.__name__) + except Exception as e: + raise RuntimeError( + f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its" + f" traceback):\n{e}" + ) from e + + def __reduce__(self): + return (self.__class__, (self._name, self.__file__, self._import_structure)) + + +class OptionalDependencyNotAvailable(BaseException): + """Internally used error class for signalling an optional dependency was not found.""" + + +def direct_transformers_import(path: str, file="__init__.py") -> ModuleType: + """Imports transformers directly + + Args: + path (`str`): The path to the source file + file (`str`, *optional*): The file to join with the path. Defaults to "__init__.py". + + Returns: + `ModuleType`: The resulting imported module + """ + name = "transformers" + location = os.path.join(path, file) + spec = importlib.util.spec_from_file_location(name, location, submodule_search_locations=[path]) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + module = sys.modules[name] + return module + + +def export(*, backends=()): + """ + This decorator enables two things: + - Attaching a `__backends` tuple to an object to see what are the necessary backends for it + to execute correctly without instantiating it + - The '@export' string is used to dynamically import objects + """ + for backend in backends: + if backend not in BACKENDS_MAPPING: + raise ValueError(f"Backend should be defined in the BACKENDS_MAPPING. Offending backend: {backend}") + + if not isinstance(backends, tuple): + raise ValueError("Backends should be a tuple.") + + def inner_fn(fun): + fun.__backends = backends + return fun + + return inner_fn + + +BASE_FILE_REQUIREMENTS = { + lambda e: "modeling_tf_" in e: ("tf",), + lambda e: "modeling_flax_" in e: ("flax",), + lambda e: "modeling_" in e: ("torch",), + lambda e: e.startswith("tokenization_") and e.endswith("_fast"): ("tokenizers",), +} + + +def fetch__all__(file_content): + """ + Returns the content of the __all__ variable in the file content. + Returns None if not defined, otherwise returns a list of strings. + """ + + if "__all__" not in file_content: + return [] + + start_index = None + lines = file_content.splitlines() + for index, line in enumerate(lines): + if line.startswith("__all__"): + start_index = index + + # There is no line starting with `__all__` + if start_index is None: + return [] + + lines = lines[start_index:] + + if not lines[0].startswith("__all__"): + raise ValueError( + "fetch__all__ accepts a list of lines, with the first line being the __all__ variable declaration" + ) + + # __all__ is defined on a single line + if lines[0].endswith("]"): + return [obj.strip("\"' ") for obj in lines[0].split("=")[1].strip(" []").split(",")] + + # __all__ is defined on multiple lines + else: + _all = [] + for __all__line_index in range(1, len(lines)): + if lines[__all__line_index].strip() == "]": + return _all + else: + _all.append(lines[__all__line_index].strip("\"', ")) + + return _all + + +@lru_cache() +def create_import_structure_from_path(module_path): + """ + This method takes the path to a file/a folder and returns the import structure. + If a file is given, it will return the import structure of the parent folder. + + Import structures are designed to be digestible by `_LazyModule` objects. They are + created from the __all__ definitions in each files as well as the `@export` decorators + above methods and objects. + + The import structure allows explicit display of the required backends for a given object. + These backends are specified in two ways: + + 1. Through their `@export`, if they are exported with that decorator. This `@export` decorator + accepts a `backend` tuple kwarg mentioning which backends are required to run this object. + + 2. If an object is defined in a file with "default" backends, it will have, at a minimum, this + backend specified. The default backends are defined according to the filename: + + - If a file is named like `modeling_*.py`, it will have a `torch` backend + - If a file is named like `modeling_tf_*.py`, it will have a `tf` backend + - If a file is named like `modeling_flax_*.py`, it will have a `flax` backend + - If a file is named like `tokenization_*_fast.py`, it will have a `tokenizers` backend + + Backends serve the purpose of displaying a clear error message to the user in case the backends are not installed. + Should an object be imported without its required backends being in the environment, any attempt to use the + object will raise an error mentioning which backend(s) should be added to the environment in order to use + that object. + + Here's an example of an input import structure at the src.transformers.models level: + + { + 'albert': { + frozenset(): { + 'configuration_albert': {'AlbertConfig', 'AlbertOnnxConfig'} + }, + frozenset({'tokenizers'}): { + 'tokenization_albert_fast': {'AlbertTokenizerFast'} + }, + }, + 'align': { + frozenset(): { + 'configuration_align': {'AlignConfig', 'AlignTextConfig', 'AlignVisionConfig'}, + 'processing_align': {'AlignProcessor'} + }, + }, + 'altclip': { + frozenset(): { + 'configuration_altclip': {'AltCLIPConfig', 'AltCLIPTextConfig', 'AltCLIPVisionConfig'}, + 'processing_altclip': {'AltCLIPProcessor'}, + } + } + } + """ + import_structure = {} + if os.path.isdir(module_path): + directory = module_path + adjacent_modules = [] + + for f in os.listdir(module_path): + if f != "__pycache__" and os.path.isdir(os.path.join(module_path, f)): + import_structure[f] = create_import_structure_from_path(os.path.join(module_path, f)) + + elif not os.path.isdir(os.path.join(directory, f)): + adjacent_modules.append(f) + + else: + directory = os.path.dirname(module_path) + adjacent_modules = [f for f in os.listdir(directory) if not os.path.isdir(os.path.join(directory, f))] + + # We're only taking a look at files different from __init__.py + # We could theoretically export things directly from the __init__.py + # files, but this is not supported at this time. + if "__init__.py" in adjacent_modules: + adjacent_modules.remove("__init__.py") + + # Modular files should not be imported + def find_substring(substring, list_): + return any(substring in x for x in list_) + + if find_substring("modular_", adjacent_modules) and find_substring("modeling_", adjacent_modules): + adjacent_modules = [module for module in adjacent_modules if "modular_" not in module] + + module_requirements = {} + for module_name in adjacent_modules: + # Only modules ending in `.py` are accepted here. + if not module_name.endswith(".py"): + continue + + with open(os.path.join(directory, module_name), encoding="utf-8") as f: + file_content = f.read() + + # Remove the .py suffix + module_name = module_name[:-3] + + previous_line = "" + previous_index = 0 + + # Some files have some requirements by default. + # For example, any file named `modeling_tf_xxx.py` + # should have TensorFlow as a required backend. + base_requirements = () + for string_check, requirements in BASE_FILE_REQUIREMENTS.items(): + if string_check(module_name): + base_requirements = requirements + break + + # Objects that have a `@export` assigned to them will get exported + # with the backends specified in the decorator as well as the file backends. + exported_objects = set() + if "@export" in file_content: + lines = file_content.split("\n") + for index, line in enumerate(lines): + # This allows exporting items with other decorators. We'll take a look + # at the line that follows at the same indentation level. + if line.startswith((" ", "\t", "@", ")")) and not line.startswith("@export"): + continue + + # Skipping line enables putting whatever we want between the + # export() call and the actual class/method definition. + # This is what enables having # Copied from statements, docs, etc. + skip_line = False + + if "@export" in previous_line: + skip_line = False + + # Backends are defined on the same line as export + if "backends" in previous_line: + backends_string = previous_line.split("backends=")[1].split("(")[1].split(")")[0] + backends = tuple(sorted([b.strip("'\",") for b in backends_string.split(", ") if b])) + + # Backends are defined in the lines following export, for example such as: + # @export( + # backends=( + # "sentencepiece", + # "torch", + # "tf", + # ) + # ) + # + # or + # + # @export( + # backends=( + # "sentencepiece", "tf" + # ) + # ) + elif "backends" in lines[previous_index + 1]: + backends = [] + for backend_line in lines[previous_index:index]: + if "backends" in backend_line: + backend_line = backend_line.split("=")[1] + if '"' in backend_line or "'" in backend_line: + if ", " in backend_line: + backends.extend(backend.strip("()\"', ") for backend in backend_line.split(", ")) + else: + backends.append(backend_line.strip("()\"', ")) + + # If the line is only a ')', then we reached the end of the backends and we break. + if backend_line.strip() == ")": + break + backends = tuple(backends) + + # No backends are registered for export + else: + backends = () + + backends = frozenset(backends + base_requirements) + if backends not in module_requirements: + module_requirements[backends] = {} + if module_name not in module_requirements[backends]: + module_requirements[backends][module_name] = set() + + if not line.startswith("class") and not line.startswith("def"): + skip_line = True + else: + start_index = 6 if line.startswith("class") else 4 + object_name = line[start_index:].split("(")[0].strip(":") + module_requirements[backends][module_name].add(object_name) + exported_objects.add(object_name) + + if not skip_line: + previous_line = line + previous_index = index + + # All objects that are in __all__ should be exported by default. + # These objects are exported with the file backends. + if "__all__" in file_content: + for _all_object in fetch__all__(file_content): + if _all_object not in exported_objects: + backends = frozenset(base_requirements) + if backends not in module_requirements: + module_requirements[backends] = {} + if module_name not in module_requirements[backends]: + module_requirements[backends][module_name] = set() + + module_requirements[backends][module_name].add(_all_object) + + import_structure = {**module_requirements, **import_structure} + return import_structure + + +def spread_import_structure(nested_import_structure): + """ + This method takes as input an unordered import structure and brings the required backends at the top-level, + aggregating modules and objects under their required backends. + + Here's an example of an input import structure at the src.transformers.models level: + + { + 'albert': { + frozenset(): { + 'configuration_albert': {'AlbertConfig', 'AlbertOnnxConfig'} + }, + frozenset({'tokenizers'}): { + 'tokenization_albert_fast': {'AlbertTokenizerFast'} + }, + }, + 'align': { + frozenset(): { + 'configuration_align': {'AlignConfig', 'AlignTextConfig', 'AlignVisionConfig'}, + 'processing_align': {'AlignProcessor'} + }, + }, + 'altclip': { + frozenset(): { + 'configuration_altclip': {'AltCLIPConfig', 'AltCLIPTextConfig', 'AltCLIPVisionConfig'}, + 'processing_altclip': {'AltCLIPProcessor'}, + } + } + } + + Here's an example of an output import structure at the src.transformers.models level: + + { + frozenset({'tokenizers'}): { + 'albert.tokenization_albert_fast': {'AlbertTokenizerFast'} + }, + frozenset(): { + 'albert.configuration_albert': {'AlbertConfig', 'AlbertOnnxConfig'}, + 'align.processing_align': {'AlignProcessor'}, + 'align.configuration_align': {'AlignConfig', 'AlignTextConfig', 'AlignVisionConfig'}, + 'altclip.configuration_altclip': {'AltCLIPConfig', 'AltCLIPTextConfig', 'AltCLIPVisionConfig'}, + 'altclip.processing_altclip': {'AltCLIPProcessor'} + } + } + + """ + + def propagate_frozenset(unordered_import_structure): + tuple_first_import_structure = {} + for _key, _value in unordered_import_structure.items(): + if not isinstance(_value, dict): + tuple_first_import_structure[_key] = _value + + elif any(isinstance(v, frozenset) for v in _value.keys()): + # Here we want to switch around key and v + for k, v in _value.items(): + if isinstance(k, frozenset): + if k not in tuple_first_import_structure: + tuple_first_import_structure[k] = {} + tuple_first_import_structure[k][_key] = v + + else: + tuple_first_import_structure[_key] = propagate_frozenset(_value) + + return tuple_first_import_structure + + def flatten_dict(_dict, previous_key=None): + items = [] + for _key, _value in _dict.items(): + _key = f"{previous_key}.{_key}" if previous_key is not None else _key + if isinstance(_value, dict): + items.extend(flatten_dict(_value, _key).items()) + else: + items.append((_key, _value)) + return dict(items) + + # The tuples contain the necessary backends. We want these first, so we propagate them up the + # import structure. + ordered_import_structure = nested_import_structure + + # 6 is a number that gives us sufficient depth to go through all files and foreseeable folder depths + # while not taking too long to parse. + for i in range(6): + ordered_import_structure = propagate_frozenset(ordered_import_structure) + + # We then flatten the dict so that it references a module path. + flattened_import_structure = {} + for key, value in ordered_import_structure.copy().items(): + if isinstance(key, str): + del ordered_import_structure[key] + else: + flattened_import_structure[key] = flatten_dict(value) + + return flattened_import_structure + + +def define_import_structure(module_path: str) -> IMPORT_STRUCTURE_T: + """ + This method takes a module_path as input and creates an import structure digestible by a _LazyModule. + + Here's an example of an output import structure at the src.transformers.models level: + + { + frozenset({'tokenizers'}): { + 'albert.tokenization_albert_fast': {'AlbertTokenizerFast'} + }, + frozenset(): { + 'albert.configuration_albert': {'AlbertConfig', 'AlbertOnnxConfig'}, + 'align.processing_align': {'AlignProcessor'}, + 'align.configuration_align': {'AlignConfig', 'AlignTextConfig', 'AlignVisionConfig'}, + 'altclip.configuration_altclip': {'AltCLIPConfig', 'AltCLIPTextConfig', 'AltCLIPVisionConfig'}, + 'altclip.processing_altclip': {'AltCLIPProcessor'} + } + } + + The import structure is a dict defined with frozensets as keys, and dicts of strings to sets of objects. + """ + import_structure = create_import_structure_from_path(module_path) + return spread_import_structure(import_structure) + + +def clear_import_cache(): + """ + Clear cached Transformers modules to allow reloading modified code. + + This is useful when actively developing/modifying Transformers code. + """ + # Get all transformers modules + transformers_modules = [mod_name for mod_name in sys.modules if mod_name.startswith("transformers.")] + + # Remove them from sys.modules + for mod_name in transformers_modules: + module = sys.modules[mod_name] + # Clear _LazyModule caches if applicable + if isinstance(module, _LazyModule): + module._objects = {} # Clear cached objects + del sys.modules[mod_name] + + # Force reload main transformers module + if "transformers" in sys.modules: + main_module = sys.modules["transformers"] + if isinstance(main_module, _LazyModule): + main_module._objects = {} # Clear cached objects + importlib.reload(main_module) diff --git a/lib/python3.12/site-packages/transformers/utils/logging.py b/lib/python3.12/site-packages/transformers/utils/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..a2915e167a0bcf9d61aeaca269608b98320c3c54 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/logging.py @@ -0,0 +1,410 @@ +# Copyright 2020 Optuna, Hugging Face +# +# 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. +"""Logging utilities.""" + +import functools +import logging +import os +import sys +import threading +from logging import ( + CRITICAL, # NOQA + DEBUG, # NOQA + ERROR, # NOQA + FATAL, # NOQA + INFO, # NOQA + NOTSET, # NOQA + WARN, # NOQA + WARNING, # NOQA +) +from logging import captureWarnings as _captureWarnings +from typing import Optional + +import huggingface_hub.utils as hf_hub_utils +from tqdm import auto as tqdm_lib + + +_lock = threading.Lock() +_default_handler: Optional[logging.Handler] = None + +log_levels = { + "detail": logging.DEBUG, # will also print filename and line number + "debug": logging.DEBUG, + "info": logging.INFO, + "warning": logging.WARNING, + "error": logging.ERROR, + "critical": logging.CRITICAL, +} + +_default_log_level = logging.WARNING + +_tqdm_active = not hf_hub_utils.are_progress_bars_disabled() + + +def _get_default_logging_level(): + """ + If TRANSFORMERS_VERBOSITY env var is set to one of the valid choices return that as the new default level. If it is + not - fall back to `_default_log_level` + """ + env_level_str = os.getenv("TRANSFORMERS_VERBOSITY", None) + if env_level_str: + if env_level_str in log_levels: + return log_levels[env_level_str] + else: + logging.getLogger().warning( + f"Unknown option TRANSFORMERS_VERBOSITY={env_level_str}, " + f"has to be one of: {', '.join(log_levels.keys())}" + ) + return _default_log_level + + +def _get_library_name() -> str: + return __name__.split(".")[0] + + +def _get_library_root_logger() -> logging.Logger: + return logging.getLogger(_get_library_name()) + + +def _configure_library_root_logger() -> None: + global _default_handler + + with _lock: + if _default_handler: + # This library has already configured the library root logger. + return + _default_handler = logging.StreamHandler() # Set sys.stderr as stream. + # set defaults based on https://github.com/pyinstaller/pyinstaller/issues/7334#issuecomment-1357447176 + if sys.stderr is None: + sys.stderr = open(os.devnull, "w") + + _default_handler.flush = sys.stderr.flush + + # Apply our default configuration to the library root logger. + library_root_logger = _get_library_root_logger() + library_root_logger.addHandler(_default_handler) + library_root_logger.setLevel(_get_default_logging_level()) + # if logging level is debug, we add pathname and lineno to formatter for easy debugging + if os.getenv("TRANSFORMERS_VERBOSITY", None) == "detail": + formatter = logging.Formatter("[%(levelname)s|%(pathname)s:%(lineno)s] %(asctime)s >> %(message)s") + _default_handler.setFormatter(formatter) + + is_ci = os.getenv("CI") is not None and os.getenv("CI").upper() in {"1", "ON", "YES", "TRUE"} + library_root_logger.propagate = True if is_ci else False + + +def _reset_library_root_logger() -> None: + global _default_handler + + with _lock: + if not _default_handler: + return + + library_root_logger = _get_library_root_logger() + library_root_logger.removeHandler(_default_handler) + library_root_logger.setLevel(logging.NOTSET) + _default_handler = None + + +def get_log_levels_dict(): + return log_levels + + +def captureWarnings(capture): + """ + Calls the `captureWarnings` method from the logging library to enable management of the warnings emitted by the + `warnings` library. + + Read more about this method here: + https://docs.python.org/3/library/logging.html#integration-with-the-warnings-module + + All warnings will be logged through the `py.warnings` logger. + + Careful: this method also adds a handler to this logger if it does not already have one, and updates the logging + level of that logger to the library's root logger. + """ + logger = get_logger("py.warnings") + + if not logger.handlers: + logger.addHandler(_default_handler) + + logger.setLevel(_get_library_root_logger().level) + + _captureWarnings(capture) + + +def get_logger(name: Optional[str] = None) -> logging.Logger: + """ + Return a logger with the specified name. + + This function is not supposed to be directly accessed unless you are writing a custom transformers module. + """ + + if name is None: + name = _get_library_name() + + _configure_library_root_logger() + return logging.getLogger(name) + + +def get_verbosity() -> int: + """ + Return the current level for the 🤗 Transformers's root logger as an int. + + Returns: + `int`: The logging level. + + + + 🤗 Transformers has following logging levels: + + - 50: `transformers.logging.CRITICAL` or `transformers.logging.FATAL` + - 40: `transformers.logging.ERROR` + - 30: `transformers.logging.WARNING` or `transformers.logging.WARN` + - 20: `transformers.logging.INFO` + - 10: `transformers.logging.DEBUG` + + """ + + _configure_library_root_logger() + return _get_library_root_logger().getEffectiveLevel() + + +def set_verbosity(verbosity: int) -> None: + """ + Set the verbosity level for the 🤗 Transformers's root logger. + + Args: + verbosity (`int`): + Logging level, e.g., one of: + + - `transformers.logging.CRITICAL` or `transformers.logging.FATAL` + - `transformers.logging.ERROR` + - `transformers.logging.WARNING` or `transformers.logging.WARN` + - `transformers.logging.INFO` + - `transformers.logging.DEBUG` + """ + + _configure_library_root_logger() + _get_library_root_logger().setLevel(verbosity) + + +def set_verbosity_info(): + """Set the verbosity to the `INFO` level.""" + return set_verbosity(INFO) + + +def set_verbosity_warning(): + """Set the verbosity to the `WARNING` level.""" + return set_verbosity(WARNING) + + +def set_verbosity_debug(): + """Set the verbosity to the `DEBUG` level.""" + return set_verbosity(DEBUG) + + +def set_verbosity_error(): + """Set the verbosity to the `ERROR` level.""" + return set_verbosity(ERROR) + + +def disable_default_handler() -> None: + """Disable the default handler of the HuggingFace Transformers's root logger.""" + + _configure_library_root_logger() + + assert _default_handler is not None + _get_library_root_logger().removeHandler(_default_handler) + + +def enable_default_handler() -> None: + """Enable the default handler of the HuggingFace Transformers's root logger.""" + + _configure_library_root_logger() + + assert _default_handler is not None + _get_library_root_logger().addHandler(_default_handler) + + +def add_handler(handler: logging.Handler) -> None: + """adds a handler to the HuggingFace Transformers's root logger.""" + + _configure_library_root_logger() + + assert handler is not None + _get_library_root_logger().addHandler(handler) + + +def remove_handler(handler: logging.Handler) -> None: + """removes given handler from the HuggingFace Transformers's root logger.""" + + _configure_library_root_logger() + + assert handler is not None and handler not in _get_library_root_logger().handlers + _get_library_root_logger().removeHandler(handler) + + +def disable_propagation() -> None: + """ + Disable propagation of the library log outputs. Note that log propagation is disabled by default. + """ + + _configure_library_root_logger() + _get_library_root_logger().propagate = False + + +def enable_propagation() -> None: + """ + Enable propagation of the library log outputs. Please disable the HuggingFace Transformers's default handler to + prevent double logging if the root logger has been configured. + """ + + _configure_library_root_logger() + _get_library_root_logger().propagate = True + + +def enable_explicit_format() -> None: + """ + Enable explicit formatting for every HuggingFace Transformers's logger. The explicit formatter is as follows: + ``` + [LEVELNAME|FILENAME|LINE NUMBER] TIME >> MESSAGE + ``` + All handlers currently bound to the root logger are affected by this method. + """ + handlers = _get_library_root_logger().handlers + + for handler in handlers: + formatter = logging.Formatter("[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s") + handler.setFormatter(formatter) + + +def reset_format() -> None: + """ + Resets the formatting for HuggingFace Transformers's loggers. + + All handlers currently bound to the root logger are affected by this method. + """ + handlers = _get_library_root_logger().handlers + + for handler in handlers: + handler.setFormatter(None) + + +def warning_advice(self, *args, **kwargs): + """ + This method is identical to `logger.warning()`, but if env var TRANSFORMERS_NO_ADVISORY_WARNINGS=1 is set, this + warning will not be printed + """ + no_advisory_warnings = os.getenv("TRANSFORMERS_NO_ADVISORY_WARNINGS", False) + if no_advisory_warnings: + return + self.warning(*args, **kwargs) + + +logging.Logger.warning_advice = warning_advice + + +@functools.lru_cache(None) +def warning_once(self, *args, **kwargs): + """ + This method is identical to `logger.warning()`, but will emit the warning with the same message only once + + Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the cache. + The assumption here is that all warning messages are unique across the code. If they aren't then need to switch to + another type of cache that includes the caller frame information in the hashing function. + """ + self.warning(*args, **kwargs) + + +logging.Logger.warning_once = warning_once + + +@functools.lru_cache(None) +def info_once(self, *args, **kwargs): + """ + This method is identical to `logger.info()`, but will emit the info with the same message only once + + Note: The cache is for the function arguments, so 2 different callers using the same arguments will hit the cache. + The assumption here is that all warning messages are unique across the code. If they aren't then need to switch to + another type of cache that includes the caller frame information in the hashing function. + """ + self.info(*args, **kwargs) + + +logging.Logger.info_once = info_once + + +class EmptyTqdm: + """Dummy tqdm which doesn't do anything.""" + + def __init__(self, *args, **kwargs): # pylint: disable=unused-argument + self._iterator = args[0] if args else None + + def __iter__(self): + return iter(self._iterator) + + def __getattr__(self, _): + """Return empty function.""" + + def empty_fn(*args, **kwargs): # pylint: disable=unused-argument + return + + return empty_fn + + def __enter__(self): + return self + + def __exit__(self, type_, value, traceback): + return + + +class _tqdm_cls: + def __call__(self, *args, **kwargs): + if _tqdm_active: + return tqdm_lib.tqdm(*args, **kwargs) + else: + return EmptyTqdm(*args, **kwargs) + + def set_lock(self, *args, **kwargs): + self._lock = None + if _tqdm_active: + return tqdm_lib.tqdm.set_lock(*args, **kwargs) + + def get_lock(self): + if _tqdm_active: + return tqdm_lib.tqdm.get_lock() + + +tqdm = _tqdm_cls() + + +def is_progress_bar_enabled() -> bool: + """Return a boolean indicating whether tqdm progress bars are enabled.""" + global _tqdm_active + return bool(_tqdm_active) + + +def enable_progress_bar(): + """Enable tqdm progress bar.""" + global _tqdm_active + _tqdm_active = True + hf_hub_utils.enable_progress_bars() + + +def disable_progress_bar(): + """Disable tqdm progress bar.""" + global _tqdm_active + _tqdm_active = False + hf_hub_utils.disable_progress_bars() diff --git a/lib/python3.12/site-packages/transformers/utils/model_parallel_utils.py b/lib/python3.12/site-packages/transformers/utils/model_parallel_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7db16b70a75ccdce7bd9c07c5dfc99f42e5c8807 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/model_parallel_utils.py @@ -0,0 +1,55 @@ +# Copyright 2020 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. + +from math import ceil + + +def assert_device_map(device_map, num_blocks): + blocks = list(range(0, num_blocks)) + + device_map_blocks = [item for sublist in list(device_map.values()) for item in sublist] + + # Duplicate check + duplicate_blocks = [] + for i in device_map_blocks: + if device_map_blocks.count(i) > 1 and i not in duplicate_blocks: + duplicate_blocks.append(i) + # Missing blocks + missing_blocks = [i for i in blocks if i not in device_map_blocks] + extra_blocks = [i for i in device_map_blocks if i not in blocks] + + if len(duplicate_blocks) != 0: + raise ValueError( + "Duplicate attention blocks specified in device_map. Attention blocks must be specified to one device." + " These attention blocks were specified more than once: " + str(duplicate_blocks) + ) + if len(missing_blocks) != 0: + raise ValueError( + "There are attention blocks for this model that are not specified in the device_map. Add these attention " + "blocks to a device on the device_map: " + str(missing_blocks) + ) + if len(extra_blocks) != 0: + raise ValueError( + "The device_map contains more attention blocks than this model has. Remove these from the device_map:" + + str(extra_blocks) + ) + + +def get_device_map(n_layers, devices): + """Returns a dictionary of layers distributed evenly across all devices.""" + layers = list(range(n_layers)) + n_blocks = int(ceil(n_layers / len(devices))) + layers_list = [layers[i : i + n_blocks] for i in range(0, n_layers, n_blocks)] + + return dict(zip(devices, layers_list)) diff --git a/lib/python3.12/site-packages/transformers/utils/notebook.py b/lib/python3.12/site-packages/transformers/utils/notebook.py new file mode 100644 index 0000000000000000000000000000000000000000..22a44d858e731f9bd4067d373165b0b376f5191b --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/notebook.py @@ -0,0 +1,383 @@ +# Copyright 2020 Hugging Face +# +# 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. + +import os +import re +import time +from typing import Optional + +import IPython.display as disp + +from ..trainer_callback import TrainerCallback +from ..trainer_utils import IntervalStrategy, has_length + + +def format_time(t): + "Format `t` (in seconds) to (h):mm:ss" + t = int(t) + h, m, s = t // 3600, (t // 60) % 60, t % 60 + return f"{h}:{m:02d}:{s:02d}" if h != 0 else f"{m:02d}:{s:02d}" + + +def html_progress_bar(value, total, prefix, label, width=300): + # docstyle-ignore + return f""" +
+ {prefix} + + {label} +
+ """ + + +def text_to_html_table(items): + "Put the texts in `items` in an HTML table." + html_code = """\n""" + html_code += """ \n \n""" + for i in items[0]: + html_code += f" \n" + html_code += " \n \n \n" + for line in items[1:]: + html_code += " \n" + for elt in line: + elt = f"{elt:.6f}" if isinstance(elt, float) else str(elt) + html_code += f" \n" + html_code += " \n" + html_code += " \n
{i}
{elt}

" + return html_code + + +class NotebookProgressBar: + """ + A progress par for display in a notebook. + + Class attributes (overridden by derived classes) + + - **warmup** (`int`) -- The number of iterations to do at the beginning while ignoring `update_every`. + - **update_every** (`float`) -- Since calling the time takes some time, we only do it every presumed + `update_every` seconds. The progress bar uses the average time passed up until now to guess the next value + for which it will call the update. + + Args: + total (`int`): + The total number of iterations to reach. + prefix (`str`, *optional*): + A prefix to add before the progress bar. + leave (`bool`, *optional*, defaults to `True`): + Whether or not to leave the progress bar once it's completed. You can always call the + [`~utils.notebook.NotebookProgressBar.close`] method to make the bar disappear. + parent ([`~notebook.NotebookTrainingTracker`], *optional*): + A parent object (like [`~utils.notebook.NotebookTrainingTracker`]) that spawns progress bars and handle + their display. If set, the object passed must have a `display()` method. + width (`int`, *optional*, defaults to 300): + The width (in pixels) that the bar will take. + + Example: + + ```python + import time + + pbar = NotebookProgressBar(100) + for val in range(100): + pbar.update(val) + time.sleep(0.07) + pbar.update(100) + ```""" + + warmup = 5 + update_every = 0.2 + + def __init__( + self, + total: int, + prefix: Optional[str] = None, + leave: bool = True, + parent: Optional["NotebookTrainingTracker"] = None, + width: int = 300, + ): + self.total = total + self.prefix = "" if prefix is None else prefix + self.leave = leave + self.parent = parent + self.width = width + self.last_value = None + self.comment = None + self.output = None + self.value = None + self.label = None + if "VSCODE_PID" in os.environ: + self.update_every = 0.5 # Adjusted for smooth updated as html rending is slow on VS Code + # This is the only adjustment required to optimize training html rending + + def update(self, value: int, force_update: bool = False, comment: Optional[str] = None): + """ + The main method to update the progress bar to `value`. + + Args: + value (`int`): + The value to use. Must be between 0 and `total`. + force_update (`bool`, *optional*, defaults to `False`): + Whether or not to force and update of the internal state and display (by default, the bar will wait for + `value` to reach the value it predicted corresponds to a time of more than the `update_every` attribute + since the last update to avoid adding boilerplate). + comment (`str`, *optional*): + A comment to add on the left of the progress bar. + """ + self.value = value + if comment is not None: + self.comment = comment + if self.last_value is None: + self.start_time = self.last_time = time.time() + self.start_value = self.last_value = value + self.elapsed_time = self.predicted_remaining = None + self.first_calls = self.warmup + self.wait_for = 1 + self.update_bar(value) + elif value <= self.last_value and not force_update: + return + elif force_update or self.first_calls > 0 or value >= min(self.last_value + self.wait_for, self.total): + if self.first_calls > 0: + self.first_calls -= 1 + current_time = time.time() + self.elapsed_time = current_time - self.start_time + # We could have value = self.start_value if the update is called twixe with the same start value. + if value > self.start_value: + self.average_time_per_item = self.elapsed_time / (value - self.start_value) + else: + self.average_time_per_item = None + if value >= self.total: + value = self.total + self.predicted_remaining = None + if not self.leave: + self.close() + elif self.average_time_per_item is not None: + self.predicted_remaining = self.average_time_per_item * (self.total - value) + self.update_bar(value) + self.last_value = value + self.last_time = current_time + if (self.average_time_per_item is None) or (self.average_time_per_item == 0): + self.wait_for = 1 + else: + self.wait_for = max(int(self.update_every / self.average_time_per_item), 1) + + def update_bar(self, value, comment=None): + spaced_value = " " * (len(str(self.total)) - len(str(value))) + str(value) + if self.elapsed_time is None: + self.label = f"[{spaced_value}/{self.total} : < :" + elif self.predicted_remaining is None: + self.label = f"[{spaced_value}/{self.total} {format_time(self.elapsed_time)}" + else: + self.label = ( + f"[{spaced_value}/{self.total} {format_time(self.elapsed_time)} <" + f" {format_time(self.predicted_remaining)}" + ) + if self.average_time_per_item == 0: + self.label += ", +inf it/s" + else: + self.label += f", {1 / self.average_time_per_item:.2f} it/s" + + self.label += "]" if self.comment is None or len(self.comment) == 0 else f", {self.comment}]" + self.display() + + def display(self): + self.html_code = html_progress_bar(self.value, self.total, self.prefix, self.label, self.width) + if self.parent is not None: + # If this is a child bar, the parent will take care of the display. + self.parent.display() + return + if self.output is None: + self.output = disp.display(disp.HTML(self.html_code), display_id=True) + else: + self.output.update(disp.HTML(self.html_code)) + + def close(self): + "Closes the progress bar." + if self.parent is None and self.output is not None: + self.output.update(disp.HTML("")) + + +class NotebookTrainingTracker(NotebookProgressBar): + """ + An object tracking the updates of an ongoing training with progress bars and a nice table reporting metrics. + + Args: + num_steps (`int`): The number of steps during training. column_names (`List[str]`, *optional*): + The list of column names for the metrics table (will be inferred from the first call to + [`~utils.notebook.NotebookTrainingTracker.write_line`] if not set). + """ + + def __init__(self, num_steps, column_names=None): + super().__init__(num_steps) + self.inner_table = None if column_names is None else [column_names] + self.child_bar = None + + def display(self): + self.html_code = html_progress_bar(self.value, self.total, self.prefix, self.label, self.width) + if self.inner_table is not None: + self.html_code += text_to_html_table(self.inner_table) + if self.child_bar is not None: + self.html_code += self.child_bar.html_code + if self.output is None: + self.output = disp.display(disp.HTML(self.html_code), display_id=True) + else: + self.output.update(disp.HTML(self.html_code)) + + def write_line(self, values): + """ + Write the values in the inner table. + + Args: + values (`Dict[str, float]`): The values to display. + """ + if self.inner_table is None: + self.inner_table = [list(values.keys()), list(values.values())] + else: + columns = self.inner_table[0] + for key in values.keys(): + if key not in columns: + columns.append(key) + self.inner_table[0] = columns + if len(self.inner_table) > 1: + last_values = self.inner_table[-1] + first_column = self.inner_table[0][0] + if last_values[0] != values[first_column]: + # write new line + self.inner_table.append([values[c] if c in values else "No Log" for c in columns]) + else: + # update last line + new_values = values + for c in columns: + if c not in new_values.keys(): + new_values[c] = last_values[columns.index(c)] + self.inner_table[-1] = [new_values[c] for c in columns] + else: + self.inner_table.append([values[c] for c in columns]) + + def add_child(self, total, prefix=None, width=300): + """ + Add a child progress bar displayed under the table of metrics. The child progress bar is returned (so it can be + easily updated). + + Args: + total (`int`): The number of iterations for the child progress bar. + prefix (`str`, *optional*): A prefix to write on the left of the progress bar. + width (`int`, *optional*, defaults to 300): The width (in pixels) of the progress bar. + """ + self.child_bar = NotebookProgressBar(total, prefix=prefix, parent=self, width=width) + return self.child_bar + + def remove_child(self): + """ + Closes the child progress bar. + """ + self.child_bar = None + self.display() + + +class NotebookProgressCallback(TrainerCallback): + """ + A [`TrainerCallback`] that displays the progress of training or evaluation, optimized for Jupyter Notebooks or + Google colab. + """ + + def __init__(self): + self.training_tracker = None + self.prediction_bar = None + self._force_next_update = False + + def on_train_begin(self, args, state, control, **kwargs): + self.first_column = "Epoch" if args.eval_strategy == IntervalStrategy.EPOCH else "Step" + self.training_loss = 0 + self.last_log = 0 + column_names = [self.first_column] + ["Training Loss"] + if args.eval_strategy != IntervalStrategy.NO: + column_names.append("Validation Loss") + self.training_tracker = NotebookTrainingTracker(state.max_steps, column_names) + + def on_step_end(self, args, state, control, **kwargs): + epoch = int(state.epoch) if int(state.epoch) == state.epoch else f"{state.epoch:.2f}" + self.training_tracker.update( + state.global_step + 1, + comment=f"Epoch {epoch}/{state.num_train_epochs}", + force_update=self._force_next_update, + ) + self._force_next_update = False + + def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs): + if not has_length(eval_dataloader): + return + if self.prediction_bar is None: + if self.training_tracker is not None: + self.prediction_bar = self.training_tracker.add_child(len(eval_dataloader)) + else: + self.prediction_bar = NotebookProgressBar(len(eval_dataloader)) + self.prediction_bar.update(1) + else: + self.prediction_bar.update(self.prediction_bar.value + 1) + + def on_predict(self, args, state, control, **kwargs): + if self.prediction_bar is not None: + self.prediction_bar.close() + self.prediction_bar = None + + def on_log(self, args, state, control, logs=None, **kwargs): + # Only for when there is no evaluation + if args.eval_strategy == IntervalStrategy.NO and "loss" in logs: + values = {"Training Loss": logs["loss"]} + # First column is necessarily Step sine we're not in epoch eval strategy + values["Step"] = state.global_step + self.training_tracker.write_line(values) + + def on_evaluate(self, args, state, control, metrics=None, **kwargs): + if self.training_tracker is not None: + values = {"Training Loss": "No log", "Validation Loss": "No log"} + for log in reversed(state.log_history): + if "loss" in log: + values["Training Loss"] = log["loss"] + break + + if self.first_column == "Epoch": + values["Epoch"] = int(state.epoch) + else: + values["Step"] = state.global_step + metric_key_prefix = "eval" + for k in metrics: + if k.endswith("_loss"): + metric_key_prefix = re.sub(r"\_loss$", "", k) + _ = metrics.pop("total_flos", None) + _ = metrics.pop("epoch", None) + _ = metrics.pop(f"{metric_key_prefix}_runtime", None) + _ = metrics.pop(f"{metric_key_prefix}_samples_per_second", None) + _ = metrics.pop(f"{metric_key_prefix}_steps_per_second", None) + _ = metrics.pop(f"{metric_key_prefix}_jit_compilation_time", None) + for k, v in metrics.items(): + splits = k.split("_") + name = " ".join([part.capitalize() for part in splits[1:]]) + if name == "Loss": + # Single dataset + name = "Validation Loss" + values[name] = v + self.training_tracker.write_line(values) + self.training_tracker.remove_child() + self.prediction_bar = None + # Evaluation takes a long time so we should force the next update. + self._force_next_update = True + + def on_train_end(self, args, state, control, **kwargs): + self.training_tracker.update( + state.global_step, + comment=f"Epoch {int(state.epoch)}/{state.num_train_epochs}", + force_update=True, + ) + self.training_tracker = None diff --git a/lib/python3.12/site-packages/transformers/utils/peft_utils.py b/lib/python3.12/site-packages/transformers/utils/peft_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3eb62a099059d4c1acd9eef75f94ee5ac22fd60f --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/peft_utils.py @@ -0,0 +1,125 @@ +# Copyright 2023 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. +import importlib +import os +from typing import Optional, Union + +from packaging import version + +from .hub import cached_file +from .import_utils import is_peft_available + + +ADAPTER_CONFIG_NAME = "adapter_config.json" +ADAPTER_WEIGHTS_NAME = "adapter_model.bin" +ADAPTER_SAFE_WEIGHTS_NAME = "adapter_model.safetensors" + + +def find_adapter_config_file( + model_id: str, + cache_dir: Optional[Union[str, os.PathLike]] = None, + force_download: bool = False, + resume_download: Optional[bool] = None, + proxies: Optional[dict[str, str]] = None, + token: Optional[Union[bool, str]] = None, + revision: Optional[str] = None, + local_files_only: bool = False, + subfolder: str = "", + _commit_hash: Optional[str] = None, +) -> Optional[str]: + r""" + Simply checks if the model stored on the Hub or locally is an adapter model or not, return the path of the adapter + config file if it is, None otherwise. + + Args: + model_id (`str`): + The identifier of the model to look for, can be either a local path or an id to the repository on the Hub. + 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. + force_download (`bool`, *optional*, defaults to `False`): + Whether or not to force to (re-)download the configuration files and override the cached versions if they + exist. + resume_download: + Deprecated and ignored. All downloads are now resumed by default when possible. + Will be removed in v5 of Transformers. + proxies (`Dict[str, str]`, *optional*): + 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. + token (`str` or *bool*, *optional*): + The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated + when running `huggingface-cli login` (stored in `~/.huggingface`). + revision (`str`, *optional*, defaults to `"main"`): + The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a + git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any + identifier allowed by git. + + + + To test a pull request you made on the Hub, you can pass `revision="refs/pr/". + + + + local_files_only (`bool`, *optional*, defaults to `False`): + If `True`, will only try to load the tokenizer configuration from local files. + 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. + """ + adapter_cached_filename = None + if model_id is None: + return None + elif os.path.isdir(model_id): + list_remote_files = os.listdir(model_id) + if ADAPTER_CONFIG_NAME in list_remote_files: + adapter_cached_filename = os.path.join(model_id, ADAPTER_CONFIG_NAME) + else: + adapter_cached_filename = cached_file( + model_id, + ADAPTER_CONFIG_NAME, + cache_dir=cache_dir, + force_download=force_download, + resume_download=resume_download, + proxies=proxies, + token=token, + revision=revision, + local_files_only=local_files_only, + subfolder=subfolder, + _commit_hash=_commit_hash, + _raise_exceptions_for_gated_repo=False, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + ) + + return adapter_cached_filename + + +def check_peft_version(min_version: str) -> None: + r""" + Checks if the version of PEFT is compatible. + + Args: + version (`str`): + The version of PEFT to check against. + """ + if not is_peft_available(): + raise ValueError("PEFT is not installed. Please install it with `pip install peft`") + + is_peft_version_compatible = version.parse(importlib.metadata.version("peft")) >= version.parse(min_version) + + if not is_peft_version_compatible: + raise ValueError( + f"The version of PEFT you are using is not compatible, please use a version that is greater" + f" than {min_version}" + ) diff --git a/lib/python3.12/site-packages/transformers/utils/quantization_config.py b/lib/python3.12/site-packages/transformers/utils/quantization_config.py new file mode 100644 index 0000000000000000000000000000000000000000..b0f119c58b67b919112884cb2e9253c1c295ebcc --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/quantization_config.py @@ -0,0 +1,1835 @@ +#!/usr/bin/env python +# coding=utf-8 + +# Copyright 2023 The HuggingFace Inc. team. All rights reserved. +# Modifications Copyright (C) 2025, Advanced Micro Devices, Inc. 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. +import copy +import dataclasses +import importlib.metadata +import json +import os +from dataclasses import dataclass +from enum import Enum +from inspect import Parameter, signature +from typing import Any, Dict, List, Optional, Tuple, Union + +from packaging import version + +from ..utils import ( + is_auto_awq_available, + is_compressed_tensors_available, + is_gptqmodel_available, + is_hqq_available, + is_quark_available, + is_torch_available, + is_torchao_available, + logging, +) +from .import_utils import is_auto_gptq_available + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class QuantizationMethod(str, Enum): + BITS_AND_BYTES = "bitsandbytes" + GPTQ = "gptq" + AWQ = "awq" + AQLM = "aqlm" + VPTQ = "vptq" + QUANTO = "quanto" + EETQ = "eetq" + HIGGS = "higgs" + HQQ = "hqq" + COMPRESSED_TENSORS = "compressed-tensors" + FBGEMM_FP8 = "fbgemm_fp8" + TORCHAO = "torchao" + BITNET = "bitnet" + SPQR = "spqr" + FP8 = "fp8" + QUARK = "quark" + + +class AWQLinearVersion(str, Enum): + GEMM = "gemm" + GEMV = "gemv" + EXLLAMA = "exllama" + IPEX = "ipex" + + @staticmethod + def from_str(version: str): + version = version.lower() + if version == "gemm": + return AWQLinearVersion.GEMM + elif version == "gemv": + return AWQLinearVersion.GEMV + elif version == "exllama": + return AWQLinearVersion.EXLLAMA + elif version == "ipex": + return AWQLinearVersion.IPEX + else: + raise ValueError(f"Unknown AWQLinearVersion {version}") + + +class AwqBackendPackingMethod(str, Enum): + AUTOAWQ = "autoawq" + LLMAWQ = "llm-awq" + + +@dataclass +class QuantizationConfigMixin: + """ + Mixin class for quantization config + """ + + quant_method: QuantizationMethod + + @classmethod + def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs): + """ + Instantiates a [`QuantizationConfigMixin`] from a Python dictionary of parameters. + + Args: + config_dict (`Dict[str, Any]`): + Dictionary that will be used to instantiate the configuration object. + return_unused_kwargs (`bool`,*optional*, defaults to `False`): + Whether or not to return a list of unused keyword arguments. Used for `from_pretrained` method in + `PreTrainedModel`. + kwargs (`Dict[str, Any]`): + Additional parameters from which to initialize the configuration object. + + Returns: + [`QuantizationConfigMixin`]: The configuration object instantiated from those parameters. + """ + config = cls(**config_dict) + + to_remove = [] + for key, value in kwargs.items(): + if hasattr(config, key): + setattr(config, key, value) + to_remove.append(key) + for key in to_remove: + kwargs.pop(key, None) + + if return_unused_kwargs: + return config, kwargs + else: + return config + + def to_json_file(self, json_file_path: Union[str, os.PathLike]): + """ + Save this instance to a JSON file. + + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file in which this configuration instance's parameters will be saved. + use_diff (`bool`, *optional*, defaults to `True`): + If set to `True`, only the difference between the config instance and the default + `QuantizationConfig()` is serialized to JSON file. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + config_dict = self.to_dict() + json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n" + + writer.write(json_string) + + def to_dict(self) -> Dict[str, Any]: + """ + Serializes this instance to a Python dictionary. Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + return copy.deepcopy(self.__dict__) + + def __iter__(self): + """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin""" + for attr, value in copy.deepcopy(self.__dict__).items(): + yield attr, value + + def __repr__(self): + return f"{self.__class__.__name__} {self.to_json_string()}" + + def to_json_string(self, use_diff: bool = True) -> str: + """ + Serializes this instance to a JSON string. + + Args: + use_diff (`bool`, *optional*, defaults to `True`): + If set to `True`, only the difference between the config instance and the default `PretrainedConfig()` + is serialized to JSON string. + + Returns: + `str`: String containing all the attributes that make up this configuration instance in JSON format. + """ + if use_diff is True: + config_dict = self.to_diff_dict() + else: + config_dict = self.to_dict() + return json.dumps(config_dict, indent=2, sort_keys=True) + "\n" + + def update(self, **kwargs): + """ + Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes, + returning all the unused kwargs. + + Args: + kwargs (`Dict[str, Any]`): + Dictionary of attributes to tentatively update this class. + + Returns: + `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance. + """ + to_remove = [] + for key, value in kwargs.items(): + if hasattr(self, key): + setattr(self, key, value) + to_remove.append(key) + + # Remove all the attributes that were updated, without modifying the input dict + unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove} + return unused_kwargs + + +@dataclass +class HqqConfig(QuantizationConfigMixin): + """ + This is wrapper around hqq's BaseQuantizeConfig. + + Args: + nbits (`int`, *optional*, defaults to 4): + Number of bits. Supported values are (8, 4, 3, 2, 1). + group_size (`int`, *optional*, defaults to 64): + Group-size value. Supported values are any value that is divisble by weight.shape[axis]). + view_as_float (`bool`, *optional*, defaults to `False`): + View the quantized weight as float (used in distributed training) if set to `True`. + axis (`Optional[int]`, *optional*): + Axis along which grouping is performed. Supported values are 0 or 1. + dynamic_config (dict, *optional*): + Parameters for dynamic configuration. The key is the name tag of the layer and the value is a quantization config. + If set, each layer specified by its id will use its dedicated quantization configuration. + skip_modules (`List[str]`, *optional*, defaults to `['lm_head']`): + List of `nn.Linear` layers to skip. + kwargs (`Dict[str, Any]`, *optional*): + Additional parameters from which to initialize the configuration object. + """ + + def __init__( + self, + nbits: int = 4, + group_size: int = 64, + view_as_float: bool = False, + axis: Optional[int] = None, + dynamic_config: Optional[dict] = None, + skip_modules: List[str] = ["lm_head"], + **kwargs, + ): + if is_hqq_available(): + from hqq.core.quantize import BaseQuantizeConfig as HQQBaseQuantizeConfig + else: + raise ImportError( + "A valid HQQ version (>=0.2.1) is not available. Please follow the instructions to install it: `https://github.com/mobiusml/hqq/`." + ) + + for deprecated_key in ["quant_zero", "quant_scale", "offload_meta"]: + if deprecated_key in kwargs: + logger.info( + deprecated_key + " is deprecated. This parameter will be ignored in quantization settings." + ) + + if axis is None: + axis = 1 + logger.info("Setting axis=1 as faster backends such as TorchAO or BitBlas are only compatible with it.") + + if axis not in [0, 1]: + raise ValueError("Invalid axis value. Only 0 and 1 are allowed.") + + if dynamic_config is not None: + self.quant_config = {} + for key in dynamic_config: + self.quant_config[key] = HQQBaseQuantizeConfig(**dynamic_config[key]) + else: + self.quant_config = HQQBaseQuantizeConfig( + **{ + "nbits": nbits, + "group_size": group_size, + "view_as_float": view_as_float, + "axis": axis, + } + ) + + self.quant_method = QuantizationMethod.HQQ + self.skip_modules = skip_modules + + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct - also replaces some NoneType arguments with their default values. + """ + pass + + @classmethod + def from_dict(cls, config: Dict[str, Any]): + """ + Override from_dict, used in AutoQuantizationConfig.from_dict in quantizers/auto.py + """ + instance = cls() + instance.quant_config = config["quant_config"] + instance.skip_modules = config["skip_modules"] + return instance + + def to_dict(self) -> Dict[str, Any]: + """ + Serializes this instance to a Python dictionary. Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + return { + "quant_config": self.quant_config, + "quant_method": self.quant_method, + "skip_modules": self.skip_modules, + } + + def __repr__(self): + config_dict = self.to_dict() + return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n" + + def to_diff_dict(self) -> Dict[str, Any]: + """ + Removes all attributes from config which correspond to the default config attributes for better readability and + serializes to a Python dictionary. + Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance, + """ + config_dict = self.to_dict() + + # get the default config dict + default_config_dict = HqqConfig().to_dict() + + serializable_config_dict = {} + + # only serialize values that differ from the default config + for key, value in config_dict.items(): + if value != default_config_dict[key]: + serializable_config_dict[key] = value + + return serializable_config_dict + + +@dataclass +class BitsAndBytesConfig(QuantizationConfigMixin): + """ + This is a wrapper class about all possible attributes and features that you can play with a model that has been + loaded using `bitsandbytes`. + + This replaces `load_in_8bit` or `load_in_4bit`therefore both options are mutually exclusive. + + Currently only supports `LLM.int8()`, `FP4`, and `NF4` quantization. If more methods are added to `bitsandbytes`, + then more arguments will be added to this class. + + Args: + load_in_8bit (`bool`, *optional*, defaults to `False`): + This flag is used to enable 8-bit quantization with LLM.int8(). + load_in_4bit (`bool`, *optional*, defaults to `False`): + This flag is used to enable 4-bit quantization by replacing the Linear layers with FP4/NF4 layers from + `bitsandbytes`. + llm_int8_threshold (`float`, *optional*, defaults to 6.0): + This corresponds to the outlier threshold for outlier detection as described in `LLM.int8() : 8-bit Matrix + Multiplication for Transformers at Scale` paper: https://arxiv.org/abs/2208.07339 Any hidden states value + that is above this threshold will be considered an outlier and the operation on those values will be done + in fp16. Values are usually normally distributed, that is, most values are in the range [-3.5, 3.5], but + there are some exceptional systematic outliers that are very differently distributed for large models. + These outliers are often in the interval [-60, -6] or [6, 60]. Int8 quantization works well for values of + magnitude ~5, but beyond that, there is a significant performance penalty. A good default threshold is 6, + but a lower threshold might be needed for more unstable models (small models, fine-tuning). + llm_int8_skip_modules (`List[str]`, *optional*): + An explicit list of the modules that we do not want to convert in 8-bit. This is useful for models such as + Jukebox that has several heads in different places and not necessarily at the last position. For example + for `CausalLM` models, the last `lm_head` is kept in its original `dtype`. + llm_int8_enable_fp32_cpu_offload (`bool`, *optional*, defaults to `False`): + This flag is used for advanced use cases and users that are aware of this feature. If you want to split + your model in different parts and run some parts in int8 on GPU and some parts in fp32 on CPU, you can use + this flag. This is useful for offloading large models such as `google/flan-t5-xxl`. Note that the int8 + operations will not be run on CPU. + llm_int8_has_fp16_weight (`bool`, *optional*, defaults to `False`): + This flag runs LLM.int8() with 16-bit main weights. This is useful for fine-tuning as the weights do not + have to be converted back and forth for the backward pass. + bnb_4bit_compute_dtype (`torch.dtype` or str, *optional*, defaults to `torch.float32`): + This sets the computational type which might be different than the input type. For example, inputs might be + fp32, but computation can be set to bf16 for speedups. + bnb_4bit_quant_type (`str`, *optional*, defaults to `"fp4"`): + This sets the quantization data type in the bnb.nn.Linear4Bit layers. Options are FP4 and NF4 data types + which are specified by `fp4` or `nf4`. + bnb_4bit_use_double_quant (`bool`, *optional*, defaults to `False`): + This flag is used for nested quantization where the quantization constants from the first quantization are + quantized again. + bnb_4bit_quant_storage (`torch.dtype` or str, *optional*, defaults to `torch.uint8`): + This sets the storage type to pack the quanitzed 4-bit prarams. + kwargs (`Dict[str, Any]`, *optional*): + Additional parameters from which to initialize the configuration object. + """ + + def __init__( + self, + load_in_8bit=False, + load_in_4bit=False, + llm_int8_threshold=6.0, + llm_int8_skip_modules=None, + llm_int8_enable_fp32_cpu_offload=False, + llm_int8_has_fp16_weight=False, + bnb_4bit_compute_dtype=None, + bnb_4bit_quant_type="fp4", + bnb_4bit_use_double_quant=False, + bnb_4bit_quant_storage=None, + **kwargs, + ): + self.quant_method = QuantizationMethod.BITS_AND_BYTES + + if load_in_4bit and load_in_8bit: + raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time") + + self._load_in_8bit = load_in_8bit + self._load_in_4bit = load_in_4bit + self.llm_int8_threshold = llm_int8_threshold + self.llm_int8_skip_modules = llm_int8_skip_modules + self.llm_int8_enable_fp32_cpu_offload = llm_int8_enable_fp32_cpu_offload + self.llm_int8_has_fp16_weight = llm_int8_has_fp16_weight + self.bnb_4bit_quant_type = bnb_4bit_quant_type + self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant + + if bnb_4bit_compute_dtype is None: + self.bnb_4bit_compute_dtype = torch.float32 + elif isinstance(bnb_4bit_compute_dtype, str): + self.bnb_4bit_compute_dtype = getattr(torch, bnb_4bit_compute_dtype) + elif isinstance(bnb_4bit_compute_dtype, torch.dtype): + self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype + else: + raise ValueError("bnb_4bit_compute_dtype must be a string or a torch.dtype") + + if bnb_4bit_quant_storage is None: + self.bnb_4bit_quant_storage = torch.uint8 + elif isinstance(bnb_4bit_quant_storage, str): + if bnb_4bit_quant_storage not in ["float16", "float32", "int8", "uint8", "float64", "bfloat16"]: + raise ValueError( + "`bnb_4bit_quant_storage` must be a valid string (one of 'float16', 'float32', 'int8', 'uint8', 'float64', 'bfloat16') " + ) + self.bnb_4bit_quant_storage = getattr(torch, bnb_4bit_quant_storage) + elif isinstance(bnb_4bit_quant_storage, torch.dtype): + self.bnb_4bit_quant_storage = bnb_4bit_quant_storage + else: + raise ValueError("bnb_4bit_quant_storage must be a string or a torch.dtype") + + if kwargs: + logger.info(f"Unused kwargs: {list(kwargs.keys())}. These kwargs are not used in {self.__class__}.") + + self.post_init() + + @property + def load_in_4bit(self): + return self._load_in_4bit + + @load_in_4bit.setter + def load_in_4bit(self, value: bool): + if not isinstance(value, bool): + raise TypeError("load_in_4bit must be a boolean") + + if self.load_in_8bit and value: + raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time") + self._load_in_4bit = value + + @property + def load_in_8bit(self): + return self._load_in_8bit + + @load_in_8bit.setter + def load_in_8bit(self, value: bool): + if not isinstance(value, bool): + raise TypeError("load_in_8bit must be a boolean") + + if self.load_in_4bit and value: + raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time") + self._load_in_8bit = value + + def post_init(self): + r""" + Safety checker that arguments are correct - also replaces some NoneType arguments with their default values. + """ + if not isinstance(self.load_in_4bit, bool): + raise TypeError("load_in_4bit must be a boolean") + + if not isinstance(self.load_in_8bit, bool): + raise TypeError("load_in_8bit must be a boolean") + + if not isinstance(self.llm_int8_threshold, float): + raise TypeError("llm_int8_threshold must be a float") + + if self.llm_int8_skip_modules is not None and not isinstance(self.llm_int8_skip_modules, list): + raise TypeError("llm_int8_skip_modules must be a list of strings") + if not isinstance(self.llm_int8_enable_fp32_cpu_offload, bool): + raise TypeError("llm_int8_enable_fp32_cpu_offload must be a boolean") + + if not isinstance(self.llm_int8_has_fp16_weight, bool): + raise TypeError("llm_int8_has_fp16_weight must be a boolean") + + if self.bnb_4bit_compute_dtype is not None and not isinstance(self.bnb_4bit_compute_dtype, torch.dtype): + raise TypeError("bnb_4bit_compute_dtype must be torch.dtype") + + if not isinstance(self.bnb_4bit_quant_type, str): + raise TypeError("bnb_4bit_quant_type must be a string") + + if not isinstance(self.bnb_4bit_use_double_quant, bool): + raise TypeError("bnb_4bit_use_double_quant must be a boolean") + + if self.load_in_4bit and not version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse( + "0.39.0" + ): + raise ValueError( + "4 bit quantization requires bitsandbytes>=0.39.0 - please upgrade your bitsandbytes version" + ) + + def is_quantizable(self): + r""" + Returns `True` if the model is quantizable, `False` otherwise. + """ + return self.load_in_8bit or self.load_in_4bit + + def quantization_method(self): + r""" + This method returns the quantization method used for the model. If the model is not quantizable, it returns + `None`. + """ + if self.load_in_8bit: + return "llm_int8" + elif self.load_in_4bit and self.bnb_4bit_quant_type == "fp4": + return "fp4" + elif self.load_in_4bit and self.bnb_4bit_quant_type == "nf4": + return "nf4" + else: + return None + + def to_dict(self) -> Dict[str, Any]: + """ + Serializes this instance to a Python dictionary. Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + output = copy.deepcopy(self.__dict__) + output["bnb_4bit_compute_dtype"] = str(output["bnb_4bit_compute_dtype"]).split(".")[1] + output["bnb_4bit_quant_storage"] = str(output["bnb_4bit_quant_storage"]).split(".")[1] + output["load_in_4bit"] = self.load_in_4bit + output["load_in_8bit"] = self.load_in_8bit + + return output + + def __repr__(self): + config_dict = self.to_dict() + return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n" + + def to_diff_dict(self) -> Dict[str, Any]: + """ + Removes all attributes from config which correspond to the default config attributes for better readability and + serializes to a Python dictionary. + + Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance, + """ + config_dict = self.to_dict() + + # get the default config dict + default_config_dict = BitsAndBytesConfig().to_dict() + + serializable_config_dict = {} + + # only serialize values that differ from the default config + for key, value in config_dict.items(): + if value != default_config_dict[key]: + serializable_config_dict[key] = value + + return serializable_config_dict + + +class ExllamaVersion(int, Enum): + ONE = 1 + TWO = 2 + + +@dataclass +class GPTQConfig(QuantizationConfigMixin): + """ + This is a wrapper class about all possible attributes and features that you can play with a model that has been + loaded using `optimum` api for gptq quantization relying on auto_gptq backend. + + Args: + bits (`int`): + The number of bits to quantize to, supported numbers are (2, 3, 4, 8). + tokenizer (`str` or `PreTrainedTokenizerBase`, *optional*): + The tokenizer used to process the dataset. You can pass either: + - A custom tokenizer object. + - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co. + - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved + using the [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`. + dataset (`Union[List[str]]`, *optional*): + The dataset used for quantization. You can provide your own dataset in a list of string or just use the + original datasets used in GPTQ paper ['wikitext2','c4','c4-new'] + group_size (`int`, *optional*, defaults to 128): + The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization. + damp_percent (`float`, *optional*, defaults to 0.1): + The percent of the average Hessian diagonal to use for dampening. Recommended value is 0.1. + desc_act (`bool`, *optional*, defaults to `False`): + Whether to quantize columns in order of decreasing activation size. Setting it to False can significantly + speed up inference but the perplexity may become slightly worse. Also known as act-order. + sym (`bool`, *optional*, defaults to `True`): + Whether to use symetric quantization. + true_sequential (`bool`, *optional*, defaults to `True`): + Whether to perform sequential quantization even within a single Transformer block. Instead of quantizing + the entire block at once, we perform layer-wise quantization. As a result, each layer undergoes + quantization using inputs that have passed through the previously quantized layers. + checkpoint_format (`str`, *optional*, defaults to `"gptq"`): + GPTQ weight format. `gptq`(v1) is supported by both gptqmodel and auto-gptq. `gptq_v2` is gptqmodel only. + meta (`Dict[str, any]`, *optional*): + Properties, such as tooling:version, that do not directly contributes to quantization or quant inference are stored in meta. + i.e. `meta.quantizer`: ["optimum:_version_", "gptqmodel:_version_"] + backend (`str`, *optional*): + Controls which gptq kernel to be used. Valid values for gptqmodel are `auto`, `auto_trainable` and more. For auto-gptq, only + valid value is None and `auto_trainable`. Ref gptqmodel backends: https://github.com/ModelCloud/GPTQModel/blob/main/gptqmodel/utils/backend.py + use_cuda_fp16 (`bool`, *optional*, defaults to `False`): + Whether or not to use optimized cuda kernel for fp16 model. Need to have model in fp16. Auto-gptq only. + model_seqlen (`int`, *optional*): + The maximum sequence length that the model can take. + block_name_to_quantize (`str`, *optional*): + The transformers block name to quantize. If None, we will infer the block name using common patterns (e.g. model.layers) + module_name_preceding_first_block (`List[str]`, *optional*): + The layers that are preceding the first Transformer block. + batch_size (`int`, *optional*, defaults to 1): + The batch size used when processing the dataset + pad_token_id (`int`, *optional*): + The pad token id. Needed to prepare the dataset when `batch_size` > 1. + use_exllama (`bool`, *optional*): + Whether to use exllama backend. Defaults to `True` if unset. Only works with `bits` = 4. + max_input_length (`int`, *optional*): + The maximum input length. This is needed to initialize a buffer that depends on the maximum expected input + length. It is specific to the exllama backend with act-order. + exllama_config (`Dict[str, Any]`, *optional*): + The exllama config. You can specify the version of the exllama kernel through the `version` key. Defaults + to `{"version": 1}` if unset. + cache_block_outputs (`bool`, *optional*, defaults to `True`): + Whether to cache block outputs to reuse as inputs for the succeeding block. + modules_in_block_to_quantize (`List[List[str]]`, *optional*): + List of list of module names to quantize in the specified block. This argument is useful to exclude certain linear modules from being quantized. + The block to quantize can be specified by setting `block_name_to_quantize`. We will quantize each list sequentially. If not set, we will quantize all linear layers. + Example: `modules_in_block_to_quantize =[["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"], ["self_attn.o_proj"]]`. + In this example, we will first quantize the q,k,v layers simultaneously since they are independent. + Then, we will quantize `self_attn.o_proj` layer with the q,k,v layers quantized. This way, we will get + better results since it reflects the real input `self_attn.o_proj` will get when the model is quantized. + """ + + def __init__( + self, + bits: int, + tokenizer: Any = None, + dataset: Optional[Union[List[str], str]] = None, + group_size: int = 128, + damp_percent: float = 0.1, + desc_act: bool = False, + sym: bool = True, + true_sequential: bool = True, + checkpoint_format: str = "gptq", + meta: Optional[Dict[str, Any]] = None, + backend: Optional[str] = None, + use_cuda_fp16: bool = False, + model_seqlen: Optional[int] = None, + block_name_to_quantize: Optional[str] = None, + module_name_preceding_first_block: Optional[List[str]] = None, + batch_size: int = 1, + pad_token_id: Optional[int] = None, + use_exllama: Optional[bool] = None, + max_input_length: Optional[int] = None, + exllama_config: Optional[Dict[str, Any]] = None, + cache_block_outputs: bool = True, + modules_in_block_to_quantize: Optional[List[List[str]]] = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.GPTQ + self.bits = bits + self.tokenizer = tokenizer + self.dataset = dataset + self.group_size = group_size + self.damp_percent = damp_percent + self.desc_act = desc_act + self.sym = sym + self.true_sequential = true_sequential + self.checkpoint_format = checkpoint_format.lower() + self.meta = meta + self.backend = backend.lower() if isinstance(backend, str) else backend + self.use_cuda_fp16 = use_cuda_fp16 + self.model_seqlen = model_seqlen + self.block_name_to_quantize = block_name_to_quantize + self.module_name_preceding_first_block = module_name_preceding_first_block + self.batch_size = batch_size + self.pad_token_id = pad_token_id + self.use_exllama = use_exllama + self.max_input_length = max_input_length + self.exllama_config = exllama_config + self.cache_block_outputs = cache_block_outputs + self.modules_in_block_to_quantize = modules_in_block_to_quantize + self.post_init() + + def get_loading_attributes(self): + attibutes_dict = copy.deepcopy(self.__dict__) + loading_attibutes = [ + "use_exllama", + "exllama_config", + "use_cuda_fp16", + "max_input_length", + "backend", + ] + loading_attibutes_dict = {i: j for i, j in attibutes_dict.items() if i in loading_attibutes} + return loading_attibutes_dict + + def post_init(self): + r""" + Safety checker that arguments are correct + """ + if self.bits not in [2, 3, 4, 8]: + raise ValueError(f"Only support quantization to [2,3,4,8] bits but found {self.bits}") + if self.group_size != -1 and self.group_size <= 0: + raise ValueError("group_size must be greater than 0 or equal to -1") + if not (0 < self.damp_percent < 1): + raise ValueError("damp_percent must between 0 and 1.") + if self.dataset is not None: + if isinstance(self.dataset, str): + if self.dataset in ["ptb", "ptb-new"]: + raise ValueError( + f"""{self.dataset} dataset was deprecated. You can only choose between + ['wikitext2','c4','c4-new']""" + ) + if self.dataset not in ["wikitext2", "c4", "c4-new"]: + raise ValueError( + f"""You have entered a string value for dataset. You can only choose between + ['wikitext2','c4','c4-new'], but we found {self.dataset}""" + ) + elif not isinstance(self.dataset, list): + raise ValueError( + f"""dataset needs to be either a list of string or a value in + ['wikitext2','c4','c4-new'], but we found {self.dataset}""" + ) + + # make sure backend is back/forward compatible with both gptqmodel (full) and auto-gptq (partial) + if is_gptqmodel_available(): + # convert auto-gptq control into gptqmodel backend + if self.backend is None: + self.backend = "auto_trainable" if self.use_exllama is not None and not self.use_exllama else "auto" + else: + # convert gptqmodel backend `auto_trainable` into auto-gptq control + if self.backend == "auto_trainable": + self.use_exllama = False + + # auto-gptq specific kernel control logic + if self.use_exllama is None: + # New default behaviour + self.use_exllama = True + + if self.exllama_config is None: + self.exllama_config = {"version": ExllamaVersion.ONE} + else: + if "version" not in self.exllama_config: + raise ValueError("`exllama_config` needs to have a `version` key.") + elif self.exllama_config["version"] not in [ExllamaVersion.ONE, ExllamaVersion.TWO]: + exllama_version = self.exllama_config["version"] + raise ValueError( + f"Only supported versions are in [ExllamaVersion.ONE, ExllamaVersion.TWO] - not recognized version {exllama_version}" + ) + + if self.bits == 4 and self.use_exllama: + if self.exllama_config["version"] == ExllamaVersion.ONE: + logger.info( + "You have activated exllama backend. Note that you can get better inference " + "speed using exllamav2 kernel by setting `exllama_config`." + ) + elif self.exllama_config["version"] == ExllamaVersion.TWO: + if is_auto_gptq_available(): + optimum_version = version.parse(importlib.metadata.version("optimum")) + autogptq_version = version.parse(importlib.metadata.version("auto_gptq")) + if optimum_version <= version.parse("1.13.2") or autogptq_version <= version.parse("0.4.2"): + raise ValueError( + f"You need optimum > 1.13.2 and auto-gptq > 0.4.2 . Make sure to have that version installed - detected version : optimum {optimum_version} and autogptq {autogptq_version}" + ) + if self.modules_in_block_to_quantize is not None: + optimum_version = version.parse(importlib.metadata.version("optimum")) + if optimum_version < version.parse("1.15.0"): + raise ValueError( + "You current version of `optimum` does not support `modules_in_block_to_quantize` quantization argument, please upgrade `optimum` package to a version superior than 1.15.0 ." + ) + + def to_dict(self): + config_dict = super().to_dict() + config_dict.pop("disable_exllama", None) + return config_dict + + def to_dict_optimum(self): + """ + Get compatible dict for optimum gptq config + """ + quant_dict = self.to_dict() + # make it compatible with optimum config + quant_dict["disable_exllama"] = not self.use_exllama + return quant_dict + + @classmethod + def from_dict_optimum(cls, config_dict): + """ + Get compatible class with optimum gptq config dict + """ + + if "disable_exllama" in config_dict: + config_dict["use_exllama"] = not config_dict["disable_exllama"] + # switch to None to not trigger the warning + config_dict.pop("disable_exllama") + + config = cls(**config_dict) + return config + + +@dataclass +class AwqConfig(QuantizationConfigMixin): + """ + This is a wrapper class about all possible attributes and features that you can play with a model that has been + loaded using `auto-awq` library awq quantization relying on auto_awq backend. + + Args: + bits (`int`, *optional*, defaults to 4): + The number of bits to quantize to. + group_size (`int`, *optional*, defaults to 128): + The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization. + zero_point (`bool`, *optional*, defaults to `True`): + Whether to use zero point quantization. + version (`AWQLinearVersion`, *optional*, defaults to `AWQLinearVersion.GEMM`): + The version of the quantization algorithm to use. GEMM is better for big batch_size (e.g. >= 8) otherwise, + GEMV is better (e.g. < 8 ). GEMM models are compatible with Exllama kernels. + backend (`AwqBackendPackingMethod`, *optional*, defaults to `AwqBackendPackingMethod.AUTOAWQ`): + The quantization backend. Some models might be quantized using `llm-awq` backend. This is useful for users + that quantize their own models using `llm-awq` library. + do_fuse (`bool`, *optional*, defaults to `False`): + Whether to fuse attention and mlp layers together for faster inference + fuse_max_seq_len (`int`, *optional*): + The Maximum sequence length to generate when using fusing. + modules_to_fuse (`dict`, *optional*, default to `None`): + Overwrite the natively supported fusing scheme with the one specified by the users. + modules_to_not_convert (`list`, *optional*, default to `None`): + The list of modules to not quantize, useful for quantizing models that explicitly require to have + some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers). + Note you cannot quantize directly with transformers, please refer to `AutoAWQ` documentation for quantizing HF models. + exllama_config (`Dict[str, Any]`, *optional*): + You can specify the version of the exllama kernel through the `version` key, the maximum sequence + length through the `max_input_len` key, and the maximum batch size through the `max_batch_size` key. + Defaults to `{"version": 2, "max_input_len": 2048, "max_batch_size": 8}` if unset. + """ + + def __init__( + self, + bits: int = 4, + group_size: int = 128, + zero_point: bool = True, + version: AWQLinearVersion = AWQLinearVersion.GEMM, + backend: AwqBackendPackingMethod = AwqBackendPackingMethod.AUTOAWQ, + do_fuse: Optional[bool] = None, + fuse_max_seq_len: Optional[int] = None, + modules_to_fuse: Optional[dict] = None, + modules_to_not_convert: Optional[List] = None, + exllama_config: Optional[Dict[str, int]] = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.AWQ + + self.bits = bits + self.group_size = group_size + self.zero_point = zero_point + self.version = version + self.backend = backend + self.fuse_max_seq_len = fuse_max_seq_len + self.modules_to_not_convert = modules_to_not_convert + self.exllama_config = exllama_config + + self.modules_to_fuse = modules_to_fuse + if do_fuse is None: + self.do_fuse = modules_to_fuse is not None and len(modules_to_fuse) > 0 + else: + self.do_fuse = do_fuse + self.fuse_max_seq_len = fuse_max_seq_len + + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct + """ + if self.backend not in [AwqBackendPackingMethod.AUTOAWQ, AwqBackendPackingMethod.LLMAWQ]: + raise ValueError( + f"Only supported quantization backends in {AwqBackendPackingMethod.AUTOAWQ} and {AwqBackendPackingMethod.LLMAWQ} - not recognized backend {self.backend}" + ) + + self.version = AWQLinearVersion.from_str(self.version) + if self.version not in [ + AWQLinearVersion.GEMM, + AWQLinearVersion.GEMV, + AWQLinearVersion.EXLLAMA, + AWQLinearVersion.IPEX, + ]: + raise ValueError( + f"Only supported versions are in [AWQLinearVersion.GEMM, AWQLinearVersion.GEMV, AWQLinearVersion.EXLLAMA, AWQLinearVersion.IPEX] - not recognized version {self.version}" + ) + + if self.backend == AwqBackendPackingMethod.LLMAWQ: + # Only cuda device can run this function + if not (torch.cuda.is_available() or torch.xpu.is_available()): + raise ValueError("LLM-AWQ backend is only supported on CUDA and XPU") + if torch.cuda.is_available(): + compute_capability = torch.cuda.get_device_capability() + major, minor = compute_capability + if major < 8: + raise ValueError("LLM-AWQ backend is only supported on CUDA GPUs with compute capability >= 8.0") + + if self.do_fuse and self.fuse_max_seq_len is None: + raise ValueError( + "You cannot enable fused modules without specifying a `fuse_max_seq_len`, make sure to pass a valid `fuse_max_seq_len` for your usecase" + ) + + if self.do_fuse: + awq_version_supports_fusing = False + MIN_AWQ_VERSION = "0.1.7" + if is_auto_awq_available(): + awq_version_supports_fusing = version.parse(importlib.metadata.version("autoawq")) >= version.parse( + MIN_AWQ_VERSION + ) + + if not awq_version_supports_fusing: + raise ValueError( + f"You current version of `autoawq` does not support module fusing, please upgrade `autoawq` package to at least {MIN_AWQ_VERSION}." + ) + + if self.modules_to_not_convert is not None: + awq_version_supports_non_conversion = False + MIN_AWQ_VERSION = "0.1.8" + if is_auto_awq_available(): + awq_version_supports_non_conversion = version.parse( + importlib.metadata.version("autoawq") + ) >= version.parse(MIN_AWQ_VERSION) + + if not awq_version_supports_non_conversion: + raise ValueError( + f"You current version of `autoawq` does not support module quantization skipping, please upgrade `autoawq` package to at least {MIN_AWQ_VERSION}." + ) + + if self.do_fuse and self.modules_to_fuse is not None: + required_keys = [ + "hidden_size", + "num_attention_heads", + "num_key_value_heads", + "mlp", + "attention", + "layernorm", + "use_alibi", + ] + if not all(key in self.modules_to_fuse for key in required_keys): + raise ValueError( + f"Required fields are missing in the fusing mapping, required fields are {required_keys}" + ) + + if self.version == AWQLinearVersion.EXLLAMA: + awq_version_supports_exllama = False + MIN_AWQ_VERSION = "0.2.0" + if is_auto_awq_available(): + awq_version_supports_exllama = version.parse(importlib.metadata.version("autoawq")) >= version.parse( + MIN_AWQ_VERSION + ) + + if not awq_version_supports_exllama: + raise ValueError( + f"You current version of `autoawq` does not support exllama backend, " + f"please upgrade `autoawq` package to at least {MIN_AWQ_VERSION}." + ) + + if self.exllama_config is None: + self.exllama_config = {"version": ExllamaVersion.TWO, "max_input_len": 2048, "max_batch_size": 8} + else: + if "version" not in self.exllama_config: + raise ValueError("`exllama_config` needs to have a `version` key.") + elif self.exllama_config["version"] not in [ExllamaVersion.ONE, ExllamaVersion.TWO]: + exllama_version = self.exllama_config["version"] + raise ValueError( + f"Only supported versions are in [ExllamaVersion.ONE, ExllamaVersion.TWO] - not recognized version {exllama_version}" + ) + + def get_loading_attributes(self): + attibutes_dict = copy.deepcopy(self.__dict__) + loading_attibutes = ["version", "do_fuse", "modules_to_fuse", "fuse_max_seq_len", "exllama_config"] + loading_attibutes_dict = {i: j for i, j in attibutes_dict.items() if i in loading_attibutes} + return loading_attibutes_dict + + +@dataclass +class AqlmConfig(QuantizationConfigMixin): + """ + This is a wrapper class about `aqlm` parameters. + + Args: + in_group_size (`int`, *optional*, defaults to 8): + The group size along the input dimension. + out_group_size (`int`, *optional*, defaults to 1): + The group size along the output dimension. It's recommended to always use 1. + num_codebooks (`int`, *optional*, defaults to 1): + Number of codebooks for the Additive Quantization procedure. + nbits_per_codebook (`int`, *optional*, defaults to 16): + Number of bits encoding a single codebook vector. Codebooks size is 2**nbits_per_codebook. + linear_weights_not_to_quantize (`Optional[List[str]]`, *optional*): + List of full paths of `nn.Linear` weight parameters that shall not be quantized. + kwargs (`Dict[str, Any]`, *optional*): + Additional parameters from which to initialize the configuration object. + """ + + def __init__( + self, + in_group_size: int = 8, + out_group_size: int = 1, + num_codebooks: int = 1, + nbits_per_codebook: int = 16, + linear_weights_not_to_quantize: Optional[List[str]] = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.AQLM + self.in_group_size = in_group_size + self.out_group_size = out_group_size + self.num_codebooks = num_codebooks + self.nbits_per_codebook = nbits_per_codebook + self.linear_weights_not_to_quantize = linear_weights_not_to_quantize + + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct - also replaces some NoneType arguments with their default values. + """ + if not isinstance(self.in_group_size, int): + raise TypeError("in_group_size must be a float") + if not isinstance(self.out_group_size, int): + raise TypeError("out_group_size must be a float") + if not isinstance(self.num_codebooks, int): + raise TypeError("num_codebooks must be a float") + if not isinstance(self.nbits_per_codebook, int): + raise TypeError("nbits_per_codebook must be a float") + + if self.linear_weights_not_to_quantize is not None and not isinstance( + self.linear_weights_not_to_quantize, list + ): + raise ValueError("linear_weights_not_to_quantize must be a list of strings") + + if self.linear_weights_not_to_quantize is None: + self.linear_weights_not_to_quantize = [] + + +@dataclass +class VptqLayerConfig(QuantizationConfigMixin): + """ + This is used to explain vptq config params for each layer + Args: + enable_norm (`bool`, *optional*, defaults to `True`): to control if we have scale/bias for fp-weight + enable_perm (`bool`, *optional*, defaults to `True`): to perm input_channel or not + group_num (`int`, *optional*, defaults to `1`): how many single groups for vector-quantization + group_size (`int`, *optional*, defaults to `-1`): depends on out-features + indices_as_float (`bool`, *optional*, defaults to `False`): for Finetuning + is_indice_packed (`bool`, *optional*, defaults to `True`): should always be True + num_centroids (`list`, *optional*, defaults to `[-1, -1]`): centriod numbers of clusters + num_res_centroids (`list`, *optional*, defaults to `[-1, -1]`): ditto for residual + outlier_size (`int`, *optional*, defaults to `1`): outliers + vector_lens (`list`, *optional*, defaults to `[-1, -1]`): centroid vector length in quantization + """ + + def __init__( + self, + enable_norm: bool = True, + enable_perm: bool = True, + group_num: int = 1, + group_size: int = -1, + in_features: int = -1, + indices_as_float: bool = False, + is_indice_packed: bool = True, + num_centroids: tuple = [-1, -1], + num_res_centroids: tuple = [-1, -1], + out_features: int = -1, + outlier_size: int = 0, + vector_lens: tuple = [-1, -1], + **kwargs, + ): + self.enable_norm = enable_norm + self.enable_perm = enable_perm + self.group_num = group_num + self.group_size = group_size + self.in_features = in_features + self.indices_as_float = indices_as_float + self.is_indice_packed = is_indice_packed + self.num_centroids = num_centroids + self.num_res_centroids = num_res_centroids + self.out_features = out_features + self.outlier_size = outlier_size + self.vector_lens = vector_lens + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct + """ + if self.is_indice_packed is False: + raise ValueError("is_indice_packed should always be True") + + +@dataclass +class VptqConfig(QuantizationConfigMixin): + """ + This is a wrapper class about `vptq` parameters. + + Args: + enable_proxy_error (`bool`, *optional*, defaults to `False`): calculate proxy error for each layer + config_for_layers (`Dict`, *optional*, defaults to `{}`): quantization params for each layer + shared_layer_config (`Dict`, *optional*, defaults to `{}`): shared quantization params among layers + modules_to_not_convert (`list`, *optional*, default to `None`): + The list of modules to not quantize, useful for quantizing models that explicitly require to have + some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers). + kwargs (`Dict[str, Any]`, *optional*): + Additional parameters from which to initialize the configuration object. + """ + + def __init__( + self, + enable_proxy_error: bool = False, + config_for_layers: Dict[str, Any] = {}, + shared_layer_config: Dict[str, Any] = {}, + modules_to_not_convert: Optional[List] = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.VPTQ + self.enable_proxy_error = enable_proxy_error + self.config_for_layers: Dict[str, Any] = config_for_layers + self.shared_layer_config: Dict[str, Any] = shared_layer_config + self.modules_to_not_convert = modules_to_not_convert + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct + """ + for layer_name, layer_param in self.config_for_layers.items(): + VptqLayerConfig(**layer_param) + if self.enable_proxy_error is True: + raise ValueError("enable_proxy_error should always be False until we support training") + + +@dataclass +class QuantoConfig(QuantizationConfigMixin): + """ + This is a wrapper class about all possible attributes and features that you can play with a model that has been + loaded using `quanto`. + + Args: + weights (`str`, *optional*, defaults to `"int8"`): + The target dtype for the weights after quantization. Supported values are ("float8","int8","int4","int2") + activations (`str`, *optional*): + The target dtype for the activations after quantization. Supported values are (None,"int8","float8") + modules_to_not_convert (`list`, *optional*, default to `None`): + The list of modules to not quantize, useful for quantizing models that explicitly require to have + some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers). + """ + + def __init__( + self, + weights="int8", + activations=None, + modules_to_not_convert: Optional[List] = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.QUANTO + self.weights = weights + self.activations = activations + self.modules_to_not_convert = modules_to_not_convert + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct + """ + accepted_weights = ["float8", "int8", "int4", "int2"] + accepted_activations = [None, "int8", "float8"] + if self.weights not in accepted_weights: + raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights}") + if self.activations not in accepted_activations: + raise ValueError(f"Only support weights in {accepted_activations} but found {self.activations}") + + +@dataclass +class EetqConfig(QuantizationConfigMixin): + """ + This is a wrapper class about all possible attributes and features that you can play with a model that has been + loaded using `eetq`. + + Args: + weights (`str`, *optional*, defaults to `"int8"`): + The target dtype for the weights. Supported value is only "int8" + modules_to_not_convert (`list`, *optional*, default to `None`): + The list of modules to not quantize, useful for quantizing models that explicitly require to have + some modules left in their original precision. + """ + + def __init__( + self, + weights: str = "int8", + modules_to_not_convert: Optional[List] = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.EETQ + self.weights = weights + self.modules_to_not_convert = modules_to_not_convert + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct + """ + accepted_weights = ["int8"] + if self.weights not in accepted_weights: + raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights}") + + +class CompressedTensorsConfig(QuantizationConfigMixin): + """ + This is a wrapper class that handles compressed-tensors quantization config options. + It is a wrapper around `compressed_tensors.QuantizationConfig` + Args: + config_groups (`typing.Dict[str, typing.Union[ForwardRef('QuantizationScheme'), typing.List[str]]]`, *optional*): + dictionary mapping group name to a quantization scheme definition + format (`str`, *optional*, defaults to `"dense"`): + format the model is represented as. Set `run_compressed` True to execute model as the + compressed format if not `dense` + quantization_status (`QuantizationStatus`, *optional*, defaults to `"initialized"`): + status of model in the quantization lifecycle, ie 'initialized', 'calibration', 'frozen' + kv_cache_scheme (`typing.Union[QuantizationArgs, NoneType]`, *optional*): + specifies quantization of the kv cache. If None, kv cache is not quantized. + global_compression_ratio (`typing.Union[float, NoneType]`, *optional*): + 0-1 float percentage of model compression + ignore (`typing.Union[typing.List[str], NoneType]`, *optional*): + layer names or types to not quantize, supports regex prefixed by 're:' + sparsity_config (`typing.Dict[str, typing.Any]`, *optional*): + configuration for sparsity compression + quant_method (`str`, *optional*, defaults to `"compressed-tensors"`): + do not override, should be compressed-tensors + run_compressed (`bool`, *optional*, defaults to `True`): alter submodules (usually linear) in order to + emulate compressed model execution if True, otherwise use default submodule + """ + + def __init__( + self, + config_groups: Dict[str, Union["QuantizationScheme", List[str]]] = None, # noqa: F821 + format: str = "dense", + quantization_status: "QuantizationStatus" = "initialized", # noqa: F821 + kv_cache_scheme: Optional["QuantizationArgs"] = None, # noqa: F821 + global_compression_ratio: Optional[float] = None, + ignore: Optional[List[str]] = None, + sparsity_config: Dict[str, Any] = None, + quant_method: str = "compressed-tensors", + run_compressed: bool = True, + **kwargs, + ): + if is_compressed_tensors_available(): + from compressed_tensors.config import SparsityCompressionConfig + from compressed_tensors.quantization import QuantizationConfig + else: + raise ImportError( + "compressed_tensors is not installed and is required for compressed-tensors quantization. Please install it with `pip install compressed-tensors`." + ) + self.quantization_config = None + self.sparsity_config = None + + self.run_compressed = run_compressed + + # parse from dict to load nested QuantizationScheme objects + if config_groups or kv_cache_scheme: + self.quantization_config = QuantizationConfig.model_validate( + { + "config_groups": config_groups, + "quant_method": quant_method, + "format": format, + "quantization_status": quantization_status, + "kv_cache_scheme": kv_cache_scheme, + "global_compression_ratio": global_compression_ratio, + "ignore": ignore, + "run_compressed": run_compressed, + **kwargs, + } + ) + + if sparsity_config: + self.sparsity_config = SparsityCompressionConfig.load_from_registry( + sparsity_config.get("format"), **sparsity_config + ) + + self.quant_method = QuantizationMethod.COMPRESSED_TENSORS + + def post_init(self): + if self.run_compressed: + if self.is_sparsification_compressed: + logger.warn( + "`run_compressed` is only supported for quantized_compressed models" + " and not for sparsified models. Setting `run_compressed=False`" + ) + self.run_compressed = False + elif not self.is_quantization_compressed: + logger.warn("`run_compressed` is only supported for compressed models. Setting `run_compressed=False`") + self.run_compressed = False + + @classmethod + def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs): + """ + Instantiates a [`CompressedTensorsConfig`] from a Python dictionary of parameters. + Optionally unwraps any args from the nested quantization_config + + Args: + config_dict (`Dict[str, Any]`): + Dictionary that will be used to instantiate the configuration object. + return_unused_kwargs (`bool`,*optional*, defaults to `False`): + Whether or not to return a list of unused keyword arguments. Used for `from_pretrained` method in + `PreTrainedModel`. + kwargs (`Dict[str, Any]`): + Additional parameters from which to initialize the configuration object. + + Returns: + [`QuantizationConfigMixin`]: The configuration object instantiated from those parameters. + + """ + + if "quantization_config" in config_dict: + config_dict = dict( + sparsity_config=config_dict.get("sparsity_config"), + **config_dict["quantization_config"], + ) + + return super().from_dict(config_dict, return_unused_kwargs=return_unused_kwargs, **kwargs) + + def to_dict(self) -> Dict[str, Any]: + """ + Quantization config to be added to config.json + + Serializes this instance to a Python dictionary. Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + quantization_config = {} + if self.quantization_config is not None: + quantization_config = self.quantization_config.dict() + else: + quantization_config["quant_method"] = QuantizationMethod.COMPRESSED_TENSORS + + if self.sparsity_config is not None: + quantization_config["sparsity_config"] = self.sparsity_config.dict() + else: + quantization_config["sparsity_config"] = {} + + return quantization_config + + def to_diff_dict(self) -> Dict[str, Any]: + """ + Removes all attributes from config which correspond to the default config attributes for better readability and + serializes to a Python dictionary. + Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance, + """ + config_dict = self.to_dict() + + # get the default config dict + default_config_dict = CompressedTensorsConfig().to_dict() + + serializable_config_dict = {} + + # only serialize values that differ from the default config + for key, value in config_dict.items(): + if key not in default_config_dict or value != default_config_dict[key]: + serializable_config_dict[key] = value + + return serializable_config_dict + + def get_loading_attributes(self): + return {"run_compressed": self.run_compressed} + + @property + def is_quantized(self): + return bool(self.quantization_config) and bool(self.quantization_config.config_groups) + + @property + def is_quantization_compressed(self): + from compressed_tensors.quantization import QuantizationStatus + + return self.is_quantized and self.quantization_config.quantization_status == QuantizationStatus.COMPRESSED + + @property + def is_sparsification_compressed(self): + from compressed_tensors.config import ( + CompressionFormat, + SparsityCompressionConfig, + ) + + return ( + isinstance(self.sparsity_config, SparsityCompressionConfig) + and self.sparsity_config.format != CompressionFormat.dense.value + ) + + +@dataclass +class FbgemmFp8Config(QuantizationConfigMixin): + """ + This is a wrapper class about all possible attributes and features that you can play with a model that has been + loaded using fbgemm fp8 quantization. + + Args: + activation_scale_ub (`float`, *optional*, defaults to 1200.0): + The activation scale upper bound. This is used when quantizing the input activation. + modules_to_not_convert (`list`, *optional*, default to `None`): + The list of modules to not quantize, useful for quantizing models that explicitly require to have + some modules left in their original precision. + """ + + def __init__( + self, + activation_scale_ub: float = 1200.0, + modules_to_not_convert: Optional[List] = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.FBGEMM_FP8 + self.activation_scale_ub = activation_scale_ub + self.modules_to_not_convert = modules_to_not_convert + + def get_loading_attributes(self): + attibutes_dict = copy.deepcopy(self.__dict__) + loading_attibutes = ["activation_scale_ub"] + loading_attibutes_dict = {i: j for i, j in attibutes_dict.items() if i in loading_attibutes} + return loading_attibutes_dict + + +@dataclass +class HiggsConfig(QuantizationConfigMixin): + """ + HiggsConfig is a configuration class for quantization using the HIGGS method. + + Args: + bits (int, *optional*, defaults to 4): + Number of bits to use for quantization. Can be 2, 3 or 4. Default is 4. + p (int, *optional*, defaults to 2): + Quantization grid dimension. 1 and 2 are supported. 2 is always better in practice. Default is 2. + modules_to_not_convert (`list`, *optional*, default to ["lm_head"]): + List of linear layers that should not be quantized. + hadamard_size (int, *optional*, defaults to 512): + Hadamard size for the HIGGS method. Default is 512. Input dimension of matrices is padded to this value. Decreasing this below 512 will reduce the quality of the quantization. + group_size (int, *optional*, defaults to 256): + Group size for the HIGGS method. Can be 64, 128 or 256. Decreasing it barely affects the performance. Default is 256. Must be a divisor of hadamard_size. + tune_metadata ('dict', *optional*, defaults to {}): + Module-wise metadata (gemm block shapes, GPU metadata, etc.) for saving the kernel tuning results. Default is an empty dictionary. Is set automatically during tuning. + """ + + def __init__( + self, + bits: int = 4, + p: int = 2, + modules_to_not_convert: Optional[List[str]] = None, + hadamard_size: int = 512, + group_size: int = 256, + tune_metadata: Optional[Dict[str, Any]] = None, + **kwargs, + ): + if tune_metadata is None: + tune_metadata = {} + self.quant_method = QuantizationMethod.HIGGS + self.bits = bits + self.p = p + self.modules_to_not_convert = modules_to_not_convert + self.hadamard_size = hadamard_size + self.group_size = group_size + self.tune_metadata = tune_metadata + + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct - also replaces some NoneType arguments with their default values. + """ + if self.bits not in [2, 3, 4]: + raise ValueError("bits must be 2, 3, or 4") + if self.p not in [1, 2]: + raise ValueError("p must be 1 or 2. 2 is always better in practice") + if self.group_size not in [64, 128, 256]: + raise ValueError("group_size must be 64, 128, or 256") + if self.hadamard_size % self.group_size != 0: + raise ValueError("hadamard_size must be divisible by group_size") + + +@dataclass +class TorchAoConfig(QuantizationConfigMixin): + quant_method: QuantizationMethod + quant_type: Union[str, "AOBaseConfig"] # noqa: F821 + modules_to_not_convert: Optional[List] + quant_type_kwargs: Dict[str, Any] + + """This is a config class for torchao quantization/sparsity techniques. + + Args: + quant_type (`Union[str, AOBaseConfig]`): + The type of quantization we want to use. Can be either: + - A string: currently supporting: `int4_weight_only`, `int8_weight_only` and `int8_dynamic_activation_int8_weight`. + - An AOBaseConfig instance: for more advanced configuration options. + modules_to_not_convert (`list`, *optional*, default to `None`): + The list of modules to not quantize, useful for quantizing models that explicitly require to have + some modules left in their original precision. + kwargs (`Dict[str, Any]`, *optional*): + The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments + `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in + https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques + + Example: + + ```python + # AOBaseConfig-based configuration + config = Int4WeightOnlyConfig(group_size=32) + quantization_config = TorchAoConfig(config) + model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config) + + # String-based configuration + quantization_config = TorchAoConfig("int4_weight_only", group_size=32) + # int4_weight_only quant is only working with *torch.bfloat16* dtype right now + model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config) + + # autoquant + # `autoquant` is a convenient way for users to search for the best quantization for each layer + # `min_sqnr` is an option to control the accuracy of the model, higher value means the model is more + # accurate, we can start with 30 and adjust it to larger or smaller (e.g. 40, 20) + # defaults to None, which means we'll try to get the best performing quantized model without + # considering accuracy + quantization_config = TorchAoConfig("autoquant", min_sqnr=30) + model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config) + # run through example inputs, quantization methods will be selected based on the shape of example input + tokenizer = AutoTokenizer.from_pretrained(model_name) + input_text = "What are we having for dinner?" + input_ids = tokenizer(input_text, return_tensors="pt").to("cuda") + MAX_NEW_TOKENS = 1000 + model.generate(**input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static") + # manually ran finalize_autoquant if needed + if hasattr(quantized_model, "finalize_autoquant"): + print("finalizing autoquant") + quantized_model.finalize_autoquant() + + ``` + """ + + def __init__( + self, + quant_type: Union[str, "AOBaseConfig"], # noqa: F821 + modules_to_not_convert: Optional[List] = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.TORCHAO + self.quant_type = quant_type + self.modules_to_not_convert = modules_to_not_convert + self.quant_type_kwargs = kwargs.get("quant_type_kwargs", kwargs) + self.post_init() + + @staticmethod + def _get_ao_version() -> version.Version: + """Centralized check for TorchAO availability and version requirements.""" + if not is_torchao_available(): + raise ValueError("TorchAoConfig requires torchao to be installed. Install with `pip install torchao`") + + return version.parse(importlib.metadata.version("torchao")) + + def post_init(self): + """Validate configuration and set defaults.""" + ao_version = self._get_ao_version() + + # Handle quant_type based on type and version + if isinstance(self.quant_type, str): + self._validate_string_quant_type() + elif ao_version > version.parse("0.9.0"): + from torchao.quantization.quant_api import AOBaseConfig + + if not isinstance(self.quant_type, AOBaseConfig): + raise ValueError( + f"quant_type must be either a string or an AOBaseConfig instance, got {type(self.quant_type)}" + ) + else: + raise ValueError( + f"In torchao <= 0.9.0, quant_type must be a string. Got {type(self.quant_type)}. " + f"Please upgrade to torchao > 0.9.0 to use AOBaseConfig instances." + ) + + def _validate_string_quant_type(self): + """Validate string quant_type and its kwargs.""" + methods = self._get_torchao_quant_type_to_method() + + if self.quant_type not in methods: + raise ValueError( + f"Unsupported string quantization type: {self.quant_type}. " + f"Supported types: {', '.join(methods.keys())}" + ) + + # Validate kwargs against method signature + method = methods[self.quant_type] + sig = signature(method) + valid_kwargs = { + param.name + for param in sig.parameters.values() + if param.kind in [Parameter.KEYWORD_ONLY, Parameter.POSITIONAL_OR_KEYWORD] + } + + invalid_kwargs = set(self.quant_type_kwargs) - valid_kwargs + if invalid_kwargs: + raise ValueError( + f"Unexpected keyword arg for {self.quant_type}: {', '.join(invalid_kwargs)}. " + f"Valid kwargs: {', '.join(valid_kwargs)}" + ) + + def _get_torchao_quant_type_to_method(self): + """Get mapping of quant_type strings to their corresponding methods.""" + from torchao.quantization import ( + autoquant, + int4_weight_only, + int8_dynamic_activation_int8_weight, + int8_weight_only, + ) + + return { + "int4_weight_only": int4_weight_only, + "int8_weight_only": int8_weight_only, + "int8_dynamic_activation_int8_weight": int8_dynamic_activation_int8_weight, + "autoquant": autoquant, + } + + def get_apply_tensor_subclass(self): + """Create the appropriate quantization method based on configuration.""" + if isinstance(self.quant_type, str): + methods = self._get_torchao_quant_type_to_method() + quant_type_kwargs = self.quant_type_kwargs.copy() + if ( + not torch.cuda.is_available() + and is_torchao_available() + and self.quant_type == "int4_weight_only" + and version.parse(importlib.metadata.version("torchao")) >= version.parse("0.8.0") + ): + from torchao.dtypes import Int4CPULayout + + quant_type_kwargs["layout"] = Int4CPULayout() + + return methods[self.quant_type](**quant_type_kwargs) + else: + return self.quant_type + + def to_dict(self): + """Convert configuration to a dictionary.""" + d = super().to_dict() + + if isinstance(self.quant_type, str): + # Handle layout serialization if present + if "quant_type_kwargs" in d and "layout" in d["quant_type_kwargs"]: + d["quant_type_kwargs"]["layout"] = dataclasses.asdict(d["quant_type_kwargs"]["layout"]) + else: + # Handle AOBaseConfig serialization + from torchao.core.config import config_to_dict + + # For now we assume there is 1 config per Transfomer, however in the future + # We may want to support a config per fqn. + d["quant_type"] = {"default": config_to_dict(self.quant_type)} + + return d + + @classmethod + def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs): + """Create configuration from a dictionary.""" + ao_verison = cls._get_ao_version() + assert ao_verison > version.parse("0.9.0"), "TorchAoConfig requires torchao > 0.9.0 for construction from dict" + config_dict = config_dict.copy() + quant_type = config_dict.pop("quant_type") + # Check if we only have one key which is "default" + # In the future we may update this + assert len(quant_type) == 1 and "default" in quant_type, ( + "Expected only one key 'default' in quant_type dictionary" + ) + quant_type = quant_type["default"] + + # Deserialize quant_type if needed + from torchao.core.config import config_from_dict + + quant_type = config_from_dict(quant_type) + + return cls(quant_type=quant_type, **config_dict) + + +@dataclass +class BitNetConfig(QuantizationConfigMixin): + def __init__( + self, + modules_to_not_convert: Optional[List] = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.BITNET + self.modules_to_not_convert = modules_to_not_convert + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct + """ + pass + + +@dataclass +class SpQRConfig(QuantizationConfigMixin): + """ + This is a wrapper class about `spqr` parameters. Refer to the original publication for more details. + + Args: + bits (`int`, *optional*, defaults to 3): + Specifies the bit count for the weights and first order zero-points and scales. + Currently only bits = 3 is supported. + beta1 (`int`, *optional*, defaults to 16): + SpQR tile width. Currently only beta1 = 16 is supported. + beta2 (`int`, *optional*, defaults to 16): + SpQR tile height. Currently only beta2 = 16 is supported. + shapes (`Optional`, *optional*): + A dictionary holding the shape of each object. We need this because it's impossible + to deduce the exact size of the parameters just from bits, beta1, beta2. + modules_to_not_convert (`Optional[List[str]]`, *optional*): + Optionally, provides a list of full paths of `nn.Linear` weight parameters that shall not be quantized. + Defaults to None. + kwargs (`Dict[str, Any]`, *optional*): + Additional parameters from which to initialize the configuration object. + """ + + def __init__( + self, + bits: int = 3, + beta1: int = 16, + beta2: int = 16, + shapes: Optional[Dict[str, int]] = None, + modules_to_not_convert: Optional[List[str]] = None, + **kwargs, + ): + if shapes is None: + shapes = {} + self.shapes = shapes + self.quant_method = QuantizationMethod.SPQR + self.bits = bits + self.beta1 = beta1 + self.beta2 = beta2 + self.modules_to_not_convert = modules_to_not_convert + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct - also replaces some NoneType arguments with their default values. + """ + if not isinstance(self.bits, int): + raise TypeError("bits must be an int") + if not isinstance(self.beta1, int): + raise TypeError("beta1 must be an int") + if not isinstance(self.beta2, int): + raise TypeError("beta2 must be an int") + + if self.bits != 3: + raise ValueError("SpQR currently only supports bits = 3") + if self.beta1 != 16: + raise ValueError("SpQR currently only supports beta1 = 16") + if self.beta2 != 16: + raise ValueError("SpQR currently only supports beta2 = 16") + if not isinstance(self.shapes, dict): + raise TypeError("shapes must be a dict") + + +@dataclass +class FineGrainedFP8Config(QuantizationConfigMixin): + """ + FineGrainedFP8Config is a configuration class for fine-grained FP8 quantization used mainly for deepseek models. + + Args: + activation_scheme (`str`, *optional*, defaults to `"dynamic"`): + The scheme used for activation, the defaults and only support scheme for now is "dynamic". + weight_block_size (`typing.Tuple[int, int]`, *optional*, defaults to `(128, 128)`): + The size of the weight blocks for quantization, default is (128, 128). + modules_to_not_convert (`list`, *optional*): + A list of module names that should not be converted during quantization. + """ + + def __init__( + self, + activation_scheme: str = "dynamic", + weight_block_size: Tuple[int, int] = (128, 128), + modules_to_not_convert: Optional[List] = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.FP8 + self.modules_to_not_convert = modules_to_not_convert + self.activation_scheme = activation_scheme + self.weight_block_size = weight_block_size + self.post_init() + + def post_init(self): + r""" + Safety checker that arguments are correct + """ + self.activation_scheme = self.activation_scheme.lower() + if self.activation_scheme not in ["dynamic"]: + raise ValueError(f"Activation scheme {self.activation_scheme} not supported") + if len(self.weight_block_size) != 2: + raise ValueError("weight_block_size must be a tuple of two integers") + if self.weight_block_size[0] <= 0 or self.weight_block_size[1] <= 0: + raise ValueError("weight_block_size must be a tuple of two positive integers") + + +class QuarkConfig(QuantizationConfigMixin): + def __init__( + self, + **kwargs, + ): + if is_torch_available() and is_quark_available(): + from quark import __version__ as quark_version + from quark.torch.export.config.config import JsonExporterConfig + from quark.torch.export.main_export.quant_config_parser import QuantConfigParser + from quark.torch.quantization.config.config import Config + + # This might be e.g. `"fp8"` or `"awq"`. + self.custom_mode = kwargs["quant_method"] + self.legacy = "export" not in kwargs + + if self.custom_mode in ["awq", "fp8"]: + # Legacy (quark<1.0) or custom export. + self.quant_config = QuantConfigParser.from_custom_config(kwargs, is_bias_quantized=False) + self.json_export_config = JsonExporterConfig() + else: + self.quant_config = Config.from_dict(kwargs) + + if "export" in kwargs: + # TODO: Remove this check once configuration version is handled natively by Quark. + if "min_kv_scale" in kwargs["export"] and version.parse(quark_version) < version.parse("0.8"): + min_kv_scale = kwargs["export"].pop("min_kv_scale") + logger.warning( + f"The parameter `min_kv_scale={min_kv_scale}` was found in the model config.json's `quantization_config.export` configuration, but this parameter is supported only for quark>=0.8. Ignoring this configuration parameter. Please update the `amd-quark` package." + ) + + self.json_export_config = JsonExporterConfig(**kwargs["export"]) + else: + # Legacy (quark<1.0) or custom export. + self.json_export_config = JsonExporterConfig() + + self.quant_method = QuantizationMethod.QUARK diff --git a/lib/python3.12/site-packages/transformers/utils/sentencepiece_model_pb2.py b/lib/python3.12/site-packages/transformers/utils/sentencepiece_model_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..b4b2992a6308c9f98f68486998d57d00bbdc1e34 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/sentencepiece_model_pb2.py @@ -0,0 +1,1511 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: sentencepiece_model.proto + +# Copyright 2022 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. +from google.protobuf import descriptor as _descriptor +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database + + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor.FileDescriptor( + name="sentencepiece_model.proto", + package="sentencepiece", + syntax="proto2", + serialized_options=b"H\003", + create_key=_descriptor._internal_create_key, + serialized_pb=( + b'\n\x19sentencepiece_model.proto\x12\rsentencepiece"\xa1\n\n\x0bTrainerSpec\x12\r\n\x05input\x18\x01' + b" \x03(\t\x12\x14\n\x0cinput_format\x18\x07 \x01(\t\x12\x14\n\x0cmodel_prefix\x18\x02" + b" \x01(\t\x12\x41\n\nmodel_type\x18\x03" + b" \x01(\x0e\x32$.sentencepiece.TrainerSpec.ModelType:\x07UNIGRAM\x12\x18\n\nvocab_size\x18\x04" + b" \x01(\x05:\x04\x38\x30\x30\x30\x12\x17\n\x0f\x61\x63\x63\x65pt_language\x18\x05 \x03(\t\x12" + b' \n\x15self_test_sample_size\x18\x06 \x01(\x05:\x01\x30\x12"\n\x12\x63haracter_coverage\x18\n' + b" \x01(\x02:\x06\x30.9995\x12\x1e\n\x13input_sentence_size\x18\x0b" + b" \x01(\x04:\x01\x30\x12$\n\x16shuffle_input_sentence\x18\x13 \x01(\x08:\x04true\x12" + b' \n\x14mining_sentence_size\x18\x0c \x01(\x05\x42\x02\x18\x01\x12"\n\x16training_sentence_size\x18\r' + b" \x01(\x05\x42\x02\x18\x01\x12(\n\x17seed_sentencepiece_size\x18\x0e" + b" \x01(\x05:\x07\x31\x30\x30\x30\x30\x30\x30\x12\x1e\n\x10shrinking_factor\x18\x0f" + b" \x01(\x02:\x04\x30.75\x12!\n\x13max_sentence_length\x18\x12" + b" \x01(\x05:\x04\x34\x31\x39\x32\x12\x17\n\x0bnum_threads\x18\x10" + b" \x01(\x05:\x02\x31\x36\x12\x1d\n\x12num_sub_iterations\x18\x11" + b" \x01(\x05:\x01\x32\x12$\n\x18max_sentencepiece_length\x18\x14" + b" \x01(\x05:\x02\x31\x36\x12%\n\x17split_by_unicode_script\x18\x15" + b" \x01(\x08:\x04true\x12\x1d\n\x0fsplit_by_number\x18\x17" + b" \x01(\x08:\x04true\x12!\n\x13split_by_whitespace\x18\x16" + b" \x01(\x08:\x04true\x12)\n\x1atreat_whitespace_as_suffix\x18\x18" + b" \x01(\x08:\x05\x66\x61lse\x12\x1b\n\x0csplit_digits\x18\x19" + b" \x01(\x08:\x05\x66\x61lse\x12\x17\n\x0f\x63ontrol_symbols\x18\x1e" + b" \x03(\t\x12\x1c\n\x14user_defined_symbols\x18\x1f \x03(\t\x12\x16\n\x0erequired_chars\x18$" + b" \x01(\t\x12\x1c\n\rbyte_fallback\x18# \x01(\x08:\x05\x66\x61lse\x12+\n\x1dvocabulary_output_piece_score\x18" + b' \x01(\x08:\x04true\x12\x1e\n\x10hard_vocab_limit\x18! \x01(\x08:\x04true\x12\x1c\n\ruse_all_vocab\x18"' + b" \x01(\x08:\x05\x66\x61lse\x12\x11\n\x06unk_id\x18( \x01(\x05:\x01\x30\x12\x11\n\x06\x62os_id\x18)" + b" \x01(\x05:\x01\x31\x12\x11\n\x06\x65os_id\x18* \x01(\x05:\x01\x32\x12\x12\n\x06pad_id\x18+" + b" \x01(\x05:\x02-1\x12\x18\n\tunk_piece\x18- \x01(\t:\x05\x12\x16\n\tbos_piece\x18." + b" \x01(\t:\x03\x12\x17\n\teos_piece\x18/ \x01(\t:\x04\x12\x18\n\tpad_piece\x18\x30" + b" \x01(\t:\x05\x12\x1a\n\x0bunk_surface\x18, \x01(\t:\x05 \xe2\x81\x87" + b" \x12+\n\x1ctrain_extremely_large_corpus\x18\x31" + b' \x01(\x08:\x05\x66\x61lse"5\n\tModelType\x12\x0b\n\x07UNIGRAM\x10\x01\x12\x07\n\x03\x42PE\x10\x02\x12\x08\n\x04WORD\x10\x03\x12\x08\n\x04\x43HAR\x10\x04*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"\xd1\x01\n\x0eNormalizerSpec\x12\x0c\n\x04name\x18\x01' + b" \x01(\t\x12\x1c\n\x14precompiled_charsmap\x18\x02 \x01(\x0c\x12\x1e\n\x10\x61\x64\x64_dummy_prefix\x18\x03" + b" \x01(\x08:\x04true\x12&\n\x18remove_extra_whitespaces\x18\x04 \x01(\x08:\x04true\x12" + b" \n\x12\x65scape_whitespaces\x18\x05 \x01(\x08:\x04true\x12\x1e\n\x16normalization_rule_tsv\x18\x06" + b' \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"y\n\x0cSelfTestData\x12\x33\n\x07samples\x18\x01' + b' \x03(\x0b\x32".sentencepiece.SelfTestData.Sample\x1a)\n\x06Sample\x12\r\n\x05input\x18\x01' + b" \x01(\t\x12\x10\n\x08\x65xpected\x18\x02" + b' \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"\xfe\x03\n\nModelProto\x12\x37\n\x06pieces\x18\x01' + b" \x03(\x0b\x32'.sentencepiece.ModelProto.SentencePiece\x12\x30\n\x0ctrainer_spec\x18\x02" + b" \x01(\x0b\x32\x1a.sentencepiece.TrainerSpec\x12\x36\n\x0fnormalizer_spec\x18\x03" + b" \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x12\x33\n\x0eself_test_data\x18\x04" + b" \x01(\x0b\x32\x1b.sentencepiece.SelfTestData\x12\x38\n\x11\x64\x65normalizer_spec\x18\x05" + b" \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x1a\xd2\x01\n\rSentencePiece\x12\r\n\x05piece\x18\x01" + b" \x01(\t\x12\r\n\x05score\x18\x02 \x01(\x02\x12\x42\n\x04type\x18\x03" + b' \x01(\x0e\x32,.sentencepiece.ModelProto.SentencePiece.Type:\x06NORMAL"T\n\x04Type\x12\n\n\x06NORMAL\x10\x01\x12\x0b\n\x07UNKNOWN\x10\x02\x12\x0b\n\x07\x43ONTROL\x10\x03\x12\x10\n\x0cUSER_DEFINED\x10\x04\x12\x08\n\x04\x42YTE\x10\x06\x12\n\n\x06UNUSED\x10\x05*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\x42\x02H\x03' + ), +) + + +_TRAINERSPEC_MODELTYPE = _descriptor.EnumDescriptor( + name="ModelType", + full_name="sentencepiece.TrainerSpec.ModelType", + filename=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + values=[ + _descriptor.EnumValueDescriptor( + name="UNIGRAM", + index=0, + number=1, + serialized_options=None, + type=None, + create_key=_descriptor._internal_create_key, + ), + _descriptor.EnumValueDescriptor( + name="BPE", + index=1, + number=2, + serialized_options=None, + type=None, + create_key=_descriptor._internal_create_key, + ), + _descriptor.EnumValueDescriptor( + name="WORD", + index=2, + number=3, + serialized_options=None, + type=None, + create_key=_descriptor._internal_create_key, + ), + _descriptor.EnumValueDescriptor( + name="CHAR", + index=3, + number=4, + serialized_options=None, + type=None, + create_key=_descriptor._internal_create_key, + ), + ], + containing_type=None, + serialized_options=None, + serialized_start=1294, + serialized_end=1347, +) +_sym_db.RegisterEnumDescriptor(_TRAINERSPEC_MODELTYPE) + +_MODELPROTO_SENTENCEPIECE_TYPE = _descriptor.EnumDescriptor( + name="Type", + full_name="sentencepiece.ModelProto.SentencePiece.Type", + filename=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + values=[ + _descriptor.EnumValueDescriptor( + name="NORMAL", + index=0, + number=1, + serialized_options=None, + type=None, + create_key=_descriptor._internal_create_key, + ), + _descriptor.EnumValueDescriptor( + name="UNKNOWN", + index=1, + number=2, + serialized_options=None, + type=None, + create_key=_descriptor._internal_create_key, + ), + _descriptor.EnumValueDescriptor( + name="CONTROL", + index=2, + number=3, + serialized_options=None, + type=None, + create_key=_descriptor._internal_create_key, + ), + _descriptor.EnumValueDescriptor( + name="USER_DEFINED", + index=3, + number=4, + serialized_options=None, + type=None, + create_key=_descriptor._internal_create_key, + ), + _descriptor.EnumValueDescriptor( + name="BYTE", + index=4, + number=6, + serialized_options=None, + type=None, + create_key=_descriptor._internal_create_key, + ), + _descriptor.EnumValueDescriptor( + name="UNUSED", + index=5, + number=5, + serialized_options=None, + type=None, + create_key=_descriptor._internal_create_key, + ), + ], + containing_type=None, + serialized_options=None, + serialized_start=2100, + serialized_end=2184, +) +_sym_db.RegisterEnumDescriptor(_MODELPROTO_SENTENCEPIECE_TYPE) + + +_TRAINERSPEC = _descriptor.Descriptor( + name="TrainerSpec", + full_name="sentencepiece.TrainerSpec", + filename=None, + file=DESCRIPTOR, + containing_type=None, + create_key=_descriptor._internal_create_key, + fields=[ + _descriptor.FieldDescriptor( + name="input", + full_name="sentencepiece.TrainerSpec.input", + index=0, + number=1, + type=9, + cpp_type=9, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="input_format", + full_name="sentencepiece.TrainerSpec.input_format", + index=1, + number=7, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="model_prefix", + full_name="sentencepiece.TrainerSpec.model_prefix", + index=2, + number=2, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="model_type", + full_name="sentencepiece.TrainerSpec.model_type", + index=3, + number=3, + type=14, + cpp_type=8, + label=1, + has_default_value=True, + default_value=1, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="vocab_size", + full_name="sentencepiece.TrainerSpec.vocab_size", + index=4, + number=4, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=8000, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="accept_language", + full_name="sentencepiece.TrainerSpec.accept_language", + index=5, + number=5, + type=9, + cpp_type=9, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="self_test_sample_size", + full_name="sentencepiece.TrainerSpec.self_test_sample_size", + index=6, + number=6, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="character_coverage", + full_name="sentencepiece.TrainerSpec.character_coverage", + index=7, + number=10, + type=2, + cpp_type=6, + label=1, + has_default_value=True, + default_value=float(0.9995), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="input_sentence_size", + full_name="sentencepiece.TrainerSpec.input_sentence_size", + index=8, + number=11, + type=4, + cpp_type=4, + label=1, + has_default_value=True, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="shuffle_input_sentence", + full_name="sentencepiece.TrainerSpec.shuffle_input_sentence", + index=9, + number=19, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=True, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="mining_sentence_size", + full_name="sentencepiece.TrainerSpec.mining_sentence_size", + index=10, + number=12, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=b"\030\001", + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="training_sentence_size", + full_name="sentencepiece.TrainerSpec.training_sentence_size", + index=11, + number=13, + type=5, + cpp_type=1, + label=1, + has_default_value=False, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=b"\030\001", + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="seed_sentencepiece_size", + full_name="sentencepiece.TrainerSpec.seed_sentencepiece_size", + index=12, + number=14, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=1000000, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="shrinking_factor", + full_name="sentencepiece.TrainerSpec.shrinking_factor", + index=13, + number=15, + type=2, + cpp_type=6, + label=1, + has_default_value=True, + default_value=float(0.75), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="max_sentence_length", + full_name="sentencepiece.TrainerSpec.max_sentence_length", + index=14, + number=18, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=4192, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="num_threads", + full_name="sentencepiece.TrainerSpec.num_threads", + index=15, + number=16, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=16, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="num_sub_iterations", + full_name="sentencepiece.TrainerSpec.num_sub_iterations", + index=16, + number=17, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=2, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="max_sentencepiece_length", + full_name="sentencepiece.TrainerSpec.max_sentencepiece_length", + index=17, + number=20, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=16, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="split_by_unicode_script", + full_name="sentencepiece.TrainerSpec.split_by_unicode_script", + index=18, + number=21, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=True, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="split_by_number", + full_name="sentencepiece.TrainerSpec.split_by_number", + index=19, + number=23, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=True, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="split_by_whitespace", + full_name="sentencepiece.TrainerSpec.split_by_whitespace", + index=20, + number=22, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=True, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="treat_whitespace_as_suffix", + full_name="sentencepiece.TrainerSpec.treat_whitespace_as_suffix", + index=21, + number=24, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=False, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="split_digits", + full_name="sentencepiece.TrainerSpec.split_digits", + index=22, + number=25, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=False, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="control_symbols", + full_name="sentencepiece.TrainerSpec.control_symbols", + index=23, + number=30, + type=9, + cpp_type=9, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="user_defined_symbols", + full_name="sentencepiece.TrainerSpec.user_defined_symbols", + index=24, + number=31, + type=9, + cpp_type=9, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="required_chars", + full_name="sentencepiece.TrainerSpec.required_chars", + index=25, + number=36, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="byte_fallback", + full_name="sentencepiece.TrainerSpec.byte_fallback", + index=26, + number=35, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=False, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="vocabulary_output_piece_score", + full_name="sentencepiece.TrainerSpec.vocabulary_output_piece_score", + index=27, + number=32, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=True, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="hard_vocab_limit", + full_name="sentencepiece.TrainerSpec.hard_vocab_limit", + index=28, + number=33, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=True, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="use_all_vocab", + full_name="sentencepiece.TrainerSpec.use_all_vocab", + index=29, + number=34, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=False, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="unk_id", + full_name="sentencepiece.TrainerSpec.unk_id", + index=30, + number=40, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=0, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="bos_id", + full_name="sentencepiece.TrainerSpec.bos_id", + index=31, + number=41, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=1, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="eos_id", + full_name="sentencepiece.TrainerSpec.eos_id", + index=32, + number=42, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=2, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="pad_id", + full_name="sentencepiece.TrainerSpec.pad_id", + index=33, + number=43, + type=5, + cpp_type=1, + label=1, + has_default_value=True, + default_value=-1, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="unk_piece", + full_name="sentencepiece.TrainerSpec.unk_piece", + index=34, + number=45, + type=9, + cpp_type=9, + label=1, + has_default_value=True, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="bos_piece", + full_name="sentencepiece.TrainerSpec.bos_piece", + index=35, + number=46, + type=9, + cpp_type=9, + label=1, + has_default_value=True, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="eos_piece", + full_name="sentencepiece.TrainerSpec.eos_piece", + index=36, + number=47, + type=9, + cpp_type=9, + label=1, + has_default_value=True, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="pad_piece", + full_name="sentencepiece.TrainerSpec.pad_piece", + index=37, + number=48, + type=9, + cpp_type=9, + label=1, + has_default_value=True, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="unk_surface", + full_name="sentencepiece.TrainerSpec.unk_surface", + index=38, + number=44, + type=9, + cpp_type=9, + label=1, + has_default_value=True, + default_value=b" \342\201\207 ".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="train_extremely_large_corpus", + full_name="sentencepiece.TrainerSpec.train_extremely_large_corpus", + index=39, + number=49, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=False, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + ], + extensions=[], + nested_types=[], + enum_types=[ + _TRAINERSPEC_MODELTYPE, + ], + serialized_options=None, + is_extendable=True, + syntax="proto2", + extension_ranges=[ + (200, 536870912), + ], + oneofs=[], + serialized_start=45, + serialized_end=1358, +) + + +_NORMALIZERSPEC = _descriptor.Descriptor( + name="NormalizerSpec", + full_name="sentencepiece.NormalizerSpec", + filename=None, + file=DESCRIPTOR, + containing_type=None, + create_key=_descriptor._internal_create_key, + fields=[ + _descriptor.FieldDescriptor( + name="name", + full_name="sentencepiece.NormalizerSpec.name", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="precompiled_charsmap", + full_name="sentencepiece.NormalizerSpec.precompiled_charsmap", + index=1, + number=2, + type=12, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"", + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="add_dummy_prefix", + full_name="sentencepiece.NormalizerSpec.add_dummy_prefix", + index=2, + number=3, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=True, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="remove_extra_whitespaces", + full_name="sentencepiece.NormalizerSpec.remove_extra_whitespaces", + index=3, + number=4, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=True, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="escape_whitespaces", + full_name="sentencepiece.NormalizerSpec.escape_whitespaces", + index=4, + number=5, + type=8, + cpp_type=7, + label=1, + has_default_value=True, + default_value=True, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="normalization_rule_tsv", + full_name="sentencepiece.NormalizerSpec.normalization_rule_tsv", + index=5, + number=6, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + serialized_options=None, + is_extendable=True, + syntax="proto2", + extension_ranges=[ + (200, 536870912), + ], + oneofs=[], + serialized_start=1361, + serialized_end=1570, +) + + +_SELFTESTDATA_SAMPLE = _descriptor.Descriptor( + name="Sample", + full_name="sentencepiece.SelfTestData.Sample", + filename=None, + file=DESCRIPTOR, + containing_type=None, + create_key=_descriptor._internal_create_key, + fields=[ + _descriptor.FieldDescriptor( + name="input", + full_name="sentencepiece.SelfTestData.Sample.input", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="expected", + full_name="sentencepiece.SelfTestData.Sample.expected", + index=1, + number=2, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + ], + extensions=[], + nested_types=[], + enum_types=[], + serialized_options=None, + is_extendable=False, + syntax="proto2", + extension_ranges=[], + oneofs=[], + serialized_start=1641, + serialized_end=1682, +) + +_SELFTESTDATA = _descriptor.Descriptor( + name="SelfTestData", + full_name="sentencepiece.SelfTestData", + filename=None, + file=DESCRIPTOR, + containing_type=None, + create_key=_descriptor._internal_create_key, + fields=[ + _descriptor.FieldDescriptor( + name="samples", + full_name="sentencepiece.SelfTestData.samples", + index=0, + number=1, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + ], + extensions=[], + nested_types=[ + _SELFTESTDATA_SAMPLE, + ], + enum_types=[], + serialized_options=None, + is_extendable=True, + syntax="proto2", + extension_ranges=[ + (200, 536870912), + ], + oneofs=[], + serialized_start=1572, + serialized_end=1693, +) + + +_MODELPROTO_SENTENCEPIECE = _descriptor.Descriptor( + name="SentencePiece", + full_name="sentencepiece.ModelProto.SentencePiece", + filename=None, + file=DESCRIPTOR, + containing_type=None, + create_key=_descriptor._internal_create_key, + fields=[ + _descriptor.FieldDescriptor( + name="piece", + full_name="sentencepiece.ModelProto.SentencePiece.piece", + index=0, + number=1, + type=9, + cpp_type=9, + label=1, + has_default_value=False, + default_value=b"".decode("utf-8"), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="score", + full_name="sentencepiece.ModelProto.SentencePiece.score", + index=1, + number=2, + type=2, + cpp_type=6, + label=1, + has_default_value=False, + default_value=float(0), + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="type", + full_name="sentencepiece.ModelProto.SentencePiece.type", + index=2, + number=3, + type=14, + cpp_type=8, + label=1, + has_default_value=True, + default_value=1, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + ], + extensions=[], + nested_types=[], + enum_types=[ + _MODELPROTO_SENTENCEPIECE_TYPE, + ], + serialized_options=None, + is_extendable=True, + syntax="proto2", + extension_ranges=[ + (200, 536870912), + ], + oneofs=[], + serialized_start=1985, + serialized_end=2195, +) + +_MODELPROTO = _descriptor.Descriptor( + name="ModelProto", + full_name="sentencepiece.ModelProto", + filename=None, + file=DESCRIPTOR, + containing_type=None, + create_key=_descriptor._internal_create_key, + fields=[ + _descriptor.FieldDescriptor( + name="pieces", + full_name="sentencepiece.ModelProto.pieces", + index=0, + number=1, + type=11, + cpp_type=10, + label=3, + has_default_value=False, + default_value=[], + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="trainer_spec", + full_name="sentencepiece.ModelProto.trainer_spec", + index=1, + number=2, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="normalizer_spec", + full_name="sentencepiece.ModelProto.normalizer_spec", + index=2, + number=3, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="self_test_data", + full_name="sentencepiece.ModelProto.self_test_data", + index=3, + number=4, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + _descriptor.FieldDescriptor( + name="denormalizer_spec", + full_name="sentencepiece.ModelProto.denormalizer_spec", + index=4, + number=5, + type=11, + cpp_type=10, + label=1, + has_default_value=False, + default_value=None, + message_type=None, + enum_type=None, + containing_type=None, + is_extension=False, + extension_scope=None, + serialized_options=None, + file=DESCRIPTOR, + create_key=_descriptor._internal_create_key, + ), + ], + extensions=[], + nested_types=[ + _MODELPROTO_SENTENCEPIECE, + ], + enum_types=[], + serialized_options=None, + is_extendable=True, + syntax="proto2", + extension_ranges=[ + (200, 536870912), + ], + oneofs=[], + serialized_start=1696, + serialized_end=2206, +) + +_TRAINERSPEC.fields_by_name["model_type"].enum_type = _TRAINERSPEC_MODELTYPE +_TRAINERSPEC_MODELTYPE.containing_type = _TRAINERSPEC +_SELFTESTDATA_SAMPLE.containing_type = _SELFTESTDATA +_SELFTESTDATA.fields_by_name["samples"].message_type = _SELFTESTDATA_SAMPLE +_MODELPROTO_SENTENCEPIECE.fields_by_name["type"].enum_type = _MODELPROTO_SENTENCEPIECE_TYPE +_MODELPROTO_SENTENCEPIECE.containing_type = _MODELPROTO +_MODELPROTO_SENTENCEPIECE_TYPE.containing_type = _MODELPROTO_SENTENCEPIECE +_MODELPROTO.fields_by_name["pieces"].message_type = _MODELPROTO_SENTENCEPIECE +_MODELPROTO.fields_by_name["trainer_spec"].message_type = _TRAINERSPEC +_MODELPROTO.fields_by_name["normalizer_spec"].message_type = _NORMALIZERSPEC +_MODELPROTO.fields_by_name["self_test_data"].message_type = _SELFTESTDATA +_MODELPROTO.fields_by_name["denormalizer_spec"].message_type = _NORMALIZERSPEC +DESCRIPTOR.message_types_by_name["TrainerSpec"] = _TRAINERSPEC +DESCRIPTOR.message_types_by_name["NormalizerSpec"] = _NORMALIZERSPEC +DESCRIPTOR.message_types_by_name["SelfTestData"] = _SELFTESTDATA +DESCRIPTOR.message_types_by_name["ModelProto"] = _MODELPROTO +_sym_db.RegisterFileDescriptor(DESCRIPTOR) + +TrainerSpec = _reflection.GeneratedProtocolMessageType( + "TrainerSpec", + (_message.Message,), + { + "DESCRIPTOR": _TRAINERSPEC, + "__module__": "sentencepiece_model_pb2", + # @@protoc_insertion_point(class_scope:sentencepiece.TrainerSpec) + }, +) +_sym_db.RegisterMessage(TrainerSpec) + +NormalizerSpec = _reflection.GeneratedProtocolMessageType( + "NormalizerSpec", + (_message.Message,), + { + "DESCRIPTOR": _NORMALIZERSPEC, + "__module__": "sentencepiece_model_pb2", + # @@protoc_insertion_point(class_scope:sentencepiece.NormalizerSpec) + }, +) +_sym_db.RegisterMessage(NormalizerSpec) + +SelfTestData = _reflection.GeneratedProtocolMessageType( + "SelfTestData", + (_message.Message,), + { + "Sample": _reflection.GeneratedProtocolMessageType( + "Sample", + (_message.Message,), + { + "DESCRIPTOR": _SELFTESTDATA_SAMPLE, + "__module__": "sentencepiece_model_pb2", + # @@protoc_insertion_point(class_scope:sentencepiece.SelfTestData.Sample) + }, + ), + "DESCRIPTOR": _SELFTESTDATA, + "__module__": "sentencepiece_model_pb2", + # @@protoc_insertion_point(class_scope:sentencepiece.SelfTestData) + }, +) +_sym_db.RegisterMessage(SelfTestData) +_sym_db.RegisterMessage(SelfTestData.Sample) + +ModelProto = _reflection.GeneratedProtocolMessageType( + "ModelProto", + (_message.Message,), + { + "SentencePiece": _reflection.GeneratedProtocolMessageType( + "SentencePiece", + (_message.Message,), + { + "DESCRIPTOR": _MODELPROTO_SENTENCEPIECE, + "__module__": "sentencepiece_model_pb2", + # @@protoc_insertion_point(class_scope:sentencepiece.ModelProto.SentencePiece) + }, + ), + "DESCRIPTOR": _MODELPROTO, + "__module__": "sentencepiece_model_pb2", + # @@protoc_insertion_point(class_scope:sentencepiece.ModelProto) + }, +) +_sym_db.RegisterMessage(ModelProto) +_sym_db.RegisterMessage(ModelProto.SentencePiece) + + +DESCRIPTOR._options = None +_TRAINERSPEC.fields_by_name["mining_sentence_size"]._options = None +_TRAINERSPEC.fields_by_name["training_sentence_size"]._options = None +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/transformers/utils/sentencepiece_model_pb2_new.py b/lib/python3.12/site-packages/transformers/utils/sentencepiece_model_pb2_new.py new file mode 100644 index 0000000000000000000000000000000000000000..2ea4f4d64ce660b3d59d60c0650f5849e4fd2251 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/sentencepiece_model_pb2_new.py @@ -0,0 +1,47 @@ +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: sentencepiece_model.proto +"""Generated protocol buffer code.""" + +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder + + +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x19sentencepiece_model.proto\x12\rsentencepiece"\x80\x0c\n\x0bTrainerSpec\x12\r\n\x05input\x18\x01 \x03(\t\x12\x14\n\x0cinput_format\x18\x07 \x01(\t\x12\x14\n\x0cmodel_prefix\x18\x02 \x01(\t\x12\x41\n\nmodel_type\x18\x03 \x01(\x0e\x32$.sentencepiece.TrainerSpec.ModelType:\x07UNIGRAM\x12\x18\n\nvocab_size\x18\x04 \x01(\x05:\x04\x38\x30\x30\x30\x12\x17\n\x0f\x61\x63\x63\x65pt_language\x18\x05 \x03(\t\x12 \n\x15self_test_sample_size\x18\x06 \x01(\x05:\x01\x30\x12*\n\x1b\x65nable_differential_privacy\x18\x32 \x01(\x08:\x05\x66\x61lse\x12+\n differential_privacy_noise_level\x18\x33 \x01(\x02:\x01\x30\x12\x32\n\'differential_privacy_clipping_threshold\x18\x34 \x01(\x04:\x01\x30\x12"\n\x12\x63haracter_coverage\x18\n \x01(\x02:\x06\x30.9995\x12\x1e\n\x13input_sentence_size\x18\x0b \x01(\x04:\x01\x30\x12$\n\x16shuffle_input_sentence\x18\x13 \x01(\x08:\x04true\x12 \n\x14mining_sentence_size\x18\x0c \x01(\x05\x42\x02\x18\x01\x12"\n\x16training_sentence_size\x18\r \x01(\x05\x42\x02\x18\x01\x12(\n\x17seed_sentencepiece_size\x18\x0e \x01(\x05:\x07\x31\x30\x30\x30\x30\x30\x30\x12\x1e\n\x10shrinking_factor\x18\x0f \x01(\x02:\x04\x30.75\x12!\n\x13max_sentence_length\x18\x12 \x01(\x05:\x04\x34\x31\x39\x32\x12\x17\n\x0bnum_threads\x18\x10 \x01(\x05:\x02\x31\x36\x12\x1d\n\x12num_sub_iterations\x18\x11 \x01(\x05:\x01\x32\x12$\n\x18max_sentencepiece_length\x18\x14 \x01(\x05:\x02\x31\x36\x12%\n\x17split_by_unicode_script\x18\x15 \x01(\x08:\x04true\x12\x1d\n\x0fsplit_by_number\x18\x17 \x01(\x08:\x04true\x12!\n\x13split_by_whitespace\x18\x16 \x01(\x08:\x04true\x12)\n\x1atreat_whitespace_as_suffix\x18\x18 \x01(\x08:\x05\x66\x61lse\x12+\n\x1c\x61llow_whitespace_only_pieces\x18\x1a \x01(\x08:\x05\x66\x61lse\x12\x1b\n\x0csplit_digits\x18\x19 \x01(\x08:\x05\x66\x61lse\x12#\n\x19pretokenization_delimiter\x18\x35 \x01(\t:\x00\x12\x17\n\x0f\x63ontrol_symbols\x18\x1e \x03(\t\x12\x1c\n\x14user_defined_symbols\x18\x1f \x03(\t\x12\x16\n\x0erequired_chars\x18$ \x01(\t\x12\x1c\n\rbyte_fallback\x18# \x01(\x08:\x05\x66\x61lse\x12+\n\x1dvocabulary_output_piece_score\x18 \x01(\x08:\x04true\x12\x1e\n\x10hard_vocab_limit\x18! \x01(\x08:\x04true\x12\x1c\n\ruse_all_vocab\x18" \x01(\x08:\x05\x66\x61lse\x12\x11\n\x06unk_id\x18( \x01(\x05:\x01\x30\x12\x11\n\x06\x62os_id\x18) \x01(\x05:\x01\x31\x12\x11\n\x06\x65os_id\x18* \x01(\x05:\x01\x32\x12\x12\n\x06pad_id\x18+ \x01(\x05:\x02-1\x12\x18\n\tunk_piece\x18- \x01(\t:\x05\x12\x16\n\tbos_piece\x18. \x01(\t:\x03\x12\x17\n\teos_piece\x18/ \x01(\t:\x04\x12\x18\n\tpad_piece\x18\x30 \x01(\t:\x05\x12\x1a\n\x0bunk_surface\x18, \x01(\t:\x05 \xe2\x81\x87 \x12+\n\x1ctrain_extremely_large_corpus\x18\x31 \x01(\x08:\x05\x66\x61lse"5\n\tModelType\x12\x0b\n\x07UNIGRAM\x10\x01\x12\x07\n\x03\x42PE\x10\x02\x12\x08\n\x04WORD\x10\x03\x12\x08\n\x04\x43HAR\x10\x04*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"\xd1\x01\n\x0eNormalizerSpec\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x1c\n\x14precompiled_charsmap\x18\x02 \x01(\x0c\x12\x1e\n\x10\x61\x64\x64_dummy_prefix\x18\x03 \x01(\x08:\x04true\x12&\n\x18remove_extra_whitespaces\x18\x04 \x01(\x08:\x04true\x12 \n\x12\x65scape_whitespaces\x18\x05 \x01(\x08:\x04true\x12\x1e\n\x16normalization_rule_tsv\x18\x06 \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"y\n\x0cSelfTestData\x12\x33\n\x07samples\x18\x01 \x03(\x0b\x32".sentencepiece.SelfTestData.Sample\x1a)\n\x06Sample\x12\r\n\x05input\x18\x01 \x01(\t\x12\x10\n\x08\x65xpected\x18\x02 \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02"\xfe\x03\n\nModelProto\x12\x37\n\x06pieces\x18\x01 \x03(\x0b\x32\'.sentencepiece.ModelProto.SentencePiece\x12\x30\n\x0ctrainer_spec\x18\x02 \x01(\x0b\x32\x1a.sentencepiece.TrainerSpec\x12\x36\n\x0fnormalizer_spec\x18\x03 \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x12\x33\n\x0eself_test_data\x18\x04 \x01(\x0b\x32\x1b.sentencepiece.SelfTestData\x12\x38\n\x11\x64\x65normalizer_spec\x18\x05 \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x1a\xd2\x01\n\rSentencePiece\x12\r\n\x05piece\x18\x01 \x01(\t\x12\r\n\x05score\x18\x02 \x01(\x02\x12\x42\n\x04type\x18\x03 \x01(\x0e\x32,.sentencepiece.ModelProto.SentencePiece.Type:\x06NORMAL"T\n\x04Type\x12\n\n\x06NORMAL\x10\x01\x12\x0b\n\x07UNKNOWN\x10\x02\x12\x0b\n\x07\x43ONTROL\x10\x03\x12\x10\n\x0cUSER_DEFINED\x10\x04\x12\x08\n\x04\x42YTE\x10\x06\x12\n\n\x06UNUSED\x10\x05*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\x42\x02H\x03' +) + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "sentencepiece_model_pb2", _globals) +if _descriptor._USE_C_DESCRIPTORS is False: + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b"H\003" + # (generated by protobuf compiler, but `_TRAINERSPEC` is not defined) + # _TRAINERSPEC.fields_by_name["mining_sentence_size"]._options = None + # _TRAINERSPEC.fields_by_name["mining_sentence_size"]._serialized_options = b"\030\001" + # _TRAINERSPEC.fields_by_name["training_sentence_size"]._options = None + # _TRAINERSPEC.fields_by_name["training_sentence_size"]._serialized_options = b"\030\001" + _globals["_TRAINERSPEC"]._serialized_start = 45 + _globals["_TRAINERSPEC"]._serialized_end = 1581 + _globals["_TRAINERSPEC_MODELTYPE"]._serialized_start = 1517 + _globals["_TRAINERSPEC_MODELTYPE"]._serialized_end = 1570 + _globals["_NORMALIZERSPEC"]._serialized_start = 1584 + _globals["_NORMALIZERSPEC"]._serialized_end = 1793 + _globals["_SELFTESTDATA"]._serialized_start = 1795 + _globals["_SELFTESTDATA"]._serialized_end = 1916 + _globals["_SELFTESTDATA_SAMPLE"]._serialized_start = 1864 + _globals["_SELFTESTDATA_SAMPLE"]._serialized_end = 1905 + _globals["_MODELPROTO"]._serialized_start = 1919 + _globals["_MODELPROTO"]._serialized_end = 2429 + _globals["_MODELPROTO_SENTENCEPIECE"]._serialized_start = 2208 + _globals["_MODELPROTO_SENTENCEPIECE"]._serialized_end = 2418 + _globals["_MODELPROTO_SENTENCEPIECE_TYPE"]._serialized_start = 2323 + _globals["_MODELPROTO_SENTENCEPIECE_TYPE"]._serialized_end = 2407 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/transformers/utils/versions.py b/lib/python3.12/site-packages/transformers/utils/versions.py new file mode 100644 index 0000000000000000000000000000000000000000..85dabec8b659022655ce570b4642febd715b6d52 --- /dev/null +++ b/lib/python3.12/site-packages/transformers/utils/versions.py @@ -0,0 +1,117 @@ +# Copyright 2020 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. +""" +Utilities for working with package versions +""" + +import importlib.metadata +import operator +import re +import sys +from typing import Optional + +from packaging import version + + +ops = { + "<": operator.lt, + "<=": operator.le, + "==": operator.eq, + "!=": operator.ne, + ">=": operator.ge, + ">": operator.gt, +} + + +def _compare_versions(op, got_ver, want_ver, requirement, pkg, hint): + if got_ver is None or want_ver is None: + raise ValueError( + f"Unable to compare versions for {requirement}: need={want_ver} found={got_ver}. This is unusual. Consider" + f" reinstalling {pkg}." + ) + if not ops[op](version.parse(got_ver), version.parse(want_ver)): + raise ImportError( + f"{requirement} is required for a normal functioning of this module, but found {pkg}=={got_ver}.{hint}" + ) + + +def require_version(requirement: str, hint: Optional[str] = None) -> None: + """ + Perform a runtime check of the dependency versions, using the exact same syntax used by pip. + + The installed module version comes from the *site-packages* dir via *importlib.metadata*. + + Args: + requirement (`str`): pip style definition, e.g., "tokenizers==0.9.4", "tqdm>=4.27", "numpy" + hint (`str`, *optional*): what suggestion to print in case of requirements not being met + + Example: + + ```python + require_version("pandas>1.1.2") + require_version("numpy>1.18.5", "this is important to have for whatever reason") + ```""" + + hint = f"\n{hint}" if hint is not None else "" + + # non-versioned check + if re.match(r"^[\w_\-\d]+$", requirement): + pkg, op, want_ver = requirement, None, None + else: + match = re.findall(r"^([^!=<>\s]+)([\s!=<>]{1,2}.+)", requirement) + if not match: + raise ValueError( + "requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23, but" + f" got {requirement}" + ) + pkg, want_full = match[0] + want_range = want_full.split(",") # there could be multiple requirements + wanted = {} + for w in want_range: + match = re.findall(r"^([\s!=<>]{1,2})(.+)", w) + if not match: + raise ValueError( + "requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23," + f" but got {requirement}" + ) + op, want_ver = match[0] + wanted[op] = want_ver + if op not in ops: + raise ValueError(f"{requirement}: need one of {list(ops.keys())}, but got {op}") + + # special case + if pkg == "python": + got_ver = ".".join([str(x) for x in sys.version_info[:3]]) + for op, want_ver in wanted.items(): + _compare_versions(op, got_ver, want_ver, requirement, pkg, hint) + return + + # check if any version is installed + try: + got_ver = importlib.metadata.version(pkg) + except importlib.metadata.PackageNotFoundError: + raise importlib.metadata.PackageNotFoundError( + f"The '{requirement}' distribution was not found and is required by this application. {hint}" + ) + + # check that the right version is installed if version number or a range was provided + if want_ver is not None: + for op, want_ver in wanted.items(): + _compare_versions(op, got_ver, want_ver, requirement, pkg, hint) + + +def require_version_core(requirement): + """require_version wrapper which emits a core-specific hint on failure""" + hint = "Try: `pip install transformers -U` or `pip install -e '.[dev]'` if you're working with git main" + return require_version(requirement, hint)