Transformers documentation

Models

Transformers

You are viewing v4.19.3 version. A newer version v4.47.1 is available.

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Models

The base classes PreTrainedModel, TFPreTrainedModel, and FlaxPreTrainedModel implement the common methods for loading/saving a model either from a local file or directory, or from a pretrained model configuration provided by the library (downloaded from HuggingFace’s AWS S3 repository).

PreTrainedModel and TFPreTrainedModel also implement a few methods which are common among all the models to:

resize the input token embeddings when new tokens are added to the vocabulary
prune the attention heads of the model.

The other methods that are common to each model are defined in ModuleUtilsMixin (for the PyTorch models) and TFModuleUtilsMixin (for the TensorFlow models) or for text generation, GenerationMixin (for the PyTorch models), TFGenerationMixin (for the TensorFlow models) and FlaxGenerationMixin (for the Flax/JAX models).

PreTrainedModel

class transformers.PreTrainedModel

( config: PretrainedConfig *inputs **kwargs )

Base class for all models.

PreTrainedModel takes care of storing the configuration of the models and handles methods for loading, downloading and saving models as well as a few methods common to all models to:

resize the input embeddings,
prune heads in the self-attention heads.

Class attributes (overridden by derived classes):

config_class (PretrainedConfig) — A subclass of PretrainedConfig to use as configuration class for this model architecture.
load_tf_weights (Callable) — A python method for loading a TensorFlow checkpoint in a PyTorch model, taking as arguments:
- model (PreTrainedModel) — An instance of the model on which to load the TensorFlow checkpoint.
- config (PreTrainedConfig) — An instance of the configuration associated to the model.
- path (str) — A path to the TensorFlow checkpoint.
base_model_prefix (str) — A string indicating the attribute associated to the base model in derived classes of the same architecture adding modules on top of the base model.
is_parallelizable (bool) — A flag indicating whether this model supports model parallelization.
main_input_name (str) — The name of the principal input to the model (often input_ids for NLP models, pixel_values for vision models and input_values for speech models).

push_to_hub

( repo_path_or_name: typing.Optional[str] = None repo_url: typing.Optional[str] = None use_temp_dir: bool = False commit_message: str = 'add model' organization: typing.Optional[str] = None private: typing.Optional[bool] = None use_auth_token: typing.Union[bool, str, NoneType] = None max_shard_size: typing.Union[int, str] = '10GB' **model_card_kwargs ) → str

Parameters

repo_path_or_name (str, optional) — Can either be a repository name for your model in the Hub or a path to a local folder (in which case the repository will have the name of that local folder). If not specified, will default to the name given by repo_url and a local directory with that name will be created.
repo_url (str, optional) — Specify this in case you want to push to an existing repository in the hub. If unspecified, a new repository will be created in your namespace (unless you specify an organization) with repo_name.
use_temp_dir (bool, optional, defaults to False) — Whether or not to clone the distant repo in a temporary directory or in repo_path_or_name inside the current working directory. This will slow things down if you are making changes in an existing repo since you will need to clone the repo before every push.
commit_message (str, optional, defaults to "add model") — Message to commit while pushing.
organization (str, optional) — Organization in which you want to push your {object} (you must be a member of this organization).
private (bool, optional) — Whether or not the repository created should be private (requires a paying subscription).
use_auth_token (bool or str, optional) — The token to use as HTTP bearer authorization for remote files. If True, will use the token generated when running transformers-cli login (stored in ~/.huggingface). Will default to True if repo_url is not specified.
max_shard_size (int or str, optional, defaults to "10GB") — The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size lower than this size. If expressed as a string, needs to be digits followed by a unit (like "5MB").

If a single weight of the model is bigger than max_shard_size, it will be in its own checkpoint shard which will be bigger than max_shard_size.

Returns

str

The url of the commit of your {object} in the given repository.

Upload the model files to the 🤗 Model Hub while synchronizing a local clone of the repo in repo_path_or_name.

Examples:

from transformers import AutoModel

model = AutoModel.from_pretrained("bert-base-cased")

# Push the model to your namespace with the name "my-finetuned-bert" and have a local clone in the
# *my-finetuned-bert* folder.
model.push_to_hub("my-finetuned-bert")

# Push the model to your namespace with the name "my-finetuned-bert" with no local clone.
model.push_to_hub("my-finetuned-bert", use_temp_dir=True)

# Push the model to an organization with the name "my-finetuned-bert" and have a local clone in the
# *my-finetuned-bert* folder.
model.push_to_hub("my-finetuned-bert", organization="huggingface")

# Make a change to an existing repo that has been cloned locally in *my-finetuned-bert*.
model.push_to_hub("my-finetuned-bert", repo_url="https://huggingface.co/sgugger/my-finetuned-bert")

from_pretrained

( pretrained_model_name_or_path: typing.Union[str, os.PathLike, NoneType] *model_args **kwargs )

