ESMC-600M蛋白质语言模型本地部署攻略

前言

之前介绍了ESMC-6B模型的网络接口调用方法,但申请token比较慢,有网友问能不能出一个本地部署ESMC小模型的攻略,遂有本文。

其实本地部署并不复杂,官方github上面也比较清楚了。

操作过程

环境配置:CUDA 12.1、torch 2.2.1+cu121、esm 3.1.1

完整的环境包列表:(因为做了些其他任务,这个里面其实不是所有都会用到,可以先把上面三个安装好,差哪些库再补哪些库)

bash 复制代码
Package                  Version
------------------------ ------------
asttokens                3.0.0
attrs                    24.3.0
biopython                1.84
biotite                  0.41.2
Brotli                   1.1.0
certifi                  2024.12.14
charset-normalizer       3.4.0
cloudpathlib             0.20.0
decorator                5.1.1
einops                   0.8.0
esm                      3.1.1
executing                2.1.0
filelock                 3.13.1
fsspec                   2024.2.0
huggingface-hub          0.27.0
idna                     3.10
ipython                  8.30.0
jedi                     0.19.2
Jinja2                   3.1.3
joblib                   1.4.2
MarkupSafe               2.1.5
matplotlib-inline        0.1.7
mpmath                   1.3.0
msgpack                  1.1.0
msgpack-numpy            0.4.8
networkx                 3.2.1
numpy                    1.26.3
nvidia-cublas-cu12       12.1.3.1
nvidia-cuda-cupti-cu12   12.1.105
nvidia-cuda-nvrtc-cu12   12.1.105
nvidia-cuda-runtime-cu12 12.1.105
nvidia-cudnn-cu12        8.9.2.26
nvidia-cufft-cu12        11.0.2.54
nvidia-curand-cu12       10.3.2.106
nvidia-cusolver-cu12     11.4.5.107
nvidia-cusparse-cu12     12.1.0.106
nvidia-nccl-cu12         2.19.3
nvidia-nvjitlink-cu12    12.1.105
nvidia-nvtx-cu12         12.1.105
packaging                24.2
pandas                   2.2.3
parso                    0.8.4
pexpect                  4.9.0
pillow                   10.2.0
pip                      24.2
prompt_toolkit           3.0.48
ptyprocess               0.7.0
pure_eval                0.2.3
Pygments                 2.18.0
python-dateutil          2.9.0.post0
pytz                     2024.2
PyYAML                   6.0.2
regex                    2024.11.6
requests                 2.32.3
safetensors              0.4.5
scikit-learn             1.6.0
scipy                    1.14.1
setuptools               75.1.0
six                      1.17.0
stack-data               0.6.3
sympy                    1.13.1
tenacity                 9.0.0
threadpoolctl            3.5.0
tokenizers               0.20.3
torch                    2.2.1+cu121
torchdata                0.7.1
torchtext                0.17.1
torchvision              0.17.1+cu121
tqdm                     4.67.1
traitlets                5.14.3
transformers             4.46.3
triton                   2.2.0
typing_extensions        4.9.0
tzdata                   2024.2
urllib3                  2.2.3
wcwidth                  0.2.13
wheel                    0.44.0

下载ESMC-600m的权重:

EvolutionaryScale/esmc-600m-2024-12 at main

下载之后把权重放在工作目录下的这个地址:data/weights

代码

和官方github上给出的例子比较类似,不过加了些修改。

python 复制代码
from esm.models.esmc import ESMC
from esm.sdk.api import *
import torch
import os
import pickle
from esm.tokenization import EsmSequenceTokenizer

# 使用预下载的参数
os.environ["INFRA_PROVIDER"] = "True"
device = torch.device("cuda:0")
client = ESMC.from_pretrained("esmc_600m",device=device)

# 读取蛋白质序列,这里需要根据自己的数据格式进行调整
def read_seq(seqfilepath):
    with open(seqfilepath,"r") as f:
        line = f.readline()
        seq = f.readline()
    return seq

# 这里沿用了上一次逆向出来的编码格式,可以替换为ESM自带的编码格式
all_amino_acid_number = {'A':5, 'C':23,'D':13,'E':9, 'F':18,
                         'G':6, 'H':21,'I':12,'K':15,'L':4,
                         'M':20,'N':17,'P':14,'Q':16,'R':10,
                         'S':8, 'T':11,'V':7, 'W':22,'Y':19,
                         '_':32}
def esm_encoder_seq(seq, pad_len):
    s = [all_amino_acid_number[x] for x in seq]
    while len(s)<pad_len:
        s.append(1)
    s.insert(0,0)
    s.append(2)
    return torch.tensor(s)

def get_esm_embedding(seq):
    protein_tensor = ESMProteinTensor(sequence=esm_encoder_seq(seq,len(seq)).to(device))
    logits_output = client.logits(protein_tensor, LogitsConfig(sequence=True, return_embeddings=True))
    esm_embedding = logits_output.embeddings
    assert isinstance(esm_embedding,torch.Tensor)
    return esm_embedding

# 这个路径设置并不重要,可以自行调整
seq_path = "seq.fasta"
seq = read_seq(seq_path)
print(seq)

# 获取序列embedding
seq_list = [seq]
emb = get_esm_embedding(seq)
with open("seq_emb.pkl","wb") as f:
    pickle.dump(emb,f)
print(emb.shape)

随便用了一个序列,得到的运行结果,tensor形状是[1,序列长度+2,1152]

相关推荐
想变成树袋熊41 分钟前
【自用】NLP算法面经(6)
人工智能·算法·自然语言处理
格林威1 小时前
Baumer工业相机堡盟工业相机如何通过YoloV8深度学习模型实现沙滩小人检测识别(C#代码UI界面版)
人工智能·深度学习·数码相机·yolo·计算机视觉
checkcheckck1 小时前
spring ai 适配 流式回答、mcp、milvus向量数据库、rag、聊天会话记忆
人工智能
Microvision维视智造1 小时前
从“人工眼”到‘智能眼’:EZ-Vision视觉系统如何重构生产线视觉检测精度?
图像处理·人工智能·重构·视觉检测
我不是大佬zvj1 小时前
解决使用vscode连接服务器出现“正在下载 VS Code 服务器...”
语言模型
巫婆理发2222 小时前
神经网络(多层感知机)(第二课第二周)
人工智能·深度学习·神经网络
lxmyzzs2 小时前
【打怪升级 - 03】YOLO11/YOLO12/YOLOv10/YOLOv8 完全指南:从理论到代码实战,新手入门必看教程
人工智能·神经网络·yolo·目标检测·计算机视觉
SEO_juper2 小时前
企业级 AI 工具选型报告:9 个技术平台的 ROI 对比与部署策略
人工智能·搜索引擎·百度·llm·工具·geo·数字营销
Coovally AI模型快速验证2 小时前
数据集分享 | 智慧农业实战数据集精选
人工智能·算法·目标检测·机器学习·计算机视觉·目标跟踪·无人机
xw33734095642 小时前
彩色转灰度的核心逻辑:三种经典方法及原理对比
人工智能·python·深度学习·opencv·计算机视觉