facebookresearch / esm

Evolutionary Scale Modeling (esm): Pretrained language models for proteins
MIT License
3.16k stars 627 forks source link

How to convert folded_positions -> PDB file format? #636

Open johnnytam100 opened 10 months ago

johnnytam100 commented 10 months ago

I followed the esmfold example here: https://huggingface.co/docs/transformers/model_doc/esm#transformers.EsmForProteinFolding

from transformers import AutoTokenizer, EsmForProteinFolding

model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1")
tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
inputs = tokenizer(["MLKNVQVQLV"], return_tensors="pt", add_special_tokens=False)  # A tiny random peptide
outputs = model(**inputs)
folded_positions = outputs.positions

which outputs a tensor

tensor([[[[[ -0.4341,   8.0899,  12.8728],
           [ -1.5881,   7.7639,  12.0398],
           [ -1.3005,   6.5501,  11.1626],
           ...,
           [ -0.0000,   0.0000,   0.0000],
           [ -0.0000,   0.0000,   0.0000],
           [ -0.0000,   0.0000,   0.0000]],

          [[ -1.1284,   6.9489,   9.9856],
           [ -1.1328,   5.6886,   9.2510],
           [ -1.7052,   5.8780,   7.8502],
           ...,
           [ -0.0000,   0.0000,   0.0000],
           [ -0.0000,   0.0000,   0.0000],
           [ -0.0000,   0.0000,   0.0000]],

          [[ -2.3477,   5.1160,   7.2933],
           [ -2.8573,   5.2312,   5.9299],
           [ -3.1533,   3.8572,   5.3357],
           ...,
           [ -0.0000,   0.0000,   0.0000],
           [ -0.0000,   0.0000,   0.0000],
           [ -0.0000,   0.0000,   0.0000]],

          ...,

May I know if there's a way that converts this tensor to a PDB file format?

ardhe-qb commented 9 months ago

There is a method for this: output_to_pdb, which you can find in the code here.