Extract contacts from fine tuned model

[nilshof01]

Hi,

I trained with an MLM approach esm2_t6_8M_UR50D on my own set of sequences and I am particularly interested in the contact maps. However, my issue is that I can not return the contacts as easy as in the tutorial (https://github.com/facebookresearch/esm/blob/main/examples/esm_structural_dataset.ipynb). I get an error which states:

TypeError: EsmModel.forward() got an unexpected keyword argument 'return_contacts'

My input_ids and attention_mask look alright, so I think it has to do something with the model setup?
That is why I want to ask if I am doing anything wrong or if you have a script which I can use to parse the attention heads and calculate the contacts on my own.

The corresponding code can be seen here:

tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
model = EsmModel.from_pretrained(DIR_TO_MODEL)
sequence = ["GSIGAIHAMGWFRQAPGSQRELVATVTTTSDSTTYGESVKGRFTISREADKNTVFLQMSSLKPEDTAVYYCNANSYCSESVCYENPWYDYWGQGTQVTSS"]
sequences = [" ".join(list(re.sub(r"[UZOB*_]", "X", sequence))) for sequence in sequence]
encoded_input = tokenizer(sequences, return_tensors='pt', padding=True)

with torch.no_grad():
    results = model(**encoded_input, return_contacts = True)

PATH_TO_MODEL contains: config.json, model.savetensors, training_args.bin
Thank you so much for your help and I am looking forward to your reply! :)
Very best regards,
Nils

[garykbrixi]

Based on the docs the huggingface module does not have a 'return_contacts' argument in the forward function, which is consistent with your error.
Try using the model.predict_contacts(tokens, attention_mask) function. See it here:
https://github.com/huggingface/transformers/blob/745bbfe4bb2b61491dedd56e1e8ee4af8ef1a9ec/src/transformers/models/esm/modeling_esm.py#L937