Parameters

pretrained_model_name_or_path (str or os.PathLike, optional) — Can be either:
- A string, the model id of a pretrained model hosted inside a model repo on huggingface.co. 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 save_pretrained(), e.g., ./my_model_directory/.
- A path or url to a tensorflow index checkpoint file (e.g, ./tf_model/model.ckpt.index). In this case, from_tf should be set to True and a configuration object should be provided as config argument. This loading path is slower than converting the TensorFlow checkpoint in a PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
- A path or url to a model folder containing a flax checkpoint file in .msgpack format (e.g, ./flax_model/ containing flax_model.msgpack). In this case, from_flax should be set to True.
- None if you are both providing the configuration and state dictionary (resp. with keyword arguments config and state_dict).
model_args (sequence of positional arguments, optional) — All remaining positional arguments will be passed to the underlying model’s __init__ method.
config (Union[PretrainedConfig, str, os.PathLike], optional) — Can be either:
- an instance of a class derived from PretrainedConfig,
- a string or path valid as input to 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 save_pretrained() and is reloaded by supplying the save directory.
- The model is loaded by supplying a local directory as pretrained_model_name_or_path and a configuration JSON file named config.json is found in the directory.
state_dict (Dict[str, torch.Tensor], optional) — A state dictionary to use instead of a state dictionary loaded from saved weights file.

