codingutf8 2021 the google flax team s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license returns the size in bytes occupied by one parameter of type dtype example py dtypebytesizenp float32 4 splits a model state dictionary in subcheckpoints so that the final size of each subcheckpoint does not exceed a given size the subcheckpoints are determined by iterating through the statedict in the order of its keys so there is no optimization made to make each subcheckpoint as close as possible to the maximum size passed for example if the limit is 10gb and we have weights of sizes 6gb 6gb 2gb 6gb 2gb 2gb they will get sharded as 6gb 62gb 622gb and not 622gb 62gb 6gb tip warningtrue if one of the model s weight is bigger that maxshardsize it will end up in its own subcheckpoint which will have a size greater than maxshardsize tip args params uniondict frozendict a pytree of model parameters maxshardsize int or str optional defaults to 10gb the maximum size of each subcheckpoint if expressed as a string needs to be digits followed by a unit like 5mb flatten the weights to chunk if this weight is going to tip up over the maximal size we split add the last block if we only have one shard we return it otherwise let s build the index add the metadata configclass none basemodelprefix maininputname inputids autoclass none missingkeys set def init self config pretrainedconfig module nn module inputshape tuple 1 1 seed int 0 dtype jnp dtype jnp float32 doinit bool true if config is none raise valueerrorconfig cannot be none if module is none raise valueerrormodule cannot be none those are private to be exposed as typed property on derived classes self config config self module module those are public as their type is generic to every derived classes self key prngkeyseed self dtype dtype self inputshape inputshape self generationconfig generationconfig frommodelconfigconfig if self cangenerate else none to check if the model was intialized automatically self isinitialized doinit if doinit randomly initialized parameters randomparams self initweightsself key inputshape paramsshapetree jax evalshapelambda params params randomparams else initfn partialself initweights inputshapeinputshape paramsshapetree jax evalshapeinitfn self key logger info model weights are not initialized as doinit is set to false fmake sure to call self class name initweights manually to initialize the weights get the shape of the parameters self paramsshapetree paramsshapetree save requiredparams as set self requiredparams setflattendictunfreezeparamsshapetree keys initialize the parameters if doinit self params randomparams def initweightsself rng jax random prngkey inputshape tuple params frozendict none dict raise notimplementederrorfinit method has to be implemented for self def enablegradientcheckpointingself raise notimplementederrorfgradient checkpointing method has to be implemented for self classmethod def fromconfigcls config kwargs return clsconfig kwargs property def frameworkself str return flax property def configself pretrainedconfig return self config property def moduleself nn module return self module property def paramsself uniondict frozendict if not self isinitialized raise valueerror params cannot be accessed from model when the model is created with doinitfalse you must call initweights manually and store the params outside of the model and pass it explicitly where needed return self params property def requiredparamsself set return self requiredparams property def paramsshapetreeself dict return self paramsshapetree params setter def paramsself params uniondict frozendict don t set params if the model is not initialized if not self isinitialized raise valueerror params cannot be set from model when the model is created with doinitfalse you store the params outside of the model if isinstanceparams frozendict params unfreezeparams paramkeys setflattendictparams keys if lenself requiredparams paramkeys 0 raise valueerror some parameters are missing make sure that params include the following fparameters self requiredparams paramkeys self params params def castfloatingtoself params uniondict frozendict dtype jnp dtype mask any none any taken from https github comdeepmindjmpblob3a8318abc3292be38582794dbf7b094e6583b192jmpsrcpolicy pyl27 def conditionalcastparam if isinstanceparam jnp ndarray and jnp issubdtypeparam dtype jnp floating param param astypedtype return param if mask is none return jax treeutil treemapconditionalcast params flatparams flattendictparams flatmask jax treeutil treeflattenmask for masked key in zipflatmask flatparams keys if masked param flatparamskey flatparamskey conditionalcastparam return unflattendictflatparams def tobf16self params uniondict frozendict mask any none r cast the floatingpoint params to jax numpy bfloat16 this returns a new params tree and does not cast the params in place this method can be used on tpu to explicitly convert the model parameters to bfloat16 precision to do full halfprecision training or to save weights in bfloat16 for inference in order to save memory and improve speed arguments params uniondict frozendict a pytree of model parameters mask uniondict frozendict a pytree with same structure as the params tree the leaves should be booleans true for params you want to cast and should be false for those you want to skip examples python from transformers import flaxbertmodel load model model flaxbertmodel frompretrainedbertbasecased by default the model parameters will be in fp32 precision to cast these to bfloat16 precision model params model tobf16model params if you want don t want to cast certain parameters for example layer norm bias and scale then pass the mask as follows from flax import traverseutil model flaxbertmodel frompretrainedbertbasecased flatparams traverseutil flattendictmodel params mask path path2 layernorm bias and path2 layernorm scale for path in flatparams mask traverseutil unflattendictmask model params model tobf16model params mask return self castfloatingtoparams jnp bfloat16 mask def tofp32self params uniondict frozendict mask any none r cast the floatingpoint parmas to jax numpy float32 this method can be used to explicitly convert the model parameters to fp32 precision this returns a new params tree and does not cast the params in place arguments params uniondict frozendict a pytree of model parameters mask uniondict frozendict a pytree with same structure as the params tree the leaves should be booleans true for params you want to cast and should be false for those you want to skip examples python from transformers import flaxbertmodel download model and configuration from huggingface co model flaxbertmodel frompretrainedbertbasecased by default the model params will be in fp32 to illustrate the use of this method we ll first cast to fp16 and back to fp32 model params model tof16model params now cast back to fp32 model params model tofp32model params return self castfloatingtoparams jnp float32 mask def tofp16self params uniondict frozendict mask any none r cast the floatingpoint parmas to jax numpy float16 this returns a new params tree and does not cast the params in place this method can be used on gpu to explicitly convert the model parameters to float16 precision to do full halfprecision training or to save weights in float16 for inference in order to save memory and improve speed arguments params uniondict frozendict a pytree of model parameters mask uniondict frozendict a pytree with same structure as the params tree the leaves should be booleans true for params you want to cast and should be false for those you want to skip examples python from transformers import flaxbertmodel load model model flaxbertmodel frompretrainedbertbasecased by default the model params will be in fp32 to cast these to float16 model params model tofp16model params if you want don t want to cast certain parameters for example layer norm bias and scale then pass the mask as follows from flax import traverseutil model flaxbertmodel frompretrainedbertbasecased flatparams traverseutil flattendictmodel params mask path path2 layernorm bias and path2 layernorm scale for path in flatparams mask traverseutil unflattendictmask model params model tofp16model params mask return self castfloatingtoparams jnp float16 mask classmethod def loadflaxweightscls resolvedarchivefile try if resolvedarchivefile endswith safetensors state safeloadfileresolvedarchivefile state unflattendictstate sep else with openresolvedarchivefile rb as statef state frombytescls statef read except unpicklingerror msgpack exceptions extradata as e try with openresolvedarchivefile as f if f read startswithversion raise oserror you seem to have cloned a repository without having gitlfs installed please install gitlfs and run git lfs install followed by git lfs pull in the folder you cloned else raise valueerror from e except unicodedecodeerror valueerror raise environmenterrorfunable to convert resolvedarchivefile to flax deserializable object return state classmethod def loadflaxshardedweightscls shardfiles load the index stateshardeddict for shardfile in shardfiles load using msgpack utils try with openshardfile rb as statef state frombytescls statef read except unpicklingerror msgpack exceptions extradata as e with openshardfile as f if f read startswithversion raise oserror you seem to have cloned a repository without having gitlfs installed please install gitlfs and run git lfs install followed by git lfs pull in the folder you cloned else raise valueerror from e except unicodedecodeerror valueerror raise environmenterrorfunable to convert shardfile to flax deserializable object state flattendictstate sep stateshardeddict updatestate del state gc collect the state dict is unflattened to the match the format of model params return unflattendictstateshardeddict sep classmethod def cangeneratecls bool detects whether prepareinputsforgeneration has been overwritten which is a requirement for generation alternativelly the model can also have a custom generate function if generationmixin in strcls prepareinputsforgeneration and generationmixin in strcls generate return false return true classmethod def frompretrained cls pretrainedmodelnameorpath unionstr os pathlike dtype jnp dtype jnp float32 modelargs config optionalunionpretrainedconfig str os pathlike none cachedir optionalunionstr os pathlike none ignoremismatchedsizes bool false forcedownload bool false localfilesonly bool false token optionalunionstr bool none revision str main kwargs r instantiate a pretrained flax model from a pretrained model configuration the warning weights from xxx not initialized from pretrained model means that the weights of xxx do not come pretrained with the rest of the model it is up to you to train those weights with a downstream finetuning task the warning weights from xxx not used in yyy means that the layer xxx is not used by yyy therefore those weights are discarded parameters pretrainedmodelnameorpath str or os pathlike can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the rootlevel like bertbaseuncased or namespaced under a user or organization name like dbmdzbertbasegermancased a path to a directory containing model weights saved using flaxpretrainedmodel savepretrained e g mymodeldirectory a path or url to a pt index checkpoint file e g tfmodelmodel ckpt index in this case frompt should be set to true dtype jax numpy dtype optional defaults to jax numpy float32 the data type of the computation can be one of jax numpy float32 jax numpy float16 on gpus and jax numpy bfloat16 on tpus this can be used to enable mixedprecision training or halfprecision inference on gpus or tpus if specified all the computation will be performed with the given dtype note that this only specifies the dtype of the computation and does not influence the dtype of model parameters if you wish to change the dtype of the model parameters see flaxpretrainedmodel tofp16 and flaxpretrainedmodel tobf16 modelargs sequence of positional arguments optional all remaining positional arguments will be passed to the underlying model s init method config unionpretrainedconfig str os pathlike optional can be either an instance of a class derived from pretrainedconfig a string or path valid as input to pretrainedconfig frompretrained configuration for the model to use instead of an automatically loaded configuration configuration can be automatically loaded when the model is a model provided by the library loaded with the model id string of a pretrained model the model was saved using pretrainedmodel savepretrained and is reloaded by supplying the save directory the model is loaded by supplying a local directory as pretrainedmodelnameorpath and a configuration json file named config json is found in the directory cachedir unionstr os pathlike optional path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used frompt bool optional defaults to false load the model weights from a pytorch checkpoint save file see docstring of pretrainedmodelnameorpath argument ignoremismatchedsizes bool optional defaults to false whether or not to raise an error if some of the weights from the checkpoint do not have the same size as the weights of the model if for instance you are instantiating a model with 10 labels from a checkpoint with 3 labels forcedownload bool optional defaults to false whether or not to force the redownload of the model weights and configuration files overriding the cached versions if they exist resumedownload bool optional defaults to false whether or not to delete incompletely received files will attempt to resume the download if such a file exists proxies dictstr str optional a dictionary of proxy servers to use by protocol or endpoint e g http foo bar 3128 http hostname foo bar 4012 the proxies are used on each request localfilesonlybool optional defaults to false whether or not to only look at local files i e do not try to download the model token str or bool optional the token to use as http bearer ization for remote files if true or not specified will use the token generated when running huggingfacecli login stored in huggingface revision str optional defaults to main the specific model version to use it can be a branch name a tag name or a commit id since we use a gitbased system for storing models and other artifacts on huggingface co so revision can be any identifier allowed by git tip to test a pull request you made on the hub you can pass revisionrefsprprnumber tip subfolder str optional defaults to in case the relevant files are located inside a subfolder of the model repo on huggingface co you can specify the folder name here kwargs remaining dictionary of keyword arguments optional can be used to update the configuration object after it being loaded and initiate the model e g outputattentionstrue behaves differently depending on whether a config is provided or automatically loaded if a configuration is provided with config kwargs will be directly passed to the underlying model s init method we assume all relevant updates to the configuration have already been done if a configuration is not provided kwargs will be first passed to the configuration class initialization function pretrainedconfig frompretrained each key of kwargs that corresponds to a configuration attribute will be used to override said attribute with the supplied kwargs value remaining keys that do not correspond to any configuration attribute will be passed to the underlying model s init function examples python from transformers import bertconfig flaxbertmodel download model and configuration from huggingface co and cache model flaxbertmodel frompretrainedbertbasecased model was saved using savepretrained testsavedmodel for example purposes not runnable model flaxbertmodel frompretrained testsavedmodel loading from a pytorch checkpoint file instead of a pytorch model slower for example purposes not runnable config bertconfig fromjsonfile ptmodelconfig json model flaxbertmodel frompretrained ptmodelpytorchmodel bin frompttrue configconfig frompt kwargs popfrompt false resumedownload kwargs popresumedownload false proxies kwargs popproxies none useauthtoken kwargs popuseauthtoken none trustremotecode kwargs poptrustremotecode none frompipeline kwargs popfrompipeline none fromautoclass kwargs popfromauto false doinit kwargs popdoinit true subfolder kwargs popsubfolder commithash kwargs popcommithash none not relevant for flax models kwargs popadapterkwargs none if useauthtoken is not none warnings warn the useauthtoken argument is deprecated and will be removed in v5 of transformers please use token instead futurewarning if token is not none raise valueerror token and useauthtoken are both specified please set only the argument token token useauthtoken if trustremotecode is true logger warning the argument trustremotecode is to be used with auto classes it has no effect here and is ignored useragent filetype model framework flax fromautoclass fromautoclass if frompipeline is not none useragentusingpipeline frompipeline if isofflinemode and not localfilesonly logger infooffline mode forcing localfilesonlytrue localfilesonly true load config if we don t provide a configuration if not isinstanceconfig pretrainedconfig configpath config if config is not none else pretrainedmodelnameorpath config modelkwargs cls configclass frompretrained configpath cachedircachedir returnunusedkwargstrue forcedownloadforcedownload resumedownloadresumedownload proxiesproxies localfilesonlylocalfilesonly tokentoken revisionrevision subfoldersubfolder fromautofromautoclass frompipelinefrompipeline commithashcommithash kwargs else modelkwargs kwargs copy if commithash is none commithash getattrconfig commithash none add the dtype to modelkwargs modelkwargsdtype dtype this variable will flag if we re loading a sharded checkpoint in this case the archive file is just the index of the files issharded false load model if pretrainedmodelnameorpath is not none pretrainedmodelnameorpath strpretrainedmodelnameorpath islocal os path isdirpretrainedmodelnameorpath if os path isdirpretrainedmodelnameorpath if os path isfileos path joinpretrainedmodelnameorpath subfolder flaxweightsname load from a flax checkpoint archivefile os path joinpretrainedmodelnameorpath subfolder flaxweightsname elif os path isfileos path joinpretrainedmodelnameorpath subfolder flaxweightsindexname load from a sharded flax checkpoint archivefile os path joinpretrainedmodelnameorpath subfolder flaxweightsindexname issharded true elif issafetensorsavailable and os path isfile os path joinpretrainedmodelnameorpath safeweightsname load from a safetensors checkpoint archivefile os path joinpretrainedmodelnameorpath safeweightsname elif frompt and os path isfileos path joinpretrainedmodelnameorpath subfolder weightsname load from a pytorch checkpoint archivefile os path joinpretrainedmodelnameorpath subfolder weightsname elif frompt and os path isfile os path joinpretrainedmodelnameorpath subfolder weightsindexname load from a sharded pytorch checkpoint archivefile os path joinpretrainedmodelnameorpath subfolder weightsindexname issharded true at this stage we don t have a weight file so we will raise an error elif issafetensorsavailable and os path isfile os path joinpretrainedmodelnameorpath safeweightsindexname load from a sharded safetensors checkpoint archivefile os path joinpretrainedmodelnameorpath safeweightsindexname issharded true raise notimplementederrorsupport for sharded checkpoints using safetensors is coming soon elif os path isfileos path joinpretrainedmodelnameorpath subfolder weightsname raise environmenterror ferror no file named flaxweightsname found in directory pretrainedmodelnameorpath but there is a file for pytorch weights use frompttrue to load this model from those weights else raise environmenterror ferror no file named flaxweightsname or weightsname found in directory fpretrainedmodelnameorpath elif os path isfileos path joinsubfolder pretrainedmodelnameorpath archivefile pretrainedmodelnameorpath islocal true elif isremoteurlpretrainedmodelnameorpath filename pretrainedmodelnameorpath resolvedarchivefile downloadurlpretrainedmodelnameorpath else if frompt filename weightsname else filename flaxweightsname try load from url or cache if already cached cachedfilekwargs cachedir cachedir forcedownload forcedownload proxies proxies resumedownload resumedownload localfilesonly localfilesonly token token useragent useragent revision revision subfolder subfolder raiseexceptionsformissingentries false commithash commithash resolvedarchivefile cachedfilepretrainedmodelnameorpath filename cachedfilekwargs maybe the checkpoint is sharded we try to grab the index name in this case if resolvedarchivefile is none and filename flaxweightsname resolvedarchivefile cachedfile pretrainedmodelnameorpath flaxweightsindexname cachedfilekwargs if resolvedarchivefile is not none issharded true maybe the checkpoint is pytorch sharded we try to grab the pytorch index name in this case if resolvedarchivefile is none and frompt resolvedarchivefile cachedfile pretrainedmodelnameorpath weightsindexname cachedfilekwargs if resolvedarchivefile is not none issharded true if we still haven t found anything look for safetensors if resolvedarchivefile is none no support for sharded safetensors yet so we ll raise an error if that s all we find filename safeweightsname resolvedarchivefile cachedfile pretrainedmodelnameorpath safeweightsname cachedfilekwargs since we set raiseexceptionsformissingentriesfalse we don t get an exception but a none result when internet is up the repo and revision exist but the file does not if resolvedarchivefile is none otherwise maybe there is a tf or torch model file we try those to give a helpful error message hasfilekwargs revision revision proxies proxies token token if hasfilepretrainedmodelnameorpath safeweightsindexname hasfilekwargs issharded true raise notimplementederror support for sharded checkpoints using safetensors is coming soon elif hasfilepretrainedmodelnameorpath weightsname hasfilekwargs raise environmenterror fpretrainedmodelnameorpath does not appear to have a file named f flaxweightsname but there is a file for pytorch weights use frompttrue to load this model from those weights elif hasfilepretrainedmodelnameorpath weightsindexname hasfilekwargs raise environmenterror fpretrainedmodelnameorpath does not appear to have a file named f flaxweightsindexname but there is a sharded file for pytorch weights use frompttrue to load this model from those weights else raise environmenterror fpretrainedmodelnameorpath does not appear to have a file named f flaxweightsname or weightsname except environmenterror raise any environment error raise by cachedfile it will have a helpful error message adapted to the original exception raise except exception for any other exception we throw a generic error raise environmenterror fcan t load the model for pretrainedmodelnameorpath if you were trying to load it from https huggingface comodels make sure you don t have a local directory with the f same name otherwise make sure pretrainedmodelnameorpath is the correct path to a f directory containing a file named flaxweightsname or weightsname if islocal logger infofloading weights file archivefile resolvedarchivefile archivefile filename resolvedarchivefile splitos path sep1 else logger infofloading weights file filename from cache at resolvedarchivefile else resolvedarchivefile none we ll need to download and cache each checkpoint shard if the checkpoint is sharded if issharded resolvedarchivefile becomes a list of files that point to the different checkpoint shards in this case resolvedarchivefile getcheckpointshardfiles pretrainedmodelnameorpath resolvedarchivefile cachedircachedir forcedownloadforcedownload proxiesproxies resumedownloadresumedownload localfilesonlylocalfilesonly tokentoken useragentuseragent revisionrevision subfoldersubfolder commithashcommithash safetensorsfrompt false if filename safeweightsname with safeopenresolvedarchivefile frameworkflax as f safetensorsmetadata f metadata if safetensorsmetadata is none or safetensorsmetadata getformat not in pt tf flax raise oserror fthe safetensors archive passed at resolvedarchivefile does not contain the valid metadata make sure you save your model with the savepretrained method safetensorsfrompt safetensorsmetadata getformat pt init random models model clsconfig modelargs doinitdoinit modelkwargs if frompt or safetensorsfrompt state loadpytorchcheckpointinflaxstatedictmodel resolvedarchivefile issharded else if issharded state cls loadflaxshardedweightsresolvedarchivefile else state cls loadflaxweightsresolvedarchivefile make sure all arrays are stored as jnp arrays note this is to prevent a bug this will be fixed in flax v0 3 4 https github comgoogleflaxissues1261 if doinit state jax treeutil treemapjnp array state else keep the params on cpu if we don t want to initialize state jax treeutil treemaplambda x jax deviceputx jax localdevicesbackendcpu0 state if batchstats in state if flax model contains batch norm layers if model is base model only use modelprefix key if cls basemodelprefix not in dictmodel paramsshapetreeparams and cls basemodelprefix in stateparams stateparams stateparamscls basemodelprefix statebatchstats statebatchstatscls basemodelprefix if model is head model and we are loading weights from base model we initialize new params dict with basemodelprefix if cls basemodelprefix in dictmodel paramsshapetreeparams and cls basemodelprefix not in stateparams state params cls basemodelprefix stateparams batchstats cls basemodelprefix statebatchstats else if model is base model only use modelprefix key if cls basemodelprefix not in dictmodel paramsshapetree and cls basemodelprefix in state state statecls basemodelprefix if model is head model and we are loading weights from base model we initialize new params dict with basemodelprefix if cls basemodelprefix in dictmodel paramsshapetree and cls basemodelprefix not in state state cls basemodelprefix state flatten dicts state flattendictstate randomstate flattendictunfreezemodel params if doinit else model paramsshapetree missingkeys model requiredparams setstate keys unexpectedkeys setstate keys model requiredparams disabling warning when porting pytorch weights to flax flax does not uses numbatchestracked for unexpectedkey in unexpectedkeys copy if numbatchestracked in unexpectedkey1 unexpectedkeys removeunexpectedkey if missingkeys and not doinit logger warning fthe checkpoint pretrainedmodelnameorpath is missing required keys missingkeys make sure to call model initweights to initialize the missing weights cls missingkeys missingkeys mistmatched keys contains tuples keyshape1shape2 of weights in the checkpoint that have a shape not matching the weights in the model mismatchedkeys for key in state keys if key in randomstate and statekey shape randomstatekey shape if ignoremismatchedsizes mismatchedkeys appendkey statekey shape randomstatekey shape statekey randomstatekey else raise valueerror ftrying to load the pretrained weight for key failed checkpoint has shape fstatekey shape which is incompatible with the model shape randomstatekey shape using ignoremismatchedsizestrue if you really want to load this checkpoint inside this model add missing keys as random parameters if we are initializing if missingkeys and doinit for missingkey in missingkeys statemissingkey randomstatemissingkey remove unexpected keys to not be saved again for unexpectedkey in unexpectedkeys del stateunexpectedkey if lenunexpectedkeys 0 logger warning fsome weights of the model checkpoint at pretrainedmodelnameorpath were not used when f initializing model class name unexpectedkeysn this is expected if you are f initializing model class name from the checkpoint of a model trained on another task or with another architecture e g initializing a bertforsequenceclassification model from a bertforpretraining model n this is not expected if you are initializing f model class name from the checkpoint of a model that you expect to be exactly identical initializing a bertforsequenceclassification model from a bertforsequenceclassification model else logger infofall model checkpoint weights were used when initializing model class name n if lenmissingkeys 0 logger warning fsome weights of model class name were not initialized from the model checkpoint at f pretrainedmodelnameorpath and are newly initialized missingkeysnyou should probably train this model on a downstream task to be able to use it for predictions and inference elif lenmismatchedkeys 0 logger info fall the weights of model class name were initialized from the model checkpoint at f pretrainedmodelnameorpath nif your task is similar to the task the model of the checkpoint f was trained on you can already use model class name for predictions without further training if lenmismatchedkeys 0 mismatchedwarning n join f key found shape shape1 in the checkpoint and shape2 in the model instantiated for key shape1 shape2 in mismatchedkeys logger warning fsome weights of model class name were not initialized from the model checkpoint at f pretrainedmodelnameorpath and are newly initialized because the shapes did not f match nmismatchedwarningnyou should probably train this model on a downstream task to be able to use it for predictions and inference dictionary of key dtypes for the model params paramdtypes jax treeutil treemaplambda x x dtype state extract keys of parameters not in jnp float32 fp16params k for k in paramdtypes if paramdtypesk jnp float16 bf16params k for k in paramdtypes if paramdtypesk jnp bfloat16 raise a warning if any of the parameters are not in jnp float32 if lenfp16params 0 logger warning fsome of the weights of model class name were initialized in float16 precision from fthe model checkpoint at pretrainedmodelnameorpath nfp16paramsn you should probably upcast the model weights to float32 if this was not intended see flaxpretrainedmodel tofp32 for further information on how to do this if lenbf16params 0 logger warning fsome of the weights of model class name were initialized in bfloat16 precision from fthe model checkpoint at pretrainedmodelnameorpath nbf16paramsn you should probably upcast the model weights to float32 if this was not intended see flaxpretrainedmodel tofp32 for further information on how to do this if it is a model with generation capabilities attempt to load the generation config if model cangenerate try model generationconfig generationconfig frompretrained pretrainedmodelnameorpath cachedircachedir forcedownloadforcedownload resumedownloadresumedownload proxiesproxies localfilesonlylocalfilesonly tokentoken revisionrevision subfoldersubfolder fromautofromautoclass frompipelinefrompipeline kwargs except oserror logger info generation config file not found using a generation config created from the model config pass if doinit set correct parameters model params unflattendictstate return model else return model unflattendictstate def savepretrained self savedirectory unionstr os pathlike paramsnone pushtohubfalse maxshardsize10gb token optionalunionstr bool none safeserialization bool false kwargs useauthtoken kwargs popuseauthtoken none if useauthtoken is not none warnings warn the useauthtoken argument is deprecated and will be removed in v5 of transformers please use token instead futurewarning if token is not none raise valueerror token and useauthtoken are both specified please set only the argument token token useauthtoken if token is not none kwargstoken token if os path isfilesavedirectory logger errorfprovided path savedirectory should be a directory not a file return os makedirssavedirectory existoktrue if pushtohub commitmessage kwargs popcommitmessage none repoid kwargs poprepoid savedirectory splitos path sep1 repoid self createreporepoid kwargs filestimestamps self getfilestimestampssavedirectory get abs dir savedirectory os path abspathsavedirectory save config as well self config architectures self class name4 if we have a custom model we copy the file defining it in the folder and set the attributes so it can be loaded from the hub if self autoclass is not none customobjectsaveself savedirectory configself config self config savepretrainedsavedirectory if self cangenerate self generationconfig savepretrainedsavedirectory save model weightsname safeweightsname if safeserialization else flaxweightsname outputmodelfile os path joinsavedirectory weightsname shards index flaxshardcheckpointparams if params is not none else self params maxshardsize clean the folder from a previous save for filename in os listdirsavedirectory fullfilename os path joinsavedirectory filename weightsnosuffix weightsname replace bin replace safetensors if filename startswithweightsnosuffix and os path isfilefullfilename and filename not in shards keys os removefullfilename if index is none if safeserialization params params if params is not none else self params flatdict flattendictparams sep safesavefileflatdict outputmodelfile metadataformat flax else with openoutputmodelfile wb as f params params if params is not none else self params modelbytes tobytesparams f writemodelbytes else saveindexfile os path joinsavedirectory flaxweightsindexname save the index as well with opensaveindexfile w encodingutf8 as f content json dumpsindex indent2 sortkeystrue n f writecontent logger info fthe model is bigger than the maximum size per checkpoint maxshardsize and is going to be fsplit in lenshards checkpoint shards you can find where each parameters has been saved in the findex located at saveindexfile for shardfile shard in shards items the shard item are unflattened to save them we need to flatten them again with openos path joinsavedirectory shardfile modewb as f params unflattendictshard sep shardbytes tobytesparams f writeshardbytes logger infofmodel weights saved in outputmodelfile if pushtohub self uploadmodifiedfiles savedirectory repoid filestimestamps commitmessagecommitmessage tokentoken classmethod def registerforautoclasscls autoclassflaxautomodel if not isinstanceautoclass str autoclass autoclass name import transformers models auto as automodule if not hasattrautomodule autoclass raise valueerrorfautoclass is not a valid auto class cls autoclass autoclass to update the docstring we need to copy the method otherwise we change the original docstring flaxpretrainedmodel pushtohub copyfuncflaxpretrainedmodel pushtohub if flaxpretrainedmodel pushtohub doc is not none flaxpretrainedmodel pushtohub doc flaxpretrainedmodel pushtohub doc format objectmodel objectclassflaxautomodel objectfilesmodel checkpoint def overwritecalldocstringmodelclass docstring copy call function to be sure docstring is changed only for this function modelclass call copyfuncmodelclass call delete existing docstring modelclass call doc none set correct docstring modelclass call addstartdocstringstomodelforwarddocstringmodelclass call def appendcallsampledocstringmodelclass checkpoint outputtype configclass masknone revisionnone modelclass call copyfuncmodelclass call modelclass call addcodesampledocstrings checkpointcheckpoint outputtypeoutputtype configclassconfigclass modelclsmodelclass name revisionrevision modelclass call def appendreplacereturndocstringsmodelclass outputtype configclass modelclass call copyfuncmodelclass call modelclass call replacereturndocstrings outputtypeoutputtype configclassconfigclass modelclass call coding utf 8 2021 the google flax team s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license returns the size in bytes occupied by one parameter of type dtype example py dtype_byte_size np float32 4 splits a model state dictionary in sub checkpoints so that the final size of each sub checkpoint does not exceed a given size the sub checkpoints are determined by iterating through the state_dict in the order of its keys so there is no optimization made to make each sub checkpoint as close as possible to the maximum size passed for example if the limit is 10gb and we have weights of sizes 6gb 6gb 2gb 6gb 2gb 2gb they will get sharded as 6gb 6 2gb 6 2 2gb and not 6 2 2gb 6 2gb 6gb tip warning true if one of the model s weight is bigger that max_shard_size it will end up in its own sub checkpoint which will have a size greater than max_shard_size tip args params union dict frozendict a pytree of model parameters max_shard_size int or str optional defaults to 10gb the maximum size of each sub checkpoint if expressed as a string needs to be digits followed by a unit like 5mb flatten the weights to chunk if this weight is going to tip up over the maximal size we split add the last block if we only have one shard we return it otherwise let s build the index add the metadata base class for all models flaxpretrainedmodel takes care of storing the configuration of the models and handles methods for loading downloading and saving models class attributes overridden by derived classes config_class pretrainedconfig a subclass of pretrainedconfig to use as configuration class for this model architecture base_model_prefix str a string indicating the attribute associated to the base model in derived classes of the same architecture adding modules on top of the base model main_input_name str the name of the principal input to the model often input_ids for nlp models pixel_values for vision models and input_values for speech models those are private to be exposed as typed property on derived classes those are public as their type is generic to every derived classes to check if the model was intialized automatically randomly initialized parameters get the shape of the parameters save required_params as set initialize the parameters all context managers that the model should be initialized under go here str identifies that this is a flax model don t set params if the model is not initialized helper method to cast floating point values of given parameter pytree to given dtype taken from https github com deepmind jmp blob 3a8318abc3292be38582794dbf7b094e6583b192 jmp _src policy py l27 cast the floating point params to jax numpy bfloat16 this returns a new params tree and does not cast the params in place this method can be used on tpu to explicitly convert the model parameters to bfloat16 precision to do full half precision training or to save weights in bfloat16 for inference in order to save memory and improve speed arguments params union dict frozendict a pytree of model parameters mask union dict frozendict a pytree with same structure as the params tree the leaves should be booleans true for params you want to cast and should be false for those you want to skip examples python from transformers import flaxbertmodel load model model flaxbertmodel from_pretrained bert base cased by default the model parameters will be in fp32 precision to cast these to bfloat16 precision model params model to_bf16 model params if you want don t want to cast certain parameters for example layer norm bias and scale then pass the mask as follows from flax import traverse_util model flaxbertmodel from_pretrained bert base cased flat_params traverse_util flatten_dict model params mask path path 2 layernorm bias and path 2 layernorm scale for path in flat_params mask traverse_util unflatten_dict mask model params model to_bf16 model params mask cast the floating point parmas to jax numpy float32 this method can be used to explicitly convert the model parameters to fp32 precision this returns a new params tree and does not cast the params in place arguments params union dict frozendict a pytree of model parameters mask union dict frozendict a pytree with same structure as the params tree the leaves should be booleans true for params you want to cast and should be false for those you want to skip examples python from transformers import flaxbertmodel download model and configuration from huggingface co model flaxbertmodel from_pretrained bert base cased by default the model params will be in fp32 to illustrate the use of this method we ll first cast to fp16 and back to fp32 model params model to_f16 model params now cast back to fp32 model params model to_fp32 model params cast the floating point parmas to jax numpy float16 this returns a new params tree and does not cast the params in place this method can be used on gpu to explicitly convert the model parameters to float16 precision to do full half precision training or to save weights in float16 for inference in order to save memory and improve speed arguments params union dict frozendict a pytree of model parameters mask union dict frozendict a pytree with same structure as the params tree the leaves should be booleans true for params you want to cast and should be false for those you want to skip examples python from transformers import flaxbertmodel load model model flaxbertmodel from_pretrained bert base cased by default the model params will be in fp32 to cast these to float16 model params model to_fp16 model params if you want don t want to cast certain parameters for example layer norm bias and scale then pass the mask as follows from flax import traverse_util model flaxbertmodel from_pretrained bert base cased flat_params traverse_util flatten_dict model params mask path path 2 layernorm bias and path 2 layernorm scale for path in flat_params mask traverse_util unflatten_dict mask model params model to_fp16 model params mask this is the same as flax serialization from_bytes https lax readthedocs io en latest _modules flax serialization html from_bytes but for a sharded checkpoint this load is performed efficiently each checkpoint shard is loaded one by one in ram and deleted after being loaded in the model args shard_files list str the list of shard files to load returns dict a nested dictionary of the model parameters in the expected format for flax models model params load the index load using msgpack utils the state dict is unflattened to the match the format of model params returns whether this model can generate sequences with generate returns bool whether this model can generate sequences with generate detects whether prepare_inputs_for_generation has been overwritten which is a requirement for generation alternativelly the model can also have a custom generate function instantiate a pretrained flax model from a pre trained model configuration the warning weights from xxx not initialized from pretrained model means that the weights of xxx do not come pretrained with the rest of the model it is up to you to train those weights with a downstream fine tuning task the warning weights from xxx not used in yyy means that the layer xxx is not used by yyy therefore those weights are discarded parameters pretrained_model_name_or_path str or os pathlike can be either a string the model id of a pretrained model hosted inside a model repo on huggingface co valid model ids can be located at the root level like bert base uncased or namespaced under a user or organization name like dbmdz bert base german cased a path to a directory containing model weights saved using flaxpretrainedmodel save_pretrained e g my_model_directory a path or url to a pt index checkpoint file e g tf_model model ckpt index in this case from_pt should be set to true dtype jax numpy dtype optional defaults to jax numpy float32 the data type of the computation can be one of jax numpy float32 jax numpy float16 on gpus and jax numpy bfloat16 on tpus this can be used to enable mixed precision training or half precision inference on gpus or tpus if specified all the computation will be performed with the given dtype note that this only specifies the dtype of the computation and does not influence the dtype of model parameters if you wish to change the dtype of the model parameters see flaxpretrainedmodel to_fp16 and flaxpretrainedmodel to_bf16 model_args sequence of positional arguments optional all remaining positional arguments will be passed to the underlying model s __init__ method config union pretrainedconfig str os pathlike optional can be either an instance of a class derived from pretrainedconfig a string or path valid as input to pretrainedconfig from_pretrained configuration for the model to use instead of an automatically loaded configuration configuration can be automatically loaded when the model is a model provided by the library loaded with the model id string of a pretrained model the model was saved using pretrainedmodel save_pretrained and is reloaded by supplying the save directory the model is loaded by supplying a local directory as pretrained_model_name_or_path and a configuration json file named config json is found in the directory cache_dir union str os pathlike optional path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used from_pt bool optional defaults to false load the model weights from a pytorch checkpoint save file see docstring of pretrained_model_name_or_path argument ignore_mismatched_sizes bool optional defaults to false whether or not to raise an error if some of the weights from the checkpoint do not have the same size as the weights of the model if for instance you are instantiating a model with 10 labels from a checkpoint with 3 labels force_download bool optional defaults to false whether or not to force the re download of the model weights and configuration files overriding the cached versions if they exist resume_download bool optional defaults to false whether or not to delete incompletely received files will attempt to resume the download if such a file exists proxies dict str str optional a dictionary of proxy servers to use by protocol or endpoint e g http foo bar 3128 http hostname foo bar 4012 the proxies are used on each request local_files_only bool optional defaults to false whether or not to only look at local files i e do not try to download the model token str or bool optional the token to use as http bearer ization for remote files if true or not specified will use the token generated when running huggingface cli login stored in huggingface revision str optional defaults to main the specific model version to use it can be a branch name a tag name or a commit id since we use a git based system for storing models and other artifacts on huggingface co so revision can be any identifier allowed by git tip to test a pull request you made on the hub you can pass revision refs pr pr_number tip subfolder str optional defaults to in case the relevant files are located inside a subfolder of the model repo on huggingface co you can specify the folder name here kwargs remaining dictionary of keyword arguments optional can be used to update the configuration object after it being loaded and initiate the model e g output_attentions true behaves differently depending on whether a config is provided or automatically loaded if a configuration is provided with config kwargs will be directly passed to the underlying model s __init__ method we assume all relevant updates to the configuration have already been done if a configuration is not provided kwargs will be first passed to the configuration class initialization function pretrainedconfig from_pretrained each key of kwargs that corresponds to a configuration attribute will be used to override said attribute with the supplied kwargs value remaining keys that do not correspond to any configuration attribute will be passed to the underlying model s __init__ function examples python from transformers import bertconfig flaxbertmodel download model and configuration from huggingface co and cache model flaxbertmodel from_pretrained bert base cased model was saved using save_pretrained test saved_model for example purposes not runnable model flaxbertmodel from_pretrained test saved_model loading from a pytorch checkpoint file instead of a pytorch model slower for example purposes not runnable config bertconfig from_json_file pt_model config json model flaxbertmodel from_pretrained pt_model pytorch_model bin from_pt true config config not relevant for flax models load config if we don t provide a configuration add the dtype to model_kwargs this variable will flag if we re loading a sharded checkpoint in this case the archive file is just the index of the files load model load from a flax checkpoint load from a sharded flax checkpoint load from a safetensors checkpoint load from a pytorch checkpoint load from a sharded pytorch checkpoint at this stage we don t have a weight file so we will raise an error load from a sharded safetensors checkpoint load from url or cache if already cached maybe the checkpoint is sharded we try to grab the index name in this case maybe the checkpoint is pytorch sharded we try to grab the pytorch index name in this case if we still haven t found anything look for safetensors no support for sharded safetensors yet so we ll raise an error if that s all we find since we set _raise_exceptions_for_missing_entries false we don t get an exception but a none result when internet is up the repo and revision exist but the file does not otherwise maybe there is a tf or torch model file we try those to give a helpful error message raise any environment error raise by cached_file it will have a helpful error message adapted to the original exception for any other exception we throw a generic error we ll need to download and cache each checkpoint shard if the checkpoint is sharded resolved_archive_file becomes a list of files that point to the different checkpoint shards in this case init random models make sure all arrays are stored as jnp arrays note this is to prevent a bug this will be fixed in flax v0 3 4 https github com google flax issues 1261 keep the params on cpu if we don t want to initialize if flax model contains batch norm layers if model is base model only use model_prefix key if model is head model and we are loading weights from base model we initialize new params dict with base_model_prefix if model is base model only use model_prefix key if model is head model and we are loading weights from base model we initialize new params dict with base_model_prefix flatten dicts disabling warning when porting pytorch weights to flax flax does not uses num_batches_tracked mistmatched keys contains tuples key shape1 shape2 of weights in the checkpoint that have a shape not matching the weights in the model add missing keys as random parameters if we are initializing remove unexpected keys to not be saved again dictionary of key dtypes for the model params extract keys of parameters not in jnp float32 raise a warning if any of the parameters are not in jnp float32 if it is a model with generation capabilities attempt to load the generation config set correct parameters save a model and its configuration file to a directory so that it can be re loaded using the flaxpretrainedmodel from_pretrained class method arguments save_directory str or os pathlike directory to which to save will be created if it doesn t exist push_to_hub bool optional defaults to false whether or not to push your model to the hugging face model hub after saving it you can specify the repository you want to push to with repo_id will default to the name of save_directory in your namespace max_shard_size int or str optional defaults to 10gb the maximum size for a checkpoint before being sharded checkpoints shard will then be each of size lower than this size if expressed as a string needs to be digits followed by a unit like 5mb tip warning true if a single weight of the model is bigger than max_shard_size it will be in its own checkpoint shard which will be bigger than max_shard_size tip token str or bool optional the token to use as http bearer ization for remote files if true or not specified will use the token generated when running huggingface cli login stored in huggingface kwargs dict str any optional additional key word arguments passed along to the utils pushtohubmixin push_to_hub method safe_serialization bool optional defaults to false whether to save the model using safetensors or through msgpack get abs dir save config as well if we have a custom model we copy the file defining it in the folder and set the attributes so it can be loaded from the hub save model clean the folder from a previous save save the index as well the shard item are unflattened to save them we need to flatten them again register this class with a given auto class this should only be used for custom models as the ones in the library are already mapped with an auto class tip warning true this api is experimental and may have some slight breaking changes in the next releases tip args auto_class str or type optional defaults to flaxautomodel the auto class to register this new model with to update the docstring we need to copy the method otherwise we change the original docstring copy __call__ function to be sure docstring is changed only for this function delete existing docstring set correct docstring
import gc
import json
import os
import re
import warnings
from functools import partial
from pickle import UnpicklingError
from typing import Any, Dict, Optional, Set, Tuple, Union

