Fix local loading of all three ESMC checkpoints

[bviggiano]

Problem

Loading any ESMC checkpoint locally currently fails for me on main. I implemented this fix to get this to work and thought I would make a PR just in case.

from esm.models.esmc import ESMC
ESMC.from_pretrained("esmc_300m")   # also esmc_600m, esmc_6b

checkpoint	error
esmc_300m	ValueError: Directory '...' does not contain a valid checkpoint
esmc_600m	ValueError: Directory '...' does not contain a valid checkpoint
esmc_6b	NotImplementedError: Cannot copy out of meta tensor; no data!

$ python cookbook/snippets/esmc.py
  File "cookbook/snippets/esmc.py", line 139, in <module>
    main(ESMC.from_pretrained("esm3_sm_open_v1"))
  File "esm/models/esmc.py", line 95, in from_pretrained
    assert isinstance(model, ESMC)
AssertionError

Three independent causes (notes from Claude):

• 300M / 600M — pretrained.py#L75 and #L89 call load_torch_model on the snapshot root, but esmc-300m-2024-12 and esmc-600m-2024-12 have a single nested data/weights/*.pth with no index file at the root.
• 6B — #L103 loads esmc-6b-2024-12, whose 808 keys are exported from the HF ESMCForMaskedLM wrapper (esmc.*, lm_head.*); ESMC expects them unprefixed with the head as sequence_head. load_torch_model defaults to strict=False, so every key is dropped silently, the parameters stay meta tensors from init_empty_weights(), and the failure only surfaces at model.to(device).
• Snippet — esmc.py#L139 passes an ESM3 checkpoint name to ESMC.from_pretrained, which resolves to the ESM3_sm_open_v0 builder and returns an ESM3: never an ESMC, so the assert fails.

Fix

esm/pretrained.py — all three ESMC builders load the state dict directly and apply it with assign=True, the pattern ESM3_sm_open_v0 and the three ESM3 decoder builders above it already use. 6B additionally walks model.safetensors.index.json and maps esmc. → `` and lm_head. → `sequence_head.`. `load_torch_model` has no remaining callers, so the import is swapped for `load_state_dict_from_file`.

assign=True keeps load_state_dict's default strict=True, so future key drift fails at the load site instead of silently.

cookbook/snippets/esmc.py — drop the esm3_sm_open_v1 call. The two lines after it already load esmc_300m and pass it to main() and raw_forward(), so it is redundant as well as incorrect.

I verified these now work in a clean micromamba env, Python 3.12 via the pip install -e . set up path.

Client returned logits with shape: torch.Size([1, 7, 64]), embeddings with shape: torch.Size([1, 7, 960]), and hidden states with shape torch.Size([30, 1, 7, 960])
Client returned hidden states with shape torch.Size([1, 1, 7, 960])
Raw model returned logits with shape: torch.Size([2, 7, 64]), embeddings with shape: torch.Size([2, 7, 960]) and hidden states with shape torch.Size([30, 2, 7, 960])

Notes

I found a few other things that appear to be out of date!

• The Hub download path has no CI coverage: data_root() returns Path("") when INFRA_PROVIDER is set, and tests/Makefile#L4 always sets it.
• CONTRIBUTIONS.md#L16 specifies Python 3.10, but pyproject.toml#L6 requires >=3.12,<3.13, so pip install -e . fails on 3.10.
• CONTRIBUTIONS.md#L20 references examples/requirements.txt; there is no examples/ directory.