Hub Python Library documentation
Chuẩn hóa
Chuẩn hóa
huggingface_hub cung cấp các hàm trợ giúp để lưu và nạp trọng số mô hình ML theo cách chuẩn hóa. Phần này của thư viện vẫn đang trong quá trình phát triển và sẽ được cải thiện trong các bản phát hành tương lai. Mục tiêu là hài hòa hóa cách lưu và nạp trọng số trên Hub, vừa để loại bỏ sự trùng lặp mã giữa các thư viện, vừa để thiết lập các quy ước nhất quán.
Định dạng tệp DDUF
DDUF là định dạng tệp được thiết kế cho các mô hình khuếch tán. Định dạng này cho phép lưu trữ toàn bộ thông tin cần thiết để chạy một mô hình trong một tệp duy nhất. Công việc này được lấy cảm hứng từ định dạng GGUF. huggingface_hub cung cấp các hàm hỗ trợ để lưu và nạp tệp DDUF, đảm bảo tuân thủ đúng định dạng tệp.
Đây là một phiên bản rất sơ khai của trình phân tích cú pháp. API và cách triển khai có thể thay đổi trong tương lai gần.
Trình phân tích hiện thực hiện rất ít kiểm tra xác thực. Để biết thêm chi tiết về định dạng tệp, hãy xem https://github.com/huggingface/huggingface.js/tree/main/packages/dduf.
Cách viết một tệp DDUF?
Dưới đây là cách xuất một thư mục chứa các phần khác nhau của một mô hình khuếch tán bằng 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")Để có nhiều linh hoạt hơn, bạn có thể sử dụng export_entries_as_dduf() và truyền vào một danh sách các tệp cần bao gồm trong tệp DDUF cuối cùng:
# 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
... ]
... )Tham số entries cũng hỗ trợ truyền vào một đối tượng lặp được chứa các đường dẫn hoặc bytes. Điều này có thể hữu ích nếu bạn đã tải một model và muốn serialize nó trực tiếp vào một tệp DDUF thay vì phải serialize từng thành phần ra đĩa trước rồi mới tạo tệp DDUF. Dưới đây là một ví dụ về cách một StableDiffusionPipeline có thể được serialize dưới dạng 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))Lưu ý: trên thực tế, diffusers cung cấp một phương thức để trực tiếp serialize một pipeline trong một tệp DDUF. Đoạn mã ở trên chỉ mang tính chất minh họa.
Cách đọc tệp 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)Trợ lý
huggingface_hub.export_entries_as_dduf
< source >( dduf_path: str | os.PathLikeentries: Iterable )
Tham số
- dduf_path (
strhoặcos.PathLike) — Đường dẫn đến tệp DDUF cần ghi. - entries (
Iterable[tuple[str, Union[str, Path, bytes]]]) — Một iterable chứa các mục cần ghi vào tệp DDUF. Mỗi mục là một tuple gồm tên tệp và nội dung. Tên tệp phải là đường dẫn đến tệp trong kho lưu trữ DDUF. Nội dung có thể là một chuỗi hoặc một pathlib.Path đại diện cho đường dẫn đến tệp trên đĩa cục bộ, hoặc trực tiếp là nội dung dưới dạng bytes.
Gây ra
-
- —
DDUFExportError: Nếu có bất kỳ sự cố nào xảy ra trong quá trình xuất (ví dụ: tên mục không hợp lệ, thiếu ‘model_index.json’, v.v.).
- —
Viết một tệp DDUF từ một đối tượng lặp được các mục nhập.
Đây là một hàm trợ giúp ở cấp độ thấp hơn so với export_folder_as_dduf(), cho phép linh hoạt hơn khi thực hiện chuỗi hóa dữ liệu. Cụ thể, bạn không cần phải lưu dữ liệu xuống đĩa trước khi xuất nó vào tệp DDUF.
Ví dụ:
# 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 )
Xuất một thư mục dưới dạng tệp DDUF.
ADùng export_entries_as_dduf() bên trong.
huggingface_hub.read_dduf_file
< source >( dduf_path: os.PathLike | str ) → dict[str, DDUFEntry]
Tham số
Trả về
dict[str, DDUFEntry]
Một từ điển của DDUFEntry được chỉ mục theo tên tệp.
Gây ra
-
- —
DDUFCorruptedFileError: Nếu tệp DDUF bị hỏng (tức là không tuân theo định dạng DDUF).
- —
Đọc tệp DDUF và trả về một từ điển các mục nhập.
Chỉ metadata được đọc, dữ liệu không được nạp vào bộ nhớ.
Ví dụ:
>>> 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)lớp huggingface_hub.DDUFEntry
< source >( filename: strlength: intoffset: intdduf_path: Path )
Đối tượng đại diện cho một mục tệp trong tệp DDUF.
Xem read_dduf_file() để biết cách đọc tệp DDUF.
Mở tệp dưới dạng tệp ánh xạ bộ nhớ.
Có ích để nạp trực tiếp safetensors từ tệp.
Đọc tệp dưới dạng văn bản.
Có ích cho các mục ‘.txt’ và ‘.json’.
Lỗi
Ngoại lệ cơ bản cho các lỗi liên quan đến định dạng DDUF.
Lỗi ngoại lệ được ném ra khi tệp DDUF bị hỏng.
Ngoại lệ cơ bản cho các lỗi trong quá trình xuất DDUF.
Ngoại lệ được ném ra khi tên mục không hợp lệ.
Lưu tensor
Hàm trợ giúp chính của module serialization nhận một nn.Module của torch làm đầu vào và lưu nó xuống đĩa. Hàm này xử lý logic để lưu các tensor dùng chung (xem giải thích về safetensors) cũng như logic để chia state dictionary thành các shard, sử dụng split_torch_state_dict_into_shards() bên trong. Hiện tại, chỉ hỗ trợ framework torch.
Nếu bạn muốn lưu một state dictionary (ví dụ: ánh xạ giữa tên các layer và các tensor liên quan) thay vì một nn.Module, bạn có thể sử dụng save_torch_state_dict() cung cấp các tính năng tương tự. Điều này hữu ích, chẳng hạn, nếu bạn muốn áp dụng logic tùy chỉnh cho state dict trước khi lưu nó.
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 )
Tham số
- model (
torch.nn.Module) — Mô hình cần lưu trên đĩa. - save_directory (
strhoặcPath) — Thư mục mà mô hình sẽ được lưu vào. - filename_pattern (
str, tùy chọn) — Mẫu để tạo tên các tệp mà mô hình sẽ được lưu vào. Mẫu phải là một chuỗi có thể được định dạng bằngfilename_pattern.format(suffix=...)và phải chứa từ khóasuffixMặc định là"model{suffix}.safetensors"hoặcpytorch_model{suffix}.bintùy thuộc vào tham sốsafe_serialization. - force_contiguous (
boolean, tùy chọn) — Buộc state_dict được lưu dưới dạng các tensor liên tục. Điều này không ảnh hưởng đến tính chính xác của mô hình, nhưng có thể thay đổi hiệu suất nếu bố cục của tensor được chọn đặc biệt vì lý do đó. Mặc định làTrue. - max_shard_size (
inthoặcstr, tùy chọn) — Kích thước tối đa của mỗi shard, tính bằng byte. Mặc định là 5GB. - metadata (
dict[str, str], tuỳ chọn) — Thông tin bổ sung để lưu cùng với mô hình. Một số metadata sẽ được thêm cho mỗi tensor bị loại bỏ. Thông tin này sẽ không đủ để khôi phục toàn bộ cấu trúc chia sẻ nhưng có thể giúp hiểu rõ hơn về các thành phần. - safe_serialization (
bool, tùy chọn) — Có lưu dưới dạng safetensors hay không, đây là hành vi mặc định. NếuFalse, các shard sẽ được lưu dưới dạng pickle. Nên sử dụng serial hóa an toàn vì lý do bảo mật. Việc lưu dưới dạng pickle đã bị ngừng hỗ trợ và sẽ bị loại bỏ trong một phiên bản tương lai. - is_main_process (
bool, tùy chọn) — Xác định xem tiến trình gọi hàm này có phải là tiến trình chính hay không. Hữu ích khi huấn luyện phân tán như trên TPU và cần gọi hàm này từ tất cả các tiến trình. Trong trường hợp này, chỉ đặtis_main_process=Truetrên tiến trình chính để tránh điều kiện đua. Mặc định là True. - shared_tensors_to_discard (
list[str], tuỳ chọn) — Danh sách tên tensor cần loại bỏ khi lưu các tensor dùng chung. Nếu không được cung cấp và phát hiện có tensor dùng chung, hệ thống sẽ loại bỏ tên đầu tiên theo thứ tự bảng chữ cái.
Lưu một mô hình torch đã cho xuống đĩa, xử lý các vấn đề về phân mảnh và tensor dùng chung.
Xem thêm save_torch_state_dict() để lưu state dict với nhiều sự linh hoạt hơn.
Để biết thêm thông tin về việc chia sẻ tensor, hãy xem hướng dẫn này.
Từ điển trạng thái mô hình được chia thành các phần (shards) sao cho mỗi phần nhỏ hơn một kích thước cho trước. Các phần này được
lưu trong save_directory với filename_pattern đã cho. Nếu mô hình quá lớn để chứa trong một phần duy nhất,
một tệp chỉ mục sẽ được lưu trong save_directory để chỉ ra vị trí lưu trữ của từng tensor. Tiện ích này sử dụng split_torch_state_dict_into_shards() bên trong. Nếu safe_serialization là True, các phần sẽ được lưu dưới dạng
safetensors (mặc định). Ngược lại, các phần sẽ được lưu dưới dạng pickle.
Trước khi lưu mô hình, save_directory sẽ được dọn sạch khỏi các tệp shard trước đó.
Nếu một trong các tensor của model có kích thước lớn hơn
max_shard_size, tensor đó sẽ nằm trong một shard riêng có kích thước lớn hơnmax_shard_size.
Nếu model của bạn là
transformers.PreTrainedModel, bạn nên truyềnmodel._tied_weights_keysvàoshared_tensors_to_discardđể xử lý việc lưu các tensor dùng chung một cách chính xác. Điều này đảm bảo các tensor trùng lặp đúng sẽ bị loại bỏ trong quá trình lưu.
Ví dụ:
>>> 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 )
Tham số
- state_dict (
dict[str, torch.Tensor]) — Từ điển trạng thái cần lưu. - save_directory (
strhoặcPath) — Thư mục mà mô hình sẽ được lưu vào. - filename_pattern (
str, tùy chọn) — Mẫu để tạo tên các tệp mà mô hình sẽ được lưu vào. Mẫu phải là một chuỗi có thể được định dạng bằngfilename_pattern.format(suffix=...)và phải chứa từ khóasuffixMặc định là"model{suffix}.safetensors"hoặcpytorch_model{suffix}.bintùy thuộc vào tham sốsafe_serialization. - force_contiguous (
boolean, tùy chọn) — Buộc state_dict được lưu dưới dạng các tensor liên tục. Điều này không ảnh hưởng đến tính chính xác của mô hình, nhưng có thể thay đổi hiệu suất nếu bố cục của tensor được chọn đặc biệt vì lý do đó. Mặc định làTrue. - max_shard_size (
inthoặcstr, tùy chọn) — Kích thước tối đa của mỗi shard, tính bằng byte. Mặc định là 5GB. - metadata (
dict[str, str], tuỳ chọn) — Thông tin bổ sung để lưu cùng với mô hình. Một số metadata sẽ được thêm cho mỗi tensor bị loại bỏ. Thông tin này sẽ không đủ để khôi phục toàn bộ cấu trúc chia sẻ nhưng có thể giúp hiểu rõ hơn về các thành phần. - safe_serialization (
bool, tùy chọn) — Có lưu dưới dạng safetensors hay không, đây là hành vi mặc định. NếuFalse, các shard sẽ được lưu dưới dạng pickle. Nên sử dụng serial hóa an toàn vì lý do bảo mật. Việc lưu dưới dạng pickle đã bị ngừng hỗ trợ và sẽ bị loại bỏ trong một phiên bản tương lai. - is_main_process (
bool, tùy chọn) — Xác định xem tiến trình gọi hàm này có phải là tiến trình chính hay không. Hữu ích khi huấn luyện phân tán như trên TPU và cần gọi hàm này từ tất cả các tiến trình. Trong trường hợp này, chỉ đặtis_main_process=Truetrên tiến trình chính để tránh điều kiện đua. Mặc định là True. - shared_tensors_to_discard (
list[str], tuỳ chọn) — Danh sách tên tensor cần loại bỏ khi lưu các tensor dùng chung. Nếu không được cung cấp và phát hiện có tensor dùng chung, hệ thống sẽ loại bỏ tên đầu tiên theo thứ tự bảng chữ cái.
Lưu từ điển trạng thái mô hình xuống đĩa, xử lý các vấn đề về phân mảnh và tensor dùng chung.
Xem thêm save_torch_model() để lưu trực tiếp một mô hình PyTorch.
Để biết thêm thông tin về việc chia sẻ tensor, hãy xem hướng dẫn này.
Từ điển trạng thái mô hình được chia thành các phần (shards) sao cho mỗi phần nhỏ hơn một kích thước cho trước. Các phần này được
lưu trong save_directory với filename_pattern đã cho. Nếu mô hình quá lớn để chứa trong một phần duy nhất,
một tệp chỉ mục sẽ được lưu trong save_directory để chỉ ra vị trí lưu trữ của từng tensor. Tiện ích này sử dụng split_torch_state_dict_into_shards() bên trong. Nếu safe_serialization là True, các phần sẽ được lưu dưới dạng
safetensors (mặc định). Ngược lại, các phần sẽ được lưu dưới dạng pickle.
Trước khi lưu mô hình, save_directory sẽ được dọn sạch khỏi các tệp shard trước đó.
Nếu một trong các tensor của model có kích thước lớn hơn
max_shard_size, tensor đó sẽ nằm trong một shard riêng có kích thước lớn hơnmax_shard_size.
Nếu model của bạn là
transformers.PreTrainedModel, bạn nên truyềnmodel._tied_weights_keysvàoshared_tensors_to_discardđể xử lý việc lưu các tensor dùng chung một cách chính xác. Điều này đảm bảo các tensor trùng lặp đúng sẽ bị loại bỏ trong quá trình lưu.
Ví dụ:
>>> 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")Module serialization cũng chứa các hàm trợ giúp cấp thấp để chia một state dictionary thành nhiều shard, đồng thời tạo một index phù hợp trong quá trình đó. Các hàm trợ giúp này có sẵn cho tensor torch và được thiết kế để dễ dàng mở rộng sang bất kỳ framework ML nào khác.
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
Tham số
- state_dict (
dict[str, torch.Tensor]) — Từ điển trạng thái cần lưu. - filename_pattern (
str, tùy chọn) — Mẫu để tạo tên các tệp mà mô hình sẽ được lưu vào. Mẫu phải là một chuỗi có thể được định dạng bằngfilename_pattern.format(suffix=...)và phải chứa từ khóasuffix. Mặc định là"model{suffix}.safetensors". - max_shard_size (
inthoặcstr, tùy chọn) — Kích thước tối đa của mỗi shard, tính bằng byte. Mặc định là 5GB.
Trả về
StateDictSplit
Một đối tượng StateDictSplit chứa các shard và chỉ mục để truy xuất chúng.
Tách từ điển trạng thái của mô hình thành các phần sao cho mỗi phần nhỏ hơn một kích thước cho trước.
Các shard được xác định bằng cách lặp qua state_dict theo thứ tự các key của nó. Không có tối ưu hóa nào được thực hiện để đảm bảo mỗi shard gần nhất có thể với kích thước tối đa được truyền vào. Ví dụ, nếu giới hạn là 10GB và chúng ta có các tensor với kích thước [6GB, 6GB, 2GB, 6GB, 2GB, 2GB], chúng sẽ được chia shard thành [6GB], [6+2GB], [6+2+2GB] chứ không phải [6+2+2GB], [6+2GB], [6GB].
Để lưu state dictionary của mô hình xuống đĩa, hãy xem save_torch_state_dict(). Hàm trợ giúp này sử dụng
split_torch_state_dict_into_shardsbên trong.
Nếu một trong các tensor của model có kích thước lớn hơn
max_shard_size, tensor đó sẽ nằm trong một shard riêng có kích thước lớn hơnmax_shard_size.
Ví dụ:
>>> 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
Đây là nhà máy cơ bản mà từ đó mỗi trình trợ giúp cụ thể cho từng framework được dẫn xuất. Trên thực tế, bạn không được kỳ vọng sử dụng trực tiếp nhà máy này, trừ khi bạn cần điều chỉnh nó cho một framework chưa được hỗ trợ. Nếu trường hợp đó xảy ra, vui lòng cho chúng tôi biết bằng cách mở một issue mới trên repo 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
Tham số
- state_dict (
dict[str, Tensor]) — Từ điển trạng thái cần lưu. - get_storage_size (
Callable[[Tensor], int]) — Một hàm trả về kích thước của một tensor khi được lưu trên đĩa, tính bằng byte. - get_storage_id (
Callable[[Tensor], Optional[Any]], tùy chọn) — Một hàm trả về định danh duy nhất cho bộ nhớ lưu trữ của một tensor. Nhiều tensor khác nhau có thể chia sẻ cùng một bộ nhớ lưu trữ cơ bản. Định danh này được đảm bảo là duy nhất và không đổi cho bộ nhớ lưu trữ của tensor này trong suốt vòng đời của nó. Hai bộ nhớ lưu trữ tensor có vòng đời không chồng lấn có thể có cùng một id. - filename_pattern (
str, tùy chọn) — Mẫu để tạo tên các tệp mà mô hình sẽ được lưu vào. Mẫu phải là một chuỗi có thể được định dạng bằngfilename_pattern.format(suffix=...)và phải chứa từ khóasuffix - max_shard_size (
inthoặcstr, tùy chọn) — Kích thước tối đa của mỗi shard, tính bằng byte. Mặc định là 5GB.
Trả về
StateDictSplit
Một đối tượng StateDictSplit chứa các shard và chỉ mục để truy xuất chúng.
Tách từ điển trạng thái của mô hình thành các phần sao cho mỗi phần nhỏ hơn một kích thước cho trước.
Các shard được xác định bằng cách lặp qua state_dict theo thứ tự các key của nó. Không có tối ưu hóa nào được thực hiện để đảm bảo mỗi shard gần nhất có thể với kích thước tối đa được truyền vào. Ví dụ, nếu giới hạn là 10GB và chúng ta có các tensor với kích thước [6GB, 6GB, 2GB, 6GB, 2GB, 2GB], chúng sẽ được chia shard thành [6GB], [6+2GB], [6+2+2GB] chứ không phải [6+2+2GB], [6+2GB], [6GB].
Nếu một trong các tensor của model có kích thước lớn hơn
max_shard_size, tensor đó sẽ nằm trong một shard riêng có kích thước lớn hơnmax_shard_size.
Tải tensor
Các hàm hỗ trợ tải hỗ trợ cả checkpoint dạng tệp đơn và dạng phân mảnh ở định dạng safetensors hoặc pickle. load_torch_model() nhận một nn.Module và đường dẫn checkpoint (một tệp đơn hoặc một thư mục) làm đầu vào và tải trọng số vào mô hình.
Cả hai hàm trợ giúp đều mặc định là
safe=True, có nghĩa là checkpoint luôn được giải chuỗi hóa bằng trình nạp safetensors — bất kể tên tệp là gì. Việc nạp checkpoint pickle yêu cầu phải chọn bật tính năng này một cách rõ ràng, vì quá trình unpickle sẽ thực thi mã tùy ý tại thời điểm nạp:>>> 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)Đường dẫn pickle mặc định thêm
weights_only=True, tức là bộ unpickler bị hạn chế của torch. Một checkpoint chứa các đối tượng không phải tensor (trạng thái optimizer, toàn bộ model, …) cầnweights_only=Falsebên cạnhsafe=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
Tham số
- model (
torch.nn.Module) — Mô hình để nạp checkpoint. - checkpoint_path (
strhoặcos.PathLike) — Đường dẫn đến tệp checkpoint hoặc thư mục chứa các checkpoint. - strict (
bool, tùy chọn, mặc định làFalse) — Liệu có thực thi nghiêm ngặt việc các khóa trong model state dict phải khớp với các khóa trong checkpoint hay không. Tương tự nhưtorch.nn.Module.load_state_dict, việc kiểm tra được thực hiện sau khi checkpoint đã được tải, do đó model có thể đã được cập nhật một phần khi lỗi được ném ra. - safe (
bool, tuỳ chọn, mặc định làTrue) — Nếusafelà True, các tệp safetensors sẽ được nạp. Nếusafelà False, hàm sẽ trước tiên cố gắng nạp các tệp safetensors nếu chúng có sẵn, nếu không sẽ chuyển sang nạp các tệp pickle. Mộtfilename_patternkhông mô tả các tệp safetensors sẽ bị từ chối khisafe=True. - weights_only (
bool, tùy chọn, mặc định làTrue) — Nếu là True, chỉ tải trọng số của mô hình mà không tải trạng thái optimizer và các metadata khác, sử dụng unpickler bị hạn chế của torch. Đặt thành False để cho phép các đối tượng Python tùy ý trong checkpoint pickle (điều này sẽ thực thi mã tùy ý khi tải). Không có tác dụng với PyTorch < 1.13, vốn không có unpickler bị hạn chế và luôn unpickle mà không có hạn chế. - map_location (
strhoặctorch.device, tùy chọn) — Một đối tượngtorch.device, chuỗi hoặc dict chỉ định cách ánh xạ lại các vị trí lưu trữ. Nó cho biết vị trí mà tất cả các tensor sẽ được nạp vào. - mmap (
bool, tùy chọn, mặc định làFalse) — Có sử dụng phương thức nạp tệp ánh xạ bộ nhớ (memory-mapped file loading) hay không. Ánh xạ bộ nhớ có thể cải thiện hiệu suất nạp cho các mô hình lớn trong PyTorch >= 2.1.0 với các checkpoint dựa trên zipfile. - filename_pattern (
str, tùy chọn) — Mẫu để tìm tệp chỉ mục. Mẫu phải là chuỗi có thể được định dạng bằngfilename_pattern.format(suffix=...)và phải chứa từ khóasuffixMặc định là"model{suffix}.safetensors".
Trả về
NamedTuple
Một named tuple có các trường missing_keys và unexpected_keys.
missing_keyslà một danh sách str chứa các khóa bị thiếu, tức là các khóa có trong model nhưng không có trong checkpoint.unexpected_keyslà một danh sách str chứa các khóa không mong đợi, tức là các khóa có trong checkpoint nhưng không có trong model.
Gây ra
FileNotFoundError hoặc ImportError hoặc ValueError
FileNotFoundError— Nếu tệp hoặc thư mục checkpoint không tồn tại.ImportError— Nếu safetensors hoặc torch chưa được cài đặt khi cố gắng tải một tệp .safetensors hoặc một checkpoint PyTorch tương ứng.ValueError— Nếu không thể xác định được định dạng checkpoint.
Tải checkpoint vào mô hình, xử lý cả checkpoint dạng phân mảnh và không phân mảnh.
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]
Tham số
- checkpoint_file (
strhoặcos.PathLike) — Đường dẫn đến tệp checkpoint cần tải. Có thể là checkpoint dạng safetensors hoặc pickle (.bin). - map_location (
strhoặctorch.device, tùy chọn) — Một đối tượngtorch.device, chuỗi hoặc dict chỉ định cách ánh xạ lại các vị trí lưu trữ. Nó cho biết vị trí mà tất cả các tensor sẽ được nạp vào. - weights_only (
bool, tùy chọn, mặc định làTrue) — Nếu là True, chỉ tải trọng số của mô hình mà không tải trạng thái optimizer và các metadata khác, sử dụng unpickler bị hạn chế của torch. Đặt thành False để cho phép các đối tượng Python tùy ý trong checkpoint pickle (điều này sẽ thực thi mã tùy ý khi tải). Không có hiệu ứng khi tải các tệp safetensors, cũng như trên PyTorch < 1.13 vốn không có unpickler bị hạn chế — các phiên bản đó luôn unpickle mà không có hạn chế và một cảnh báo sẽ được ghi log. - mmap (
bool, tùy chọn, mặc định làFalse) — Có sử dụng phương thức nạp tệp ánh xạ bộ nhớ (memory-mapped file loading) hay không. Ánh xạ bộ nhớ có thể cải thiện hiệu suất nạp cho các mô hình lớn trong PyTorch >= 2.1.0 với các checkpoint dựa trên zipfile. Không có tác dụng khi nạp các tệp safetensors, vì thư việnsafetensorssử dụng ánh xạ bộ nhớ theo mặc định. - safe (
bool, tùy chọn, mặc định làTrue) — Nếu là True, checkpoint luôn được nạp dưới dạng safetensors, bất kể tên của nó là gì. Mọi định dạng khác (ví dụ một tệp pickle.bin) sẽ gây raValueErrorthay vì được khử chuỗi hóa: các checkpoint pickle có thể thực thi mã tùy ý khi nạp. Nếu là False, tệp sẽ được nạp dưới dạng safetensors nếu tên của nó cho thấy điều đó, và bằngtorch.loadtrong các trường hợp còn lại.
Trả về
Union[dict[str, "torch.Tensor"], Any]
Checkpoint đã được tải.
- Đối với các tệp safetensors: luôn trả về một từ điển ánh xạ tên tham số tới các tensor.
- Đối với tệp pickle: trả về bất kỳ đối tượng Python nào đã được pickle (thường là state dict, nhưng cũng có thể là toàn bộ model, trạng thái optimizer, hoặc bất kỳ đối tượng Python nào khác).
Gây ra
FileNotFoundError hoặc ImportError hoặc OSError hoặc ValueError
FileNotFoundError— Nếu tệp checkpoint không tồn tại.ImportError— Nếu safetensors hoặc torch chưa được cài đặt khi cố gắng tải một tệp .safetensors hoặc một checkpoint PyTorch tương ứng.OSError— Nếu định dạng tệp checkpoint không hợp lệ hoặc các tệp git-lfs không được tải xuống đúng cách.ValueError— Nếu đường dẫn tệp checkpoint trống hoặc không hợp lệ, hoặc nếu tệp không thể được khử chuỗi hóa dưới dạng safetensors khisafe=True(truyềnsafe=Falseđể cho phép các checkpoint pickle).
Tải tệp checkpoint, xử lý cả định dạng checkpoint safetensors và pickle.
Ví dụ:
>>> 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)Trợ lý Tensors
get_torch_storage_id
Trả về định danh duy nhất cho bộ nhớ tensor.
Nhiều tensor khác nhau có thể chia sẻ cùng một bộ nhớ lưu trữ (storage) cơ bản. Định danh này được đảm bảo là duy nhất và không đổi đối với storage của tensor này trong suốt vòng đời của nó. Hai storage của tensor có vòng đời không chồng lấn có thể có cùng id. Đối với meta tensor, chúng tôi trả về None vì không thể xác định xem chúng có chia sẻ cùng một storage hay không.