import flax.linen as nn
import jax
import jax.numpy as jnp
import msgpack.exceptions
from flax.core.frozen_dict import FrozenDict, unfreeze
from flax.serialization import from_bytes, to_bytes
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from .configuration_utils import PretrainedConfig
from .dynamic_module_utils import custom_object_save
from .generation import FlaxGenerationMixin, GenerationConfig
from .modeling_flax_pytorch_utils import load_pytorch_checkpoint_in_flax_state_dict
from .utils import (
    FLAX_WEIGHTS_INDEX_NAME,
    FLAX_WEIGHTS_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
    PushToHubMixin,
    add_code_sample_docstrings,
    add_start_docstrings_to_model_forward,
    cached_file,
    copy_func,
    download_url,
    has_file,
    is_offline_mode,
    is_remote_url,
    logging,
    replace_return_docstrings,
)
from .utils.hub import convert_file_size_to_int, get_checkpoint_shard_files
from .utils.import_utils import is_safetensors_available


if is_safetensors_available():
    from safetensors import safe_open
    from safetensors.flax import load_file as safe_load_file
    from safetensors.flax import save_file as safe_save_file

logger = logging.get_logger(__name__)


def quick_gelu(x):
    return x * jax.nn.sigmoid(1.702 * x)


ACT2FN = {
    "gelu": partial(nn.gelu, approximate=False),
    "relu": nn.relu,
    "silu": nn.swish,
    "swish": nn.swish,
    "gelu_new": partial(nn.gelu, approximate=True),
    "quick_gelu": quick_gelu,
}


