Hub Python Library documentation
Serialização
Serialização
huggingface_hub fornece auxiliares para salvar e carregar pesos de modelos de ML de uma maneira padronizada. Esta parte da biblioteca ainda está em desenvolvimento e será aprimorada em versões futuras. O objetivo é harmonizar como os pesos são salvos e carregados no Hub, tanto para eliminar a duplicação de código entre as bibliotecas quanto para estabelecer convenções consistentes.
Formato de arquivo DDUF
DDUF é um formato de arquivo projetado para modelos de difusão. Ele permite salvar todas as informações necessárias para executar um modelo em um único arquivo. Este trabalho é inspirado no formato GGUF. O huggingface_hub fornece auxiliares para salvar e carregar arquivos DDUF, garantindo que o formato do arquivo seja respeitado.
Esta é uma versão muito inicial do parser. A API e a implementação podem evoluir em um futuro próximo.
O parser atualmente realiza muito pouca validação. Para mais detalhes sobre o formato do arquivo, consulte https://github.com/huggingface/huggingface.js/tree/main/packages/dduf.
Como escrever um arquivo DDUF?
Aqui está como exportar uma pasta contendo diferentes partes de um modelo de difusão usando export_folder_as_dduf():
# Export a folder as a DDUF file
>>> from huggingface_hub import export_folder_as_dduf
>>> export_folder_as_dduf("FLUX.1-dev.dduf", folder_path="path/to/FLUX.1-dev")Para mais flexibilidade, você pode usar export_entries_as_dduf() e passar uma lista de arquivos para incluir no arquivo DDUF final:
# Export specific files from the local disk.
>>> from huggingface_hub import export_entries_as_dduf
>>> export_entries_as_dduf(
... dduf_path="stable-diffusion-v1-4-FP16.dduf",
... entries=[ # List entries to add to the DDUF file (here, only FP16 weights)
... ("model_index.json", "path/to/model_index.json"),
... ("vae/config.json", "path/to/vae/config.json"),
... ("vae/diffusion_pytorch_model.fp16.safetensors", "path/to/vae/diffusion_pytorch_model.fp16.safetensors"),
... ("text_encoder/config.json", "path/to/text_encoder/config.json"),
... ("text_encoder/model.fp16.safetensors", "path/to/text_encoder/model.fp16.safetensors"),
... # ... add more entries here
... ]
... )O parâmetro entries também suporta a passagem de um iterável de caminhos ou bytes. Isso pode ser útil se você tiver um modelo carregado e quiser serializá-lo diretamente em um arquivo DDUF, em vez de precisar serializar cada componente para o disco primeiro e, em seguida, como um arquivo DDUF. Aqui está um exemplo de como um StableDiffusionPipeline pode ser serializado como DDUF:
# Export state_dicts one by one from a loaded pipeline
>>> from diffusers import DiffusionPipeline
>>> from typing import Generator, Tuple
>>> import safetensors.torch
>>> from huggingface_hub import export_entries_as_dduf
>>> pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
... # ... do some work with the pipeline
>>> def as_entries(pipe: DiffusionPipeline) -> Generator[Tuple[str, bytes], None, None]:
... # Build a generator that yields the entries to add to the DDUF file.
... # The first element of the tuple is the filename in the DDUF archive (must use UNIX separator!). The second element is the content of the file.
... # Entries will be evaluated lazily when the DDUF file is created (only 1 entry is loaded in memory at a time)
... yield "vae/config.json", pipe.vae.to_json_string().encode()
... yield "vae/diffusion_pytorch_model.safetensors", safetensors.torch.save(pipe.vae.state_dict())
... yield "text_encoder/config.json", pipe.text_encoder.config.to_json_string().encode()
... yield "text_encoder/model.safetensors", safetensors.torch.save(pipe.text_encoder.state_dict())
... # ... add more entries here
>>> export_entries_as_dduf(dduf_path="stable-diffusion-v1-4.dduf", entries=as_entries(pipe))Observação: na prática, o diffusers fornece um método para serializar diretamente um pipeline em um arquivo DDUF. O trecho acima é apenas um exemplo.
Como ler um arquivo DDUF?
>>> import json
>>> import safetensors.torch
>>> from huggingface_hub import read_dduf_file
# Read DDUF metadata
>>> dduf_entries = read_dduf_file("FLUX.1-dev.dduf")
# Returns a mapping filename <> DDUFEntry
>>> dduf_entries["model_index.json"]
DDUFEntry(filename='model_index.json', offset=66, length=587)
# Load model index as JSON
>>> json.loads(dduf_entries["model_index.json"].read_text())
{'_class_name': 'FluxPipeline', '_diffusers_version': '0.32.0.dev0', '_name_or_path': 'black-forest-labs/FLUX.1-dev', 'scheduler': ['diffusers', 'FlowMatchEulerDiscreteScheduler'], 'text_encoder': ['transformers', 'CLIPTextModel'], 'text_encoder_2': ['transformers', 'T5EncoderModel'], 'tokenizer': ['transformers', 'CLIPTokenizer'], 'tokenizer_2': ['transformers', 'T5TokenizerFast'], 'transformer': ['diffusers', 'FluxTransformer2DModel'], 'vae': ['diffusers', 'AutoencoderKL']}
# Load VAE weights using safetensors
>>> with dduf_entries["vae/diffusion_pytorch_model.safetensors"].as_mmap() as mm:
... state_dict = safetensors.torch.load(mm)Auxiliares
huggingface_hub.export_entries_as_dduf
< source >( dduf_path: str | os.PathLikeentries: Iterable )
Parâmetros
- dduf_path (
strouos.PathLike) — O caminho para o arquivo DDUF a ser gravado. - entries (
Iterable[tuple[str, Union[str, Path, bytes]]]) — Um iterável de entradas para gravar no arquivo DDUF. Cada entrada é uma tupla com o nome do arquivo e o conteúdo. O nome do arquivo deve ser o caminho do arquivo no arquivo DDUF. O conteúdo pode ser uma string ou um pathlib.Path representando um caminho para um arquivo no disco local ou diretamente o conteúdo como bytes.
Exceções
-
- —
DDUFExportError: Se algo der errado durante a exportação (por exemplo, nome de entrada inválido, ausência de ‘model_index.json’, etc.).
- —
Escreva um arquivo DDUF a partir de um iterável de entradas.
Este é um auxiliar de nível inferior em relação ao export_folder_as_dduf() que permite mais flexibilidade ao serializar dados. Em particular, você não precisa salvar os dados em disco antes de exportá-los no arquivo DDUF.
Exemplo:
# Export specific files from the local disk.
>>> from huggingface_hub import export_entries_as_dduf
>>> export_entries_as_dduf(
... dduf_path="stable-diffusion-v1-4-FP16.dduf",
... entries=[ # List entries to add to the DDUF file (here, only FP16 weights)
... ("model_index.json", "path/to/model_index.json"),
... ("vae/config.json", "path/to/vae/config.json"),
... ("vae/diffusion_pytorch_model.fp16.safetensors", "path/to/vae/diffusion_pytorch_model.fp16.safetensors"),
... ("text_encoder/config.json", "path/to/text_encoder/config.json"),
... ("text_encoder/model.fp16.safetensors", "path/to/text_encoder/model.fp16.safetensors"),
... # ... add more entries here
... ]
... )# Export state_dicts one by one from a loaded pipeline
>>> from diffusers import DiffusionPipeline
>>> from typing import Generator, Tuple
>>> import safetensors.torch
>>> from huggingface_hub import export_entries_as_dduf
>>> pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
... # ... do some work with the pipeline
>>> def as_entries(pipe: DiffusionPipeline) -> Generator[tuple[str, bytes], None, None]:
... # Build a generator that yields the entries to add to the DDUF file.
... # The first element of the tuple is the filename in the DDUF archive (must use UNIX separator!). The second element is the content of the file.
... # Entries will be evaluated lazily when the DDUF file is created (only 1 entry is loaded in memory at a time)
... yield "vae/config.json", pipe.vae.to_json_string().encode()
... yield "vae/diffusion_pytorch_model.safetensors", safetensors.torch.save(pipe.vae.state_dict())
... yield "text_encoder/config.json", pipe.text_encoder.config.to_json_string().encode()
... yield "text_encoder/model.safetensors", safetensors.torch.save(pipe.text_encoder.state_dict())
... # ... add more entries here
>>> export_entries_as_dduf(dduf_path="stable-diffusion-v1-4.dduf", entries=as_entries(pipe))huggingface_hub.export_folder_as_dduf
< source >( dduf_path: str | os.PathLikefolder_path: str | os.PathLike )
Exporte uma pasta como um arquivo DDUF.
AUtiliza export_entries_as_dduf() por baixo dos panos.
huggingface_hub.read_dduf_file
< source >( dduf_path: os.PathLike | str ) → dict[str, DDUFEntry]
Parâmetros
Retorna
dict[str, DDUFEntry]
Um dicionário de DDUFEntry indexado pelo nome do arquivo.
Exceções
-
- —
DDUFCorruptedFileError: Se o arquivo DDUF estiver corrompido (ou seja, não seguir o formato DDUF).
- —
Lê um arquivo DDUF e retorna um dicionário de entradas.
Apenas os metadados são lidos; os dados não são carregados na memória.
Exemplo:
>>> import json
>>> import safetensors.torch
>>> from huggingface_hub import read_dduf_file
# Read DDUF metadata
>>> dduf_entries = read_dduf_file("FLUX.1-dev.dduf")
# Returns a mapping filename <> DDUFEntry
>>> dduf_entries["model_index.json"]
DDUFEntry(filename='model_index.json', offset=66, length=587)
# Load model index as JSON
>>> json.loads(dduf_entries["model_index.json"].read_text())
{'_class_name': 'FluxPipeline', '_diffusers_version': '0.32.0.dev0', '_name_or_path': 'black-forest-labs/FLUX.1-dev', ...
# Load VAE weights using safetensors
>>> with dduf_entries["vae/diffusion_pytorch_model.safetensors"].as_mmap() as mm:
... state_dict = safetensors.torch.load(mm)classe huggingface_hub.DDUFEntry
< source >( filename: strlength: intoffset: intdduf_path: Path )
Objeto que representa uma entrada de arquivo em um arquivo DDUF.
Veja read_dduf_file() para saber como ler um arquivo DDUF.
Abra o arquivo como um arquivo mapeado em memória.
Útil para carregar safetensors diretamente do arquivo.
Leia o arquivo como texto.
Útil para entradas ‘.txt’ e ‘.json’.
Erros
Exceção base para erros relacionados ao formato DDUF.
Exceção lançada quando o arquivo DDUF está corrompido.
Exceção base para erros durante a exportação do DDUF.
Exceção lançada quando o nome da entrada é inválido.
Salvando tensores
O principal auxiliar do módulo serialization recebe um nn.Module do torch como entrada e o salva em disco. Ele trata a lógica para salvar tensores compartilhados (veja explicação do safetensors) e a lógica para dividir o dicionário de estado em fragmentos, usando split_torch_state_dict_into_shards() por baixo dos panos. No momento, apenas o framework torch é suportado.
Se você deseja salvar um dicionário de estado (por exemplo, um mapeamento entre nomes de camadas e os tensores relacionados) em vez de um nn.Module, você pode usar save_torch_state_dict(), que oferece os mesmos recursos. Isso é útil, por exemplo, se você quiser aplicar uma lógica personalizada ao dicionário de estado antes de salvá-lo.
save_torch_model
huggingface_hub.save_torch_model
< source >( model: torch.nn.Modulesave_directory: str | pathlib.Pathfilename_pattern: str | None = Noneforce_contiguous: bool = Truemax_shard_size: int | str = '5GB'metadata: dict[str, str] | None = Nonesafe_serialization: bool = Trueis_main_process: bool = Trueshared_tensors_to_discard: list[str] | None = None )
Parâmetros
- model (
torch.nn.Module) — O modelo a ser salvo em disco. - save_directory (
strouPath) — O diretório no qual o modelo será salvo. - filename_pattern (
str, opcional) — O padrão para gerar os nomes dos arquivos nos quais o modelo será salvo. O padrão deve ser uma string que pode ser formatada comfilename_pattern.format(suffix=...)e deve conter a palavra-chavesuffixPadrão:"model{suffix}.safetensors"oupytorch_model{suffix}.bin, dependendo do parâmetrosafe_serialization. - force_contiguous (
boolean, opcional) — Força o state_dict a ser salvo como tensores contíguos. Isso não afeta a correção do modelo, mas pode potencialmente alterar o desempenho se o layout do tensor foi escolhido especificamente por esse motivo. O valor padrão éTrue. - max_shard_size (
intoustr, opcional) — O tamanho máximo de cada shard, em bytes. O padrão é 5GB. - metadata (
dict[str, str], opcional) — Informações extras para salvar junto com o modelo. Algumas metadados serão adicionados para cada tensor descartado. Essas informações não serão suficientes para recuperar toda a estrutura compartilhada, mas podem ajudar a entender as coisas. - safe_serialization (
bool, opcional) — Se deve salvar como safetensors, que é o comportamento padrão. SeFalse, os shards são salvos como pickle. A serialização segura é recomendada por motivos de segurança. Salvar como pickle está obsoleto e será removido em uma versão futura. - is_main_process (
bool, opcional) — Indica se o processo que está chamando esta função é o processo principal ou não. Útil em treinamento distribuído, como em TPUs, quando é necessário chamar esta função a partir de todos os processos. Nesse caso, definais_main_process=Trueapenas no processo principal para evitar condições de corrida. O valor padrão é True. - shared_tensors_to_discard (
list[str], opcional) — Lista de nomes de tensores a serem descartados ao salvar tensores compartilhados. Se não for fornecida e tensores compartilhados forem detectados, o primeiro nome em ordem alfabética será descartado.
Salva um modelo torch específico no disco, lidando com problemas de sharding e tensores compartilhados.
Veja também save_torch_state_dict() para salvar um state dict com mais flexibilidade.
Para mais informações sobre compartilhamento de tensores, consulte este guia.
O dicionário de estado do modelo é dividido em fragmentos para que cada fragmento seja menor que um tamanho dado. Os fragmentos são
salvos no save_directory com o filename_pattern fornecido. Se o modelo for grande demais para caber em um único fragmento,
um arquivo de índice é salvo no save_directory para indicar onde cada tensor é salvo. Este helper usa split_torch_state_dict_into_shards() por baixo dos panos. Se safe_serialization for True, os fragmentos são salvos como
safetensors (o padrão). Caso contrário, os fragmentos são salvos como pickle.
Antes de salvar o modelo, o save_directory é limpo de quaisquer arquivos de shard anteriores.
Se um dos tensores do modelo for maior que
max_shard_size, ele ficará em seu próprio shard, que terá um tamanho maior quemax_shard_size.
Se o seu modelo for um
transformers.PreTrainedModel, você deve passarmodel._tied_weights_keyscomoshared_tensors_to_discardpara lidar corretamente com o salvamento de tensores compartilhados. Isso garante que os tensores duplicados corretos sejam descartados durante o salvamento.
Exemplo:
>>> from huggingface_hub import save_torch_model
>>> model = ... # A PyTorch model
# Save state dict to "path/to/folder". The model will be split into shards of 5GB each and saved as safetensors.
>>> save_torch_model(model, "path/to/folder")
# Load model back
>>> from huggingface_hub import load_torch_model # TODO
>>> load_torch_model(model, "path/to/folder")
>>>save_torch_state_dict
huggingface_hub.save_torch_state_dict
< source >( state_dict: dictsave_directory: str | pathlib.Pathfilename_pattern: str | None = Noneforce_contiguous: bool = Truemax_shard_size: int | str = '5GB'metadata: dict[str, str] | None = Nonesafe_serialization: bool = Trueis_main_process: bool = Trueshared_tensors_to_discard: list[str] | None = None )
Parâmetros
- state_dict (
dict[str, torch.Tensor]) — O dicionário de estado a ser salvo. - save_directory (
strouPath) — O diretório no qual o modelo será salvo. - filename_pattern (
str, opcional) — O padrão para gerar os nomes dos arquivos nos quais o modelo será salvo. O padrão deve ser uma string que pode ser formatada comfilename_pattern.format(suffix=...)e deve conter a palavra-chavesuffixPadrão:"model{suffix}.safetensors"oupytorch_model{suffix}.bin, dependendo do parâmetrosafe_serialization. - force_contiguous (
boolean, opcional) — Força o state_dict a ser salvo como tensores contíguos. Isso não afeta a correção do modelo, mas pode potencialmente alterar o desempenho se o layout do tensor foi escolhido especificamente por esse motivo. O valor padrão éTrue. - max_shard_size (
intoustr, opcional) — O tamanho máximo de cada shard, em bytes. O padrão é 5GB. - metadata (
dict[str, str], opcional) — Informações extras para salvar junto com o modelo. Algumas metadados serão adicionados para cada tensor descartado. Essas informações não serão suficientes para recuperar toda a estrutura compartilhada, mas podem ajudar a entender as coisas. - safe_serialization (
bool, opcional) — Se deve salvar como safetensors, que é o comportamento padrão. SeFalse, os shards são salvos como pickle. A serialização segura é recomendada por motivos de segurança. Salvar como pickle está obsoleto e será removido em uma versão futura. - is_main_process (
bool, opcional) — Indica se o processo que está chamando esta função é o processo principal ou não. Útil em treinamento distribuído, como em TPUs, quando é necessário chamar esta função a partir de todos os processos. Nesse caso, definais_main_process=Trueapenas no processo principal para evitar condições de corrida. O valor padrão é True. - shared_tensors_to_discard (
list[str], opcional) — Lista de nomes de tensores a serem descartados ao salvar tensores compartilhados. Se não for fornecida e tensores compartilhados forem detectados, o primeiro nome em ordem alfabética será descartado.
Salve um dicionário de estado de modelo no disco, lidando com problemas de sharding e tensores compartilhados.
Veja também save_torch_model() para salvar diretamente um modelo do PyTorch.
Para mais informações sobre compartilhamento de tensores, consulte este guia.
O dicionário de estado do modelo é dividido em fragmentos para que cada fragmento seja menor que um tamanho dado. Os fragmentos são
salvos no save_directory com o filename_pattern fornecido. Se o modelo for grande demais para caber em um único fragmento,
um arquivo de índice é salvo no save_directory para indicar onde cada tensor é salvo. Este helper usa split_torch_state_dict_into_shards() por baixo dos panos. Se safe_serialization for True, os fragmentos são salvos como
safetensors (o padrão). Caso contrário, os fragmentos são salvos como pickle.
Antes de salvar o modelo, o save_directory é limpo de quaisquer arquivos de shard anteriores.
Se um dos tensores do modelo for maior que
max_shard_size, ele ficará em seu próprio shard, que terá um tamanho maior quemax_shard_size.
Se o seu modelo for um
transformers.PreTrainedModel, você deve passarmodel._tied_weights_keyscomoshared_tensors_to_discardpara lidar corretamente com o salvamento de tensores compartilhados. Isso garante que os tensores duplicados corretos sejam descartados durante o salvamento.
Exemplo:
>>> from huggingface_hub import save_torch_state_dict
>>> model = ... # A PyTorch model
# Save state dict to "path/to/folder". The model will be split into shards of 5GB each and saved as safetensors.
>>> state_dict = model_to_save.state_dict()
>>> save_torch_state_dict(state_dict, "path/to/folder")O módulo serialization também contém auxiliares de baixo nível para dividir um dicionário de estado em vários shards, criando um índice adequado no processo. Esses auxiliares estão disponíveis para tensores torch e foram projetados para serem facilmente estendidos a qualquer outro framework de ML.
split_torch_state_dict_into_shards
huggingface_hub.split_torch_state_dict_into_shards
< source >( state_dict: dictfilename_pattern: str = 'model{suffix}.safetensors'max_shard_size: int | str = '5GB' ) → StateDictSplit
Parâmetros
- state_dict (
dict[str, torch.Tensor]) — O dicionário de estado a ser salvo. - filename_pattern (
str, opcional) — O padrão para gerar os nomes dos arquivos nos quais o modelo será salvo. O padrão deve ser uma string que pode ser formatada comfilename_pattern.format(suffix=...)e deve conter a palavra-chavesuffixPadrão:"model{suffix}.safetensors". - max_shard_size (
intoustr, opcional) — O tamanho máximo de cada shard, em bytes. O padrão é 5GB.
Retorna
StateDictSplit
Um objeto StateDictSplit contendo os shards e o índice para recuperá-los.
Divida o dicionário de estado de um modelo em fragmentos para que cada fragmento seja menor que um tamanho dado.
Os shards são determinados iterando pelo state_dict na ordem de suas chaves. Não há otimização
para tornar cada shard o mais próximo possível do tamanho máximo passado. Por exemplo, se o limite for 10GB e
tivermos tensores de tamanhos [6GB, 6GB, 2GB, 6GB, 2GB, 2GB], eles serão shardados como [6GB], [6+2GB], [6+2+2GB] e não
[6+2+2GB], [6+2GB], [6GB].
Para salvar um dicionário de estado de modelo no disco, consulte save_torch_state_dict(). Este helper usa
split_torch_state_dict_into_shardspor baixo dos panos.
Se um dos tensores do modelo for maior que
max_shard_size, ele ficará em seu próprio shard, que terá um tamanho maior quemax_shard_size.
Exemplo:
>>> import json
>>> import os
>>> from safetensors.torch import save_file as safe_save_file
>>> from huggingface_hub import split_torch_state_dict_into_shards
>>> def save_state_dict(state_dict: dict[str, torch.Tensor], save_directory: str):
... state_dict_split = split_torch_state_dict_into_shards(state_dict)
... for filename, tensors in state_dict_split.filename_to_tensors.items():
... shard = {tensor: state_dict[tensor] for tensor in tensors}
... safe_save_file(
... shard,
... os.path.join(save_directory, filename),
... metadata={"format": "pt"},
... )
... if state_dict_split.is_sharded:
... index = {
... "metadata": state_dict_split.metadata,
... "weight_map": state_dict_split.tensor_to_filename,
... }
... with open(os.path.join(save_directory, "model.safetensors.index.json"), "w") as f:
... f.write(json.dumps(index, indent=2))split_state_dict_into_shards_factory
Esta é a fábrica subjacente a partir da qual cada helper específico de framework é derivado. Na prática, não se espera que você use esta fábrica diretamente, exceto se precisar adaptá-la a um framework que ainda não é suportado. Se esse for o caso, por favor, nos avise abrindo uma nova issue no repositório huggingface_hub.
huggingface_hub.split_state_dict_into_shards_factory
< source >( state_dict: dictget_storage_size: Callablefilename_pattern: strget_storage_id: Callable = <function <lambda> at 0x7f9673968670>max_shard_size: int | str = '5GB' ) → StateDictSplit
Parâmetros
- state_dict (
dict[str, Tensor]) — O dicionário de estado a ser salvo. - get_storage_size (
Callable[[Tensor], int]) — Uma função que retorna o tamanho de um tensor quando salvo em disco, em bytes. - get_storage_id (
Callable[[Tensor], Optional[Any]], opcional) — Uma função que retorna um identificador único para o armazenamento de um tensor. Vários tensores diferentes podem compartilhar o mesmo armazenamento subjacente. Esse identificador é garantido como único e constante para o armazenamento deste tensor durante seu ciclo de vida. Dois armazenamentos de tensor com ciclos de vida não sobrepostos podem ter o mesmo id. - filename_pattern (
str, opcional) — O padrão para gerar os nomes dos arquivos nos quais o modelo será salvo. O padrão deve ser uma string que pode ser formatada comfilename_pattern.format(suffix=...)e deve conter a palavra-chavesuffix - max_shard_size (
intoustr, opcional) — O tamanho máximo de cada shard, em bytes. O padrão é 5GB.
Retorna
StateDictSplit
Um objeto StateDictSplit contendo os shards e o índice para recuperá-los.
Divida o dicionário de estado de um modelo em fragmentos para que cada fragmento seja menor que um tamanho dado.
Os shards são determinados iterando pelo state_dict na ordem de suas chaves. Não há otimização
para tornar cada shard o mais próximo possível do tamanho máximo passado. Por exemplo, se o limite for 10GB e
tivermos tensores de tamanhos [6GB, 6GB, 2GB, 6GB, 2GB, 2GB], eles serão shardados como [6GB], [6+2GB], [6+2+2GB] e não
[6+2+2GB], [6+2GB], [6GB].
Se um dos tensores do modelo for maior que
max_shard_size, ele ficará em seu próprio shard, que terá um tamanho maior quemax_shard_size.
Carregando tensores
Os helpers de carregamento suportam checkpoints em arquivo único e fragmentados, no formato safetensors ou pickle. load_torch_model() recebe um nn.Module e um caminho de checkpoint (seja um arquivo único ou um diretório) como entrada e carrega os pesos no modelo.
Os dois helpers têm como padrão
safe=True, o que significa que o checkpoint é sempre desserializado com o carregador de safetensors — independentemente do nome do arquivo. Carregar um checkpoint em pickle requer opt-in explícito, pois o unpickling executa código arbitrário no momento do carregamento:>>> from huggingface_hub import load_state_dict_from_file, load_torch_model # Raises: the file is not safetensors >>> load_state_dict_from_file("path/to/pytorch_model.bin") >>> load_torch_model(model, "path/to/pytorch_model.bin") # Explicit opt-in >>> load_state_dict_from_file("path/to/pytorch_model.bin", safe=False) >>> load_torch_model(model, "path/to/pytorch_model.bin", safe=False)O caminho do pickle também usa
weights_only=Truepor padrão, ou seja, o descompactador restrito do torch. Um checkpoint que contém objetos não tensor (estados do otimizador, um modelo inteiro, …) precisa deweights_only=Falseem conjunto comsafe=False.
load_torch_model
huggingface_hub.load_torch_model
< source >( model: torch.nn.Modulecheckpoint_path: str | os.PathLikestrict: bool = Falsesafe: bool = Trueweights_only: bool = Truemap_location: typing.Union[str, ForwardRef('torch.device'), NoneType] = Nonemmap: bool = Falsefilename_pattern: str | None = None ) → NamedTuple
Parâmetros
- model (
torch.nn.Module) — O modelo no qual o checkpoint será carregado. - checkpoint_path (
strouos.PathLike) — Caminho para o arquivo de checkpoint ou para o diretório contendo os checkpoints. - strict (
bool, opcional, padrãoFalse) — Se as chaves no state dict do modelo devem corresponder estritamente às chaves no checkpoint. Assim como emtorch.nn.Module.load_state_dict, a verificação ocorre após o carregamento do checkpoint, portanto o modelo pode ter sido atualizado parcialmente quando o erro é lançado. - safe (
bool, opcional, padrãoTrue) — Sesafefor True, os arquivos safetensors serão carregados. Sesafefor False, a função primeiro tentará carregar arquivos safetensors, se disponíveis; caso contrário, usará arquivos pickle. Umfilename_patternque não descreva arquivos safetensors é rejeitado quandosafe=True. - weights_only (
bool, opcional, padrãoTrue) — Se True, carrega apenas os pesos do modelo, sem os estados do otimizador e outros metadados, usando o unpickler restrito do torch. Defina como False para permitir objetos Python arbitrários em um checkpoint pickle (isso executa código arbitrário no momento da carga). Não tem efeito no PyTorch < 1.13, que não possui um unpickler restrito e sempre faz o unpickle sem restrições. - map_location (
stroutorch.device, opcional) — Um objetotorch.device, uma string ou um dicionário especificando como remapear os locais de armazenamento. Indica o local onde todos os tensors devem ser carregados. - mmap (
bool, opcional, padrãoFalse) — Se deve usar o carregamento de arquivo mapeado em memória. O mapeamento de memória pode melhorar o desempenho de carregamento para modelos grandes no PyTorch >= 2.1.0 com checkpoints baseados em zipfile. - filename_pattern (
str, opcional) — O padrão para localizar o arquivo de índice. O padrão deve ser uma string que pode ser formatada comfilename_pattern.format(suffix=...)e deve conter a palavra-chavesuffixPadrão:"model{suffix}.safetensors".
Retorna
NamedTuple
Uma tupla nomeada com os campos missing_keys e unexpected_keys.
missing_keysé uma lista de str contendo as chaves ausentes, ou seja, chaves que estão no modelo mas não no checkpoint.unexpected_keysé uma lista de str contendo as chaves inesperadas, ou seja, chaves que estão no checkpoint mas não no modelo.
Exceções
FileNotFoundError ou ImportError ou ValueError
FileNotFoundError— Se o arquivo ou diretório de checkpoint não existir.ImportError— Se safetensors ou torch não estiverem instalados ao tentar carregar um arquivo .safetensors ou um checkpoint do PyTorch, respectivamente.ValueError— Se o formato do checkpoint não puder ser determinado.
Carrega um checkpoint em um modelo, tratando checkpoints fragmentados e não fragmentados.
load_state_dict_from_file
huggingface_hub.load_state_dict_from_file
< source >( checkpoint_file: str | os.PathLikemap_location: typing.Union[str, ForwardRef('torch.device'), NoneType] = Noneweights_only: bool = Truemmap: bool = Falsesafe: bool = True ) → Union[dict[str, "torch.Tensor"], Any]
Parâmetros
- checkpoint_file (
strouos.PathLike) — Caminho para o arquivo de checkpoint a ser carregado. Pode ser um checkpoint em safetensors ou pickle (.bin). - map_location (
stroutorch.device, opcional) — Um objetotorch.device, uma string ou um dicionário especificando como remapear os locais de armazenamento. Indica o local onde todos os tensors devem ser carregados. - weights_only (
bool, opcional, padrãoTrue) — Se True, carrega apenas os pesos do modelo, sem os estados do otimizador e outros metadados, usando o unpickler restrito do torch. Defina como False para permitir objetos Python arbitrários em um checkpoint pickle (isso executa código arbitrário no momento da carga). Não tem efeito ao carregar arquivos safetensors, nem no PyTorch < 1.13, que não possui um unpickler restrito — essas versões sempre fazem unpickle sem restrição e um aviso é registrado. - mmap (
bool, opcional, padrãoFalse) — Se deve usar carregamento de arquivo mapeado em memória. O mapeamento de memória pode melhorar o desempenho de carregamento para modelos grandes no PyTorch >= 2.1.0 com checkpoints baseados em zipfile. Não tem efeito ao carregar arquivos safetensors, pois a bibliotecasafetensorsusa mapeamento de memória por padrão. - safe (
bool, opcional, padrãoTrue) — Se True, o checkpoint é sempre carregado como safetensors, independentemente do seu nome. Qualquer outro formato (por exemplo, um arquivo pickle.bin) levanta umValueErrorem vez de ser desserializado: checkpoints pickle podem executar código arbitrário no momento do carregamento. Se False, o arquivo é carregado como safetensors quando seu nome indica isso, e comtorch.loadcaso contrário.
Retorna
Union[dict[str, "torch.Tensor"], Any]
O checkpoint carregado.
- Para arquivos safetensors: sempre retorna um dicionário que mapeia nomes de parâmetros para tensores.
- Para arquivos pickle: retorna qualquer objeto Python que foi pickled (comumente um state dict, mas pode ser um modelo inteiro, o estado do otimizador ou qualquer outro objeto Python).
Exceções
FileNotFoundError ou ImportError ou OSError ou ValueError
FileNotFoundError— Se o arquivo de checkpoint não existir.ImportError— Se safetensors ou torch não estiverem instalados ao tentar carregar um arquivo .safetensors ou um checkpoint do PyTorch, respectivamente.OSError— Se o formato do arquivo de checkpoint for inválido ou se os arquivos git-lfs não forem baixados corretamente.ValueError— Se o caminho do arquivo de checkpoint estiver vazio ou inválido, ou se ele não puder ser desserializado como safetensors enquantosafe=True(passarsafe=Falsepara permitir checkpoints pickle).
Carrega um arquivo de checkpoint, tratando tanto o formato de checkpoint safetensors quanto o formato pickle.
Exemplo:
>>> from huggingface_hub import load_state_dict_from_file
# Load a safetensors checkpoint (safe by default)
>>> state_dict = load_state_dict_from_file("path/to/model.safetensors", safe=True)
>>> model.load_state_dict(state_dict)
# Load a pickle checkpoint. `safe=False` is required: pickle files can execute arbitrary code.
>>> state_dict = load_state_dict_from_file("path/to/model.bin", safe=False, map_location="cpu")
>>> model.load_state_dict(state_dict)Helpers de tensores
get_torch_storage_id
Retorna um identificador único para o armazenamento de um tensor.
Vários tensores diferentes podem compartilhar o mesmo armazenamento subjacente. Este identificador é garantido como único e constante para o armazenamento deste tensor durante seu tempo de vida. Dois armazenamentos de tensor com tempos de vida não sobrepostos podem ter o mesmo id. No caso de tensores meta, retornamos None, pois não podemos determinar se eles compartilham o mesmo armazenamento.