This option can be used if you want to create a model from a pretrained configuration but load your own weights. In this case though, you should check if using save_pretrained() and from_pretrained() is not a simpler option.
cache_dir (Union[str, os.PathLike], optional) — Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used.
from_tf (bool, optional, defaults to False) — Load the model weights from a TensorFlow checkpoint save file (see docstring of pretrained_model_name_or_path argument).
from_flax (bool, optional, defaults to False) — Load the model weights from a Flax checkpoint save file (see docstring of pretrained_model_name_or_path argument).
ignore_mismatched_sizes (bool, optional, defaults to False) — Whether or not to raise an error if some of the weights from the checkpoint do not have the same size as the weights of the model (if for instance, you are instantiating a model with 10 labels from a checkpoint with 3 labels).
force_download (bool, optional, defaults to False) — Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist.
resume_download (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.
output_loading_info(bool, optional, defaults to False) — Whether ot not to also return a dictionary containing missing keys, unexpected keys and error messages.
local_files_only(bool, optional, defaults to False) — Whether or not to only look at local files (i.e., do not try to download the model).
use_auth_token (str or bool, optional) — The token to use as HTTP bearer authorization for remote files. If True, will use the token generated when running transformers-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.
mirror (str, optional) — Mirror source to accelerate downloads in China. If you are from China and have an accessibility problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety. Please refer to the mirror site for more information.
_fast_init(bool, optional, defaults to True) — Whether or not to disable fast initialization.

One should only disable _fast_init to ensure backwards compatibility with transformers.__version__ < 4.6.0 for seeded model initialization. This argument will be removed at the next major version. See pull request 11471 for more information.
low_cpu_mem_usage(bool, optional, defaults to False) — Tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model. This is an experimental feature and a subject to change at any moment.
torch_dtype (str or torch.dtype, optional) — Override the default torch.dtype and load the model under this dtype. If "auto" is passed the dtype will be automatically derived from the model’s weights.
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 (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.

Instantiate a pretrained pytorch model from a pre-trained model configuration.

The model is set in evaluation mode by default using model.eval() (Dropout modules are deactivated). To train the model, you should first set it back in training mode with model.train().

The warning Weights from XXX not initialized from pretrained model means that the weights of XXX do not come pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning task.

The warning Weights from XXX not used in YYY means that the layer XXX is not used by YYY, therefore those weights are discarded.

Passing `use_auth_token=True“ is required when you want to use a private model.

Activate the special “offline-mode” to use this method in a firewalled environment.

Examples:

>>> from transformers import BertConfig, BertModel

>>> # Download model and configuration from huggingface.co and cache.
>>> model = BertModel.from_pretrained("bert-base-uncased")
>>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable).
>>> model = BertModel.from_pretrained("./test/saved_model/")
>>> # Update configuration during loading.
>>> model = BertModel.from_pretrained("bert-base-uncased", output_attentions=True)
>>> assert model.config.output_attentions == True
>>> # Loading from a TF checkpoint file instead of a PyTorch model (slower, for example purposes, not runnable).
>>> config = BertConfig.from_json_file("./tf_model/my_tf_model_config.json")
>>> model = BertModel.from_pretrained("./tf_model/my_tf_checkpoint.ckpt.index", from_tf=True, config=config)
>>> # Loading from a Flax checkpoint file instead of a PyTorch model (slower)
>>> model = BertModel.from_pretrained("bert-base-uncased", from_flax=True)

low_cpu_mem_usage algorithm:

This is an experimental function that loads the model using ~1x model size CPU memory

Here is how it works:

save which state_dict keys we have
drop state_dict before the model is created, since the latter takes 1x model size CPU memory
after the model has been instantiated switch to the meta device all params/buffers that are going to be replaced from the loaded state_dict
load state_dict 2nd time
replace the params/buffers from the state_dict

Currently, it can’t handle deepspeed ZeRO stage 3 and ignores loading errors

get_input_embeddings

( ) → nn.Module

Returns

nn.Module

A torch module mapping vocabulary to hidden states.

Returns the model’s input embeddings.

get_output_embeddings

( ) → nn.Module

Returns

nn.Module

A torch module mapping hidden states to vocabulary.

Returns the model’s output embeddings.

gradient_checkpointing_disable

( )

Deactivates gradient checkpointing for the current model.

Note that in other frameworks this feature can be referred to as “activation checkpointing” or “checkpoint activations”.

gradient_checkpointing_enable

( )

Activates gradient checkpointing for the current model.

Note that in other frameworks this feature can be referred to as “activation checkpointing” or “checkpoint activations”.

init_weights

( )

If needed prunes and maybe initializes weights.

post_init

( )

A method executed at the end of each Transformer model initialization, to execute code that needs the model’s modules properly initialized (such as weight initialization).

prune_heads

( heads_to_prune: typing.Dict[int, typing.List[int]] )

Parameters

heads_to_prune (Dict[int, List[int]]) — Dictionary with keys being selected layer indices (int) and associated values being the list of heads to prune in said layer (list of int). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on layer 1 and heads 2 and 3 on layer 2.

Prunes heads of the base model.

register_for_auto_class

( auto_class = 'AutoModel' )

Parameters

auto_class (str or type, optional, defaults to "AutoModel") — The auto class to register this new model with.

Register this class with a given auto class. This should only be used for custom models as the ones in the library are already mapped with an auto class.

This API is experimental and may have some slight breaking changes in the next releases.

resize_token_embeddings

( new_num_tokens: typing.Optional[int] = None ) → torch.nn.Embedding

Parameters

new_num_tokens (int, optional) — The number of new tokens in the embedding matrix. Increasing the size will add newly initialized vectors at the end. Reducing the size will remove vectors from the end. If not provided or None, just returns a pointer to the input tokens torch.nn.Embedding module of the model without doing anything.

Returns

torch.nn.Embedding

Pointer to the input tokens Embeddings Module of the model.

Resizes input token embeddings matrix of the model if new_num_tokens != config.vocab_size.

Takes care of tying weights embeddings afterwards if the model class has a tie_weights() method.

save_pretrained

( save_directory: typing.Union[str, os.PathLike] is_main_process: bool = True state_dict: typing.Optional[dict] = None save_function: typing.Callable = <function save at 0x7f1b109b1430> push_to_hub: bool = False max_shard_size: typing.Union[int, str] = '10GB' **kwargs )

Parameters

save_directory (str or os.PathLike) — Directory to which to save. Will be created if it doesn’t exist.
is_main_process (bool, optional, defaults to True) — Whether the process calling this is the main process or not. Useful when in distributed training like TPUs and need to call this function on all processes. In this case, set is_main_process=True only on the main process to avoid race conditions.
state_dict (nested dictionary of torch.Tensor) — The state dictionary of the model to save. Will default to self.state_dict(), but can be used to only save parts of the model or if special precautions need to be taken when recovering the state dictionary of a model (like when using model parallelism).
save_function (Callable) — The function to use to save the state dictionary. Useful on distributed training like TPUs when one need to replace torch.save by another method.
push_to_hub (bool, optional, defaults to False) — Whether or not to push your model to the Hugging Face model hub after saving it.

Using push_to_hub=True will synchronize the repository you are pushing to with save_directory, which requires save_directory to be a local clone of the repo you are pushing to if it’s an existing folder. Pass along temp_dir=True to use a temporary directory instead.
max_shard_size (int or str, optional, defaults to "10GB") — The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size lower than this size. If expressed as a string, needs to be digits followed by a unit (like "5MB").

If a single weight of the model is bigger than max_shard_size, it will be in its own checkpoint shard which will be bigger than max_shard_size.

kwargs — Additional key word arguments passed along to the push_to_hub() method.

Save a model and its configuration file to a directory, so that it can be re-loaded using the [from_pretrained()](/docs/transformers/v4.19.3/en/main_classes/model#transformers.PreTrainedModel.from_pretrained) class method.

set_input_embeddings

( value: Module )

Parameters

value (nn.Module) — A module mapping vocabulary to hidden states.

Set model’s input embeddings.

tie_weights

( )

Tie the weights between the input embeddings and the output embeddings.

If the torchscript flag is set in the configuration, can’t handle parameter sharing so we are cloning the weights instead.

Model Instantiation dtype

Under Pytorch a model normally gets instantiated with torch.float32 format. This can be an issue if one tries to load a model whose weights are in fp16, since it’d require twice as much memory. To overcome this limitation, you can either explicitly pass the desired dtype using torch_dtype argument:

model = T5ForConditionalGeneration.from_pretrained("t5", torch_dtype=torch.float16)

or, if you want the model to always load in the most optimal memory pattern, you can use the special value "auto", and then dtype will be automatically derived from the model’s weights:

model = T5ForConditionalGeneration.from_pretrained("t5", torch_dtype="auto")

Models instantiated from scratch can also be told which dtype to use with:

config = T5Config.from_pretrained("t5")
model = AutoModel.from_config(config)

Due to Pytorch design, this functionality is only available for floating dtypes.

ModuleUtilsMixin

class transformers.modeling_utils.ModuleUtilsMixin

( )

A few utilities for torch.nn.Modules, to be used as a mixin.

add_memory_hooks

( )

Add a memory hook before and after each sub-module forward pass to record increase in memory consumption.

Increase in memory consumption is stored in a mem_rss_diff attribute for each module and can be reset to zero with model.reset_memory_hooks_state().

estimate_tokens

( input_dict: typing.Dict[str, typing.Union[torch.Tensor, typing.Any]] ) → int

Parameters

inputs (dict) — The model inputs.

Returns

int

The total number of tokens.

Helper function to estimate the total number of tokens from the model inputs.

floating_point_ops

( input_dict: typing.Dict[str, typing.Union[torch.Tensor, typing.Any]] exclude_embeddings: bool = True ) → int

Parameters

batch_size (int) — The batch size for the forward pass.
sequence_length (int) — The number of tokens in each line of the batch.
exclude_embeddings (bool, optional, defaults to True) — Whether or not to count embedding and softmax operations.

Returns

int

The number of floating-point operations.

Get number of (optionally, non-embeddings) floating-point operations for the forward and backward passes of a batch with this transformer model. Default approximation neglects the quadratic dependency on the number of tokens (valid if 12 * d_model << sequence_length) as laid out in this paper section 2.1. Should be overridden for transformers with parameter re-use e.g. Albert or Universal Transformers, or if doing long-range modeling with very high sequence lengths.

get_extended_attention_mask

( attention_mask: Tensor input_shape: typing.Tuple[int] device: <property object at 0x7f1aa7d4e270> = None )

Parameters

attention_mask (torch.Tensor) — Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
input_shape (Tuple[int]) — The shape of the input to the model.

Makes broadcastable attention and causal masks so that future and masked tokens are ignored.

get_head_mask

( head_mask: typing.Optional[torch.Tensor] num_hidden_layers: int is_attention_chunked: bool = False )

Parameters

head_mask (torch.Tensor with shape [num_heads] or [num_hidden_layers x num_heads], optional) — The mask indicating if we should keep the heads or not (1.0 for keep, 0.0 for discard).
num_hidden_layers (int) — The number of hidden layers in the model. is_attention_chunked — (bool, optional, defaults to False): Whether or not the attentions scores are computed by chunks or not.

Prepare the head mask if needed.

invert_attention_mask

( encoder_attention_mask: Tensor ) → torch.Tensor

Parameters

encoder_attention_mask (torch.Tensor) — An attention mask.

Returns

torch.Tensor

The inverted attention mask.

Invert an attention mask (e.g., switches 0. and 1.).

num_parameters

( only_trainable: bool = False exclude_embeddings: bool = False ) → int

Parameters

only_trainable (bool, optional, defaults to False) — Whether or not to return only the number of trainable parameters
exclude_embeddings (bool, optional, defaults to False) — Whether or not to return only the number of non-embeddings parameters

Returns

int

The number of parameters.

Get number of (optionally, trainable or non-embeddings) parameters in the module.

reset_memory_hooks_state

( )

Reset the mem_rss_diff attribute of each module (see add_memory_hooks()).

TFPreTrainedModel

class transformers.TFPreTrainedModel

( *args **kwargs )

Base class for all TF models.

TFPreTrainedModel takes care of storing the configuration of the models and handles methods for loading, downloading and saving models as well as a few methods common to all models to:

resize the input embeddings,
prune heads in the self-attention heads.

Class attributes (overridden by derived classes):

config_class (PretrainedConfig) — A subclass of PretrainedConfig to use as configuration class for this model architecture.
base_model_prefix (str) — A string indicating the attribute associated to the base model in derived classes of the same architecture adding modules on top of the base model.
main_input_name (str) — The name of the principal input to the model (often input_ids for NLP models, pixel_values for vision models and input_values for speech models).

push_to_hub

( repo_path_or_name: typing.Optional[str] = None repo_url: typing.Optional[str] = None use_temp_dir: bool = False commit_message: typing.Optional[str] = None organization: typing.Optional[str] = None private: typing.Optional[bool] = None use_auth_token: typing.Union[bool, str, NoneType] = None **model_card_kwargs ) → str

Parameters

repo_path_or_name (str, optional) — Can either be a repository name for your model in the Hub or a path to a local folder (in which case the repository will have the name of that local folder). If not specified, will default to the name given by repo_url and a local directory with that name will be created.
repo_url (str, optional) — Specify this in case you want to push to an existing repository in the hub. If unspecified, a new repository will be created in your namespace (unless you specify an organization) with repo_name.
use_temp_dir (bool, optional, defaults to False) — Whether or not to clone the distant repo in a temporary directory or in repo_path_or_name inside the current working directory. This will slow things down if you are making changes in an existing repo since you will need to clone the repo before every push.
commit_message (str, optional) — Message to commit while pushing. Will default to "add model".
organization (str, optional) — Organization in which you want to push your model (you must be a member of this organization).
private (bool, optional) — Whether or not the repository created should be private (requires a paying subscription).
use_auth_token (bool or str, optional) — The token to use as HTTP bearer authorization for remote files. If True, will use the token generated when running transformers-cli login (stored in ~/.huggingface). Will default to True if repo_url is not specified.

Returns

str

The url of the commit of your model in the given repository.

Upload the model checkpoint to the 🤗 Model Hub while synchronizing a local clone of the repo in repo_path_or_name.

Examples:

from transformers import TFAutoModel

model = TFAutoModel.from_pretrained("bert-base-cased")

# Push the model to your namespace with the name "my-finetuned-bert" and have a local clone in the
# *my-finetuned-bert* folder.
model.push_to_hub("my-finetuned-bert")

# Push the model to your namespace with the name "my-finetuned-bert" with no local clone.
model.push_to_hub("my-finetuned-bert", use_temp_dir=True)

# Push the model to an organization with the name "my-finetuned-bert" and have a local clone in the
# *my-finetuned-bert* folder.
model.push_to_hub("my-finetuned-bert", organization="huggingface")

# Make a change to an existing repo that has been cloned locally in *my-finetuned-bert*.
model.push_to_hub("my-finetuned-bert", repo_url="https://huggingface.co/sgugger/my-finetuned-bert")

compile

( optimizer = 'rmsprop' loss = 'passthrough' metrics = None loss_weights = None weighted_metrics = None run_eagerly = None steps_per_execution = None **kwargs )

This is a thin wrapper that sets the model’s loss output head as the loss if the user does not specify a loss function themselves.

from_pretrained

( pretrained_model_name_or_path *model_args **kwargs )

Parameters

pretrained_model_name_or_path (str, optional) — Can be either:
- A string, the model id of a pretrained model hosted inside a model repo on huggingface.co. 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 save_pretrained(), e.g., ./my_model_directory/.
- A path or url to a PyTorch state_dict save file (e.g, ./pt_model/pytorch_model.bin). In this case, from_pt should be set to True and a configuration object should be provided as config argument. This loading path is slower than converting the PyTorch model in a TensorFlow model using the provided conversion scripts and loading the TensorFlow model afterwards.
- None if you are both providing the configuration and state dictionary (resp. with keyword arguments config and state_dict).
model_args (sequence of positional arguments, optional) — All remaining positional arguments will be passed to the underlying model’s __init__ method.
config (Union[PretrainedConfig, str], optional) — Can be either:
- an instance of a class derived from PretrainedConfig,
- a string valid as input to 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 save_pretrained() and is reloaded by supplying the save directory.
- The model is loaded by supplying a local directory as pretrained_model_name_or_path and a configuration JSON file named config.json is found in the directory. from_pt — (bool, optional, defaults to False): Load the model weights from a PyTorch state_dict save file (see docstring of pretrained_model_name_or_path argument).
ignore_mismatched_sizes (bool, optional, defaults to False) — Whether or not to raise an error if some of the weights from the checkpoint do not have the same size as the weights of the model (if for instance, you are instantiating a model with 10 labels from a checkpoint with 3 labels).
cache_dir (str, optional) — Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used.
force_download (bool, optional, defaults to False) — Whether or not to force the (re-)download of the model weights and configuration files, overriding the cached versions if they exist.
resume_download (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. output_loading_info(bool, *optional*, defaults to False`): Whether ot not to also return a dictionary containing missing keys, unexpected keys and error messages.
local_files_only(bool, optional, defaults to False) — Whether or not to only look at local files (e.g., not try doanloading the model).
use_auth_token (str or bool, optional) — The token to use as HTTP bearer authorization for remote files. If True, will use the token generated when running transformers-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.
mirror (str, optional) — Mirror source to accelerate downloads in China. If you are from China and have an accessibility problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety. Please refer to the mirror site for more information.
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 (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.

Instantiate a pretrained TF 2.0 model from a pre-trained model configuration.

The warning Weights from XXX not initialized from pretrained model means that the weights of XXX do not come pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning task.

The warning Weights from XXX not used in YYY means that the layer XXX is not used by YYY, therefore those weights are discarded.

Passing use_auth_token=True is required when you want to use a private model.

Examples:

>>> from transformers import BertConfig, TFBertModel

>>> # Download model and configuration from huggingface.co and cache.
>>> model = TFBertModel.from_pretrained("bert-base-uncased")
>>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable).
>>> model = TFBertModel.from_pretrained("./test/saved_model/")
>>> # Update configuration during loading.
>>> model = TFBertModel.from_pretrained("bert-base-uncased", output_attentions=True)
>>> assert model.config.output_attentions == True
>>> # Loading from a Pytorch model file instead of a TensorFlow checkpoint (slower, for example purposes, not runnable).
>>> config = BertConfig.from_json_file("./pt_model/my_pt_model_config.json")
>>> model = TFBertModel.from_pretrained("./pt_model/my_pytorch_model.bin", from_pt=True, config=config)

get_bias

( ) → tf.Variable

Returns

tf.Variable

The weights representing the bias, None if not an LM model.

Dict of bias attached to an LM head. The key represents the name of the bias attribute.

get_input_embeddings

( ) → tf.Variable

Returns

tf.Variable

The embeddings layer mapping vocabulary to hidden states.

Returns the model’s input embeddings layer.

get_lm_head

( ) → tf.keras.layers.Layer

Returns

tf.keras.layers.Layer

The LM head layer if the model has one, None if not.

The LM Head layer. This method must be overwritten by all the models that have a lm head.

get_output_embeddings

( ) → tf.Variable

Returns

tf.Variable

The new weights mapping vocabulary to hidden states.

Returns the model’s output embeddings

get_output_layer_with_bias

( ) → tf.keras.layers.Layer

Returns

tf.keras.layers.Layer

The layer that handles the bias, None if not an LM model.

Get the layer that handles a bias attribute in case the model has an LM head with weights tied to the embeddings

get_prefix_bias_name

( ) → str

Returns

str

The _prefix name of the bias.

Get the concatenated _prefix name of the bias from the model name to the parent layer

load_repo_checkpoint

( repo_path_or_name ) → dict

Parameters

repo_path_or_name (str) — Can either be a repository name for your {object} in the Hub or a path to a local folder (in which case the repository will have the name of that local folder).

Returns

dict

A dictionary of extra metadata from the checkpoint, most commonly an “epoch” count.

Loads a saved checkpoint (model weights and optimizer state) from a repo. Returns the current epoch count when the checkpoint was made.

prune_heads

( heads_to_prune )

Parameters

heads_to_prune (Dict[int, List[int]]) — Dictionary with keys being selected layer indices (int) and associated values being the list of heads to prune in said layer (list of int). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on layer 1 and heads 2 and 3 on layer 2.

Prunes heads of the base model.

resize_token_embeddings

( new_num_tokens = None ) → tf.Variable

Parameters

new_num_tokens (int, optional) — The number of new tokens in the embedding matrix. Increasing the size will add newly initialized vectors at the end. Reducing the size will remove vectors from the end. If not provided or None, just returns a pointer to the input tokens tf.Variable module of the model without doing anything.

Returns

tf.Variable

Pointer to the input tokens Embeddings Module of the model.

Resizes input token embeddings matrix of the model if new_num_tokens != config.vocab_size.

Takes care of tying weights embeddings afterwards if the model class has a tie_weights() method.

save_pretrained

( save_directory saved_model = False version = 1 push_to_hub = False **kwargs )

Parameters

save_directory (str) — Directory to which to save. Will be created if it doesn’t exist.
saved_model (bool, optional, defaults to False) — If the model has to be saved in saved model format as well or not.
version (int, optional, defaults to 1) — The version of the saved model. A saved model needs to be versioned in order to be properly loaded by TensorFlow Serving as detailed in the official documentation https://www.tensorflow.org/tfx/serving/serving_basic
push_to_hub (bool, optional, defaults to False) — Whether or not to push your model to the Hugging Face model hub after saving it.

Using push_to_hub=True will synchronize the repository you are pushing to with save_directory, which requires save_directory to be a local clone of the repo you are pushing to if it’s an existing folder. Pass along temp_dir=True to use a temporary directory instead.

kwargs — Additional key word arguments passed along to the push_to_hub() method.

Save a model and its configuration file to a directory, so that it can be re-loaded using the from_pretrained() class method.

serving

( inputs )

Parameters

inputs (Dict[str, tf.Tensor]) — The input of the saved model as a dictionary of tensors.

Method used for serving the model.

serving_output

( output )

Parameters

output (TFBaseModelOutput) — The output returned by the model.

Prepare the output of the saved model. Each model must implement this function.

set_bias

( value )

Parameters

value (Dict[tf.Variable]) — All the new bias attached to an LM head.

Set all the bias in the LM head.

set_input_embeddings

( value )

Parameters

value (tf.Variable) — The new weights mapping hidden states to vocabulary.

Set model’s input embeddings

set_output_embeddings

( value )

Parameters

value (tf.Variable) — The new weights mapping hidden states to vocabulary.

Set model’s output embeddings

test_step

( data )

A modification of Keras’s default test_step that cleans up the printed metrics when we use a dummy loss. If a user specifies a loss at model compile time, this function behaves as the original Keras test_step.

When the model is compiled without specifying the loss, our overridden compile function can set a simple dummy loss that just reads the loss output head of the model. When using this dummy loss, inputs can be passed either as keys in the input dictionary, or as normal Keras labels.

train_step

( data )

A modification of Keras’s default train_step that cleans up the printed metrics when we use a dummy loss. If a user specifies a loss at model compile time, this function behaves as the original Keras train_step.

When the model is compiled without specifying the loss, our overridden compile function can set a simple dummy loss that just reads the loss output head of the model. When using this dummy loss, inputs can be passed either as keys in the input dictionary, or as normal Keras labels.

TFModelUtilsMixin

class transformers.modeling_tf_utils.TFModelUtilsMixin

( )

A few utilities for tf.keras.Model, to be used as a mixin.

num_parameters

( only_trainable: bool = False ) → int

Parameters

only_trainable (bool, optional, defaults to False) — Whether or not to return only the number of trainable parameters

Returns

int

The number of parameters.

Get the number of (optionally, trainable) parameters in the model.

FlaxPreTrainedModel

class transformers.FlaxPreTrainedModel

( config: PretrainedConfig module: Module input_shape: typing.Tuple = (1, 1) seed: int = 0 dtype: dtype = <class 'jax._src.numpy.lax_numpy.float32'> _do_init: bool = True )

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).