def dtype_byte_size(dtype):
    
    if dtype == bool:
        return 1 / 8
    bit_search = re.search(r"[^\d](\d+)$", dtype.name)
    if bit_search is None:
        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
    bit_size = int(bit_search.groups()[0])
    return bit_size // 8


def flax_shard_checkpoint(params, max_shard_size="10GB"):
    
    max_shard_size = convert_file_size_to_int(max_shard_size)

    sharded_state_dicts = []
    current_block = {}
    current_block_size = 0
    total_size = 0

    
    weights = flatten_dict(params, sep="/")
    for item in weights:
        weight_size = weights[item].size * dtype_byte_size(weights[item].dtype)

        
        if current_block_size + weight_size > max_shard_size:
            sharded_state_dicts.append(current_block)
            current_block = {}
            current_block_size = 0

        current_block[item] = weights[item]
        current_block_size += weight_size
        total_size += weight_size

    
    sharded_state_dicts.append(current_block)

    
    if len(sharded_state_dicts) == 1:
        return {FLAX_WEIGHTS_NAME: sharded_state_dicts[0]}, None

    
    weight_map = {}
    shards = {}
    for idx, shard in enumerate(sharded_state_dicts):
        shard_file = FLAX_WEIGHTS_NAME.replace(".msgpack", f"-{idx+1:05d}-of-{len(sharded_state_dicts):05d}.msgpack")
        shards[shard_file] = shard
        for weight_name in shard.keys():
            weight_map[weight_name] = shard_file

    
    metadata = {"total_size": total_size}
    index = {"metadata": metadata, "weight_map": weight_map}
    return shards, index


