Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.3# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16"""TF general model utils."""17 18from __future__ import annotations19 20import functools21import gc22import inspect23import json24import os25import pickle26import re27import warnings28from collections.abc import Mapping29from pathlib import Path30from typing import TYPE_CHECKING, Any, Callable, Union31 32import h5py33import numpy as np34import tensorflow as tf35from packaging.version import parse36 37from . import DataCollatorWithPadding, DefaultDataCollator38from .activations_tf import get_tf_activation39from .configuration_utils import PretrainedConfig40from .dynamic_module_utils import custom_object_save41from .generation import GenerationConfig, TFGenerationMixin42from .tf_utils import (43 convert_batch_encoding,44 expand_1d,45 load_attributes_from_hdf5_group,46 save_attributes_to_hdf5_group,47 shape_list,48)49from .utils import (50 SAFE_WEIGHTS_INDEX_NAME,51 SAFE_WEIGHTS_NAME,52 TF2_WEIGHTS_INDEX_NAME,53 TF2_WEIGHTS_NAME,54 TF_WEIGHTS_NAME,55 WEIGHTS_INDEX_NAME,56 WEIGHTS_NAME,57 ModelOutput,58 PushToHubMixin,59 cached_file,60 download_url,61 find_labels,62 has_file,63 is_offline_mode,64 is_remote_url,65 is_safetensors_available,66 is_tf_symbolic_tensor,67 logging,68 requires_backends,69 working_or_temp_dir,70)71from .utils.hub import convert_file_size_to_int, get_checkpoint_shard_files72 73 74if is_safetensors_available():75 from safetensors import safe_open76 from safetensors.tensorflow import save_file as safe_save_file77 78if TYPE_CHECKING:79 from . import PreTrainedTokenizerBase80 81logger = logging.get_logger(__name__)82 83if "TF_USE_LEGACY_KERAS" not in os.environ:84 os.environ["TF_USE_LEGACY_KERAS"] = "1" # Compatibility fix to make sure tf.keras stays at Keras 285elif os.environ["TF_USE_LEGACY_KERAS"] != "1":86 logger.warning(87 "Transformers is only compatible with Keras 2, but you have explicitly set `TF_USE_LEGACY_KERAS` to `0`. "88 "This may result in unexpected behaviour or errors if Keras 3 objects are passed to Transformers models."89 )90 91try:92 import tf_keras as keras93 from tf_keras import backend as K94except (ModuleNotFoundError, ImportError):95 import keras96 from keras import backend as K97 98 if parse(keras.__version__).major > 2:99 raise ValueError(100 "Your currently installed version of Keras is Keras 3, but this is not yet supported in "101 "Transformers. Please install the backwards-compatible tf-keras package with "102 "`pip install tf-keras`."103 )104 105 106tf_logger = tf.get_logger()107 108TFModelInputType = Union[109 list[tf.Tensor],110 list[np.ndarray],111 dict[str, tf.Tensor],112 dict[str, np.ndarray],113 tf.Tensor,114 np.ndarray,115]116 117 118def dummy_loss(y_true, y_pred):119 if y_pred.shape.rank <= 1:120 return y_pred121 else:122 reduction_axes = list(range(1, y_pred.shape.rank))123 return tf.reduce_mean(y_pred, axis=reduction_axes)124 125 126class TFModelUtilsMixin:127 """128 A few utilities for `keras.Model`, to be used as a mixin.129 """130 131 def num_parameters(self, only_trainable: bool = False) -> int:132 """133 Get the number of (optionally, trainable) parameters in the model.134 135 Args:136 only_trainable (`bool`, *optional*, defaults to `False`):137 Whether or not to return only the number of trainable parameters138 139 Returns:140 `int`: The number of parameters.141 """142 if only_trainable:143 return int(sum(np.prod(w.shape.as_list()) for w in self.trainable_variables))144 else:145 return self.count_params()146 147 148def keras_serializable(cls):149 """150 Decorate a Keras Layer class to support Keras serialization.151 152 This is done by:153 154 1. Adding a `transformers_config` dict to the Keras config dictionary in `get_config` (called by Keras at155 serialization time.156 2. Wrapping `__init__` to accept that `transformers_config` dict (passed by Keras at deserialization time) and157 convert it to a config object for the actual layer initializer.158 3. Registering the class as a custom object in Keras (if the Tensorflow version supports this), so that it does not159 need to be supplied in `custom_objects` in the call to `keras.models.load_model`.160 161 Args:162 cls (a `keras.layers.Layers subclass`):163 Typically a `TF.MainLayer` class in this project, in general must accept a `config` argument to its164 initializer.165 166 Returns:167 The same class object, with modifications for Keras deserialization.168 """169 initializer = cls.__init__170 171 config_class = getattr(cls, "config_class", None)172 if config_class is None:173 raise AttributeError("Must set `config_class` to use @keras_serializable")174 175 @functools.wraps(initializer)176 def wrapped_init(self, *args, **kwargs):177 config = args[0] if args and isinstance(args[0], PretrainedConfig) else kwargs.pop("config", None)178 179 if isinstance(config, dict):180 config = config_class.from_dict(config)181 initializer(self, config, *args, **kwargs)182 elif isinstance(config, PretrainedConfig):183 if len(args) > 0:184 initializer(self, *args, **kwargs)185 else:186 initializer(self, config, *args, **kwargs)187 else:188 raise TypeError("Must pass either `config` (PretrainedConfig) or `config` (dict)")189 190 self._config = config191 self._kwargs = kwargs192 193 cls.__init__ = wrapped_init194 195 if not hasattr(cls, "get_config"):196 raise TypeError("Only use @keras_serializable on keras.layers.Layer subclasses")197 if hasattr(cls.get_config, "_is_default"):198 199 def get_config(self):200 cfg = super(cls, self).get_config()201 cfg["config"] = self._config.to_dict()202 cfg.update(self._kwargs)203 return cfg204 205 cls.get_config = get_config206 207 cls._keras_serializable = True208 if hasattr(keras.utils, "register_keras_serializable"):209 cls = keras.utils.register_keras_serializable()(cls)210 return cls211 212 213class TFCausalLanguageModelingLoss:214 """215 Loss function suitable for causal language modeling (CLM), that is, the task of guessing the next token.216 217 <Tip>218 219 Any label of -100 will be ignored (along with the corresponding logits) in the loss computation.220 221 </Tip>222 """223 224 def hf_compute_loss(self, labels, logits):225 loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)226 if self.config.tf_legacy_loss:227 # make sure only labels that are not equal to -100 affect the loss228 active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)229 reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)230 labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)231 return loss_fn(labels, reduced_logits)232 233 # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway234 unmasked_loss = loss_fn(tf.nn.relu(labels), logits)235 # make sure only labels that are not equal to -100 affect the loss236 loss_mask = tf.cast(labels != -100, dtype=unmasked_loss.dtype)237 masked_loss = unmasked_loss * loss_mask238 reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)239 return tf.reshape(reduced_masked_loss, (1,))240 241 242class TFQuestionAnsweringLoss:243 """244 Loss function suitable for question answering.245 """246 247 def hf_compute_loss(self, labels, logits):248 loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)249 start_loss = loss_fn(labels["start_position"], logits[0])250 end_loss = loss_fn(labels["end_position"], logits[1])251 252 return (start_loss + end_loss) / 2.0253 254 255class TFTokenClassificationLoss:256 """257 Loss function suitable for token classification.258 259 <Tip>260 261 Any label of -100 will be ignored (along with the corresponding logits) in the loss computation.262 263 </Tip>264 """265 266 def hf_compute_loss(self, labels, logits):267 loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)268 if tf.executing_eagerly(): # Data-dependent conditionals are forbidden in XLA269 if tf.math.reduce_any(labels == -1):270 tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")271 272 if self.config.tf_legacy_loss:273 # make sure only labels that are not equal to -100274 # are taken into account as loss275 if tf.math.reduce_any(labels == -1):276 tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")277 active_loss = tf.reshape(labels, (-1,)) != -1278 else:279 active_loss = tf.reshape(labels, (-1,)) != -100280 reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)281 labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)282 283 return loss_fn(labels, reduced_logits)284 285 # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway286 unmasked_loss = loss_fn(tf.nn.relu(labels), logits)287 # make sure only labels that are not equal to -100 or -1288 # are taken into account as loss289 loss_mask = tf.cast(labels >= 0, dtype=unmasked_loss.dtype)290 # Avoid possible division by zero later291 # Masked positions will have a loss of NaN because -100 and -1 are not valid labels292 masked_loss = unmasked_loss * loss_mask293 reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)294 return tf.reshape(reduced_masked_loss, (1,))295 296 297class TFSequenceClassificationLoss:298 """299 Loss function suitable for sequence classification.300 """301 302 def hf_compute_loss(self, labels, logits):303 if logits.shape.rank == 1 or logits.shape[1] == 1:304 loss_fn = keras.losses.MeanSquaredError(reduction=keras.losses.Reduction.NONE)305 if labels.shape.rank == 1:306 # MeanSquaredError returns a scalar loss if the labels are 1D, so avoid that307 labels = tf.expand_dims(labels, axis=-1)308 else:309 loss_fn = keras.losses.SparseCategoricalCrossentropy(310 from_logits=True, reduction=keras.losses.Reduction.NONE311 )312 313 return loss_fn(labels, logits)314 315 316class TFMultipleChoiceLoss:317 """Loss function suitable for multiple choice tasks."""318 319 def hf_compute_loss(self, labels, logits):320 loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)321 return loss_fn(labels, logits)322 323 324class TFMaskedLanguageModelingLoss(TFCausalLanguageModelingLoss):325 """326 Loss function suitable for masked language modeling (MLM), that is, the task of guessing the masked tokens.327 328 <Tip>329 330 Any label of -100 will be ignored (along with the corresponding logits) in the loss computation.331 332 </Tip>333 """334 335 336class TFNextSentencePredictionLoss:337 """338 Loss function suitable for next sentence prediction (NSP), that is, the task of guessing the next sentence.339 340 <Tip>341 342 Any label of -100 will be ignored (along with the corresponding logits) in the loss computation.343 344 </Tip>345 """346 347 def hf_compute_loss(self, labels, logits):348 loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=keras.losses.Reduction.NONE)349 if self.config.tf_legacy_loss:350 # make sure only labels that are not equal to -100351 # are taken into account as loss352 next_sentence_active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)353 next_sentence_reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, 2)), next_sentence_active_loss)354 next_sentence_label = tf.boolean_mask(tf.reshape(labels, (-1,)), next_sentence_active_loss)355 356 return loss_fn(next_sentence_label, next_sentence_reduced_logits)357 358 # make sure only labels that are not equal to -100359 # are taken into account as loss360 361 # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway362 unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels), y_pred=logits)363 ns_loss_mask = tf.cast(labels != -100, dtype=unmasked_ns_loss.dtype)364 # Just zero out samples where label is -100, no reduction365 masked_ns_loss = unmasked_ns_loss * ns_loss_mask366 367 return masked_ns_loss368 369 370def booleans_processing(config, **kwargs):371 """372 Process the input booleans of each model.373 374 Args:375 config ([`PretrainedConfig`]):376 The config of the running model.377 **kwargs:378 The boolean parameters379 380 Returns:381 A dictionary with the proper values for each boolean382 """383 final_booleans = {}384 385 # Pure conv models (such as ConvNext) do not have `output_attentions`. If the signature has386 # `output_attentions`, it will be present here in `kwargs`, even if unset (in that case, as `None`)387 if "output_attentions" in kwargs:388 final_booleans["output_attentions"] = (389 kwargs["output_attentions"] if kwargs["output_attentions"] is not None else config.output_attentions390 )391 final_booleans["output_hidden_states"] = (392 kwargs["output_hidden_states"] if kwargs["output_hidden_states"] is not None else config.output_hidden_states393 )394 final_booleans["return_dict"] = kwargs["return_dict"] if kwargs["return_dict"] is not None else config.return_dict395 396 if "use_cache" in kwargs:397 final_booleans["use_cache"] = (398 kwargs["use_cache"] if kwargs["use_cache"] is not None else getattr(config, "use_cache", None)399 )400 return final_booleans401 402 403def unpack_inputs(func):404 """405 Decorator that processes the inputs to a Keras layer, passing them to the layer as keyword arguments. This enables406 downstream use of the inputs by their variable name, even if they arrive packed as a dictionary in the first input407 (common case in Keras).408 409 Args:410 func (`callable`):411 The callable function of the TensorFlow model.412 413 414 Returns:415 A callable that wraps the original `func` with the behavior described above.416 """417 418 original_signature = inspect.signature(func)419 420 @functools.wraps(func)421 def run_call_with_unpacked_inputs(self, *args, **kwargs):422 # isolates the actual `**kwargs` for the decorated function423 kwargs_call = {key: val for key, val in kwargs.items() if key not in dict(original_signature.parameters)}424 fn_args_and_kwargs = {key: val for key, val in kwargs.items() if key not in kwargs_call}425 fn_args_and_kwargs.update({"kwargs_call": kwargs_call})426 427 # move any arg into kwargs, if they exist428 fn_args_and_kwargs.update(dict(zip(func.__code__.co_varnames[1:], args)))429 430 # Encoder Decoder models delegate the application of the configuration options to their inner models.431 if "EncoderDecoder" in self.__class__.__name__:432 config = None433 else:434 config = self.config435 436 unpacked_inputs = input_processing(func, config, **fn_args_and_kwargs)437 return func(self, **unpacked_inputs)438 439 # Keras enforces the first layer argument to be passed, and checks it through `inspect.getfullargspec()`. This440 # function does not follow wrapper chains (i.e. ignores `functools.wraps()`), meaning that without the line below441 # Keras would attempt to check the first argument against the literal signature of the wrapper.442 run_call_with_unpacked_inputs.__signature__ = original_signature443 444 return run_call_with_unpacked_inputs445 446 447def input_processing(func, config, **kwargs):448 """449 Process the input of each TensorFlow model including the booleans. In case of a list of symbolic inputs, each input450 has to be named accordingly to the parameters name, i.e. `input_ids = keras.Input(shape=(128,), dtype='int32',451 name="input_ids")` otherwise the order of the tensors will not be guaranteed during the training.452 453 Args:454 func (`callable`):455 The callable function of the TensorFlow model.456 config ([`PretrainedConfig`]):457 The config of the running model.458 **kwargs:459 The inputs of the model.460 461 Returns:462 Two lists, one for the missing layers, and another one for the unexpected layers.463 """464 signature = dict(inspect.signature(func).parameters)465 has_kwargs = bool(signature.pop("kwargs", None))466 signature.pop("self", None)467 parameter_names = list(signature.keys())468 main_input_name = parameter_names[0]469 main_input = kwargs.pop(main_input_name, None)470 output = {}471 allowed_types = (tf.Tensor, bool, int, ModelOutput, tuple, list, dict, np.ndarray)472 473 if "inputs" in kwargs["kwargs_call"]:474 warnings.warn(475 "The `inputs` argument is deprecated and will be removed in a future version, use `input_ids` instead.",476 FutureWarning,477 )478 479 output["input_ids"] = kwargs["kwargs_call"].pop("inputs")480 481 if "decoder_cached_states" in kwargs["kwargs_call"]:482 warnings.warn(483 "The `decoder_cached_states` argument is deprecated and will be removed in a future version, use"484 " `past_key_values` instead.",485 FutureWarning,486 )487 output["past_key_values"] = kwargs["kwargs_call"].pop("decoder_cached_states")488 489 if "past" in kwargs["kwargs_call"] and "past_key_values" in parameter_names:490 warnings.warn(491 "The `past` argument is deprecated and will be removed in a future version, use `past_key_values`"492 " instead.",493 FutureWarning,494 )495 kwargs["past_key_values"] = kwargs["kwargs_call"].pop("past")496 elif "past_key_values" in kwargs["kwargs_call"] and "past" in parameter_names:497 kwargs["past"] = kwargs["kwargs_call"].pop("past_key_values")498 499 if has_kwargs:500 output["kwargs"] = kwargs.pop("kwargs_call", {})501 else:502 if len(kwargs["kwargs_call"]) > 0:503 raise ValueError(504 "The following keyword arguments are not supported by this model:"505 f" {list(kwargs['kwargs_call'].keys())}."506 )507 kwargs.pop("kwargs_call")508 509 for k, v in kwargs.items():510 if isinstance(v, allowed_types) or tf.is_tensor(v) or v is None:511 output[k] = v512 else:513 raise ValueError(f"Data of type {type(v)} is not allowed only {allowed_types} is accepted for {k}.")514 515 if isinstance(main_input, (tuple, list)):516 for i, input in enumerate(main_input):517 # EagerTensors don't allow to use the .name property so we check for a real Tensor518 if is_tf_symbolic_tensor(input):519 # Tensor names have always the pattern `name:id` then we check only the520 # `name` part521 tensor_name = input.name.split(":")[0]522 523 if tensor_name in parameter_names:524 output[tensor_name] = input525 else:526 output[parameter_names[i]] = input527 elif isinstance(input, allowed_types) or input is None:528 output[parameter_names[i]] = input529 else:530 raise ValueError(531 f"Data of type {type(input)} is not allowed only {allowed_types} is accepted for"532 f" {parameter_names[i]}."533 )534 elif isinstance(main_input, Mapping):535 if "inputs" in main_input:536 warnings.warn(537 "The `inputs` argument is deprecated and will be removed in a future version, use `input_ids`"538 " instead.",539 FutureWarning,540 )541 542 output["input_ids"] = main_input.pop("inputs")543 544 if "decoder_cached_states" in main_input:545 warnings.warn(546 "The `decoder_cached_states` argument is deprecated and will be removed in a future version, use"547 " `past_key_values` instead.",548 FutureWarning,549 )550 output["past_key_values"] = main_input.pop("decoder_cached_states")551 552 for k, v in dict(main_input).items():553 if isinstance(v, allowed_types) or v is None:554 output[k] = v555 elif k not in parameter_names and "args" not in parameter_names:556 logger.warning(557 f"The parameter {k} does not belongs to the parameter list {parameter_names} and will be ignored."558 )559 continue560 else:561 raise ValueError(f"Data of type {type(v)} is not allowed only {allowed_types} is accepted for {k}.")562 else:563 if tf.is_tensor(main_input) or main_input is None:564 output[main_input_name] = main_input565 else:566 raise ValueError(567 f"Data of type {type(main_input)} is not allowed only {allowed_types} is accepted for"568 f" {main_input_name}."569 )570 571 # Populates any unspecified argument with their default value, according to the signature.572 for name in parameter_names:573 if name not in list(output.keys()) and name != "args":574 output[name] = kwargs.pop(name, signature[name].default)575 576 # When creating a SavedModel TF calls the method with LayerCall.__call__(args, **kwargs)577 # So to respect the proper output we have to add this exception578 if "args" in output:579 if output["args"] is not None and is_tf_symbolic_tensor(output["args"]):580 tensor_name = output["args"].name.split(":")[0]581 output[tensor_name] = output["args"]582 else:583 # `args` in this case is always the first parameter, then `input_ids`584 output["input_ids"] = output["args"]585 586 del output["args"]587 588 if "kwargs" in output:589 del output["kwargs"]590 591 cast_output = {}592 for key, val in output.items():593 if isinstance(val, tf.Tensor) and val.dtype == tf.int64:594 cast_output[key] = tf.cast(val, tf.int32)595 elif isinstance(val, np.ndarray) and val.dtype == np.int64:596 cast_output[key] = val.astype(np.int32)597 else:598 cast_output[key] = val599 600 output = cast_output601 del cast_output602 603 if config is not None:604 boolean_dict = {605 k: v606 for k, v in output.items()607 if k in ["return_dict", "output_attentions", "output_hidden_states", "use_cache"]608 }609 610 output.update(611 booleans_processing(612 config=config,613 **boolean_dict,614 )615 )616 617 return output618 619 620def strip_model_name_and_prefix(name, _prefix=None):621 if _prefix is not None and name.startswith(_prefix):622 name = name[len(_prefix) :]623 if name.startswith("/"):624 name = name[1:]625 if "model." not in name and len(name.split("/")) > 1:626 name = "/".join(name.split("/")[1:])627 return name628 629 630def tf_shard_checkpoint(weights, max_shard_size="10GB", weights_name: str = TF2_WEIGHTS_NAME):631 """632 Splits a model state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a633 given size.634 635 The sub-checkpoints are determined by iterating through the `state_dict` in the order of its keys, so there is no636 optimization made to make each sub-checkpoint as close as possible to the maximum size passed. For example, if the637 limit is 10GB and we have weights of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB],638 [6+2+2GB] and not [6+2+2GB], [6+2GB], [6GB].639 640 <Tip warning={true}>641 642 If one of the model's weight is bigger that `max_shard_size`, it will end up in its own sub-checkpoint which will643 have a size greater than `max_shard_size`.644 645 </Tip>646 647 Args:648 weights (`dict[str, tf.RessourceVariable]`): The list of tf.RessourceVariable of a model to save.649 max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):650 The maximum size of each sub-checkpoint. If expressed as a string, needs to be digits followed by a unit651 (like `"5MB"`).652 """653 max_shard_size = convert_file_size_to_int(max_shard_size)654 655 sharded_state_dicts = []656 current_block = []657 current_block_size = 0658 total_size = 0659 660 for item in weights:661 weight_size = item.numpy().size * item.dtype.size662 663 # If this weight is going to tip up over the maximal size, we split.664 if current_block_size + weight_size > max_shard_size:665 sharded_state_dicts.append(current_block)666 current_block = []667 current_block_size = 0668 669 current_block.append(item)670 current_block_size += weight_size671 total_size += weight_size672 673 # Add the last block674 sharded_state_dicts.append(current_block)675 676 # If we only have one shard, we return it677 if len(sharded_state_dicts) == 1:678 return {weights_name: sharded_state_dicts[0]}, None679 680 # Otherwise, let's build the index681 weight_map = {}682 shards = {}683 for idx, shard in enumerate(sharded_state_dicts):684 shard_file = weights_name.replace(".h5", f"-{idx + 1:05d}-of-{len(sharded_state_dicts):05d}.h5")685 shard_file = shard_file.replace(686 ".safetensors", f"-{idx + 1:05d}-of-{len(sharded_state_dicts):05d}.safetensors"687 )688 shards[shard_file] = shard689 for weight in shard:690 weight_name = weight.name691 weight_map[weight_name] = shard_file692 693 # Add the metadata694 metadata = {"total_size": total_size}695 index = {"metadata": metadata, "weight_map": weight_map}696 return shards, index697 698 699def load_tf_sharded_weights(model, shard_files, ignore_mismatched_sizes=False, strict=False, _prefix=None):700 """701 This is the same as `load_tf_weights` but for a sharded checkpoint. Detect missing and unexpected layers and load702 the TF weights from the shard file accordingly to their names and shapes.703 704 This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being705 loaded in the model.706 707 Args:708 model (`keras.models.Model`): The model in which to load the checkpoint.709 shard_files (`str` or `os.PathLike`): A list containing the sharded checkpoint names.710 ignore_mismatched_sizes`bool`, *optional`, defaults to `True`):711 Whether or not to ignore the mismatch between the sizes712 strict (`bool`, *optional*, defaults to `True`):713 Whether to strictly enforce that the keys in the model state dict match the keys in the sharded checkpoint.714 715 Returns:716 Three lists, one for the missing layers, another one for the unexpected layers, and a last one for the717 mismatched layers.718 """719 720 # Load the index721 unexpected_keys = set()722 saved_keys = set()723 mismatched_keys = set()724 725 # Since TF adds the name of the class to its weights, and uses the index and not the name of the layer to load726 # the weight, we have to get rid of the first prefix of the name of the layer.727 model_keys = set()728 model_layer_map = {}729 for i, k in enumerate(model.weights):730 layer_name = k.name731 if _prefix is not None and layer_name.startswith(_prefix):732 layer_name = layer_name[len(_prefix) :]733 layer_name = layer_name.lstrip("/")734 if not ("model." in layer_name or len(layer_name.split("/")) == 1):735 layer_name = "/".join(layer_name.split("/")[1:])736 model_keys.add(layer_name)737 model_layer_map[layer_name] = i738 739 for shard_file in shard_files:740 saved_weight_names_set, unexpected_keys_set, mismatched_keys_set = load_tf_shard(741 model,742 model_layer_map,743 shard_file,744 ignore_mismatched_sizes=ignore_mismatched_sizes,745 _prefix=_prefix,746 )747 saved_keys.update(saved_weight_names_set)748 unexpected_keys.update(unexpected_keys_set)749 mismatched_keys.update(mismatched_keys_set)750 gc.collect()751 752 missing_keys = model_keys - saved_keys753 if strict and (len(missing_keys) > 0 or len(unexpected_keys) > 0):754 error_message = f"Error(s) in loading state_dict for {model.__class__.__name__}"755 if len(missing_keys) > 0:756 str_missing_keys = ",".join([f'"{k}"' for k in missing_keys])757 error_message += f"\nMissing key(s): {str_missing_keys}."758 if len(unexpected_keys) > 0:759 str_unexpected_keys = ",".join([f'"{k}"' for k in unexpected_keys])760 error_message += f"\nMissing key(s): {str_unexpected_keys}."761 raise RuntimeError(error_message)762 763 return missing_keys, unexpected_keys, mismatched_keys764 765 766def load_tf_shard(model, model_layer_map, resolved_archive_file, ignore_mismatched_sizes=False, _prefix=None):767 """768 Loads a shard from a sharded checkpoint file. Can be either H5 or Safetensors.769 Handles missing keys and unexpected keys.770 771 Args:772 model (`keras.models.Model`): Model in which the weights are loaded773 model_layer_map (`Dict`): A dictionary mapping the layer name to the index of the layer in the model.774 resolved_archive_file (`str`): Path to the checkpoint file from which the weights will be loaded775 ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`): Whether to ignore the mismatched keys776 777 Returns:778 `keras.models.Model`: Three lists, one for the layers that were found and successfully restored (from the779 shard file), one for the mismatched layers, and another one for the unexpected layers.780 """781 saved_weight_names_set = set()782 saved_weights = {}783 mismatched_keys = set()784 unexpected_keys = set()785 # Read the H5 file786 try:787 with h5py.File(resolved_archive_file, "r") as sharded_checkpoint_file:788 # Retrieve the name of each layer from the H5 file789 saved_h5_model_layers_name = set(load_attributes_from_hdf5_group(sharded_checkpoint_file, "layer_names"))790 weight_value_tuples = []791 792 # Compute missing and unexpected sub layers793 # Store the weights in list of tuples that looks like [(weight_object, value_of_weight),...]794 for layer_name in saved_h5_model_layers_name:795 h5_layer_object = sharded_checkpoint_file[layer_name]796 saved_weights[layer_name] = np.asarray(h5_layer_object)797 798 saved_weight_names_set.add(layer_name)799 800 if layer_name not in model_layer_map:801 unexpected_keys.add(layer_name)802 else:803 symbolic_weight = model.weights[model_layer_map[layer_name]]804 805 saved_weight_value = saved_weights[layer_name]806 # If the current weight is found807 if saved_weight_value is not None:808 # Check if the shape of the current weight and the one from the H5 file are different809 if K.int_shape(symbolic_weight) != saved_weight_value.shape:810 # If yes we reshape the weight from the H5 file accordingly to the current weight811 # If the two shapes are not compatible we raise an issue812 try:813 array = np.reshape(saved_weight_value, K.int_shape(symbolic_weight))814 except ValueError as e:815 if ignore_mismatched_sizes:816 mismatched_keys.add(817 (layer_name, saved_weight_value.shape, K.int_shape(symbolic_weight))818 )819 continue820 else:821 raise e822 else:823 array = saved_weight_value824 825 # We create the tuple that will be loaded and add it to the final list826 weight_value_tuples.append((symbolic_weight, array))827 828 K.batch_set_value(weight_value_tuples)829 830 return saved_weight_names_set, unexpected_keys, mismatched_keys831 832 except Exception as e:833 try:834 with open(resolved_archive_file) as f:835 if f.read().startswith("version"):836 raise OSError(837 "You seem to have cloned a repository without having git-lfs installed. Please install "838 "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder "839 "you cloned."840 )841 else:842 raise ValueError(843 f"Unable to locate the file {resolved_archive_file} which is necessary to load this pretrained"844 " model. Make sure you have saved the model properly."845 ) from e846 except (UnicodeDecodeError, ValueError):847 raise OSError(848 f"Unable to load weights from TF checkpoint file for '{resolved_archive_file}' "849 f"at '{resolved_archive_file}'. "850 "If you tried to load a TF model from a sharded checkpoint, you should try converting the model "851 "by loading it in pytorch and saving it locally. A conversion script should be released soon."852 )853 854 855def load_tf_sharded_weights_from_safetensors(856 model, shard_files, ignore_mismatched_sizes=False, strict=False, _prefix=None857):858 """859 This is the same as `load_tf_weights_from_safetensors` but for a sharded TF-format safetensors checkpoint.860 Detect missing and unexpected layers and load the TF weights from the shard file accordingly to their names and861 shapes.862 863 This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being864 loaded in the model.865 866 Args:867 model (`keras.models.Model`): The model in which to load the checkpoint.868 shard_files (`str` or `os.PathLike`): A list containing the sharded checkpoint names.869 ignore_mismatched_sizes`bool`, *optional`, defaults to `True`):870 Whether or not to ignore the mismatch between the sizes871 strict (`bool`, *optional*, defaults to `True`):872 Whether to strictly enforce that the keys in the model state dict match the keys in the sharded checkpoint.873 874 Returns:875 Three lists, one for the missing layers, another one for the unexpected layers, and a last one for the876 mismatched layers.877 """878 879 # Load the index880 unexpected_keys = set()881 all_missing_keys = []882 mismatched_keys = set()883 884 for shard_file in shard_files:885 missing_layers, unexpected_layers, mismatched_layers = load_tf_weights_from_safetensors(886 model,887 shard_file,888 ignore_mismatched_sizes=ignore_mismatched_sizes,889 _prefix=_prefix,890 )891 all_missing_keys.append(set(missing_layers))892 unexpected_keys.update(unexpected_layers)893 mismatched_keys.update(mismatched_layers)894 gc.collect()895 missing_keys = set.intersection(*all_missing_keys)896 897 if strict and (len(missing_keys) > 0 or len(unexpected_keys) > 0):898 error_message = f"Error(s) in loading state_dict for {model.__class__.__name__}"899 if len(missing_keys) > 0:900 str_missing_keys = ",".join([f'"{k}"' for k in missing_keys])901 error_message += f"\nMissing key(s): {str_missing_keys}."902 if len(unexpected_keys) > 0:903 str_unexpected_keys = ",".join([f'"{k}"' for k in unexpected_keys])904 error_message += f"\nMissing key(s): {str_unexpected_keys}."905 raise RuntimeError(error_message)906 907 return missing_keys, unexpected_keys, mismatched_keys908 909 910def load_tf_weights(model, resolved_archive_file, ignore_mismatched_sizes=False, _prefix=None):911 """912 Detect missing and unexpected layers and load the TF weights from the shard file accordingly to their names and913 shapes.914 915 Args:916 model (`keras.models.Model`):917 The model to load the weights into.918 resolved_archive_file (`str`):919 The location of the H5 file.920 ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`):921 Whether or not to ignore weights with shapes that don't match between the checkpoint of the model.922 923 Returns:924 Three lists, one for the missing layers, another one for the unexpected layers, and a last one for the925 mismatched layers.926 """927 if resolved_archive_file.endswith(".safetensors"):928 load_function = load_tf_weights_from_safetensors929 else:930 load_function = load_tf_weights_from_h5931 932 return load_function(933 model, resolved_archive_file, ignore_mismatched_sizes=ignore_mismatched_sizes, _prefix=_prefix934 )935 936 937def load_tf_weights_from_h5(model, resolved_archive_file, ignore_mismatched_sizes=False, _prefix=None):938 mismatched_layers = []939 940 # Read the H5 file941 with h5py.File(resolved_archive_file, "r") as sharded_checkpoint_file:942 # Retrieve the name of each layer from the H5 file943 saved_h5_model_layers_name = set(load_attributes_from_hdf5_group(sharded_checkpoint_file, "layer_names"))944 945 # Find the missing layers from the high level list of layers946 missing_layers = list({layer.name for layer in model.layers} - saved_h5_model_layers_name)947 948 # Find the unexpected layers from the high level list of layers949 unexpected_layers = list(saved_h5_model_layers_name - {layer.name for layer in model.layers})950 saved_weight_names_set = set()951 symbolic_weights_names = set()952 weight_value_tuples = []953 954 # Compute missing and unexpected sub layers955 # Store the weights in list of tuples that looks like [(weight_object, value_of_weight),...]956 for layer in model.layers:957 # if layer_name from the H5 file belongs to the layers from the instantiated model958 if layer.name in saved_h5_model_layers_name:959 # Get the H5 layer object from its name960 h5_layer_object = sharded_checkpoint_file[layer.name]961 # Get all the weights as a list from the layer object962 symbolic_weights = layer.trainable_weights + layer.non_trainable_weights963 saved_weights = {}964 965 # Create a dict from the H5 saved model that looks like {"weight_name": weight_value}966 # And a set with only the names967 for weight_name in load_attributes_from_hdf5_group(h5_layer_object, "weight_names"):968 # TF names always start with the model name so we ignore it969 name = "/".join(weight_name.split("/")[1:])970 971 if _prefix is not None:972 name = _prefix + "/" + name973 974 saved_weights[name] = np.asarray(h5_layer_object[weight_name])975 976 # Add the updated name to the final list for computing missing/unexpected values977 saved_weight_names_set.add(name)978 979 # Loop over each weights from the instantiated model and compare with the weights from the H5 file980 for symbolic_weight in symbolic_weights:981 # TF names always start with the model name so we ignore it982 if _prefix is not None:983 delimiter = len(_prefix.split("/"))984 symbolic_weight_name = "/".join(985 symbolic_weight.name.split("/")[:delimiter]986 + symbolic_weight.name.split("/")[delimiter + 1 :]987 )988 else:989 symbolic_weight_name = "/".join(symbolic_weight.name.split("/")[1:])990 991 # here we check if the current weight is among the weights from the H5 file992 # If yes, get the weight_value of the corresponding weight from the H5 file993 # If not, make the value to None994 saved_weight_value = saved_weights.get(symbolic_weight_name)995 996 # Retrocompatibility patch: some embeddings are stored with the weights name (e.g. Bart's997 # `model.shared/embeddings:0` are stored as `model.shared/weights:0`)998 if saved_weight_value is None and symbolic_weight_name.endswith("embeddings:0"):999 symbolic_weight_name = symbolic_weight_name[:-12] + "weight:0"1000 saved_weight_value = saved_weights.get(symbolic_weight_name)1001 1002 # Add the updated name to the final list for computing missing/unexpected values1003 symbolic_weights_names.add(symbolic_weight_name)1004 1005 # If the current weight is found1006 if saved_weight_value is not None:1007 # Check if the shape of the current weight and the one from the H5 file are different1008 if K.int_shape(symbolic_weight) != saved_weight_value.shape:1009 # If yes we reshape the weight from the H5 file accordingly to the current weight1010 # If the two shapes are not compatible we raise an issue1011 try:1012 array = np.reshape(saved_weight_value, K.int_shape(symbolic_weight))1013 except ValueError as e:1014 if ignore_mismatched_sizes:1015 mismatched_layers.append(1016 (symbolic_weight_name, saved_weight_value.shape, K.int_shape(symbolic_weight))1017 )1018 continue1019 else:1020 raise e1021 else:1022 array = saved_weight_value1023 1024 # We create the tuple that will be loaded and add it to the final list1025 weight_value_tuples.append((symbolic_weight, array))1026 1027 # Load all the weights1028 K.batch_set_value(weight_value_tuples)1029 1030 # Compute the missing and unexpected layers1031 missing_layers.extend(list(symbolic_weights_names - saved_weight_names_set))1032 unexpected_layers.extend(list(saved_weight_names_set - symbolic_weights_names))1033 1034 return missing_layers, unexpected_layers, mismatched_layers1035 1036 1037def load_tf_weights_from_safetensors(model, resolved_archive_file, ignore_mismatched_sizes=False, _prefix=None):1038 # Read the safetensors file1039 with safe_open(resolved_archive_file, framework="tf") as safetensors_archive:1040 mismatched_layers = []1041 weight_names = [strip_model_name_and_prefix(w.name, _prefix=_prefix) for w in model.weights]1042 loaded_weight_names = list(safetensors_archive.keys())1043 # Find the missing layers from the high level list of layers1044 missing_layers = list(set(weight_names) - set(loaded_weight_names))1045 # Find the unexpected layers from the high level list of layers1046 unexpected_layers = list(set(loaded_weight_names) - set(weight_names))1047 1048 for weight in model.weights:1049 weight_name = strip_model_name_and_prefix(weight.name, _prefix=_prefix)1050 if weight_name in loaded_weight_names:1051 weight_value = safetensors_archive.get_tensor(weight_name)1052 # Check if the shape of the current weight and the one from the H5 file are different1053 if K.int_shape(weight) != weight_value.shape:1054 # If yes we reshape the weight from the H5 file accordingly to the current weight1055 # If the two shapes are not compatible we raise an issue1056 try:1057 weight_value = tf.reshape(weight_value, K.int_shape(weight))1058 except (ValueError, tf.errors.InvalidArgumentError) as e:1059 if ignore_mismatched_sizes:1060 mismatched_layers.append((weight_name, weight_value.shape, K.int_shape(weight)))1061 continue1062 else:1063 raise e1064 1065 K.set_value(weight, weight_value) # weight.assign() might break if weight is a DTensor1066 return missing_layers, unexpected_layers, mismatched_layers1067 1068 1069def init_copy_embeddings(old_embeddings, new_num_tokens):1070 r"""1071 This function aims to reduce the embeddings in case new_num_tokens < old_num_tokens or to pad with -1 in case1072 new_num_tokens > old_num_tokens. A mask is also computed in order to know which weight in the embeddings should be1073 kept or not. Example:1074 1075 - if new_num_tokens=5 and old_num_tokens=4 and old_embeddings=[w1,w2,w3,w4]1076 1077 - mask=[True,True,True,True,False] and current_weights=[w1,w2,w3,w4,-1]1078 - if new_num_tokens=4 and old_num_tokens=5 and old_embeddings=[w1,w2,w3,w4,w5]1079 1080 - mask=[True,True,True,True] and current_weights=[w1,w2,w3,w4]1081 """1082 old_num_tokens, old_embedding_dim = shape_list(old_embeddings)1083 size_diff = new_num_tokens - old_num_tokens1084 1085 # initialize new embeddings1086 # Copy token embeddings from the previous ones1087 if tf.math.greater(size_diff, 0):1088 # if the new size is greater than the old one, we extend the current embeddings with a padding until getting new size1089 # and we create a mask to properly identify the padded values and be replaced by the values of the newly created1090 # embeddings1091 current_weights = tf.pad(1092 old_embeddings.value(), tf.convert_to_tensor([[0, size_diff], [0, 0]]), constant_values=-11093 )1094 num_tokens_to_copy = min(old_num_tokens, new_num_tokens)1095 mask = tf.fill(tf.convert_to_tensor([num_tokens_to_copy, 1]), True)1096 mask = tf.pad(mask, tf.convert_to_tensor([[0, size_diff], [0, 0]]), constant_values=False)1097 else:1098 # if the new size if lower than the old one, we take the current embeddings until the new size1099 current_weights = tf.slice(1100 old_embeddings.value(),1101 tf.convert_to_tensor([0, 0]),1102 tf.convert_to_tensor([new_num_tokens, old_embedding_dim]),1103 )1104 mask = tf.fill(tf.convert_to_tensor([new_num_tokens, 1]), True)1105 1106 return mask, current_weights1107 1108 1109class TFPreTrainedModel(keras.Model, TFModelUtilsMixin, TFGenerationMixin, PushToHubMixin):1110 r"""1111 Base class for all TF models.1112 1113 [`TFPreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading,1114 downloading and saving models as well as a few methods common to all models to:1115 1116 - resize the input embeddings,1117 - prune heads in the self-attention heads.1118 1119 Class attributes (overridden by derived classes):1120 1121 - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class1122 for this model architecture.1123 - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived1124 classes of the same architecture adding modules on top of the base model.1125 - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP1126 models, `pixel_values` for vision models and `input_values` for speech models).1127 """1128 1129 config_class = None1130 base_model_prefix = ""1131 main_input_name = "input_ids"1132 _auto_class = None1133 _using_dummy_loss = None1134 _label_to_output_map = None1135 1136 # a list of re pattern of tensor names to ignore from the model when loading the model weights1137 # (and avoid unnecessary warnings).1138 _keys_to_ignore_on_load_missing = None1139 # a list of re pattern of tensor names to ignore from the weights when loading the model weights1140 # (and avoid unnecessary warnings).1141 _keys_to_ignore_on_load_unexpected = None1142 _requires_load_weight_prefix = False1143 1144 @property1145 def dummy_inputs(self) -> dict[str, tf.Tensor]:1146 """1147 Dummy inputs to build the network.1148 1149 Returns:1150 `dict[str, tf.Tensor]`: The dummy inputs.1151 """1152 dummies = {}1153 for key, spec in self.input_signature.items():1154 # 2 is the most correct arbitrary size. I will not be taking questions1155 dummy_shape = [dim if dim is not None else 2 for dim in spec.shape]1156 if spec.shape[0] is None:1157 # But let's make the batch size 1 to save memory anyway1158 dummy_shape[0] = 11159 dummies[key] = tf.ones(shape=dummy_shape, dtype=spec.dtype)1160 if key == "token_type_ids":1161 # Some models have token_type_ids but with a vocab_size of 11162 dummies[key] = tf.zeros_like(dummies[key])1163 if self.config.add_cross_attention and "encoder_hidden_states" in inspect.signature(self.call).parameters:1164 if "encoder_hidden_states" not in dummies:1165 if self.main_input_name == "input_ids":1166 dummies["encoder_hidden_states"] = tf.ones(1167 shape=(1, 2, self.config.hidden_size), dtype=tf.float32, name="encoder_hidden_states"1168 )1169 else:1170 raise NotImplementedError(1171 "Model has cross-attention but we couldn't infer the shape for the encoder hidden states. Please manually override dummy_inputs!"1172 )1173 return dummies1174 1175 def build_in_name_scope(self):1176 with tf.name_scope(self.name):1177 self.build(input_shape=None)1178 1179 @property1180 def framework(self) -> str:1181 """1182 :str: Identifies that this is a TensorFlow model.1183 """1184 return "tf"1185 1186 def build(self, input_shape=None):1187 pass # This is just here to make sure we don't call the superclass build()1188 1189 def __init__(self, config, *inputs, **kwargs):1190 super().__init__(*inputs, **kwargs)1191 if not isinstance(config, PretrainedConfig):1192 raise TypeError(1193 f"Parameter config in `{self.__class__.__name__}(config)` should be an instance of class "1194 "`PretrainedConfig`. To create a model from a pretrained model use "1195 f"`model = {self.__class__.__name__}.from_pretrained(PRETRAINED_MODEL_NAME)`"1196 )1197 # Save config and origin of the pretrained weights if given in model1198 self.config = config1199 self.name_or_path = config.name_or_path1200 self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None