push_to_hub

( repo_path_or_name: typing.Optional[str] = None repo_url: typing.Optional[str] = None use_temp_dir: bool = False commit_message: typing.Optional[str] = None organization: typing.Optional[str] = None private: typing.Optional[bool] = None use_auth_token: typing.Union[bool, str, NoneType] = None **model_card_kwargs ) → str

Parameters

repo_path_or_name (str, optional) — Can either be a repository name for your model in the Hub or a path to a local folder (in which case the repository will have the name of that local folder). If not specified, will default to the name given by repo_url and a local directory with that name will be created.
repo_url (str, optional) — Specify this in case you want to push to an existing repository in the hub. If unspecified, a new repository will be created in your namespace (unless you specify an organization) with repo_name.
use_temp_dir (bool, optional, defaults to False) — Whether or not to clone the distant repo in a temporary directory or in repo_path_or_name inside the current working directory. This will slow things down if you are making changes in an existing repo since you will need to clone the repo before every push.
commit_message (str, optional) — Message to commit while pushing. Will default to "add model".
organization (str, optional) — Organization in which you want to push your model (you must be a member of this organization).
private (bool, optional) — Whether or not the repository created should be private (requires a paying subscription).
use_auth_token (bool or str, optional) — The token to use as HTTP bearer authorization for remote files. If True, will use the token generated when running transformers-cli login (stored in ~/.huggingface). Will default to True if repo_url is not specified.

