Locally Installed v1 Pretrained Model for Structure Prediction

[SkwisgaarSkwigelf]

Maybe a simple issue, but I'm trying to reproduce the simple structure prediction using the v1 model from the readme example (see bottom for that unedited code). However, I get SSL: CERTIFICATE_VERIFY_FAILED error at this line model = esm.pretrained.esmfold_v1(). Seems like it's trying to download the pretrained model on the fly and my setup doesn't like that. So instead I just downloaded v1 locally like this:

cd checkpoints/
curl -s https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v1.pt -o esmfold_3B_v1.pt

And now I'm just trying to point it to the locally installed v1 model. However when I do this it complains:

input:

model = esm.pretrained.load_model_and_alphabet_local('/path/to/esm/checkpoints/esmfold_3B_v1.pt')
model = model.eval()

error:

/path/to/esm/esm/pretrained.py in load_model_and_alphabet_local(model_location)
     75         regression_data = None
     76     return load_model_and_alphabet_core(model_name, model_data, regression_data)
---> 77 
     78 
     79 def has_emb_layer_norm_before(model_state):

/path/to/esm/esm/pretrained.py in load_model_and_alphabet_core(model_name, model_data, regression_data)
    191     else:
    192         model, alphabet, model_state = _load_model_and_alphabet_core_v1(model_data)
--> 193 
    194     expected_keys = set(model.state_dict().keys())
    195     found_keys = set(model_state.keys())

/path/to/esm/esm/pretrained.py in _load_model_and_alphabet_core_v1(model_data)
     86 
     87     alphabet = esm.Alphabet.from_architecture(model_data["args"].arch)
---> 88 
     89     if model_data["args"].arch == "roberta_large":
     90         # upgrade state dict

KeyError: 'args'

Not sure what the issue is. Is it still looking for a contact-regression.pt file? My understanding is the v1 model doesn't come with this. There must be a way to point it to the locally installed v1 pretrained model that I'm just missing?

Here is the full code from the README for reference:

import torch
import esm

model = esm.pretrained.esmfold_v1()
model = model.eval().cuda()

# Optionally, uncomment to set a chunk size for axial attention. This can help reduce memory.
# Lower sizes will have lower memory requirements at the cost of increased speed.
# model.set_chunk_size(128)

sequence = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
# Multimer prediction can be done with chains separated by ':'

with torch.no_grad():
    output = model.infer_pdb(sequence)

with open("result.pdb", "w") as f:
    f.write(output)

import biotite.structure.io as bsio
struct = bsio.load_structure("result.pdb", extra_fields=["b_factor"])
print(struct.b_factor.mean())  # this will be the pLDDT
# 88.3