用途:fair-esm2 技能基于 Meta AI 的 ESM-2 模型(fair-esm 包,MIT 许可,github.com/facebookresearch/esm),用于为蛋白质序列提取嵌入,供下游机器学习使用。
主要能力:可提取逐残基或逐序列嵌入;进行掩码语言模型打分(如突变评分);在 return_contacts=True 时输出接触预测矩阵。
使用方式:安装 pip install fair-esm 后 import esm,通过 esm.pretrained.esm2_t33_650M_UR50D() 等加载模型。示例:model, alphabet = esm.pretrained.esm2_t33_650M_UR50D();model = model.eval().cuda();bc = alphabet.get_batch_converter();toks 经 bc 转换后调用 model(toks.cuda(), repr_layers=[33])。输出 out["representations"][layer] 形状为 (B, L+2, D),需用 [:, 1:-1, :] 去掉 BOS/EOS;逐序列均值可 emb[0, 1:-1].mean(0)。接触预测输出 out["contacts"] 形状 (B, L, L)。
可用模型:esm2_t6_8M_UR50D(6 层,320 维,8M 参数,快速测试)、esm2_t33_650M_UR50D(33 层,1280 维,650M 参数,默认嵌入模型)、esm2_t36_3B_UR50D(36 层,2560 维,3B 参数,需 24GB+ 显存)。
环境要求:Python 3.8+(推荐 3.11),CUDA 11.7+(推荐 12.x),显存 8GB(8M 模型)、16GB(650M 模型),推荐 24GB+。
远程计算:650M 模型需 ≥16GB 显存,且需预缓存 .pt 检查点或可访问 dl.fbaipublicfiles.com。可通过 host.compute.create(provider) 创建计算并 submit_job 提交任务,命令为 python3 embed_esm2.py,输出 embeddings.pt;脚本内应将 TORCH_HOME 指向提供方的 torch-hub 缓存挂载路径。任务完成后调用 wait_for_notification,收到 compute_done 后 save_artifacts(payload["featured_files"]),或通过 c.attach_job(job_id).result() 获取完整结果。
限制与注意:本技能针对 Meta 的 fair-esm 包(使用 esm.pretrained.*);Biohub 的 esm 分支(from esm.models.esmfold2 import ...)是不同库,见 esmfold2 技能。首次调用若通过 torch.hub 下载权重会较慢,建议设置 TORCH_HOME 到缓存位置。结构预测请使用 esmfold2。