Returns

str

The url of the commit of your model in the given repository.

Upload the model checkpoint to the 🤗 Model Hub while synchronizing a local clone of the repo in repo_path_or_name.

Examples:

from transformers import FlaxAutoModel

model = FlaxAutoModel.from_pretrained("bert-base-cased")

# Push the model to your namespace with the name "my-finetuned-bert" and have a local clone in the
# *my-finetuned-bert* folder.
model.push_to_hub("my-finetuned-bert")

# Push the model to your namespace with the name "my-finetuned-bert" with no local clone.
model.push_to_hub("my-finetuned-bert", use_temp_dir=True)

# Push the model to an organization with the name "my-finetuned-bert" and have a local clone in the
# *my-finetuned-bert* folder.
model.push_to_hub("my-finetuned-bert", organization="huggingface")

# Make a change to an existing repo that has been cloned locally in *my-finetuned-bert*.
model.push_to_hub("my-finetuned-bert", repo_url="https://huggingface.co/sgugger/my-finetuned-bert")

from_pretrained

( pretrained_model_name_or_path: typing.Union[str, os.PathLike] dtype: dtype = <class 'jax._src.numpy.lax_numpy.float32'> *model_args **kwargs )

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 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 to_fp16() and 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 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 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).
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.
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 (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.

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.

Examples:

>>> 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)

