* Studio: fail closed on index-referenced nested pickle shards in the offline embedding gate
The offline embedding security gate (HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE)
only scanned the direct files of each SentenceTransformer load root and never
parsed local weight indexes, so a cached snapshot whose pytorch_model.bin.index.json
maps a weight to a nested shard (e.g. shards/pytorch_model-00001-of-00001.bin) was
treated as inert and allowed. The loader then follows the index into the subdir and
unpickles the shard. The online gate already blocks index-referenced subdir pickles,
so the offline path was strictly weaker.
Parse each local weight index in a load root and follow weight_map into nested dirs,
flagging any referenced pickle-extension shard. Paths resolve lexically (normpath),
never Path.resolve(), since HF cache snapshot files symlink into blobs/ and resolving
would leave the snapshot dir and false-block every sharded model offline. An absolute
path, a .. traversal that escapes the snapshot, or an unreadable/invalid index fails
closed. The existing safetensors-sibling suppression is kept.
* Studio: classify offline indexed shards by torch.load path, not pickle extension
load_state_dict picks safetensors vs torch.load per shard by the shard's own
suffix, so two offline-gate gaps remained:
- A model.safetensors.index.json whose weight_map points at a .bin shard was
suppressed by has_base_safetensors (the index file itself matches the base
safetensors regex), yet Transformers still torch.loads that shard. Only the
pytorch index is superseded by a base safetensors now; a safetensors index is
the chosen archive, so its non-safetensors targets are always flagged.
- A pytorch index can map weights to arbitrary names (shards/payload,
weights.data); the loader torch.loads any target not ending in .safetensors.
Flag indexed shards by that rule instead of a pickle-extension allowlist.
Restrict the scan to the two torch-family indexes (tf/flax load via non-pickle
loaders). Add regression tests for both cases.
* Studio: match offline weight-index filenames case-insensitively
The index-name check compared the on-disk filename exactly, while the
surrounding weight and safetensors matches use case-insensitive rules. On a
case-insensitive volume (Windows or macOS) from_pretrained opens an oddly-cased
cache file such as PYTORCH_MODEL.BIN.INDEX.JSON when it requests the canonical
lowercase name, so the exact-case check skipped it and a nested pickle shard it
referenced was allowed through. Lower-case the index name before matching, as
the rest of the gate does, and add a regression test.
* Studio: match load_state_dict format/selection exactly in the offline index scan
Two edge cases in the offline weight-index scan:
- load_state_dict decides safetensors vs torch.load with a case-sensitive
endswith(".safetensors"), so a shard named payload.SAFETENSORS still
deserializes via torch.load. Classify indexed shard suffixes case-sensitively
to match, instead of lower-casing (which treated such a shard as inert).
- A complete direct model.safetensors is selected before either sharded index,
so a stale model.safetensors.index.json referencing a .bin shard never loads.
Skip both indexes when a direct model.safetensors is present, so an otherwise
loadable model is not over-blocked.
Add regression tests for both.
* Studio: read the offline weight index as UTF-8
Path.read_text() uses the locale default, which is cp1252 on Windows, so a
UTF-8 weight index with non-ASCII bytes raised UnicodeDecodeError and the gate
blocked an otherwise loadable model. JSON is UTF-8 by spec (and how the loader
reads it), so pin the encoding.
* Studio: resolve safetensors alternatives via the loader's own filename lookup
The offline gate decided a safetensors alternative existed by case-folding the
directory listing. On a case-sensitive filesystem that let an uppercase decoy
such as MODEL.SAFETENSORS suppress the pickle scan, yet from_pretrained asks for
the canonical lowercase model.safetensors, does not find the decoy, and selects
the pickle (a direct pytorch_model.bin or the pytorch index) and deserializes it.
Probe each alternative with (root / name).is_file() instead, mirroring the
loader: is_file() honors the platform's case rules, so a decoy suppresses only
where the loader would truly open it. Suppression must never fail open; detection
stays case-insensitive (fail closed). Add regression tests for the direct and
indexed pickle decoys (skipped on case-insensitive volumes, where no bypass
exists).
* Studio: resolve indexes and shards exactly as from_pretrained does
Two more loader-fidelity gaps in the offline index scan:
- Shard lookup normalized backslashes to forward slashes. On POSIX a backslash
is a literal filename character, so an index naming dir\payload.bin matches a
real pickle of that exact name that Transformers joins and deserializes, while
the normalized dir/payload.bin missed it. Join the raw weight_map value with
os.path.join so the probe mirrors the loader on each platform.
- Index detection case-folded the directory listing, so on a case-sensitive
filesystem an uppercase PYTORCH_MODEL.BIN.INDEX.JSON artifact the loader never
opens was treated as live and its shard blocked. Probe the canonical name with
the loader's own is_file lookup instead, so an index counts only where
from_pretrained would actually load it.
Update the uppercase-index tests to assert the correct per-filesystem behavior
and add a POSIX backslash-shard regression test.
---------
Co-authored-by: danielhanchen <unslothai@gmail.com>
c2114d64dd