File size: 3,590 Bytes
ebcab53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | import os
import polars as pl
from datasets import load_dataset
import requests
from tqdm import tqdm
# Dataset loading and conversion to Polars DataFrame
ds = load_dataset("Datasets/eve-bio")
train_ds = ds["train"]
df = pl.from_pandas(train_ds.to_pandas())
# Directory to save PDB files
pdb_dir = "uniprot_pdbs"
os.makedirs(pdb_dir, exist_ok=True)
from transformers import EsmForProteinFolding, AutoTokenizer
import torch
# 1. 加载模型 (第一次会自动下载权重，约 3GB)
print("Loading Local Model...")
tokenizer = AutoTokenizer.from_pretrained("huggingface/esmfold_v1")
model = EsmForProteinFolding.from_pretrained("huggingface/esmfold_v1")
# 如果有 GPU，转到 GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
# 开启半精度推理，省显存
model.esm = model.esm.half()
def fold_local(sequence):
inputs = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids']
tokenized_input = inputs.to(device)
with torch.no_grad():
output = model(tokenized_input)
print(output)
return convert_outputs_to_pdb(output)
from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein
from transformers.models.esm.openfold_utils.feats import atom14_to_atom37
def convert_outputs_to_pdb(outputs):
final_atom_positions = atom14_to_atom37(outputs["positions"][-1], outputs)
outputs = {k: v.to("cpu").numpy() for k, v in outputs.items()}
final_atom_positions = final_atom_positions.cpu().numpy()
final_atom_mask = outputs["atom37_atom_exists"]
pdbs = []
for i in range(outputs["aatype"].shape[0]):
aa = outputs["aatype"][i]
pred_pos = final_atom_positions[i]
mask = final_atom_mask[i]
resid = outputs["residue_index"][i] + 1
pred = OFProtein(
aatype=aa,
atom_positions=pred_pos,
atom_mask=mask,
residue_index=resid,
b_factors=outputs["plddt"][i],
chain_index=outputs["chain_index"][i] if "chain_index" in outputs else None,
)
pdbs.append(to_pdb(pred))
return pdbs
# Function to download FASTA and fold to PDB
def download_and_fold_uniprot(uniprot_id, pdb_path):
url = f"https://rest.uniprot.org/uniprotkb/{uniprot_id}.fasta"
resp = requests.get(url)
print("analysing", uniprot_id)
if resp.status_code == 200:
print(resp.text)
lines = resp.text.splitlines()
sequence = "".join(lines[1:])
print("\n提取出的纯序列：")
print(sequence[:50] + "...")
else:
print("下载失败")
test_seq = sequence
output_file = pdb_path
pdb = fold_local(test_seq)
# import py3Dmol
# view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js', width=800, height=400)
# view.addModel("".join(pdb), 'pdb')
# view.setStyle({'model': -1}, {"cartoon": {'color': 'spectrum'}})
with open(pdb_path, "w") as f:
f.write("".join(pdb))
# Collect all unique UniProt IDs in the dataframe, skip null/missing
uniprot_ids = df["target__uniprot_id"].unique()
uniprot_ids = [x for x in uniprot_ids if isinstance(x, str) and x.strip()]
for uniprot_id in tqdm(uniprot_ids, desc="Processing UniProt IDs"):
try:
pdb_path = os.path.join(pdb_dir, f"{uniprot_id}.pdb")
# Robust restart: skip already-completed
if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 500:
continue
print(uniprot_id)
download_and_fold_uniprot(uniprot_id, pdb_path)
except:
pass
|