register_for_auto_class

( auto_class = 'FlaxAutoModel' )

Parameters

auto_class (str or type, optional, defaults to "FlaxAutoModel") — The auto class to register this new model with.

Register this class with a given auto class. This should only be used for custom models as the ones in the library are already mapped with an auto class.

This API is experimental and may have some slight breaking changes in the next releases.

save_pretrained

( save_directory: typing.Union[str, os.PathLike] params = None push_to_hub = False **kwargs )

Parameters

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.

Using push_to_hub=True will synchronize the repository you are pushing to with save_directory, which requires save_directory to be a local clone of the repo you are pushing to if it’s an existing folder. Pass along temp_dir=True to use a temporary directory instead.

kwargs — Additional key word arguments passed along to the push_to_hub() method.

Save a model and its configuration file to a directory, so that it can be re-loaded using the [from_pretrained()](/docs/transformers/v4.19.3/en/main_classes/model#transformers.FlaxPreTrainedModel.from_pretrained) class method

to_bf16

( params: typing.Union[typing.Dict, flax.core.frozen_dict.FrozenDict] mask: typing.Any = None )

Parameters

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.

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.

Examples:

>>> 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)

to_fp16

( params: typing.Union[typing.Dict, flax.core.frozen_dict.FrozenDict] mask: typing.Any = None )

