ESM-PEFT Functional Peptide Classification

基于蛋白质语言模型 ESM-2 与参数高效微调（PEFT）的功能肽分类项目。当前公开版本聚焦 IA3 微调路径，包含模型结构、特征融合、分类器训练和 Streamlit 推理界面，用于抗菌肽（AMP）识别及抗病毒肽（AVP）迁移实验。

项目概览

模型以 facebook/esm2_t30_150M_UR50D 为基础，通过 IA3 对注意力投影层与前馈网络进行参数高效适配；随后融合一维卷积局部特征、Transformer 全局上下文和残差特征，使用增强型 MLP 完成二分类。

核心设计

• ESM-2 序列表征：使用 30 层、150M 参数的 ESM-2 提取 640 维残基表示。
• IA3 参数高效微调：对 query、key、value 和 output.dense 等模块注入可训练缩放参数。
• 局部与全局特征融合：一维卷积提取局部模式，Transformer Encoder 建模长程依赖。
• 残差分类器：通过特征注意力、残差全连接块和增强型 MLP 输出功能肽概率。
• 批量推理界面：Streamlit 支持单序列预测、CSV 批量预测、置信度展示和结果下载。
• 资源自适应：推理自动选择 CPU/CUDA，分类器训练根据显存容量调整批量和隐藏层规模。

模型流程

flowchart TB
    A["蛋白质序列"] --> B["ESM-2 + IA3"]
    B --> C["残基级表示"]
    C --> D["1D CNN 局部特征"]
    C --> E["Transformer 全局特征"]
    D --> F["残差特征融合"]
    E --> F
    F --> G["自适应池化"]
    G --> H["增强型 MLP"]
    H --> I["AMP / Non-AMP"]

Loading

实验结果

以下为项目最终实验记录：

任务	Accuracy	F1	AUC
AMP 识别	0.8011	0.8103	0.8554
AVP 迁移实验	0.8259	0.8240	0.8750

IA3 实验中的可训练参数占比约为 0.38%。这些指标来自本地最终实验；原始数据和大型 checkpoint 未随仓库分发，因此完整复现实验还需要对应数据划分与模型文件。

目录结构

ESM_PEFT/
├── app.py                  # Streamlit 入口
├── home_page.py            # 项目介绍页面
├── prediction_page.py      # 单条与批量预测页面
├── predictor.py            # 推理组件加载与预测
├── model.py                # ESM-2、IA3、CNN/Transformer 与 MLP
├── train_classifier.py     # 特征分类器训练与评估
├── configs/                # IA3 结构配置快照
├── assets/                 # 模型结构图
├── data/                   # 数据格式说明与输入样例
└── models/                 # checkpoint 放置说明

环境安装

推荐 Python 3.10 或 3.11：

python -m venv .venv

Windows PowerShell：

.\.venv\Scripts\Activate.ps1
pip install -r requirements.txt

PyTorch GPU 版本应根据目标机器的 CUDA 环境从 PyTorch 官方安装页选择。

模型准备

复制配置模板：

Copy-Item .env.example .env

将 IA3 adapter、AIMP checkpoint 和分类器 checkpoint 放入 models/ 指定位置。也可以通过以下环境变量覆盖：

ESM_BASE_MODEL=facebook/esm2_t30_150M_UR50D
ESM_PEFT_ADAPTER_DIR=models/finetuned_esm2_ia3_150M
ESM_PEFT_AIMP_CHECKPOINT=models/net_best_aimp_model_ia3.pth
ESM_PEFT_CLASSIFIER_CHECKPOINT=models/stable_ia3_model.pth

启动

streamlit run app.py

默认访问地址：http://localhost:8501

没有本地研究数据时，首页仍可打开；数据下载按钮会处于禁用状态。进入预测页面前需要按说明准备模型 checkpoint。

批量输入格式

CSV 文件必须包含 Seq 列：

Seq
SVLGKSVAKHLPHVVPVIAEKTG
MVSKGEELFTGVVPILVELDGDVNGQK

示例文件见 data/sample_sequences.csv。

公开范围

仓库公开模型定义、推理逻辑、分类器训练代码和可视化界面，不包含：

• ESM-2 基础模型权重；
• PEFT 训练过程中的中间 checkpoint 与优化器状态；
• 最终 AIMP/MLP 大型 checkpoint；
• 需要单独确认再分发许可的原始 AMP/AVP 数据集。

基础模型来源：facebook/esm2_t30_150M_UR50D。

IA3 目标层配置见 configs/ia3_adapter_config.json。

使用边界

该项目用于功能肽计算筛选与研究验证，模型输出不能替代湿实验、生物安全评估或临床判断。