Safetensors

Safetensors là một định dạng mới, đơn giản để lưu trữ tensor một cách an toàn (khác với pickle) và vẫn nhanh (zero-copy). Safetensors thực sự nhanh 🚀.

Cài đặt

với pip:

pip install safetensors

với conda:

conda install -c conda-forge safetensors

Cách sử dụng

Tải tensor

from safetensors import safe_open

tensors = {}
with safe_open("model.safetensors", framework="pt", device=0) as f:
    for k in f.keys():
        tensors[k] = f.get_tensor(k)

Chỉ tải một phần các tensor (hữu ích khi chạy trên nhiều GPU)

from safetensors import safe_open

tensors = {}
with safe_open("model.safetensors", framework="pt", device=0) as f:
    tensor_slice = f.get_slice("embedding")
    vocab_size, hidden_dim = tensor_slice.get_shape()
    tensor = tensor_slice[:, :hidden_dim]

Lưu tensor

import torch
from safetensors.torch import save_file

tensors = {
    "embedding": torch.zeros((2, 2)),
    "attention": torch.zeros((2, 3))
}
save_file(tensors, "model.safetensors")

Định dạng

Giả sử bạn có tệp safetensors tên là model.safetensors, thì model.safetensors sẽ có định dạng bên trong như sau:

Dự án nổi bật

Safetensors đang được sử dụng rộng rãi tại các doanh nghiệp AI hàng đầu, chẳng hạn như Hugging Face, EleutherAI và StabilityAI. Dưới đây là danh sách không đầy đủ các dự án đang sử dụng safetensors:

Update on GitHub