Parameters

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

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.

Examples:

>>> 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)

to_fp32

( params: typing.Union[typing.Dict, flax.core.frozen_dict.FrozenDict] mask: typing.Any = None )

Parameters

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

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.

Examples:

>>> 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)

Pushing to the Hub

class transformers.utils.PushToHubMixin

( )

A Mixin containing the functionality to push a model or tokenizer to the hub.

push_to_hub

( repo_path_or_name: typing.Optional[str] = None repo_url: typing.Optional[str] = None use_temp_dir: bool = False commit_message: typing.Optional[str] = None organization: typing.Optional[str] = None private: typing.Optional[bool] = None use_auth_token: typing.Union[bool, str, NoneType] = None **model_card_kwargs ) → str

Parameters

repo_path_or_name (str, optional) — Can either be a repository name for your {object} in the Hub or a path to a local folder (in which case the repository will have the name of that local folder). If not specified, will default to the name given by repo_url and a local directory with that name will be created.
repo_url (str, optional) — Specify this in case you want to push to an existing repository in the hub. If unspecified, a new repository will be created in your namespace (unless you specify an organization) with repo_name.
use_temp_dir (bool, optional, defaults to False) — Whether or not to clone the distant repo in a temporary directory or in repo_path_or_name inside the current working directory. This will slow things down if you are making changes in an existing repo since you will need to clone the repo before every push.
commit_message (str, optional) — Message to commit while pushing. Will default to "add {object}".
organization (str, optional) — Organization in which you want to push your {object} (you must be a member of this organization).
private (bool, optional) — Whether or not the repository created should be private (requires a paying subscription).
use_auth_token (bool or str, optional) — The token to use as HTTP bearer authorization for remote files. If True, will use the token generated when running transformers-cli login (stored in ~/.huggingface). Will default to True if repo_url is not specified.

