Skip to content

Add safetensors to the pytorch extra - #75

Open
manjunathshiva wants to merge 1 commit into
google-research:mainfrom
manjunathshiva:add-safetensors-to-pytorch-extra
Open

Add safetensors to the pytorch extra#75
manjunathshiva wants to merge 1 commit into
google-research:mainfrom
manjunathshiva:add-safetensors-to-pytorch-extra

Conversation

@manjunathshiva

@manjunathshiva manjunathshiva commented Jul 20, 2026

Copy link
Copy Markdown

A fresh pip install tabfm[pytorch] cannot load the released checkpoint. On tabfm 1.0.1 (cb6ba46), Python 3.13, torch 2.13.0:

from tabfm.src.pytorch import tabfm_v1_0_0
model = tabfm_v1_0_0.load(model_type="classification")
File ".../huggingface_hub/hub_mixin.py", line 813, in _load_as_safetensor
    if packaging.version.parse(safetensors.__version__) < packaging.version.parse("0.4.3"):
NameError: name 'safetensors' is not defined

Cause: huggingface_hub declares safetensors as an optional dependency and imports it lazily, but the PyTorch loader always goes through PyTorchModelHubMixin, and since 1.0.1 the Hugging Face checkpoint ships only model.safetensors (#33 / the 1.0.1 weight-loading fix dropped the pytorch_model.bin path). So safetensors is a hard runtime dependency of the pytorch extra — and torch does not pull it in transitively. Anyone installing per the README's PyTorch instructions hits the NameError on their first load().

It likely went unnoticed because CI's pip install -e .[dev,jax,pytorch] never loads the released checkpoint (tests build small models directly), and any environment that has ever installed transformers/accelerate already has safetensors present.

Change: add safetensors to the pytorch extra in pyproject.toml, and a matching pinned entry (safetensors==0.8.0, current release) to the requirements.txt lock in pip-compile format/ordering. The JAX path is unaffected — its weights restore via orbax.

Verified in a clean venv: pip install -e .[pytorch] on this branch pulls safetensors and tabfm_v1_0_0.load(model_type="classification") succeeds; on main the same steps reproduce the NameError above.

Note on overlap with #49: the MLX backend PR bundles this same dependency fix (both the safetensors line in the pytorch extra and the safetensors==0.8.0 pin in requirements.txt). This PR extracts it as a standalone fix so it isn't gated on the larger feature review; if this lands first, the corresponding hunks in #49 resolve trivially on its rebase, and if #49 lands first this PR can simply be closed.

pip install tabfm[pytorch] cannot load the released checkpoint:
tabfm_v1_0_0.load() fails with

    NameError: name 'safetensors' is not defined

inside huggingface_hub's _load_as_safetensor. huggingface_hub declares
safetensors optional and imports it lazily, but the PyTorch loader always
goes through PyTorchModelHubMixin -> model.safetensors (since 1.0.1
dropped pytorch_model.bin), so safetensors is a hard runtime dependency
of the pytorch extra. Torch does not pull it in transitively.

Adds it to the pytorch extra and to the pinned requirements.txt lock.
The JAX path is unaffected (weights restore via orbax).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant