Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2021 The Google Flax Team Authors and The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16 17import gc18import json19import os20import warnings21from functools import partial22from pickle import UnpicklingError23from typing import Any, Optional, Union24 25import flax.linen as nn26import jax27import jax.numpy as jnp28import msgpack.exceptions29from flax.core.frozen_dict import FrozenDict, unfreeze30from flax.serialization import from_bytes, to_bytes31from flax.traverse_util import flatten_dict, unflatten_dict32from jax.random import PRNGKey33 34from .configuration_utils import PretrainedConfig35from .dynamic_module_utils import custom_object_save36from .generation import FlaxGenerationMixin, GenerationConfig37from .modeling_flax_pytorch_utils import load_pytorch_checkpoint_in_flax_state_dict38from .utils import (39 FLAX_WEIGHTS_INDEX_NAME,40 FLAX_WEIGHTS_NAME,41 SAFE_WEIGHTS_INDEX_NAME,42 SAFE_WEIGHTS_NAME,43 WEIGHTS_INDEX_NAME,44 WEIGHTS_NAME,45 PushToHubMixin,46 add_code_sample_docstrings,47 add_start_docstrings_to_model_forward,48 cached_file,49 copy_func,50 download_url,51 has_file,52 is_offline_mode,53 is_remote_url,54 logging,55 replace_return_docstrings,56)57from .utils.hub import convert_file_size_to_int, get_checkpoint_shard_files58from .utils.import_utils import is_safetensors_available59 60 61if is_safetensors_available():62 from safetensors import safe_open63 from safetensors.flax import load_file as safe_load_file64 from safetensors.flax import save_file as safe_save_file65 66logger = logging.get_logger(__name__)67 68 69def quick_gelu(x):70 return x * jax.nn.sigmoid(1.702 * x)71 72 73ACT2FN = {74 "gelu": partial(nn.gelu, approximate=False),75 "relu": nn.relu,76 "silu": nn.swish,77 "swish": nn.swish,78 "gelu_new": partial(nn.gelu, approximate=True),79 "quick_gelu": quick_gelu,80 "gelu_pytorch_tanh": partial(nn.gelu, approximate=True),81 "tanh": nn.tanh,82}83 84 85def flax_shard_checkpoint(params, max_shard_size="10GB"):86 """87 Splits a model state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a88 given size. The sub-checkpoints are determined by iterating through the `state_dict` in the order of its keys, so89 there is no optimization made to make each sub-checkpoint as close as possible to the maximum size passed. For90 example, if the limit is 10GB and we have weights of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as91 [6GB], [6+2GB], [6+2+2GB] and not [6+2+2GB], [6+2GB], [6GB].92 93 <Tip warning={true}>94 95 If one of the model's weight is bigger that `max_shard_size`, it will end up in its own sub-checkpoint which will96 have a size greater than `max_shard_size`.97 98 </Tip>99 100 Args:101 params (`Union[Dict, FrozenDict]`): A `PyTree` of model parameters.102 max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):103 The maximum size of each sub-checkpoint. If expressed as a string, needs to be digits followed by a unit104 (like `"5MB"`).105 """106 max_shard_size = convert_file_size_to_int(max_shard_size)107 108 sharded_state_dicts = []109 current_block = {}110 current_block_size = 0111 total_size = 0112 113 # flatten the weights to chunk114 weights = flatten_dict(params, sep="/")115 for item in weights:116 weight_size = weights[item].size * weights[item].dtype.itemsize117 118 # If this weight is going to tip up over the maximal size, we split.119 if current_block_size + weight_size > max_shard_size:120 sharded_state_dicts.append(current_block)121 current_block = {}122 current_block_size = 0123 124 current_block[item] = weights[item]125 current_block_size += weight_size126 total_size += weight_size127 128 # Add the last block129 sharded_state_dicts.append(current_block)130 131 # If we only have one shard, we return it132 if len(sharded_state_dicts) == 1:133 return {FLAX_WEIGHTS_NAME: sharded_state_dicts[0]}, None134 135 # Otherwise, let's build the index136 weight_map = {}137 shards = {}138 for idx, shard in enumerate(sharded_state_dicts):139 shard_file = FLAX_WEIGHTS_NAME.replace(".msgpack", f"-{idx + 1:05d}-of-{len(sharded_state_dicts):05d}.msgpack")140 shards[shard_file] = shard141 for weight_name in shard:142 weight_map[weight_name] = shard_file143 144 # Add the metadata145 metadata = {"total_size": total_size}146 index = {"metadata": metadata, "weight_map": weight_map}147 return shards, index148 149 150class FlaxPreTrainedModel(PushToHubMixin, FlaxGenerationMixin):151 r"""152 Base class for all models.153 154 [`FlaxPreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading,155 downloading and saving models.156 157 Class attributes (overridden by derived classes):158 159 - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class160 for this model architecture.161 - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived162 classes of the same architecture adding modules on top of the base model.163 - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP164 models, `pixel_values` for vision models and `input_values` for speech models).165 """166 167 config_class = None168 base_model_prefix = ""169 main_input_name = "input_ids"170 _auto_class = None171 _missing_keys = set()172 173 def __init__(174 self,175 config: PretrainedConfig,176 module: nn.Module,177 input_shape: tuple = (1, 1),178 seed: int = 0,179 dtype: jnp.dtype = jnp.float32,180 _do_init: bool = True,181 ):182 logger.warning_once(183 "TensorFlow and JAX classes are deprecated and will be removed in Transformers v5. We "184 "recommend migrating to PyTorch classes or pinning your version of Transformers."185 )186 if config is None:187 raise ValueError("config cannot be None")188 189 if module is None:190 raise ValueError("module cannot be None")191 192 # Those are private to be exposed as typed property on derived classes.193 self._config = config194 self._module = module195 196 # Those are public as their type is generic to every derived classes.197 self.key = PRNGKey(seed)198 self.dtype = dtype199 self.input_shape = input_shape200 self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None201 202 # To check if the model was initialized automatically.203 self._is_initialized = _do_init204 205 if _do_init:206 # randomly initialized parameters207 random_params = self.init_weights(self.key, input_shape)208 params_shape_tree = jax.eval_shape(lambda params: params, random_params)209 else:210 init_fn = partial(self.init_weights, input_shape=input_shape)211 params_shape_tree = jax.eval_shape(init_fn, self.key)212 213 logger.info(214 "Model weights are not initialized as `_do_init` is set to `False`. "215 f"Make sure to call `{self.__class__.__name__}.init_weights` manually to initialize the weights."216 )217 218 # get the shape of the parameters219 self._params_shape_tree = params_shape_tree220 221 # save required_params as set222 self._required_params = set(flatten_dict(unfreeze(params_shape_tree)).keys())223 224 # initialize the parameters225 if _do_init:226 self.params = random_params227 228 def init_weights(self, rng: jax.random.PRNGKey, input_shape: tuple, params: FrozenDict = None) -> dict:229 raise NotImplementedError(f"init method has to be implemented for {self}")230 231 def enable_gradient_checkpointing(self):232 raise NotImplementedError(f"gradient checkpointing method has to be implemented for {self}")233 234 @classmethod235 def _from_config(cls, config, **kwargs):236 """237 All context managers that the model should be initialized under go here.238 """239 return cls(config, **kwargs)240 241 @property242 def framework(self) -> str:243 """244 :str: Identifies that this is a Flax model.245 """246 return "flax"247 248 @property249 def config(self) -> PretrainedConfig:250 return self._config251 252 @property253 def module(self) -> nn.Module:254 return self._module255 256 @property257 def params(self) -> Union[dict, FrozenDict]:258 if not self._is_initialized:259 raise ValueError(260 "`params` cannot be accessed from model when the model is created with `_do_init=False`. "261 "You must call `init_weights` manually and store the params outside of the model and "262 "pass it explicitly where needed."263 )264 return self._params265 266 @property267 def required_params(self) -> set:268 return self._required_params269 270 @property271 def params_shape_tree(self) -> dict:272 return self._params_shape_tree273 274 @params.setter275 def params(self, params: Union[dict, FrozenDict]):276 # don't set params if the model is not initialized277 if not self._is_initialized:278 raise ValueError(279 "`params` cannot be set from model when the model is created with `_do_init=False`. "280 "You store the params outside of the model."281 )282 283 if isinstance(params, FrozenDict):284 params = unfreeze(params)285 param_keys = set(flatten_dict(params).keys())286 if len(self.required_params - param_keys) > 0:287 raise ValueError(288 "Some parameters are missing. Make sure that `params` include the following "289 f"parameters {self.required_params - param_keys}"290 )291 self._params = params292 293 def _cast_floating_to(self, params: Union[dict, FrozenDict], dtype: jnp.dtype, mask: Any = None) -> Any:294 """295 Helper method to cast floating-point values of given parameter `PyTree` to given `dtype`.296 """297 298 # taken from https://github.com/deepmind/jmp/blob/3a8318abc3292be38582794dbf7b094e6583b192/jmp/_src/policy.py#L27299 def conditional_cast(param):300 if isinstance(param, jnp.ndarray) and jnp.issubdtype(param.dtype, jnp.floating):301 param = param.astype(dtype)302 return param303 304 if mask is None:305 return jax.tree_util.tree_map(conditional_cast, params)306 307 flat_params = flatten_dict(params)308 flat_mask, _ = jax.tree_util.tree_flatten(mask)309 310 for masked, key in zip(flat_mask, sorted(flat_params.keys())):311 if masked:312 flat_params[key] = conditional_cast(flat_params[key])313 314 return unflatten_dict(flat_params)315 316 def to_bf16(self, params: Union[dict, FrozenDict], mask: Any = None):317 r"""318 Cast the floating-point `params` to `jax.numpy.bfloat16`. This returns a new `params` tree and does not cast319 the `params` in place.320 321 This method can be used on TPU to explicitly convert the model parameters to bfloat16 precision to do full322 half-precision training or to save weights in bfloat16 for inference in order to save memory and improve speed.323 324 Arguments:325 params (`Union[Dict, FrozenDict]`):326 A `PyTree` of model parameters.327 mask (`Union[Dict, FrozenDict]`):328 A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params329 you want to cast, and should be `False` for those you want to skip.330 331 Examples:332 333 ```python334 >>> from transformers import FlaxBertModel335 336 >>> # load model337 >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")338 >>> # By default, the model parameters will be in fp32 precision, to cast these to bfloat16 precision339 >>> model.params = model.to_bf16(model.params)340 >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale)341 >>> # then pass the mask as follows342 >>> from flax import traverse_util343 344 >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")345 >>> flat_params = traverse_util.flatten_dict(model.params)346 >>> mask = {347 ... path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale"))348 ... for path in flat_params349 ... }350 >>> mask = traverse_util.unflatten_dict(mask)351 >>> model.params = model.to_bf16(model.params, mask)352 ```"""353 return self._cast_floating_to(params, jnp.bfloat16, mask)354 355 def to_fp32(self, params: Union[dict, FrozenDict], mask: Any = None):356 r"""357 Cast the floating-point `params` to `jax.numpy.float32`. This method can be used to explicitly convert the358 model parameters to fp32 precision. This returns a new `params` tree and does not cast the `params` in place.359 360 Arguments:361 params (`Union[Dict, FrozenDict]`):362 A `PyTree` of model parameters.363 mask (`Union[Dict, FrozenDict]`):364 A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params365 you want to cast, and should be `False` for those you want to skip366 367 Examples:368 369 ```python370 >>> from transformers import FlaxBertModel371 372 >>> # Download model and configuration from huggingface.co373 >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")374 >>> # By default, the model params will be in fp32, to illustrate the use of this method,375 >>> # we'll first cast to fp16 and back to fp32376 >>> model.params = model.to_f16(model.params)377 >>> # now cast back to fp32378 >>> model.params = model.to_fp32(model.params)379 ```"""380 return self._cast_floating_to(params, jnp.float32, mask)381 382 def to_fp16(self, params: Union[dict, FrozenDict], mask: Any = None):383 r"""384 Cast the floating-point `params` to `jax.numpy.float16`. This returns a new `params` tree and does not cast the385 `params` in place.386 387 This method can be used on GPU to explicitly convert the model parameters to float16 precision to do full388 half-precision training or to save weights in float16 for inference in order to save memory and improve speed.389 390 Arguments:391 params (`Union[Dict, FrozenDict]`):392 A `PyTree` of model parameters.393 mask (`Union[Dict, FrozenDict]`):394 A `PyTree` with same structure as the `params` tree. The leaves should be booleans, `True` for params395 you want to cast, and should be `False` for those you want to skip396 397 Examples:398 399 ```python400 >>> from transformers import FlaxBertModel401 402 >>> # load model403 >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")404 >>> # By default, the model params will be in fp32, to cast these to float16405 >>> model.params = model.to_fp16(model.params)406 >>> # If you want don't want to cast certain parameters (for example layer norm bias and scale)407 >>> # then pass the mask as follows408 >>> from flax import traverse_util409 410 >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")411 >>> flat_params = traverse_util.flatten_dict(model.params)412 >>> mask = {413 ... path: (path[-2] != ("LayerNorm", "bias") and path[-2:] != ("LayerNorm", "scale"))414 ... for path in flat_params415 ... }416 >>> mask = traverse_util.unflatten_dict(mask)417 >>> model.params = model.to_fp16(model.params, mask)418 ```"""419 return self._cast_floating_to(params, jnp.float16, mask)420 421 @classmethod422 def load_flax_weights(cls, resolved_archive_file):423 try:424 if resolved_archive_file.endswith(".safetensors"):425 state = safe_load_file(resolved_archive_file)426 state = unflatten_dict(state, sep=".")427 else:428 with open(resolved_archive_file, "rb") as state_f:429 state = from_bytes(cls, state_f.read())430 except (UnpicklingError, msgpack.exceptions.ExtraData) as e:431 try:432 with open(resolved_archive_file) as f:433 if f.read().startswith("version"):434 raise OSError(435 "You seem to have cloned a repository without having git-lfs installed. Please"436 " install git-lfs and run `git lfs install` followed by `git lfs pull` in the"437 " folder you cloned."438 )439 else:440 raise ValueError from e441 except (UnicodeDecodeError, ValueError):442 raise OSError(f"Unable to convert {resolved_archive_file} to Flax deserializable object. ")443 444 return state445 446 @classmethod447 def load_flax_sharded_weights(cls, shard_files):448 """449 This is the same as [`flax.serialization.from_bytes`]450 (https:lax.readthedocs.io/en/latest/_modules/flax/serialization.html#from_bytes) but for a sharded checkpoint.451 452 This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being453 loaded in the model.454 455 Args:456 shard_files (`list[str]`:457 The list of shard files to load.458 459 Returns:460 `Dict`: A nested dictionary of the model parameters, in the expected format for flax models : `{'model':461 {'params': {'...'}}}`.462 """463 464 # Load the index465 state_sharded_dict = {}466 467 for shard_file in shard_files:468 # load using msgpack utils469 try:470 with open(shard_file, "rb") as state_f:471 state = from_bytes(cls, state_f.read())472 except (UnpicklingError, msgpack.exceptions.ExtraData) as e:473 with open(shard_file) as f:474 if f.read().startswith("version"):475 raise OSError(476 "You seem to have cloned a repository without having git-lfs installed. Please"477 " install git-lfs and run `git lfs install` followed by `git lfs pull` in the"478 " folder you cloned."479 )480 else:481 raise ValueError from e482 except (UnicodeDecodeError, ValueError):483 raise OSError(f"Unable to convert {shard_file} to Flax deserializable object. ")484 485 state = flatten_dict(state, sep="/")486 state_sharded_dict.update(state)487 del state488 gc.collect()489 490 # the state dict is unflattened to the match the format of model.params491 return unflatten_dict(state_sharded_dict, sep="/")492 493 @classmethod494 def can_generate(cls) -> bool:495 """496 Returns whether this model can generate sequences with `.generate()`. Returns:497 `bool`: Whether this model can generate sequences with `.generate()`.498 """499 # Detects whether `prepare_inputs_for_generation` has been overwritten, which is a requirement for generation.500 # Alternatively, the model can also have a custom `generate` function.501 if "GenerationMixin" in str(cls.prepare_inputs_for_generation) and "GenerationMixin" in str(cls.generate):502 return False503 return True504 505 @classmethod506 def from_pretrained(507 cls,508 pretrained_model_name_or_path: Union[str, os.PathLike],509 dtype: jnp.dtype = jnp.float32,510 *model_args,511 config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,512 cache_dir: Optional[Union[str, os.PathLike]] = None,513 ignore_mismatched_sizes: bool = False,514 force_download: bool = False,515 local_files_only: bool = False,516 token: Optional[Union[str, bool]] = None,517 revision: str = "main",518 **kwargs,519 ):520 r"""521 Instantiate a pretrained flax model from a pre-trained model configuration.522 523 The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come524 pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning525 task.526 527 The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those528 weights are discarded.529 530 Parameters:531 pretrained_model_name_or_path (`str` or `os.PathLike`):532 Can be either:533 534 - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.535 - A path to a *directory* containing model weights saved using536 [`~FlaxPreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.537 - A path or url to a *pt index checkpoint file* (e.g, `./tf_model/model.ckpt.index`). In this case,538 `from_pt` should be set to `True`.539 dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`):540 The data type of the computation. Can be one of `jax.numpy.float32`, `jax.numpy.float16` (on GPUs) and541 `jax.numpy.bfloat16` (on TPUs).542 543 This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If544 specified all the computation will be performed with the given `dtype`.545 546 **Note that this only specifies the dtype of the computation and does not influence the dtype of model547 parameters.**548 549 If you wish to change the dtype of the model parameters, see [`~FlaxPreTrainedModel.to_fp16`] and550 [`~FlaxPreTrainedModel.to_bf16`].551 model_args (sequence of positional arguments, *optional*):552 All remaining positional arguments will be passed to the underlying model's `__init__` method.553 config (`Union[PretrainedConfig, str, os.PathLike]`, *optional*):554 Can be either:555 556 - an instance of a class derived from [`PretrainedConfig`],557 - a string or path valid as input to [`~PretrainedConfig.from_pretrained`].558 559 Configuration for the model to use instead of an automatically loaded configuration. Configuration can560 be automatically loaded when:561 562 - The model is a model provided by the library (loaded with the *model id* string of a pretrained563 model).564 - The model was saved using [`~PreTrainedModel.save_pretrained`] and is reloaded by supplying the565 save directory.566 - The model is loaded by supplying a local directory as `pretrained_model_name_or_path` and a567 configuration JSON file named *config.json* is found in the directory.568 cache_dir (`Union[str, os.PathLike]`, *optional*):569 Path to a directory in which a downloaded pretrained model configuration should be cached if the570 standard cache should not be used.571 from_pt (`bool`, *optional*, defaults to `False`):572 Load the model weights from a PyTorch checkpoint save file (see docstring of573 `pretrained_model_name_or_path` argument).574 ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`):575 Whether or not to raise an error if some of the weights from the checkpoint do not have the same size576 as the weights of the model (if for instance, you are instantiating a model with 10 labels from a577 checkpoint with 3 labels).578 force_download (`bool`, *optional*, defaults to `False`):579 Whether or not to force the (re-)download of the model weights and configuration files, overriding the580 cached versions if they exist.581 resume_download:582 Deprecated and ignored. All downloads are now resumed by default when possible.583 Will be removed in v5 of Transformers.584 proxies (`dict[str, str]`, *optional*):585 A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',586 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.587 local_files_only(`bool`, *optional*, defaults to `False`):588 Whether or not to only look at local files (i.e., do not try to download the model).589 token (`str` or `bool`, *optional*):590 The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use591 the token generated when running `hf auth login` (stored in `~/.huggingface`).592 revision (`str`, *optional*, defaults to `"main"`):593 The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a594 git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any595 identifier allowed by git.596 597 598 <Tip>599 600 To test a pull request you made on the Hub, you can pass `revision="refs/pr/<pr_number>"`.601 602 </Tip>603 604 subfolder (`str`, *optional*, defaults to `""`):605 In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can606 specify the folder name here.607 kwargs (remaining dictionary of keyword arguments, *optional*):608 Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,609 `output_attentions=True`). Behaves differently depending on whether a `config` is provided or610 automatically loaded:611 612 - If a configuration is provided with `config`, `**kwargs` will be directly passed to the613 underlying model's `__init__` method (we assume all relevant updates to the configuration have614 already been done)615 - If a configuration is not provided, `kwargs` will be first passed to the configuration class616 initialization function ([`~PretrainedConfig.from_pretrained`]). Each key of `kwargs` that617 corresponds to a configuration attribute will be used to override said attribute with the618 supplied `kwargs` value. Remaining keys that do not correspond to any configuration attribute619 will be passed to the underlying model's `__init__` function.620 621 Examples:622 623 ```python624 >>> from transformers import BertConfig, FlaxBertModel625 626 >>> # Download model and configuration from huggingface.co and cache.627 >>> model = FlaxBertModel.from_pretrained("google-bert/bert-base-cased")628 >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable).629 >>> model = FlaxBertModel.from_pretrained("./test/saved_model/")630 >>> # Loading from a PyTorch checkpoint file instead of a PyTorch model (slower, for example purposes, not runnable).631 >>> config = BertConfig.from_json_file("./pt_model/config.json")632 >>> model = FlaxBertModel.from_pretrained("./pt_model/pytorch_model.bin", from_pt=True, config=config)633 ```"""634 from_pt = kwargs.pop("from_pt", False)635 resume_download = kwargs.pop("resume_download", None)636 proxies = kwargs.pop("proxies", None)637 use_auth_token = kwargs.pop("use_auth_token", None)638 trust_remote_code = kwargs.pop("trust_remote_code", None)639 from_pipeline = kwargs.pop("_from_pipeline", None)640 from_auto_class = kwargs.pop("_from_auto", False)641 _do_init = kwargs.pop("_do_init", True)642 subfolder = kwargs.pop("subfolder", "")643 commit_hash = kwargs.pop("_commit_hash", None)644 645 # Not relevant for Flax Models646 _ = kwargs.pop("adapter_kwargs", None)647 648 if use_auth_token is not None:649 warnings.warn(650 "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",651 FutureWarning,652 )653 if token is not None:654 raise ValueError(655 "`token` and `use_auth_token` are both specified. Please set only the argument `token`."656 )657 token = use_auth_token658 659 if trust_remote_code is True:660 logger.warning(661 "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is"662 " ignored."663 )664 665 user_agent = {"file_type": "model", "framework": "flax", "from_auto_class": from_auto_class}666 if from_pipeline is not None:667 user_agent["using_pipeline"] = from_pipeline668 669 if is_offline_mode() and not local_files_only:670 logger.info("Offline mode: forcing local_files_only=True")671 local_files_only = True672 673 # Load config if we don't provide a configuration674 if not isinstance(config, PretrainedConfig):675 config_path = config if config is not None else pretrained_model_name_or_path676 config, model_kwargs = cls.config_class.from_pretrained(677 config_path,678 cache_dir=cache_dir,679 return_unused_kwargs=True,680 force_download=force_download,681 resume_download=resume_download,682 proxies=proxies,683 local_files_only=local_files_only,684 token=token,685 revision=revision,686 subfolder=subfolder,687 _from_auto=from_auto_class,688 _from_pipeline=from_pipeline,689 _commit_hash=commit_hash,690 **kwargs,691 )692 else:693 model_kwargs = kwargs.copy()694 695 if commit_hash is None:696 commit_hash = getattr(config, "_commit_hash", None)697 698 # Add the dtype to model_kwargs699 model_kwargs["dtype"] = dtype700 701 # This variable will flag if we're loading a sharded checkpoint. In this case the archive file is just the702 # index of the files.703 is_sharded = False704 705 # Load model706 if pretrained_model_name_or_path is not None:707 pretrained_model_name_or_path = str(pretrained_model_name_or_path)708 is_local = os.path.isdir(pretrained_model_name_or_path)709 if is_local:710 if os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME)):711 # Load from a Flax checkpoint712 archive_file = os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME)713 elif os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_INDEX_NAME)):714 # Load from a sharded Flax checkpoint715 archive_file = os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_INDEX_NAME)716 is_sharded = True717 elif is_safetensors_available() and os.path.isfile(718 os.path.join(pretrained_model_name_or_path, subfolder, SAFE_WEIGHTS_NAME)719 ):720 # Load from a safetensors checkpoint721 archive_file = os.path.join(pretrained_model_name_or_path, subfolder, SAFE_WEIGHTS_NAME)722 elif is_safetensors_available() and os.path.isfile(723 os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME)724 ):725 # Load from a safetensors checkpoint726 archive_file = os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME)727 elif from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)):728 # Load from a PyTorch checkpoint729 archive_file = os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)730 elif from_pt and os.path.isfile(731 os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_INDEX_NAME)732 ):733 # Load from a sharded pytorch checkpoint734 archive_file = os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_INDEX_NAME)735 is_sharded = True736 # At this stage we don't have a weight file so we will raise an error.737 elif is_safetensors_available() and os.path.isfile(738 os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME)739 ):740 # Load from a sharded safetensors checkpoint741 archive_file = os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME)742 is_sharded = True743 raise NotImplementedError("Support for sharded checkpoints using safetensors is coming soon!")744 elif os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)):745 raise OSError(746 f"Error no file named {FLAX_WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} "747 "but there is a file for PyTorch weights. Use `from_pt=True` to load this model from those "748 "weights."749 )750 else:751 raise OSError(752 f"Error no file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME} found in directory "753 f"{pretrained_model_name_or_path}."754 )755 elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)):756 archive_file = pretrained_model_name_or_path757 is_local = True758 elif is_remote_url(pretrained_model_name_or_path):759 filename = pretrained_model_name_or_path760 resolved_archive_file = download_url(pretrained_model_name_or_path)761 else:762 if from_pt:763 filename = WEIGHTS_NAME764 else:765 filename = FLAX_WEIGHTS_NAME766 767 try:768 # Load from URL or cache if already cached769 cached_file_kwargs = {770 "cache_dir": cache_dir,771 "force_download": force_download,772 "proxies": proxies,773 "resume_download": resume_download,774 "local_files_only": local_files_only,775 "token": token,776 "user_agent": user_agent,777 "revision": revision,778 "subfolder": subfolder,779 "_raise_exceptions_for_gated_repo": False,780 "_raise_exceptions_for_missing_entries": False,781 "_commit_hash": commit_hash,782 }783 resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)784 785 # Maybe the checkpoint is sharded, we try to grab the index name in this case.786 if resolved_archive_file is None and filename == FLAX_WEIGHTS_NAME:787 resolved_archive_file = cached_file(788 pretrained_model_name_or_path, FLAX_WEIGHTS_INDEX_NAME, **cached_file_kwargs789 )790 if resolved_archive_file is not None:791 is_sharded = True792 793 # Maybe the checkpoint is pytorch sharded, we try to grab the pytorch index name in this case.794 if resolved_archive_file is None and from_pt:795 resolved_archive_file = cached_file(796 pretrained_model_name_or_path, WEIGHTS_INDEX_NAME, **cached_file_kwargs797 )798 if resolved_archive_file is not None:799 is_sharded = True800 801 # If we still haven't found anything, look for `safetensors`.802 if resolved_archive_file is None:803 # No support for sharded safetensors yet, so we'll raise an error if that's all we find.804 filename = SAFE_WEIGHTS_NAME805 resolved_archive_file = cached_file(806 pretrained_model_name_or_path, SAFE_WEIGHTS_NAME, **cached_file_kwargs807 )808 809 # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None810 # result when internet is up, the repo and revision exist, but the file does not.811 if resolved_archive_file is None:812 # Otherwise, maybe there is a TF or Torch model file. We try those to give a helpful error813 # message.814 has_file_kwargs = {815 "revision": revision,816 "proxies": proxies,817 "token": token,818 "cache_dir": cache_dir,819 "local_files_only": local_files_only,820 }821 if has_file(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME, **has_file_kwargs):822 is_sharded = True823 raise NotImplementedError(824 "Support for sharded checkpoints using safetensors is coming soon!"825 )826 elif has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs):827 raise OSError(828 f"{pretrained_model_name_or_path} does not appear to have a file named"829 f" {FLAX_WEIGHTS_NAME} but there is a file for PyTorch weights. Use `from_pt=True` to"830 " load this model from those weights."831 )832 elif has_file(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME, **has_file_kwargs):833 raise OSError(834 f"{pretrained_model_name_or_path} does not appear to have a file named"835 f" {FLAX_WEIGHTS_INDEX_NAME} but there is a sharded file for PyTorch weights. Use"836 " `from_pt=True` to load this model from those weights."837 )838 else:839 raise OSError(840 f"{pretrained_model_name_or_path} does not appear to have a file named"841 f" {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}."842 )843 except OSError:844 # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted845 # to the original exception.846 raise847 except Exception:848 # For any other exception, we throw a generic error.849 raise OSError(850 f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it"851 " from 'https://huggingface.co/models', make sure you don't have a local directory with the"852 f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"853 f" directory containing a file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}."854 )855 856 if is_local:857 logger.info(f"loading weights file {archive_file}")858 resolved_archive_file = archive_file859 filename = resolved_archive_file.split(os.path.sep)[-1]860 else:861 logger.info(f"loading weights file {filename} from cache at {resolved_archive_file}")862 else:863 resolved_archive_file = None864 865 # We'll need to download and cache each checkpoint shard if the checkpoint is sharded.866 if is_sharded:867 # resolved_archive_file becomes a list of files that point to the different checkpoint shards in this case.868 resolved_archive_file, _ = get_checkpoint_shard_files(869 pretrained_model_name_or_path,870 resolved_archive_file,871 cache_dir=cache_dir,872 force_download=force_download,873 proxies=proxies,874 resume_download=resume_download,875 local_files_only=local_files_only,876 token=token,877 user_agent=user_agent,878 revision=revision,879 subfolder=subfolder,880 _commit_hash=commit_hash,881 )882 883 safetensors_from_pt = False884 if filename == SAFE_WEIGHTS_NAME:885 with safe_open(resolved_archive_file, framework="flax") as f:886 safetensors_metadata = f.metadata()887 if safetensors_metadata is None or safetensors_metadata.get("format") not in ["pt", "tf", "flax"]:888 raise OSError(889 f"The safetensors archive passed at {resolved_archive_file} does not contain the valid metadata."890 " Make sure you save your model with the `save_pretrained` method."891 )892 safetensors_from_pt = safetensors_metadata.get("format") == "pt"893 894 # init random models895 model = cls(config, *model_args, _do_init=_do_init, **model_kwargs)896 897 if from_pt or safetensors_from_pt:898 state = load_pytorch_checkpoint_in_flax_state_dict(model, resolved_archive_file, is_sharded)899 else:900 if is_sharded:901 state = cls.load_flax_sharded_weights(resolved_archive_file)902 else:903 state = cls.load_flax_weights(resolved_archive_file)904 # make sure all arrays are stored as jnp.arrays905 # NOTE: This is to prevent a bug this will be fixed in Flax >= v0.3.4:906 # https://github.com/google/flax/issues/1261907 if _do_init:908 state = jax.tree_util.tree_map(jnp.array, state)909 else:910 # keep the params on CPU if we don't want to initialize911 state = jax.tree_util.tree_map(lambda x: jax.device_put(x, jax.local_devices(backend="cpu")[0]), state)912 913 if "batch_stats" in state: # if flax model contains batch norm layers914 # if model is base model only use model_prefix key915 if (916 cls.base_model_prefix not in dict(model.params_shape_tree["params"])917 and cls.base_model_prefix in state["params"]918 ):919 state["params"] = state["params"][cls.base_model_prefix]920 state["batch_stats"] = state["batch_stats"][cls.base_model_prefix]921 922 # if model is head model and we are loading weights from base model923 # we initialize new params dict with base_model_prefix924 if (925 cls.base_model_prefix in dict(model.params_shape_tree["params"])926 and cls.base_model_prefix not in state["params"]927 ):928 state = {929 "params": {cls.base_model_prefix: state["params"]},930 "batch_stats": {cls.base_model_prefix: state["batch_stats"]},931 }932 933 else:934 # if model is base model only use model_prefix key935 if cls.base_model_prefix not in dict(model.params_shape_tree) and cls.base_model_prefix in state:936 state = state[cls.base_model_prefix]937 938 # if model is head model and we are loading weights from base model939 # we initialize new params dict with base_model_prefix940 if cls.base_model_prefix in dict(model.params_shape_tree) and cls.base_model_prefix not in state:941 state = {cls.base_model_prefix: state}942 943 # flatten dicts944 state = flatten_dict(state)945 946 random_state = flatten_dict(unfreeze(model.params if _do_init else model.params_shape_tree))947 948 missing_keys = model.required_params - set(state.keys())949 unexpected_keys = set(state.keys()) - model.required_params950 951 # Disabling warning when porting pytorch weights to flax, flax does not uses num_batches_tracked952 for unexpected_key in unexpected_keys.copy():953 if "num_batches_tracked" in unexpected_key[-1]:954 unexpected_keys.remove(unexpected_key)955 956 if missing_keys and not _do_init:957 logger.warning(958 f"The checkpoint {pretrained_model_name_or_path} is missing required keys: {missing_keys}. "959 "Make sure to call model.init_weights to initialize the missing weights."960 )961 cls._missing_keys = missing_keys962 963 # Mismatched keys contains tuples key/shape1/shape2 of weights in the checkpoint that have a shape not964 # matching the weights in the model.965 mismatched_keys = []966 for key in state:967 if key in random_state and state[key].shape != random_state[key].shape:968 if ignore_mismatched_sizes:969 mismatched_keys.append((key, state[key].shape, random_state[key].shape))970 state[key] = random_state[key]971 else:972 raise ValueError(973 f"Trying to load the pretrained weight for {key} failed: checkpoint has shape "974 f"{state[key].shape} which is incompatible with the model shape {random_state[key].shape}. "975 "Using `ignore_mismatched_sizes=True` if you really want to load this checkpoint inside this "976 "model."977 )978 979 # add missing keys as random parameters if we are initializing980 if missing_keys and _do_init:981 for missing_key in missing_keys:982 state[missing_key] = random_state[missing_key]983 984 # remove unexpected keys to not be saved again985 for unexpected_key in unexpected_keys:986 del state[unexpected_key]987 988 if len(unexpected_keys) > 0:989 logger.warning(990 f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"991 f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"992 f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or"993 " with another architecture (e.g. initializing a BertForSequenceClassification model from a"994 " BertForPreTraining model).\n- This IS NOT expected if you are initializing"995 f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical"996 " (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."997 )998 else:999 logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")1000 1001 if len(missing_keys) > 0:1002 logger.warning(1003 f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"1004 f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"1005 " TRAIN this model on a down-stream task to be able to use it for predictions and inference."1006 )1007 elif len(mismatched_keys) == 0:1008 logger.info(1009 f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"1010 f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint"1011 f" was trained on, you can already use {model.__class__.__name__} for predictions without further"1012 " training."1013 )1014 if len(mismatched_keys) > 0:1015 mismatched_warning = "\n".join(1016 [1017 f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"1018 for key, shape1, shape2 in mismatched_keys1019 ]1020 )1021 logger.warning(1022 f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"1023 f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not"1024 f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able"1025 " to use it for predictions and inference."1026 )1027 1028 # dictionary of key: dtypes for the model params1029 param_dtypes = jax.tree_util.tree_map(lambda x: x.dtype, state)1030 # extract keys of parameters not in jnp.float321031 fp16_params = [k for k in param_dtypes if param_dtypes[k] == jnp.float16]1032 bf16_params = [k for k in param_dtypes if param_dtypes[k] == jnp.bfloat16]1033 1034 # raise a warning if any of the parameters are not in jnp.float321035 if len(fp16_params) > 0:1036 logger.warning(1037 f"Some of the weights of {model.__class__.__name__} were initialized in float16 precision from "1038 f"the model checkpoint at {pretrained_model_name_or_path}:\n{fp16_params}\n"1039 "You should probably UPCAST the model weights to float32 if this was not intended. "1040 "See [`~FlaxPreTrainedModel.to_fp32`] for further information on how to do this."1041 )1042 1043 if len(bf16_params) > 0:1044 logger.warning(1045 f"Some of the weights of {model.__class__.__name__} were initialized in bfloat16 precision from "1046 f"the model checkpoint at {pretrained_model_name_or_path}:\n{bf16_params}\n"1047 "You should probably UPCAST the model weights to float32 if this was not intended. "1048 "See [`~FlaxPreTrainedModel.to_fp32`] for further information on how to do this."1049 )1050 1051 # If it is a model with generation capabilities, attempt to load the generation config1052 if model.can_generate():1053 try:1054 model.generation_config = GenerationConfig.from_pretrained(1055 pretrained_model_name_or_path,1056 cache_dir=cache_dir,1057 force_download=force_download,1058 resume_download=resume_download,1059 proxies=proxies,1060 local_files_only=local_files_only,1061 token=token,1062 revision=revision,1063 subfolder=subfolder,1064 _from_auto=from_auto_class,1065 _from_pipeline=from_pipeline,1066 **kwargs,1067 )1068 except OSError:1069 logger.info(1070 "Generation config file not found, using a generation config created from the model config."1071 )1072 pass1073 1074 if _do_init:1075 # set correct parameters1076 model.params = unflatten_dict(state)1077 return model1078 else:1079 return model, unflatten_dict(state)1080 1081 def save_pretrained(1082 self,1083 save_directory: Union[str, os.PathLike],1084 params=None,1085 push_to_hub=False,1086 max_shard_size="10GB",1087 token: Optional[Union[str, bool]] = None,1088 safe_serialization: bool = False,1089 **kwargs,1090 ):1091 """1092 Save a model and its configuration file to a directory, so that it can be re-loaded using the1093 `[`~FlaxPreTrainedModel.from_pretrained`]` class method1094 1095 Arguments:1096 save_directory (`str` or `os.PathLike`):1097 Directory to which to save. Will be created if it doesn't exist.1098 push_to_hub (`bool`, *optional*, defaults to `False`):1099 Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the1100 repository you want to push to with `repo_id` (will default to the name of `save_directory` in your1101 namespace).1102 max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):1103 The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size1104 lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).1105 1106 <Tip warning={true}>1107 1108 If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard1109 which will be bigger than `max_shard_size`.1110 1111 </Tip>1112 1113 token (`str` or `bool`, *optional*):1114 The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use1115 the token generated when running `hf auth login` (stored in `~/.huggingface`).1116 kwargs (`dict[str, Any]`, *optional*):1117 Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.1118 safe_serialization (`bool`, *optional*, defaults to `False`):1119 Whether to save the model using `safetensors` or through msgpack.1120 """1121 use_auth_token = kwargs.pop("use_auth_token", None)1122 1123 if use_auth_token is not None:1124 warnings.warn(1125 "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",1126 FutureWarning,1127 )1128 if token is not None:1129 raise ValueError(1130 "`token` and `use_auth_token` are both specified. Please set only the argument `token`."1131 )1132 token = use_auth_token1133 1134 if token is not None:1135 kwargs["token"] = token1136 1137 if os.path.isfile(save_directory):1138 logger.error(f"Provided path ({save_directory}) should be a directory, not a file")1139 return1140 1141 os.makedirs(save_directory, exist_ok=True)1142 1143 if push_to_hub:1144 commit_message = kwargs.pop("commit_message", None)1145 repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])1146 repo_id = self._create_repo(repo_id, **kwargs)1147 files_timestamps = self._get_files_timestamps(save_directory)1148 1149 # get abs dir1150 save_directory = os.path.abspath(save_directory)1151 # save config as well1152 self.config.architectures = [self.__class__.__name__[4:]]1153 1154 # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be1155 # loaded from the Hub.1156 if self._auto_class is not None:1157 custom_object_save(self, save_directory, config=self.config)1158 1159 self.config.save_pretrained(save_directory)1160 if self.can_generate():1161 self.generation_config.save_pretrained(save_directory)1162 1163 # save model1164 weights_name = SAFE_WEIGHTS_NAME if safe_serialization else FLAX_WEIGHTS_NAME1165 output_model_file = os.path.join(save_directory, weights_name)1166 1167 shards, index = flax_shard_checkpoint(params if params is not None else self.params, max_shard_size)1168 # Clean the folder from a previous save1169 for filename in os.listdir(save_directory):1170 full_filename = os.path.join(save_directory, filename)1171 weights_no_suffix = weights_name.replace(".bin", "").replace(".safetensors", "")1172 if filename.startswith(weights_no_suffix) and os.path.isfile(full_filename) and filename not in shards:1173 os.remove(full_filename)1174 1175 if index is None:1176 if safe_serialization:1177 params = params if params is not None else self.params1178 flat_dict = flatten_dict(params, sep=".")1179 safe_save_file(flat_dict, output_model_file, metadata={"format": "flax"})1180 else:1181 with open(output_model_file, "wb") as f:1182 params = params if params is not None else self.params1183 model_bytes = to_bytes(params)1184 f.write(model_bytes)1185 1186 else:1187 save_index_file = os.path.join(save_directory, FLAX_WEIGHTS_INDEX_NAME)1188 # Save the index as well1189 with open(save_index_file, "w", encoding="utf-8") as f:1190 content = json.dumps(index, indent=2, sort_keys=True) + "\n"1191 f.write(content)1192 logger.info(1193 f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "1194 f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the "1195 f"index located at {save_index_file}."1196 )1197 for shard_file, shard in shards.items():1198 # the shard item are unflattened, to save them we need to flatten them again1199 with open(os.path.join(save_directory, shard_file), mode="wb") as f:1200 params = unflatten_dict(shard, sep="/")