Returns

str

The url of the commit of your {object} in the given repository.

Upload the {object_files} to the 🤗 Model Hub while synchronizing a local clone of the repo in repo_path_or_name.

Examples:

from transformers import {object_class}

{object} = {object_class}.from_pretrained("bert-base-cased")

# Push the {object} to your namespace with the name "my-finetuned-bert" and have a local clone in the
# *my-finetuned-bert* folder.
{object}.push_to_hub("my-finetuned-bert")

# Push the {object} to your namespace with the name "my-finetuned-bert" with no local clone.
{object}.push_to_hub("my-finetuned-bert", use_temp_dir=True)

# Push the {object} to an organization with the name "my-finetuned-bert" and have a local clone in the
# *my-finetuned-bert* folder.
{object}.push_to_hub("my-finetuned-bert", organization="huggingface")

# Make a change to an existing repo that has been cloned locally in *my-finetuned-bert*.
{object}.push_to_hub("my-finetuned-bert", repo_url="https://huggingface.co/sgugger/my-finetuned-bert")

Sharded checkpoints

transformers.modeling_utils.load_sharded_checkpoint

( model folder strict = True ) → NamedTuple

Parameters

model (torch.nn.Module) — The model in which to load the checkpoint.
folder (str or os.PathLike) — A path to a folder containing the sharded checkpoint.
strict (bool, *optional, defaults to True`) — Whether to strictly enforce that the keys in the model state dict match the keys in the sharded checkpoint.

Returns

NamedTuple

A named tuple with missing_keys and unexpected_keys fields

missing_keys is a list of str containing the missing keys
unexpected_keys is a list of str containing the unexpected keys

This is the same as torch.nn.Module.load_state_dict but for a sharded checkpoint.

This load is performed efficiently: each checkpoint shard is loaded one by one in RAM and deleted after being loaded in the model.

←Logging Text Generation→