class FlaxPreTrainedModel(PushToHubMixin, FlaxGenerationMixin):
    r

    config_class = None
    base_model_prefix = ""
    main_input_name = "input_ids"
    _auto_class = None
    _missing_keys = set()

    def __init__(
        self,
        config: PretrainedConfig,
        module: nn.Module,
        input_shape: Tuple = (1, 1),
        seed: int = 0,
        dtype: jnp.dtype = jnp.float32,
        _do_init: bool = True,
    ):
        if config is None:
            raise ValueError("config cannot be None")

        if module is None:
            raise ValueError("module cannot be None")

        
        self._config = config
        self._module = module

        
        self.key = PRNGKey(seed)
        self.dtype = dtype
        self.input_shape = input_shape
        self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None

        
        self._is_initialized = _do_init

        if _do_init:
            
            random_params = self.init_weights(self.key, input_shape)
            params_shape_tree = jax.eval_shape(lambda params: params, random_params)
        else:
            init_fn = partial(self.init_weights, input_shape=input_shape)
            params_shape_tree = jax.eval_shape(init_fn, self.key)

            logger.info(
                "Model weights are not initialized as `_do_init` is set to `False`. "
                f"Make sure to call `{self.__class__.__name__}.init_weights` manually to initialize the weights."
            )

        
        self._params_shape_tree = params_shape_tree

        
        self._required_params = set(flatten_dict(unfreeze(params_shape_tree)).keys())

        
        if _do_init:
            self.params = random_params

    def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple, params: FrozenDict = None) -> Dict:
        raise NotImplementedError(f"init method has to be implemented for {self}")

    def enable_gradient_checkpointing(self):
        raise NotImplementedError(f"gradient checkpointing method has to be implemented for {self}")

    @classmethod
    def _from_config(cls, config, **kwargs):
        
        return cls(config, **kwargs)

    @property
    def framework(self) -> str:
        
        return "flax"

    @property
    def config(self) -> PretrainedConfig:
        return self._config

    @property
    def module(self) -> nn.Module:
        return self._module

    @property
    def params(self) -> Union[Dict, FrozenDict]:
        if not self._is_initialized:
            raise ValueError(
                "`params` cannot be accessed from model when the model is created with `_do_init=False`. "
                "You must call `init_weights` manually and store the params outside of the model and "
                "pass it explicitly where needed."
            )
        return self._params

    @property
    def required_params(self) -> Set:
        return self._required_params

    @property
    def params_shape_tree(self) -> Dict:
        return self._params_shape_tree

    @params.setter
    def params(self, params: Union[Dict, FrozenDict]):
        
        if not self._is_initialized:
            raise ValueError(
                "`params` cannot be set from model when the model is created with `_do_init=False`. "
                "You store the params outside of the model."
            )

        if isinstance(params, FrozenDict):
            params = unfreeze(params)
        param_keys = set(flatten_dict(params).keys())
        if len(self.required_params - param_keys) > 0:
            raise ValueError(
                "Some parameters are missing. Make sure that `params` include the following "
                f"parameters {self.required_params - param_keys}"
            )
        self._params = params

    def _cast_floating_to(self, params: Union[Dict, FrozenDict], dtype: jnp.dtype, mask: Any = None) -> Any:
        

        
        def conditional_cast(param):
            if isinstance(param, jnp.ndarray) and jnp.issubdtype(param.dtype, jnp.floating):
                param = param.astype(dtype)
            return param

        if mask is None:
            return jax.tree_util.tree_map(conditional_cast, params)

        flat_params = flatten_dict(params)
        flat_mask, _ = jax.tree_util.tree_flatten(mask)

        for masked, key in zip(flat_mask, flat_params.keys()):
            if masked:
                param = flat_params[key]
                flat_params[key] = conditional_cast(param)

        return unflatten_dict(flat_params)

    def to_bf16(self, params: Union[Dict, FrozenDict], mask: Any = None):
        r
        return self._cast_floating_to(params, jnp.bfloat16, mask)

    def to_fp32(self, params: Union[Dict, FrozenDict], mask: Any = None):
        r
        return self._cast_floating_to(params, jnp.float32, mask)

    def to_fp16(self, params: Union[Dict, FrozenDict], mask: Any = None):
        r
        return self._cast_floating_to(params, jnp.float16, mask)

    @classmethod
    def load_flax_weights(cls, resolved_archive_file):
        try:
            if resolved_archive_file.endswith(".safetensors"):
                state = safe_load_file(resolved_archive_file)
                state = unflatten_dict(state, sep=".")
            else:
                with open(resolved_archive_file, "rb") as state_f:
                    state = from_bytes(cls, state_f.read())
        except (UnpicklingError, msgpack.exceptions.ExtraData) as e:
            try:
                with open(resolved_archive_file) as f:
                    if f.read().startswith("version"):
                        raise OSError(
                            "You seem to have cloned a repository without having git-lfs installed. Please"
                            " install git-lfs and run `git lfs install` followed by `git lfs pull` in the"
                            " folder you cloned."
                        )
                    else:
                        raise ValueError from e
            except (UnicodeDecodeError, ValueError):
                raise EnvironmentError(f"Unable to convert {resolved_archive_file} to Flax deserializable object. ")

        return state

    @classmethod
    def load_flax_sharded_weights(cls, shard_files):
        

        
        state_sharded_dict = {}

        for shard_file in shard_files:
            
            try:
                with open(shard_file, "rb") as state_f:
                    state = from_bytes(cls, state_f.read())
            except (UnpicklingError, msgpack.exceptions.ExtraData) as e:
                with open(shard_file) as f:
                    if f.read().startswith("version"):
                        raise OSError(
                            "You seem to have cloned a repository without having git-lfs installed. Please"
                            " install git-lfs and run `git lfs install` followed by `git lfs pull` in the"
                            " folder you cloned."
                        )
                    else:
                        raise ValueError from e
            except (UnicodeDecodeError, ValueError):
                raise EnvironmentError(f"Unable to convert {shard_file} to Flax deserializable object. ")

            state = flatten_dict(state, sep="/")
            state_sharded_dict.update(state)
            del state
            gc.collect()

        
        return unflatten_dict(state_sharded_dict, sep="/")

    @classmethod
    def can_generate(cls) -> bool:
        
        
        
        if "GenerationMixin" in str(cls.prepare_inputs_for_generation) and "GenerationMixin" in str(cls.generate):
            return False
        return True

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Union[str, os.PathLike],
        dtype: jnp.dtype = jnp.float32,
        *model_args,
        config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
        cache_dir: Optional[Union[str, os.PathLike]] = None,
        ignore_mismatched_sizes: bool = False,
        force_download: bool = False,
        local_files_only: bool = False,
        token: Optional[Union[str, bool]] = None,
        revision: str = "main",
        **kwargs,
    ):
        r
        from_pt = kwargs.pop("from_pt", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        use_auth_token = kwargs.pop("use_auth_token", None)
        trust_remote_code = kwargs.pop("trust_remote_code", None)
        from_pipeline = kwargs.pop("_from_pipeline", None)
        from_auto_class = kwargs.pop("_from_auto", False)
        _do_init = kwargs.pop("_do_init", True)
        subfolder = kwargs.pop("subfolder", "")
        commit_hash = kwargs.pop("_commit_hash", None)

        
        _ = kwargs.pop("adapter_kwargs", None)

        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            token = use_auth_token

        if trust_remote_code is True:
            logger.warning(
                "The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is"
                " ignored."
            )

        user_agent = {"file_type": "model", "framework": "flax", "from_auto_class": from_auto_class}
        if from_pipeline is not None:
            user_agent["using_pipeline"] = from_pipeline

        if is_offline_mode() and not local_files_only:
            logger.info("Offline mode: forcing local_files_only=True")
            local_files_only = True

        
        if not isinstance(config, PretrainedConfig):
            config_path = config if config is not None else pretrained_model_name_or_path
            config, model_kwargs = cls.config_class.from_pretrained(
                config_path,
                cache_dir=cache_dir,
                return_unused_kwargs=True,
                force_download=force_download,
                resume_download=resume_download,
                proxies=proxies,
                local_files_only=local_files_only,
                token=token,
                revision=revision,
                subfolder=subfolder,
                _from_auto=from_auto_class,
                _from_pipeline=from_pipeline,
                _commit_hash=commit_hash,
                **kwargs,
            )
        else:
            model_kwargs = kwargs.copy()

        if commit_hash is None:
            commit_hash = getattr(config, "_commit_hash", None)

        
        model_kwargs["dtype"] = dtype

        
        
        is_sharded = False

        
        if pretrained_model_name_or_path is not None:
            pretrained_model_name_or_path = str(pretrained_model_name_or_path)
            is_local = os.path.isdir(pretrained_model_name_or_path)
            if os.path.isdir(pretrained_model_name_or_path):
                if os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME)):
                    
                    archive_file = os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_NAME)
                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_INDEX_NAME)):
                    
                    archive_file = os.path.join(pretrained_model_name_or_path, subfolder, FLAX_WEIGHTS_INDEX_NAME)
                    is_sharded = True
                elif is_safetensors_available() and os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME)
                ):
                    
                    archive_file = os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_NAME)
                elif from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)):
                    
                    archive_file = os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)
                elif from_pt and os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_INDEX_NAME)
                ):
                    
                    archive_file = os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_INDEX_NAME)
                    is_sharded = True
                
                elif is_safetensors_available() and os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME)
                ):
                    
                    archive_file = os.path.join(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME)
                    is_sharded = True
                    raise NotImplementedError("Support for sharded checkpoints using safetensors is coming soon!")
                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, subfolder, WEIGHTS_NAME)):
                    raise EnvironmentError(
                        f"Error no file named {FLAX_WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} "
                        "but there is a file for PyTorch weights. Use `from_pt=True` to load this model from those "
                        "weights."
                    )
                else:
                    raise EnvironmentError(
                        f"Error no file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME} found in directory "
                        f"{pretrained_model_name_or_path}."
                    )
            elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)):
                archive_file = pretrained_model_name_or_path
                is_local = True
            elif is_remote_url(pretrained_model_name_or_path):
                filename = pretrained_model_name_or_path
                resolved_archive_file = download_url(pretrained_model_name_or_path)
            else:
                if from_pt:
                    filename = WEIGHTS_NAME
                else:
                    filename = FLAX_WEIGHTS_NAME

                try:
                    
                    cached_file_kwargs = {
                        "cache_dir": cache_dir,
                        "force_download": force_download,
                        "proxies": proxies,
                        "resume_download": resume_download,
                        "local_files_only": local_files_only,
                        "token": token,
                        "user_agent": user_agent,
                        "revision": revision,
                        "subfolder": subfolder,
                        "_raise_exceptions_for_missing_entries": False,
                        "_commit_hash": commit_hash,
                    }
                    resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)

                    
                    if resolved_archive_file is None and filename == FLAX_WEIGHTS_NAME:
                        resolved_archive_file = cached_file(
                            pretrained_model_name_or_path, FLAX_WEIGHTS_INDEX_NAME, **cached_file_kwargs
                        )
                        if resolved_archive_file is not None:
                            is_sharded = True

                    
                    if resolved_archive_file is None and from_pt:
                        resolved_archive_file = cached_file(
                            pretrained_model_name_or_path, WEIGHTS_INDEX_NAME, **cached_file_kwargs
                        )
                        if resolved_archive_file is not None:
                            is_sharded = True

                    
                    if resolved_archive_file is None:
                        
                        filename = SAFE_WEIGHTS_NAME
                        resolved_archive_file = cached_file(
                            pretrained_model_name_or_path, SAFE_WEIGHTS_NAME, **cached_file_kwargs
                        )

                    
                    
                    if resolved_archive_file is None:
                        
                        
                        has_file_kwargs = {
                            "revision": revision,
                            "proxies": proxies,
                            "token": token,
                        }
                        if has_file(pretrained_model_name_or_path, SAFE_WEIGHTS_INDEX_NAME, **has_file_kwargs):
                            is_sharded = True
                            raise NotImplementedError(
                                "Support for sharded checkpoints using safetensors is coming soon!"
                            )
                        elif has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs):
                            raise EnvironmentError(
                                f"{pretrained_model_name_or_path} does not appear to have a file named"
                                f" {FLAX_WEIGHTS_NAME} but there is a file for PyTorch weights. Use `from_pt=True` to"
                                " load this model from those weights."
                            )
                        elif has_file(pretrained_model_name_or_path, WEIGHTS_INDEX_NAME, **has_file_kwargs):
                            raise EnvironmentError(
                                f"{pretrained_model_name_or_path} does not appear to have a file named"
                                f" {FLAX_WEIGHTS_INDEX_NAME} but there is a sharded file for PyTorch weights. Use"
                                " `from_pt=True` to load this model from those weights."
                            )
                        else:
                            raise EnvironmentError(
                                f"{pretrained_model_name_or_path} does not appear to have a file named"
                                f" {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}."
                            )
                except EnvironmentError:
                    
                    
                    raise
                except Exception:
                    
                    raise EnvironmentError(
                        f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it"
                        " from 'https://huggingface.co/models', make sure you don't have a local directory with the"
                        f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"
                        f" directory containing a file named {FLAX_WEIGHTS_NAME} or {WEIGHTS_NAME}."
                    )

            if is_local:
                logger.info(f"loading weights file {archive_file}")
                resolved_archive_file = archive_file
                filename = resolved_archive_file.split(os.path.sep)[-1]
            else:
                logger.info(f"loading weights file {filename} from cache at {resolved_archive_file}")
        else:
            resolved_archive_file = None

        
        if is_sharded:
            
            resolved_archive_file, _ = get_checkpoint_shard_files(
                pretrained_model_name_or_path,
                resolved_archive_file,
                cache_dir=cache_dir,
                force_download=force_download,
                proxies=proxies,
                resume_download=resume_download,
                local_files_only=local_files_only,
                token=token,
                user_agent=user_agent,
                revision=revision,
                subfolder=subfolder,
                _commit_hash=commit_hash,
            )

        safetensors_from_pt = False
        if filename == SAFE_WEIGHTS_NAME:
            with safe_open(resolved_archive_file, framework="flax") as f:
                safetensors_metadata = f.metadata()
            if safetensors_metadata is None or safetensors_metadata.get("format") not in ["pt", "tf", "flax"]:
                raise OSError(
                    f"The safetensors archive passed at {resolved_archive_file} does not contain the valid metadata."
                    " Make sure you save your model with the `save_pretrained` method."
                )
            safetensors_from_pt = safetensors_metadata.get("format") == "pt"

        
        model = cls(config, *model_args, _do_init=_do_init, **model_kwargs)

        if from_pt or safetensors_from_pt:
            state = load_pytorch_checkpoint_in_flax_state_dict(model, resolved_archive_file, is_sharded)
        else:
            if is_sharded:
                state = cls.load_flax_sharded_weights(resolved_archive_file)
            else:
                state = cls.load_flax_weights(resolved_archive_file)
            
            
            
            if _do_init:
                state = jax.tree_util.tree_map(jnp.array, state)
            else:
                
                state = jax.tree_util.tree_map(lambda x: jax.device_put(x, jax.local_devices(backend="cpu")[0]), state)

        if "batch_stats" in state:  
            
            if (
                cls.base_model_prefix not in dict(model.params_shape_tree["params"])
                and cls.base_model_prefix in state["params"]
            ):
                state["params"] = state["params"][cls.base_model_prefix]
                state["batch_stats"] = state["batch_stats"][cls.base_model_prefix]

            
            
            if (
                cls.base_model_prefix in dict(model.params_shape_tree["params"])
                and cls.base_model_prefix not in state["params"]
            ):
                state = {
                    "params": {cls.base_model_prefix: state["params"]},
                    "batch_stats": {cls.base_model_prefix: state["batch_stats"]},
                }

        else:
            
            if cls.base_model_prefix not in dict(model.params_shape_tree) and cls.base_model_prefix in state:
                state = state[cls.base_model_prefix]

            
            
            if cls.base_model_prefix in dict(model.params_shape_tree) and cls.base_model_prefix not in state:
                state = {cls.base_model_prefix: state}

        
        state = flatten_dict(state)

        random_state = flatten_dict(unfreeze(model.params if _do_init else model.params_shape_tree))

        missing_keys = model.required_params - set(state.keys())
        unexpected_keys = set(state.keys()) - model.required_params

        
        for unexpected_key in unexpected_keys.copy():
            if "num_batches_tracked" in unexpected_key[-1]:
                unexpected_keys.remove(unexpected_key)

        if missing_keys and not _do_init:
            logger.warning(
                f"The checkpoint {pretrained_model_name_or_path} is missing required keys: {missing_keys}. "
                "Make sure to call model.init_weights to initialize the missing weights."
            )
            cls._missing_keys = missing_keys

        
        
        mismatched_keys = []
        for key in state.keys():
            if key in random_state and state[key].shape != random_state[key].shape:
                if ignore_mismatched_sizes:
                    mismatched_keys.append((key, state[key].shape, random_state[key].shape))
                    state[key] = random_state[key]
                else:
                    raise ValueError(
                        f"Trying to load the pretrained weight for {key} failed: checkpoint has shape "
                        f"{state[key].shape} which is incompatible with the model shape {random_state[key].shape}. "
                        "Using `ignore_mismatched_sizes=True` if you really want to load this checkpoint inside this "
                        "model."
                    )

        
        if missing_keys and _do_init:
            for missing_key in missing_keys:
                state[missing_key] = random_state[missing_key]

        
        for unexpected_key in unexpected_keys:
            del state[unexpected_key]

        if len(unexpected_keys) > 0:
            logger.warning(
                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
                f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
                f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or"
                " with another architecture (e.g. initializing a BertForSequenceClassification model from a"
                " BertForPreTraining model).\n- This IS NOT expected if you are initializing"
                f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical"
                " (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
            )
        else:
            logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")

        if len(missing_keys) > 0:
            logger.warning(
                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
                " TRAIN this model on a down-stream task to be able to use it for predictions and inference."
            )
        elif len(mismatched_keys) == 0:
            logger.info(
                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint"
                f" was trained on, you can already use {model.__class__.__name__} for predictions without further"
                " training."
            )
        if len(mismatched_keys) > 0:
            mismatched_warning = "\n".join(
                [
                    f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
                    for key, shape1, shape2 in mismatched_keys
                ]
            )
            logger.warning(
                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path} and are newly initialized because the shapes did not"
                f" match:\n{mismatched_warning}\nYou should probably TRAIN this model on a down-stream task to be able"
                " to use it for predictions and inference."
            )

        
        param_dtypes = jax.tree_util.tree_map(lambda x: x.dtype, state)
        
        fp16_params = [k for k in param_dtypes if param_dtypes[k] == jnp.float16]
        bf16_params = [k for k in param_dtypes if param_dtypes[k] == jnp.bfloat16]

        
        if len(fp16_params) > 0:
            logger.warning(
                f"Some of the weights of {model.__class__.__name__} were initialized in float16 precision from "
                f"the model checkpoint at {pretrained_model_name_or_path}:\n{fp16_params}\n"
                "You should probably UPCAST the model weights to float32 if this was not intended. "
                "See [`~FlaxPreTrainedModel.to_fp32`] for further information on how to do this."
            )

        if len(bf16_params) > 0:
            logger.warning(
                f"Some of the weights of {model.__class__.__name__} were initialized in bfloat16 precision from "
                f"the model checkpoint at {pretrained_model_name_or_path}:\n{bf16_params}\n"
                "You should probably UPCAST the model weights to float32 if this was not intended. "
                "See [`~FlaxPreTrainedModel.to_fp32`] for further information on how to do this."
            )

        
        if model.can_generate():
            try:
                model.generation_config = GenerationConfig.from_pretrained(
                    pretrained_model_name_or_path,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    resume_download=resume_download,
                    proxies=proxies,
                    local_files_only=local_files_only,
                    token=token,
                    revision=revision,
                    subfolder=subfolder,
                    _from_auto=from_auto_class,
                    _from_pipeline=from_pipeline,
                    **kwargs,
                )
            except OSError:
                logger.info(
                    "Generation config file not found, using a generation config created from the model config."
                )
                pass

        if _do_init:
            
            model.params = unflatten_dict(state)
            return model
        else:
            return model, unflatten_dict(state)

    def save_pretrained(
        self,
        save_directory: Union[str, os.PathLike],
        params=None,
        push_to_hub=False,
        max_shard_size="10GB",
        token: Optional[Union[str, bool]] = None,
        safe_serialization: bool = False,
        **kwargs,
    ):
        
        use_auth_token = kwargs.pop("use_auth_token", None)

        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
                )
            token = use_auth_token

        if token is not None:
            kwargs["token"] = token

        if os.path.isfile(save_directory):
            logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
            return

        os.makedirs(save_directory, exist_ok=True)

        if push_to_hub:
            commit_message = kwargs.pop("commit_message", None)
            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
            repo_id = self._create_repo(repo_id, **kwargs)
            files_timestamps = self._get_files_timestamps(save_directory)

        
        save_directory = os.path.abspath(save_directory)
        
        self.config.architectures = [self.__class__.__name__[4:]]

        
        
        if self._auto_class is not None:
            custom_object_save(self, save_directory, config=self.config)

        self.config.save_pretrained(save_directory)
        if self.can_generate():
            self.generation_config.save_pretrained(save_directory)

        
        weights_name = SAFE_WEIGHTS_NAME if safe_serialization else FLAX_WEIGHTS_NAME
        output_model_file = os.path.join(save_directory, weights_name)

        shards, index = flax_shard_checkpoint(params if params is not None else self.params, max_shard_size)
        
        for filename in os.listdir(save_directory):
            full_filename = os.path.join(save_directory, filename)
            weights_no_suffix = weights_name.replace(".bin", "").replace(".safetensors", "")
            if (
                filename.startswith(weights_no_suffix)
                and os.path.isfile(full_filename)
                and filename not in shards.keys()
            ):
                os.remove(full_filename)

        if index is None:
            if safe_serialization:
                params = params if params is not None else self.params
                flat_dict = flatten_dict(params, sep=".")
                safe_save_file(flat_dict, output_model_file, metadata={"format": "flax"})
            else:
                with open(output_model_file, "wb") as f:
                    params = params if params is not None else self.params
                    model_bytes = to_bytes(params)
                    f.write(model_bytes)

        else:
            save_index_file = os.path.join(save_directory, FLAX_WEIGHTS_INDEX_NAME)
            
            with open(save_index_file, "w", encoding="utf-8") as f:
                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
                f.write(content)
            logger.info(
                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
                f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the "
                f"index located at {save_index_file}."
            )
            for shard_file, shard in shards.items():
                
                with open(os.path.join(save_directory, shard_file), mode="wb") as f:
                    params = unflatten_dict(shard, sep="/")
                    shard_bytes = to_bytes(params)
                    f.write(shard_bytes)

        logger.info(f"Model weights saved in {output_model_file}")

        if push_to_hub:
            self._upload_modified_files(
                save_directory,
                repo_id,
                files_timestamps,
                commit_message=commit_message,
                token=token,
            )

    @classmethod
    def register_for_auto_class(cls, auto_class="FlaxAutoModel"):
        
        if not isinstance(auto_class, str):
            auto_class = auto_class.__name__

        import transformers.models.auto as auto_module

        if not hasattr(auto_module, auto_class):
            raise ValueError(f"{auto_class} is not a valid auto class.")

        cls._auto_class = auto_class



