Aluode/PerceptionLabPortable
0
1# Copyright 2021 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14import copy15import dataclasses16import warnings17from abc import ABC, abstractmethod18from collections import OrderedDict19from collections.abc import Iterable, Mapping20from typing import TYPE_CHECKING, Any, Callable, Optional, Union21 22import numpy as np23from packaging import version24 25from ..utils import TensorType, is_torch_available, is_vision_available, logging26from .utils import ParameterFormat, compute_effective_axis_dimension, compute_serialized_parameters_size27 28 29if TYPE_CHECKING:30 from ..configuration_utils import PretrainedConfig31 from ..feature_extraction_utils import FeatureExtractionMixin32 from ..image_processing_utils import ImageProcessingMixin33 from ..tokenization_utils_base import PreTrainedTokenizerBase34 35 36if is_vision_available():37 from PIL import Image38 39logger = logging.get_logger(__name__)40 41 42DEFAULT_ONNX_OPSET = 1143 44# 2 Gb45EXTERNAL_DATA_FORMAT_SIZE_LIMIT = 2 * 1024 * 1024 * 102446 47 48@dataclasses.dataclass49class PatchingSpec:50 """51 Data class that holds patching specifications.52 53 Args:54 o: Module / object where the op to patch is located55 name: Name of the op to monkey patch56 custom_op: Custom op that patches the original op57 orig_op: Original op that is being patched58 op_wrapper: Wrapper (optional) that wraps both the original and custom ops.59 It is useful for ops that are class or static methods for instance.60 """61 62 o: Any63 name: str64 custom_op: Callable65 orig_op: Optional[Callable] = None66 op_wrapper: Optional[Callable] = None67 68 69class OnnxConfig(ABC):70 """71 Base class for ONNX exportable model describing metadata on how to export the model through the ONNX format.72 """73 74 default_fixed_batch = 275 default_fixed_sequence = 876 default_fixed_num_choices = 477 torch_onnx_minimum_version = version.parse("1.8")78 _tasks_to_common_outputs = {79 "causal-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),80 "default": OrderedDict({"last_hidden_state": {0: "batch", 1: "sequence"}}),81 "image-classification": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),82 "image-segmentation": OrderedDict(83 {84 "logits": {0: "batch", 1: "sequence"},85 "pred_boxes": {0: "batch", 1: "sequence"},86 "pred_masks": {0: "batch", 1: "sequence"},87 }88 ),89 "masked-im": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),90 "masked-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),91 "multiple-choice": OrderedDict({"logits": {0: "batch"}}),92 "object-detection": OrderedDict(93 {94 "logits": {0: "batch", 1: "sequence"},95 "pred_boxes": {0: "batch", 1: "sequence"},96 }97 ),98 "question-answering": OrderedDict(99 {100 "start_logits": {0: "batch", 1: "sequence"},101 "end_logits": {0: "batch", 1: "sequence"},102 }103 ),104 "semantic-segmentation": OrderedDict({"logits": {0: "batch", 1: "num_labels", 2: "height", 3: "width"}}),105 "seq2seq-lm": OrderedDict({"logits": {0: "batch", 1: "decoder_sequence"}}),106 "sequence-classification": OrderedDict({"logits": {0: "batch"}}),107 "token-classification": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),108 "vision2seq-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),109 "speech2seq-lm": OrderedDict({"logits": {0: "batch", 1: "sequence"}}),110 }111 112 def __init__(113 self, config: "PretrainedConfig", task: str = "default", patching_specs: Optional[list[PatchingSpec]] = None114 ):115 self._config = config116 117 if task not in self._tasks_to_common_outputs:118 raise ValueError(119 f"{task} is not a supported task, supported tasks: {self._tasks_to_common_outputs.keys()}"120 )121 self.task = task122 123 self._patching_specs = []124 for spec in patching_specs if patching_specs is not None else []:125 final_spec = spec126 if spec.orig_op is None:127 final_spec = dataclasses.replace(spec, orig_op=getattr(spec.o, spec.name))128 self._patching_specs.append(final_spec)129 130 @classmethod131 def from_model_config(cls, config: "PretrainedConfig", task: str = "default") -> "OnnxConfig":132 """133 Instantiate a OnnxConfig for a specific model134 135 Args:136 config: The model's configuration to use when exporting to ONNX137 138 Returns:139 OnnxConfig for this model140 """141 return cls(config, task=task)142 143 @property144 @abstractmethod145 def inputs(self) -> Mapping[str, Mapping[int, str]]:146 """147 Mapping containing the axis definition of the input tensors to provide to the model148 149 Returns:150 For each input: its name associated to the axes symbolic name and the axis position within the tensor151 """152 raise NotImplementedError()153 154 @property155 def outputs(self) -> Mapping[str, Mapping[int, str]]:156 """157 Mapping containing the axis definition of the output tensors to provide to the model158 159 Returns:160 For each output: its name associated to the axes symbolic name and the axis position within the tensor161 """162 common_outputs = self._tasks_to_common_outputs[self.task]163 return copy.deepcopy(common_outputs)164 165 @property166 def values_override(self) -> Optional[Mapping[str, Any]]:167 """168 Dictionary of keys to override in the model's config before exporting169 170 Returns:171 Dictionary with the keys (and their corresponding values) to override172 """173 if hasattr(self._config, "use_cache"):174 return {"use_cache": False}175 176 return None177 178 @property179 def default_batch_size(self) -> int:180 """181 The default batch size to use if no other indication182 183 Returns:184 Integer > 0185 """186 # Using 2 avoid ONNX making assumption about single sample batch187 return OnnxConfig.default_fixed_batch188 189 @property190 def default_sequence_length(self) -> int:191 """192 The default sequence length to use if no other indication193 194 Returns:195 Integer > 0196 """197 return OnnxConfig.default_fixed_sequence198 199 @property200 def default_num_choices(self) -> int:201 """202 The default number of choices to use if no other indication203 204 Returns:205 Integer > 0206 """207 return OnnxConfig.default_fixed_num_choices208 209 @property210 def default_onnx_opset(self) -> int:211 """212 Which onnx opset to use when exporting the model213 214 Returns:215 Integer ONNX Opset version216 """217 return DEFAULT_ONNX_OPSET218 219 @property220 def atol_for_validation(self) -> float:221 """222 What absolute tolerance value to use during model conversion validation.223 224 Returns:225 Float absolute tolerance value.226 """227 return 1e-5228 229 @property230 def is_torch_support_available(self) -> bool:231 """232 The minimum PyTorch version required to export the model.233 234 Returns:235 `bool`: Whether the installed version of PyTorch is compatible with the model.236 """237 if is_torch_available():238 from transformers.utils import get_torch_version239 240 return version.parse(get_torch_version()) >= self.torch_onnx_minimum_version241 else:242 return False243 244 @staticmethod245 def use_external_data_format(num_parameters: int) -> bool:246 """247 Flag indicating if the model requires using external data format248 249 Args:250 num_parameters: Number of parameter on the model251 252 Returns:253 True if model.num_parameters() * size_of(float32) >= 2Gb False otherwise254 """255 256 return (257 compute_serialized_parameters_size(num_parameters, ParameterFormat.Float)258 >= EXTERNAL_DATA_FORMAT_SIZE_LIMIT259 )260 261 def _generate_dummy_images(262 self, batch_size: int = 2, num_channels: int = 3, image_height: int = 40, image_width: int = 40263 ):264 images = []265 for _ in range(batch_size):266 data = np.random.rand(image_height, image_width, num_channels) * 255267 images.append(Image.fromarray(data.astype("uint8")).convert("RGB"))268 return images269 270 def _generate_dummy_audio(271 self, batch_size: int = 2, sampling_rate: int = 22050, time_duration: float = 5.0, frequency: int = 220272 ):273 audio_data = []274 for _ in range(batch_size):275 # time variable276 t = np.linspace(0, time_duration, int(time_duration * sampling_rate), endpoint=False)277 278 # generate pure sine wave at `frequency` Hz279 audio_data.append(0.5 * np.sin(2 * np.pi * frequency * t))280 281 return audio_data282 283 def generate_dummy_inputs(284 self,285 preprocessor: Union["PreTrainedTokenizerBase", "FeatureExtractionMixin", "ImageProcessingMixin"],286 batch_size: int = -1,287 seq_length: int = -1,288 num_choices: int = -1,289 is_pair: bool = False,290 framework: Optional[TensorType] = None,291 num_channels: int = 3,292 image_width: int = 40,293 image_height: int = 40,294 sampling_rate: int = 22050,295 time_duration: float = 5.0,296 frequency: int = 220,297 tokenizer: Optional["PreTrainedTokenizerBase"] = None,298 ) -> Mapping[str, Any]:299 """300 Generate inputs to provide to the ONNX exporter for the specific framework301 302 Args:303 preprocessor: ([`PreTrainedTokenizerBase`], [`FeatureExtractionMixin`], or [`ImageProcessingMixin`]):304 The preprocessor associated with this model configuration.305 batch_size (`int`, *optional*, defaults to -1):306 The batch size to export the model for (-1 means dynamic axis).307 num_choices (`int`, *optional*, defaults to -1):308 The number of candidate answers provided for multiple choice task (-1 means dynamic axis).309 seq_length (`int`, *optional*, defaults to -1):310 The sequence length to export the model for (-1 means dynamic axis).311 is_pair (`bool`, *optional*, defaults to `False`):312 Indicate if the input is a pair (sentence 1, sentence 2)313 framework (`TensorType`, *optional*, defaults to `None`):314 The framework (PyTorch or TensorFlow) that the tokenizer will generate tensors for.315 num_channels (`int`, *optional*, defaults to 3):316 The number of channels of the generated images.317 image_width (`int`, *optional*, defaults to 40):318 The width of the generated images.319 image_height (`int`, *optional*, defaults to 40):320 The height of the generated images.321 sampling_rate (`int`, *optional* defaults to 22050)322 The sampling rate for audio data generation.323 time_duration (`float`, *optional* defaults to 5.0)324 Total seconds of sampling for audio data generation.325 frequency (`int`, *optional* defaults to 220)326 The desired natural frequency of generated audio.327 328 Returns:329 Mapping[str, Tensor] holding the kwargs to provide to the model's forward function330 """331 from ..feature_extraction_utils import FeatureExtractionMixin332 from ..image_processing_utils import ImageProcessingMixin333 from ..tokenization_utils_base import PreTrainedTokenizerBase334 335 if isinstance(preprocessor, PreTrainedTokenizerBase) and tokenizer is not None:336 raise ValueError("You cannot provide both a tokenizer and a preprocessor to generate dummy inputs.")337 if tokenizer is not None:338 warnings.warn(339 "The `tokenizer` argument is deprecated and will be removed in version 5 of Transformers. Use"340 " `preprocessor` instead.",341 FutureWarning,342 )343 logger.warning("Overwriting the `preprocessor` argument with `tokenizer` to generate dummy inputs.")344 preprocessor = tokenizer345 if isinstance(preprocessor, PreTrainedTokenizerBase):346 # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX347 batch_size = compute_effective_axis_dimension(348 batch_size, fixed_dimension=OnnxConfig.default_fixed_batch, num_token_to_add=0349 )350 # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX351 token_to_add = preprocessor.num_special_tokens_to_add(is_pair)352 seq_length = compute_effective_axis_dimension(353 seq_length, fixed_dimension=OnnxConfig.default_fixed_sequence, num_token_to_add=token_to_add354 )355 # Generate dummy inputs according to compute batch and sequence356 input_token = (357 preprocessor.unk_token358 if (preprocessor.unk_token is not None and len(preprocessor.unk_token) > 0)359 else "0"360 )361 dummy_input = [" ".join([input_token]) * seq_length] * batch_size362 if self.task == "multiple-choice":363 # If dynamic axis (-1) we forward with a fixed dimension of 4 candidate answers to avoid optimizations364 # made by ONNX365 num_choices = compute_effective_axis_dimension(366 num_choices, fixed_dimension=OnnxConfig.default_fixed_num_choices, num_token_to_add=0367 )368 dummy_input = dummy_input * num_choices369 # The shape of the tokenized inputs values is [batch_size * num_choices, seq_length]370 tokenized_input = preprocessor(dummy_input, text_pair=dummy_input)371 # Unflatten the tokenized inputs values expanding it to the shape [batch_size, num_choices, seq_length]372 for k, v in tokenized_input.items():373 tokenized_input[k] = [v[i : i + num_choices] for i in range(0, len(v), num_choices)]374 return dict(tokenized_input.convert_to_tensors(tensor_type=framework))375 return dict(preprocessor(dummy_input, return_tensors=framework))376 elif isinstance(preprocessor, ImageProcessingMixin):377 if preprocessor.model_input_names[0] != "pixel_values":378 raise ValueError(379 f"The `preprocessor` is an image processor ({preprocessor.__class__.__name__}) and expects"380 f' `model_input_names[0]` to be "pixel_values", but got {preprocessor.model_input_names[0]}'381 )382 # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX383 batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch)384 dummy_input = self._generate_dummy_images(batch_size, num_channels, image_height, image_width)385 return dict(preprocessor(images=dummy_input, return_tensors=framework))386 elif isinstance(preprocessor, FeatureExtractionMixin) and preprocessor.model_input_names[0] == "pixel_values":387 # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX388 batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch)389 dummy_input = self._generate_dummy_images(batch_size, num_channels, image_height, image_width)390 return dict(preprocessor(images=dummy_input, return_tensors=framework))391 elif (392 isinstance(preprocessor, FeatureExtractionMixin) and preprocessor.model_input_names[0] == "input_features"393 ):394 # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX395 batch_size = compute_effective_axis_dimension(batch_size, fixed_dimension=OnnxConfig.default_fixed_batch)396 dummy_input = self._generate_dummy_audio(batch_size, sampling_rate, time_duration, frequency)397 return dict(preprocessor(dummy_input, return_tensors=framework))398 else:399 raise ValueError(400 "Unable to generate dummy inputs for the model. Please provide a tokenizer or a preprocessor."401 )402 403 def generate_dummy_inputs_onnxruntime(self, reference_model_inputs: Mapping[str, Any]) -> Mapping[str, Any]:404 """405 Generate inputs for ONNX Runtime using the reference model inputs. Override this to run inference with seq2seq406 models which have the encoder and decoder exported as separate ONNX files.407 408 Args:409 reference_model_inputs ([`Mapping[str, Tensor]`):410 Reference inputs for the model.411 412 Returns:413 `Mapping[str, Tensor]`: The mapping holding the kwargs to provide to the model's forward function414 """415 return reference_model_inputs416 417 def patch_ops(self):418 for spec in self._patching_specs:419 custom_op = spec.custom_op if spec.op_wrapper is None else spec.op_wrapper(spec.custom_op)420 setattr(spec.o, spec.name, custom_op)421 422 def restore_ops(self):423 for spec in self._patching_specs:424 orig_op = spec.orig_op if spec.op_wrapper is None else spec.op_wrapper(spec.orig_op)425 setattr(spec.o, spec.name, orig_op)426 427 @classmethod428 def flatten_output_collection_property(cls, name: str, field: Iterable[Any]) -> dict[str, Any]:429 """430 Flatten any potential nested structure expanding the name of the field with the index of the element within the431 structure.432 433 Args:434 name: The name of the nested structure435 field: The structure to, potentially, be flattened436 437 Returns:438 (dict[str, Any]): Outputs with flattened structure and key mapping this new structure.439 440 """441 from itertools import chain442 443 return {f"{name}.{idx}": item for idx, item in enumerate(chain.from_iterable(field))}444 445 446class OnnxConfigWithPast(OnnxConfig, ABC):447 def __init__(448 self,449 config: "PretrainedConfig",450 task: str = "default",451 patching_specs: Optional[list[PatchingSpec]] = None,452 use_past: bool = False,453 ):454 super().__init__(config, task=task, patching_specs=patching_specs)455 self.use_past = use_past456 457 @classmethod458 def with_past(cls, config: "PretrainedConfig", task: str = "default") -> "OnnxConfigWithPast":459 """460 Instantiate a OnnxConfig with `use_past` attribute set to True461 462 Args:463 config: The underlying model's config to use when exporting to ONNX464 465 Returns:466 OnnxConfig with `.use_past = True`467 """468 return cls(config, task=task, use_past=True)469 470 @property471 def outputs(self) -> Mapping[str, Mapping[int, str]]:472 common_outputs = super().outputs473 if self.use_past:474 self.fill_with_past_key_values_(common_outputs, direction="outputs")475 476 return common_outputs477 478 @property479 def values_override(self) -> Optional[Mapping[str, Any]]:480 if hasattr(self._config, "use_cache"):481 return {"use_cache": self.use_past}482 483 return None484 485 @property486 def num_layers(self) -> int:487 """488 The number of layers attribute retrieved from the model config. Override this for model configs where the489 number of layers attribute is not called `num_layers`.490 """491 if not hasattr(self._config, "num_layers"):492 raise AttributeError(493 "could not find the number of layers attribute in the model configuration, override the num_layers"494 " property of the model OnnxConfig to solve this"495 )496 return self._config.num_layers497 498 @property499 def num_attention_heads(self) -> int:500 """501 The number of attention heads attribute retrieved from the model config. Override this for model configs where502 the number of attention heads attribute is not called `num_attention_heads`.503 """504 if not hasattr(self._config, "num_attention_heads"):505 raise AttributeError(506 "could not find the number of attention heads attribute in the model configuration, override the"507 " num_attention_heads property of the model OnnxConfig to solve this"508 )509 return self._config.num_attention_heads510 511 def generate_dummy_inputs(512 self,513 tokenizer: "PreTrainedTokenizerBase",514 batch_size: int = -1,515 seq_length: int = -1,516 is_pair: bool = False,517 framework: Optional[TensorType] = None,518 ) -> Mapping[str, Any]:519 # TODO: should we set seq_length = 1 when self.use_past = True?520 common_inputs = super().generate_dummy_inputs(521 tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework522 )523 524 if self.use_past:525 if not is_torch_available():526 raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")527 else:528 import torch529 530 batch, seqlen = common_inputs["input_ids"].shape531 # Not using the same length for past_key_values532 past_key_values_length = seqlen + 2533 shape = (534 batch,535 self.num_attention_heads,536 past_key_values_length,537 self._config.hidden_size // self.num_attention_heads,538 )539 540 if "attention_mask" in common_inputs:541 mask_dtype = common_inputs["attention_mask"].dtype542 common_inputs["attention_mask"] = torch.cat(543 [common_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)],544 dim=1,545 )546 547 common_inputs["past_key_values"] = []548 for _ in range(self.num_layers):549 common_inputs["past_key_values"].append((torch.zeros(shape), torch.zeros(shape)))550 551 return common_inputs552 553 def fill_with_past_key_values_(554 self, inputs_or_outputs: Mapping[str, Mapping[int, str]], direction: str, inverted_values_shape: bool = False555 ):556 """557 Fill the input_or_outputs mapping with past_key_values dynamic axes considering.558 559 Args:560 inputs_or_outputs: The mapping to fill.561 direction: either "inputs" or "outputs", it specifies whether input_or_outputs is the input mapping or the562 output mapping, this is important for axes naming.563 inverted_values_shape:564 If `True`, store values on dynamic axis 1, else on axis 2.565 566 """567 if direction not in ["inputs", "outputs"]:568 raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given')569 570 name = "past_key_values" if direction == "inputs" else "present"571 for i in range(self.num_layers):572 inputs_or_outputs[f"{name}.{i}.key"] = {0: "batch", 2: "past_sequence + sequence"}573 if inverted_values_shape:574 inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch", 1: "past_sequence + sequence"}575 else:576 inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch", 2: "past_sequence + sequence"}577 578 def _flatten_past_key_values_(self, flattened_output, name, idx, t):579 flattened_output[f"{name}.{idx}.key"] = t[0]580 flattened_output[f"{name}.{idx}.value"] = t[1]581 582 def flatten_output_collection_property(self, name: str, field: Iterable[Any]) -> dict[str, Any]:583 flattened_output = {}584 if name in ["present", "past_key_values"]:585 for idx, t in enumerate(field):586 self._flatten_past_key_values_(flattened_output, name, idx, t)587 else:588 flattened_output = super().flatten_output_collection_property(name, field)589 590 return flattened_output591 592 593class OnnxSeq2SeqConfigWithPast(OnnxConfigWithPast):594 @property595 def outputs(self) -> Mapping[str, Mapping[int, str]]:596 common_outputs = super(OnnxConfigWithPast, self).outputs597 # Renaming the outputs axes properly.598 for name, axes_names in common_outputs.items():599 sequence_name = "encoder_sequence" if "encoder" in name else "decoder_sequence"600 for axis_idx, name in axes_names.items():601 if "sequence" in name:602 axes_names[axis_idx] = sequence_name603 # We reset the value as the order in common_outputs (OrderedDict) is lost otherwise604 else:605 axes_names[axis_idx] = name606 if self.use_past:607 self.fill_with_past_key_values_(common_outputs, direction="outputs")608 609 return common_outputs610 611 @property612 def num_layers(self) -> tuple[int, ...]:613 try:614 num_layers = super().num_layers615 num_layers = (num_layers, num_layers)616 except AttributeError:617 if hasattr(self._config, "encoder_layers") and hasattr(self._config, "decoder_layers"):618 num_layers = (self._config.encoder_layers, self._config.decoder_layers)619 else:620 raise AttributeError(621 "could not find the number of encoder and decoder layers attributes in the model configuration,"622 " override the num_layers property of the model OnnxConfig to solve this"623 )624 625 return num_layers626 627 @property628 def num_attention_heads(self) -> tuple[int, ...]:629 try:630 num_attention_heads = super().num_attention_heads631 num_attention_heads = (num_attention_heads, num_attention_heads)632 except AttributeError:633 if hasattr(self._config, "encoder_attention_heads") and hasattr(self._config, "decoder_attention_heads"):634 num_attention_heads = (self._config.encoder_attention_heads, self._config.decoder_attention_heads)635 else:636 raise AttributeError(637 "could not find the number of attention heads for the encoder and the decoder attributes in the"638 " model configuration, override the num_attention_heads property of the model OnnxConfig to solve"639 " this"640 )641 return num_attention_heads642 643 def generate_dummy_inputs(644 self,645 tokenizer: Optional["PreTrainedTokenizerBase"],646 batch_size: int = -1,647 seq_length: int = -1,648 is_pair: bool = False,649 framework: Optional[TensorType] = None,650 ) -> Mapping[str, Any]:651 encoder_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(652 tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework653 )654 655 # Generate decoder inputs656 decoder_seq_length = seq_length if not self.use_past else 1657 decoder_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(658 tokenizer, batch_size=batch_size, seq_length=decoder_seq_length, is_pair=is_pair, framework=framework659 )660 decoder_inputs = {f"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}661 common_inputs = dict(**encoder_inputs, **decoder_inputs)662 663 if self.use_past:664 if not is_torch_available():665 raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")666 else:667 import torch668 batch = common_inputs["input_ids"].shape[0]669 encoder_seq_length = common_inputs["input_ids"].shape[1]670 decoder_seq_length = common_inputs["decoder_input_ids"].shape[1]671 num_encoder_attention_heads, num_decoder_attention_heads = self.num_attention_heads672 encoder_shape = (673 batch,674 num_encoder_attention_heads,675 encoder_seq_length,676 self._config.hidden_size // num_encoder_attention_heads,677 )678 decoder_shape = (679 batch,680 num_decoder_attention_heads,681 # Not using the same length for past_key_values682 decoder_seq_length + 3,683 self._config.hidden_size // num_decoder_attention_heads,684 )685 686 common_inputs["past_key_values"] = []687 # If the number of encoder and decoder layers are present in the model configuration, both are considered688 num_encoder_layers, num_decoder_layers = self.num_layers689 min_num_layers = min(num_encoder_layers, num_decoder_layers)690 max_num_layers = max(num_encoder_layers, num_decoder_layers) - min_num_layers691 remaining_side_name = "encoder" if num_encoder_layers > num_decoder_layers else "decoder"692 693 for _ in range(min_num_layers):694 # For encoder-decoder models, past_key_values contains pre-computed values for both the encoder and the695 # decoder layers, hence a tuple of 4 tensors instead of 2696 common_inputs["past_key_values"].append(697 (698 torch.zeros(decoder_shape),699 torch.zeros(decoder_shape),700 torch.zeros(encoder_shape),701 torch.zeros(encoder_shape),702 )703 )704 705 # TODO: test this.706 shape = encoder_shape if remaining_side_name == "encoder" else decoder_shape707 for _ in range(min_num_layers, max_num_layers):708 common_inputs["past_key_values"].append((torch.zeros(shape), torch.zeros(shape)))709 710 return common_inputs711 712 def fill_with_past_key_values_(self, inputs_or_outputs: Mapping[str, Mapping[int, str]], direction: str):713 if direction not in ["inputs", "outputs"]:714 raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given')715 716 name = "past_key_values" if direction == "inputs" else "present"717 718 # If the number of encoder and decoder layers are present in the model configuration, both are considered719 num_encoder_layers, num_decoder_layers = self.num_layers720 min_num_layers = min(num_encoder_layers, num_decoder_layers)721 max_num_layers = max(num_encoder_layers, num_decoder_layers) - min_num_layers722 remaining_side_name = "encoder" if num_encoder_layers > num_decoder_layers else "decoder"723 724 encoder_sequence = "past_encoder_sequence"725 decoder_sequence = "past_decoder_sequence" if direction == "inputs" else "past_decoder_sequence + sequence"726 727 for i in range(min_num_layers):728 inputs_or_outputs[f"{name}.{i}.decoder.key"] = {0: "batch", 2: decoder_sequence}729 inputs_or_outputs[f"{name}.{i}.decoder.value"] = {0: "batch", 2: decoder_sequence}730 inputs_or_outputs[f"{name}.{i}.encoder.key"] = {0: "batch", 2: encoder_sequence}731 inputs_or_outputs[f"{name}.{i}.encoder.value"] = {0: "batch", 2: encoder_sequence}732 733 for i in range(min_num_layers, max_num_layers):734 if remaining_side_name == "encoder":735 axes_info = {0: "batch", 2: encoder_sequence}736 else:737 axes_info = {0: "batch", 2: decoder_sequence}738 inputs_or_outputs[f"{name}.{i}.{remaining_side_name}.key"] = axes_info739 740 def _flatten_past_key_values_(self, flattened_output, name, idx, t):741 flattened_output[f"{name}.{idx}.decoder.key"] = t[0]742 flattened_output[f"{name}.{idx}.decoder.value"] = t[1]743 flattened_output[f"{name}.{idx}.encoder.key"] = t[2]744 flattened_output[f"{name}.{idx}.encoder.value"] = t[3]745