FlaxPreTrainedModel.push_to_hub = copy_func(FlaxPreTrainedModel.push_to_hub)
if FlaxPreTrainedModel.push_to_hub.__doc__ is not None:
    FlaxPreTrainedModel.push_to_hub.__doc__ = FlaxPreTrainedModel.push_to_hub.__doc__.format(
        object="model", object_class="FlaxAutoModel", object_files="model checkpoint"
    )


def overwrite_call_docstring(model_class, docstring):
    
    model_class.__call__ = copy_func(model_class.__call__)
    
    model_class.__call__.__doc__ = None
    
    model_class.__call__ = add_start_docstrings_to_model_forward(docstring)(model_class.__call__)


def append_call_sample_docstring(model_class, checkpoint, output_type, config_class, mask=None, revision=None):
    model_class.__call__ = copy_func(model_class.__call__)
    model_class.__call__ = add_code_sample_docstrings(
        checkpoint=checkpoint,
        output_type=output_type,
        config_class=config_class,
        model_cls=model_class.__name__,
        revision=revision,
    )(model_class.__call__)


def append_replace_return_docstrings(model_class, output_type, config_class):
    model_class.__call__ = copy_func(model_class.__call__)
    model_class.__call__ = replace_return_docstrings(
        output_type=output_type,
        config_class=config_class,
    )(model_class.__call__)