mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-07-21 21:22:05 +08:00
Add LLM molecular encoder branch (BioT5/MolT5/MoleculeSTM/SELFIES-TED)
This commit is contained in:
Dicong Li
parent
804bda5576
commit
fb4f90b68a
12
.gitignore
vendored
12
.gitignore
vendored
@ -188,4 +188,14 @@ cython_debug/
|
||||
.ruff_cache/
|
||||
|
||||
# PyPI configuration file
|
||||
.pypirc
|
||||
.pypirc
|
||||
|
||||
# ============ LLM encoder branch additions ============
|
||||
# 大文件 / 模型权重（不进仓库）
|
||||
*.pt
|
||||
*.pth
|
||||
*.ckpt
|
||||
*.safetensors
|
||||
|
||||
# 数据
|
||||
*.parquet
|
||||
|
||||
205
README_llm_encoders.md
Normal file
205
README_llm_encoders.md
Normal file
@ -0,0 +1,205 @@
|
||||
# LLM 分子编码器分支（LLM Encoder Branch）
|
||||
|
||||
本模块在 LNP 多任务预测模型上新增了一个**可选的 LLM 分子编码器分支**。它把可电离脂质的 SMILES 送入一个冻结的预训练分子模型，提取语义向量，作为一个额外的 token（V5）追加到模型的 fusion 序列中参与多任务预测。
|
||||
|
||||
支持四种分子编码器，可通过参数自由切换，用于横向对比：
|
||||
|
||||
| llm_type | 模型 | 年份 / 发表 | 参数量 | hidden | 输入 |
|
||||
|---|---|---|---|---|---|
|
||||
| `biot5` | BioT5-plus | 2024, arXiv | ~250M | 768 | SMILES |
|
||||
| `molt5` | MolT5-base | 2022, EMNLP | 109M | 768 | SMILES |
|
||||
| `moleculestm` | MoleculeSTM GNN | 2023, Nature MI | ~2M (GNN) | 300 | 分子图 |
|
||||
| `selfiested` | SELFIES-TED | 2025, IBM | 358M | 1024 | SELFIES |
|
||||
|
||||
---
|
||||
|
||||
## 1. 设计说明
|
||||
|
||||
- **非侵入式**：默认 `use_llm=False`，此时模型行为与原版完全一致，不影响既有功能。只有显式开启才会启用 LLM 分支。
|
||||
- **统一接口**：四个 encoder 的 `forward(List[str]) -> [B, d_model]` 接口完全一致，通过 `models.build_llm_encoder()` 工厂函数按 `llm_type` 实例化。
|
||||
- **冻结主体**：LLM / GNN 主体全程冻结，只有内部一层 `projection`（LayerNorm + Linear + ReLU + Dropout）参与训练，参数量增加极小。
|
||||
- **惰性导入**：`encoders/__init__.py` 对四个 LLM encoder 采用惰性导入，未安装 `torch_geometric` / `selfies` 等可选依赖时，只要不使用对应 encoder，导入不会报错。
|
||||
|
||||
---
|
||||
|
||||
## 2. 文件清单
|
||||
|
||||
新增文件（均在 `lnp_ml/modeling/encoders/` 下）：
|
||||
|
||||
```
|
||||
encoders/
|
||||
├── llm_encoder.py # BioT5-plus (类名 LLMEncoder)
|
||||
├── molt5_encoder.py # MolT5 (类名 MolT5Encoder)
|
||||
├── moleculestm_encoder.py # MoleculeSTM (类名 MoleculeSTMEncoder)
|
||||
├── molecule_gnn_model.py # MoleculeSTM 的 GNN 网络定义（被上一个文件 import）
|
||||
└── selfiested_encoder.py # SELFIES-TED (类名 SELFIESTEDEncoder)
|
||||
```
|
||||
|
||||
修改文件：
|
||||
|
||||
```
|
||||
lnp_ml/modeling/models.py # 新增 use_llm / llm_type 等参数 + build_llm_encoder 工厂 + V5 追加逻辑；desc 维度 210 -> 217
|
||||
lnp_ml/modeling/encoders/__init__.py # 惰性导出四个 LLM encoder
|
||||
verify_llm_encoder.py # LLM 模块验证脚本（四项测试）
|
||||
```
|
||||
|
||||
> 注：`models.py` 中 `DEFAULT_INPUT_DIMS["desc"]` 已从 210 更新为 **217**，以匹配新版 RDKit 的描述符维度。
|
||||
|
||||
---
|
||||
|
||||
## 3. 依赖安装
|
||||
|
||||
**基础依赖（所有模型都需要）：**
|
||||
```bash
|
||||
pip install rdkit loguru transformers sentencepiece optuna
|
||||
```
|
||||
|
||||
**按模型追加：**
|
||||
|
||||
| 模型 | 额外依赖 |
|
||||
|---|---|
|
||||
| biot5 | 无（基础依赖即可） |
|
||||
| molt5 | 无 |
|
||||
| moleculestm | `pip install torch_geometric torch_scatter torch_sparse ogb` |
|
||||
| selfiested | `pip install selfies` |
|
||||
|
||||
---
|
||||
|
||||
## 4. 预训练权重
|
||||
|
||||
| 模型 | 是否需手动下载 | 说明 |
|
||||
|---|---|---|
|
||||
| **molt5** | 否 | 训练时自动从 HuggingFace 拉取 `laituan245/molt5-base` |
|
||||
| **selfiested** | 否 | 训练时自动从 HuggingFace 拉取 `ibm/materials.selfies-ted` |
|
||||
| **biot5** | 是 | 需下载 BioT5-plus 权重文件夹，见下方 |
|
||||
| **moleculestm** | 是 | 需下载 MoleculeSTM GNN 权重，见下方 |
|
||||
|
||||
**BioT5-plus 权重下载：**
|
||||
从 HuggingFace `QizhiPei/biot5-plus-base` 下载整个权重文件夹（含 `config.json`、`tokenizer_config.json` 等），放到本地任意路径，训练时用 `--llm-model-path <该路径>` 指定。也可直接传 HuggingFace ID `QizhiPei/biot5-plus-base` 在线加载。
|
||||
|
||||
**MoleculeSTM GNN 权重下载：**
|
||||
```python
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download(
|
||||
repo_id="chao1224/MoleculeSTM",
|
||||
repo_type="model",
|
||||
local_dir="<本地目录>/moleculestm_weights",
|
||||
allow_patterns="*Graph*",
|
||||
)
|
||||
```
|
||||
下载后使用的 GNN 权重路径为：
|
||||
```
|
||||
<本地目录>/moleculestm_weights/pretrained_MoleculeSTM/SciBERT-Graph-3e-5-1-1e-4-1-InfoNCE-0.1-32-32/molecule_model.pth
|
||||
```
|
||||
训练时用 `--llm-model-path <该 .pth 路径>` 指定。
|
||||
|
||||
---
|
||||
|
||||
## 5. 用法
|
||||
|
||||
### 5.1 代码中直接构造模型
|
||||
|
||||
```python
|
||||
from lnp_ml.modeling.models import LNPModelWithoutMPNN
|
||||
|
||||
# 不启用 LLM（与原版行为一致）
|
||||
model = LNPModelWithoutMPNN(d_model=256)
|
||||
|
||||
# 启用 LLM 分支（以 MolT5 为例）
|
||||
model = LNPModelWithoutMPNN(
|
||||
d_model=256,
|
||||
use_llm=True,
|
||||
llm_type="molt5", # biot5 / molt5 / moleculestm / selfiested
|
||||
llm_model_path="laituan245/molt5-base", # 见上表，不同模型路径不同
|
||||
llm_device="cuda",
|
||||
)
|
||||
```
|
||||
|
||||
各模型对应的 `llm_type` 和 `llm_model_path`：
|
||||
|
||||
| llm_type | llm_model_path 示例 |
|
||||
|---|---|
|
||||
| `biot5` | `QizhiPei/biot5-plus-base` 或本地 BioT5 权重文件夹路径 |
|
||||
| `molt5` | `laituan245/molt5-base` |
|
||||
| `moleculestm` | `.../moleculestm_weights/.../molecule_model.pth` |
|
||||
| `selfiested` | `ibm/materials.selfies-ted` |
|
||||
|
||||
### 5.2 验证 LLM 模块（训练前建议先跑）
|
||||
|
||||
```bash
|
||||
python verify_llm_encoder.py \
|
||||
--model_path <llm_model_path> \
|
||||
--lnp_repo_path <仓库根目录>
|
||||
```
|
||||
四项测试（输出维度、维度对齐、可复现性、过拟合 sanity）全部 PASS 再训练。
|
||||
|
||||
### 5.3 训练
|
||||
|
||||
> 训练脚本（`final_train_optuna_cv.py` / `pretrain.py`）需要支持 `--use-llm` / `--llm-type` / `--llm-model-path` / `--llm-device` 参数。各模型的完整 Colab 操作步骤见 `docs/` 下四份 quickstart 指南。
|
||||
|
||||
---
|
||||
|
||||
## 6. 实验结果
|
||||
|
||||
四个模型在 LNP 测试集上的对比（核心指标 delivery R²，baseline 为不含 LLM 的版本）：
|
||||
|
||||
| 任务 | Baseline | MolT5 | MoleculeSTM | BioT5-plus | SELFIES-TED |
|
||||
|---|---|---|---|---|---|
|
||||
| delivery R² | 0.373 | 0.602 | 0.552 | **0.633** | 0.618 |
|
||||
| size R² | 0.215 | 0.440 | 0.408 | 0.423 | 0.381 |
|
||||
| pdi acc | 0.786 | 0.748 | 0.710 | 0.695 | 0.664 |
|
||||
| ee acc | 0.679 | 0.702 | 0.679 | 0.679 | 0.679 |
|
||||
| toxic acc | 0.980 | 0.980 | 0.951 | 0.960 | 0.941 |
|
||||
| biodist KL | 0.293 | 0.727 | 0.646 | 0.703 | 0.643 |
|
||||
|
||||
**主要结论：**
|
||||
1. 四个 LLM 编码器都显著提升了核心回归任务，delivery R² 从 baseline 的 0.37 最高提升到 0.63。
|
||||
2. 领域专用预训练比堆参数更重要：BioT5-plus（针对生物分子预训练）效果最好，而参数最大、最新的 SELFIES-TED 仅排第二。
|
||||
3. 序列模型整体优于图模型：三个基于 SMILES/SELFIES 的模型均优于基于图的 MoleculeSTM。
|
||||
4. MolT5 性价比最高：仅 109M 参数，delivery R² 排第三且与第一仅差 0.03，分类任务表现最均衡。
|
||||
5. 共同局限：biodist（分布预测）四个 LLM 均略低于 baseline，属多任务学习中的 trade-off。
|
||||
|
||||
> 各模型每个任务的详细指标见 `results/<模型名>/` 目录。
|
||||
|
||||
---
|
||||
|
||||
## 附：训练脚本改动
|
||||
|
||||
本 PR 已经更新了两个训练脚本以支持 LLM 分支（MoE 相关代码未引入）：
|
||||
|
||||
- `lnp_ml/modeling/final_train_optuna_cv.py`
|
||||
- `lnp_ml/modeling/pretrain.py`
|
||||
|
||||
两者都新增了四个命令行参数，并透传给模型：
|
||||
|
||||
```
|
||||
--use-llm # 是否启用 LLM 分支
|
||||
--llm-type {biot5,molt5,moleculestm,selfiested} # 选择哪个 encoder
|
||||
--llm-model-path <路径或HF_ID> # 见第 4/5 节
|
||||
--llm-device cuda # LLM 推理设备
|
||||
```
|
||||
|
||||
预训练保存的 checkpoint 的 config 中也会记录 LLM 设置，恢复模型时自动读取。
|
||||
|
||||
### 训练示例（以 MolT5 为例）
|
||||
|
||||
```bash
|
||||
# 预训练
|
||||
python -m lnp_ml.modeling.pretrain main \
|
||||
--train-path data/processed/train_pretrain.parquet \
|
||||
--val-path data/processed/val_pretrain.parquet \
|
||||
--use-llm --llm-type molt5 \
|
||||
--llm-model-path laituan245/molt5-base \
|
||||
--llm-device cuda --epochs 50 --lr 1e-4 --device cuda
|
||||
|
||||
# 正式训练
|
||||
python lnp_ml/modeling/final_train_optuna_cv.py \
|
||||
--init-from-pretrain models/pretrain_delivery.pt \
|
||||
--use-llm --llm-type molt5 \
|
||||
--llm-model-path laituan245/molt5-base \
|
||||
--llm-device cuda \
|
||||
--n-trials 20 --epochs-per-trial 30 --seed 42 \
|
||||
--device cuda --output-dir models/final
|
||||
```
|
||||
|
||||
切换其他模型只需改 `--llm-type` 和 `--llm-model-path`（见第 5.1 节对照表）。
|
||||
42
RESULTS.md
Normal file
42
RESULTS.md
Normal file
@ -0,0 +1,42 @@
|
||||
# LLM 编码器分支 — 实验结果
|
||||
|
||||
LNP 多任务预测模型在引入不同 LLM 分子编码器（V5 token）后的测试集表现。
|
||||
所有实验：`LNPModelWithoutMPNN`（无 MPNN、无 MoE），固定 seed=42，先在外部数据上预训练 delivery，再用 Optuna 做 3-fold 调参 + 全量训练。
|
||||
|
||||
## 总览对比
|
||||
|
||||
| 任务 | 指标 | Baseline (无 LLM) | MolT5 | MoleculeSTM | BioT5-plus | SELFIES-TED |
|
||||
|---|---|---|---|---|---|---|
|
||||
| delivery | R² | 0.373 | 0.602 | 0.552 | **0.633** | 0.618 |
|
||||
| size | R² | 0.215 | **0.440** | 0.408 | 0.423 | 0.381 |
|
||||
| pdi | acc | **0.786** | 0.748 | 0.710 | 0.695 | 0.664 |
|
||||
| ee | acc | 0.679 | **0.702** | 0.679 | 0.679 | 0.679 |
|
||||
| toxic | acc | **0.980** | 0.980 | 0.951 | 0.960 | 0.941 |
|
||||
| biodist | KL↓ | **0.293** | 0.727 | 0.646 | 0.703 | 0.643 |
|
||||
|
||||
（粗体为该行最优；KL 越低越好，其余越高越好。）
|
||||
|
||||
## 模型信息
|
||||
|
||||
| llm_type | 模型 | 年份 | 参数量 | hidden | 输入表示 |
|
||||
|---|---|---|---|---|---|
|
||||
| biot5 | BioT5-plus | 2024 | ~250M | 768 | SMILES |
|
||||
| molt5 | MolT5-base | 2022 | 109M | 768 | SMILES |
|
||||
| moleculestm | MoleculeSTM GNN | 2023 | ~2M (GNN) | 300 | 分子图 |
|
||||
| selfiested | SELFIES-TED | 2025 | 358M | 1024 | SELFIES |
|
||||
|
||||
## 主要结论
|
||||
|
||||
1. **四个 LLM 编码器都显著提升核心回归任务**：delivery R² 从 baseline 的 0.373 最高提升到 0.633（BioT5-plus），size R² 也从 0.215 提升到 0.38–0.44。
|
||||
|
||||
2. **领域专用预训练 > 堆参数**：BioT5-plus（针对生物分子预训练）效果最好；参数最大、最新的 SELFIES-TED（358M）仅排第二，说明针对生物分子的预训练比单纯扩大规模更有效。
|
||||
|
||||
3. **序列模型 > 图模型**：三个基于序列（SMILES/SELFIES）的语言模型都优于基于图的 MoleculeSTM GNN。
|
||||
|
||||
4. **MolT5 性价比最高**：仅 109M 参数，delivery R² 排第三（0.602，与第一仅差 0.03），且三个分类任务表现最均衡。
|
||||
|
||||
5. **共同局限 — biodist**：四个 LLM 在 biodist（组织分布预测）上都比 baseline 差，属于多任务学习中的 trade-off：模型容量向回归任务倾斜，牺牲了分布预测。后续可通过调整各任务损失权重缓解。
|
||||
|
||||
## 复现方式
|
||||
|
||||
各模型的完整 Colab 操作步骤见 `docs/` 下的四份 quickstart 指南；模型构造与参数说明见 `README_llm_encoders.md`。
|
||||
209
docs/colab_biot5_quickstart.md
Normal file
209
docs/colab_biot5_quickstart.md
Normal file
@ -0,0 +1,209 @@
|
||||
# LNP-ML + BioT5-plus 快速上手指南
|
||||
|
||||
适用于：拿到代码包、从未运行过任何代码的新组员。
|
||||
|
||||
---
|
||||
|
||||
## 前提条件（收到代码包前确认）
|
||||
|
||||
- Google Drive 的 `MyDrive/lnp_project/` 下已有以下两个文件夹：
|
||||
- `lnp_ml/`：项目代码
|
||||
- `biot5-plus-base/`：BioT5-plus 权重（无需重新下载）
|
||||
- `lnp_ml/data/processed/` 下已有处理好的数据文件
|
||||
|
||||
---
|
||||
|
||||
## Cell 1：挂载 Drive
|
||||
|
||||
```python
|
||||
from google.colab import drive
|
||||
drive.mount('/content/drive')
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cell 2：安装依赖
|
||||
|
||||
> 每次重新打开 Colab 都需要重跑这一步，依赖不会自动保留。
|
||||
|
||||
```python
|
||||
!pip install rdkit loguru transformers sentencepiece optuna -q
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cell 3：配置路径 + 确认文件完整
|
||||
|
||||
> 所有 cell 都依赖这里定义的 `LNP_PATH` 和 `MODEL_PATH`，必须先运行。
|
||||
|
||||
```python
|
||||
import os, sys
|
||||
|
||||
LNP_PATH = "/content/drive/MyDrive/lnp_project/lnp_ml"
|
||||
MODEL_PATH = "/content/drive/MyDrive/lnp_project/biot5-plus-base"
|
||||
|
||||
sys.path.insert(0, LNP_PATH)
|
||||
os.chdir(LNP_PATH)
|
||||
|
||||
checks = [
|
||||
f"{LNP_PATH}/lnp_ml/modeling/models.py",
|
||||
f"{LNP_PATH}/lnp_ml/modeling/encoders/llm_encoder.py",
|
||||
f"{LNP_PATH}/lnp_ml/modeling/encoders/__init__.py",
|
||||
f"{LNP_PATH}/lnp_ml/featurization/smiles.py",
|
||||
f"{MODEL_PATH}/config.json",
|
||||
f"{MODEL_PATH}/tokenizer_config.json",
|
||||
f"{LNP_PATH}/data/processed/train.parquet",
|
||||
f"{LNP_PATH}/data/processed/val.parquet",
|
||||
f"{LNP_PATH}/data/processed/test.parquet",
|
||||
f"{LNP_PATH}/data/processed/train_pretrain.parquet",
|
||||
f"{LNP_PATH}/data/processed/val_pretrain.parquet",
|
||||
]
|
||||
for f in checks:
|
||||
print(f"{'✓' if os.path.exists(f) else '✗ 缺失'} {f}")
|
||||
```
|
||||
|
||||
**所有文件都显示 ✓ 再继续。如果有 ✗，先补齐对应文件。**
|
||||
|
||||
---
|
||||
|
||||
## Cell 4：验证 LLM 模块（正式训练前必做）
|
||||
|
||||
> 验证数据流、维度对齐、梯度链路、可复现性，四项全部通过再训练。
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python {LNP_PATH}/verify_llm_encoder.py \
|
||||
--model_path "{MODEL_PATH}" \
|
||||
--lnp_repo_path "{LNP_PATH}"
|
||||
```
|
||||
|
||||
**正常输出：**
|
||||
```
|
||||
TEST 1 PASSED ✓
|
||||
TEST 2 PASSED ✓
|
||||
TEST 3 PASSED ✓
|
||||
TEST 4 PASSED ✓
|
||||
ALL TESTS PASSED ✓
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cell 5：预训练（用外部 LiON 数据）
|
||||
|
||||
> 在约 9000 条外部数据上预训练 delivery 任务，产出 `models/pretrain_delivery.pt`。
|
||||
> 约需 10-15 分钟。
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python -m lnp_ml.modeling.pretrain main \
|
||||
--train-path data/processed/train_pretrain.parquet \
|
||||
--val-path data/processed/val_pretrain.parquet \
|
||||
--epochs 50 \
|
||||
--lr 1e-4 \
|
||||
--device cuda
|
||||
```
|
||||
|
||||
完成后确认：
|
||||
```python
|
||||
import os
|
||||
print("✓ 预训练权重存在" if os.path.exists(f"{LNP_PATH}/models/pretrain_delivery.pt")
|
||||
else "✗ 预训练权重未生成，检查上面的输出")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cell 6：正式训练（含 LLM，约 20-25 分钟）
|
||||
|
||||
> 使用 Optuna 做 3-fold 超参搜索（20 trials），然后全量数据训练。
|
||||
> 产出 `models/final/model.pt`。
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python lnp_ml/modeling/final_train_optuna_cv.py \
|
||||
--init-from-pretrain models/pretrain_delivery.pt \
|
||||
--use-llm \
|
||||
--llm-model-path "{MODEL_PATH}" \
|
||||
--llm-device cuda \
|
||||
--n-trials 20 \
|
||||
--epochs-per-trial 30 \
|
||||
--seed 42 \
|
||||
--device cuda \
|
||||
--output-dir models/final
|
||||
```
|
||||
|
||||
完成后确认：
|
||||
```python
|
||||
print("✓ 模型权重存在" if os.path.exists(f"{LNP_PATH}/models/final/model.pt")
|
||||
else "✗ 模型未生成，检查上面的输出")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Cell 7：测试评估
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python lnp_ml/modeling/predict.py test \
|
||||
--test-path data/processed/test.parquet \
|
||||
--model-path models/final/model.pt \
|
||||
--output-path models/final/test_results.json
|
||||
|
||||
import json
|
||||
with open(f'{LNP_PATH}/models/final/test_results.json') as f:
|
||||
results = json.load(f)
|
||||
|
||||
print("=== 分类任务 ===")
|
||||
for task in ['pdi', 'ee', 'toxic']:
|
||||
m = results['detailed_metrics'][task]
|
||||
print(f" {task}: acc={m['accuracy']:.4f}, f1={m['f1']:.4f}")
|
||||
|
||||
print("\n=== 回归任务 ===")
|
||||
for task in ['size', 'delivery']:
|
||||
m = results['detailed_metrics'][task]
|
||||
print(f" {task}: R²={m['r2']:.4f}, RMSE={m['rmse']:.4f}")
|
||||
|
||||
print("\n=== 分布任务 ===")
|
||||
m = results['detailed_metrics']['biodist']
|
||||
print(f" biodist: KL={m['kl_divergence']:.4f}, JS={m['js_divergence']:.4f}")
|
||||
```
|
||||
|
||||
**参考指标（基于已有实验结果）：**
|
||||
|
||||
| 任务 | 指标 | 参考值 |
|
||||
|---|---|---|
|
||||
| delivery | R² | ~0.63 |
|
||||
| size | R² | ~0.42 |
|
||||
| pdi | acc | ~0.69 |
|
||||
| ee | acc | ~0.68 |
|
||||
| toxic | acc | ~0.96 |
|
||||
| biodist | KL | ~0.70 |
|
||||
|
||||
---
|
||||
|
||||
## Cell 8：备份模型到 Drive
|
||||
|
||||
```python
|
||||
!cp -r {LNP_PATH}/models/final \
|
||||
{LNP_PATH}/models/final_backup
|
||||
print("备份完成 ✓")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 常见问题
|
||||
|
||||
| 报错 | 原因 | 解决方法 |
|
||||
|---|---|---|
|
||||
| `No module named 'rdkit'` | 依赖未安装 | 重新运行 Cell 2 |
|
||||
| `No module named 'lnp_ml'` | 路径未配置 | 重新运行 Cell 3 |
|
||||
| `No such option '--use-llm'` | Drive 里的 `final_train_optuna_cv.py` 是旧版本 | 确认代码包里的文件是最新版本 |
|
||||
| `running_mean should contain 217 elements not 210` | `models.py` 里 `desc` 维度是旧值 210 | 把 `DEFAULT_INPUT_DIMS` 里的 `"desc": 210` 改为 `"desc": 217` |
|
||||
| `KeyError: attribute 'projection' already exists` | `llm_encoder.py` 是旧版本 | 确认代码包里的文件是最新版本 |
|
||||
| `RuntimeError: Unexpected key(s) in state_dict` | 加载模型时 strict=True | 确认 `predict.py` 里用的是 `strict=False` |
|
||||
| Colab 断开后重连 | 环境变量丢失 | 从 Cell 1 重跑，Cell 5/6 已有权重可跳过 |
|
||||
|
||||
---
|
||||
|
||||
## 如果已有训练好的模型，直接从 Cell 7 开始
|
||||
|
||||
只需运行 Cell 1 → Cell 2 → Cell 3 → Cell 7，跳过 Cell 4、5、6。
|
||||
334
docs/colab_moleculestm_quickstart.md
Normal file
334
docs/colab_moleculestm_quickstart.md
Normal file
@ -0,0 +1,334 @@
|
||||
# LNP-ML + MoleculeSTM GNN 快速上手指南
|
||||
|
||||
适用于：拿到代码包、从未运行过任何代码的新组员。
|
||||
|
||||
---
|
||||
|
||||
## 前提条件（开始前确认）
|
||||
|
||||
- Google Drive 的 `MyDrive/lnp_project_moleculestm/` 下已有以下内容：
|
||||
- `lnp_ml/`：项目代码（从 BioT5 版本复制来的，已修改）
|
||||
- `moleculestm_weights/`：GNN 预训练权重
|
||||
|
||||
如果以上文件夹不存在，先看文末的【从零开始】章节。
|
||||
|
||||
---
|
||||
|
||||
## 每次重新打开 Colab 都需要运行
|
||||
|
||||
### Cell 1：挂载 Drive
|
||||
|
||||
```python
|
||||
from google.colab import drive
|
||||
drive.mount('/content/drive')
|
||||
```
|
||||
|
||||
> 弹出授权页面时，确认选择有 `lnp_project_moleculestm` 的 Google 账号。
|
||||
|
||||
---
|
||||
|
||||
### Cell 2：安装依赖
|
||||
|
||||
```python
|
||||
!pip install rdkit loguru transformers sentencepiece optuna \
|
||||
torch_geometric torch_scatter torch_sparse ogb -q
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Cell 3：配置路径
|
||||
|
||||
```python
|
||||
import os, sys
|
||||
|
||||
LNP_PATH = "/content/drive/MyDrive/lnp_project_moleculestm/lnp_ml"
|
||||
MODEL_PATH = "/content/drive/MyDrive/lnp_project_moleculestm/moleculestm_weights/pretrained_MoleculeSTM/SciBERT-Graph-3e-5-1-1e-4-1-InfoNCE-0.1-32-32/molecule_model.pth"
|
||||
|
||||
sys.path.insert(0, LNP_PATH)
|
||||
os.chdir(LNP_PATH)
|
||||
print("路径配置完成 ✓")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Cell 4：确认文件完整
|
||||
|
||||
```python
|
||||
checks = [
|
||||
f"{LNP_PATH}/lnp_ml/modeling/models.py",
|
||||
f"{LNP_PATH}/lnp_ml/modeling/encoders/moleculestm_encoder.py",
|
||||
f"{LNP_PATH}/lnp_ml/modeling/encoders/molecule_gnn_model.py",
|
||||
f"{LNP_PATH}/lnp_ml/modeling/encoders/__init__.py",
|
||||
f"{MODEL_PATH}",
|
||||
f"{LNP_PATH}/data/processed/train.parquet",
|
||||
f"{LNP_PATH}/data/processed/val.parquet",
|
||||
f"{LNP_PATH}/data/processed/test.parquet",
|
||||
f"{LNP_PATH}/data/processed/train_pretrain.parquet",
|
||||
f"{LNP_PATH}/data/processed/val_pretrain.parquet",
|
||||
]
|
||||
for f in checks:
|
||||
print(f"{'✓' if os.path.exists(f) else '✗ 缺失'} {f}")
|
||||
```
|
||||
|
||||
**所有文件都显示 ✓ 再继续。**
|
||||
|
||||
---
|
||||
|
||||
## 正式流程
|
||||
|
||||
### Cell 5：验证 MoleculeSTM 模块
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python {LNP_PATH}/verify_llm_encoder.py \
|
||||
--model_path "{MODEL_PATH}" \
|
||||
--lnp_repo_path "{LNP_PATH}"
|
||||
```
|
||||
|
||||
**正常输出：**
|
||||
```
|
||||
TEST 1 PASSED ✓
|
||||
TEST 2 PASSED ✓
|
||||
TEST 3 PASSED ✓
|
||||
TEST 4 PASSED ✓
|
||||
ALL TESTS PASSED ✓
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Cell 6：预训练（约 10 分钟）
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python -m lnp_ml.modeling.pretrain main \
|
||||
--train-path data/processed/train_pretrain.parquet \
|
||||
--val-path data/processed/val_pretrain.parquet \
|
||||
--epochs 50 \
|
||||
--lr 1e-4 \
|
||||
--device cuda
|
||||
```
|
||||
|
||||
完成后确认权重存在：
|
||||
```python
|
||||
print("✓ 预训练权重存在" if os.path.exists(f"{LNP_PATH}/models/pretrain_delivery.pt")
|
||||
else "✗ 预训练权重未生成")
|
||||
```
|
||||
|
||||
参考结果：`Best val_loss: 0.5667`
|
||||
|
||||
---
|
||||
|
||||
### Cell 7：正式训练（约 20 分钟）
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python lnp_ml/modeling/final_train_optuna_cv.py \
|
||||
--init-from-pretrain models/pretrain_delivery.pt \
|
||||
--use-llm \
|
||||
--llm-model-path "{MODEL_PATH}" \
|
||||
--llm-device cuda \
|
||||
--n-trials 20 \
|
||||
--epochs-per-trial 30 \
|
||||
--seed 42 \
|
||||
--device cuda \
|
||||
--output-dir models/final
|
||||
```
|
||||
|
||||
完成后确认：
|
||||
```python
|
||||
print("✓ 模型权重存在" if os.path.exists(f"{LNP_PATH}/models/final/model.pt")
|
||||
else "✗ 模型未生成")
|
||||
```
|
||||
|
||||
参考结果：`Best val_loss: 2.4433`
|
||||
|
||||
---
|
||||
|
||||
### Cell 8：测试评估
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python lnp_ml/modeling/predict.py test \
|
||||
--test-path data/processed/test.parquet \
|
||||
--model-path models/final/model.pt \
|
||||
--output-path models/final/test_results.json
|
||||
|
||||
import json
|
||||
with open(f'{LNP_PATH}/models/final/test_results.json') as f:
|
||||
results = json.load(f)
|
||||
|
||||
print("=== 分类任务 ===")
|
||||
for task in ['pdi', 'ee', 'toxic']:
|
||||
m = results['detailed_metrics'][task]
|
||||
print(f" {task}: acc={m['accuracy']:.4f}, f1={m['f1']:.4f}")
|
||||
|
||||
print("\n=== 回归任务 ===")
|
||||
for task in ['size', 'delivery']:
|
||||
m = results['detailed_metrics'][task]
|
||||
print(f" {task}: R²={m['r2']:.4f}, RMSE={m['rmse']:.4f}")
|
||||
|
||||
print("\n=== 分布任务 ===")
|
||||
m = results['detailed_metrics']['biodist']
|
||||
print(f" biodist: KL={m['kl_divergence']:.4f}, JS={m['js_divergence']:.4f}")
|
||||
```
|
||||
|
||||
**参考指标：**
|
||||
|
||||
| 任务 | 指标 | 参考值 |
|
||||
|---|---|---|
|
||||
| delivery | R² | ~0.55 |
|
||||
| size | R² | ~0.41 |
|
||||
| pdi | acc | ~0.71 |
|
||||
| ee | acc | ~0.68 |
|
||||
| toxic | acc | ~0.95 |
|
||||
| biodist | KL | ~0.65 |
|
||||
|
||||
---
|
||||
|
||||
### Cell 9：备份模型
|
||||
|
||||
```python
|
||||
!cp -r {LNP_PATH}/models/final \
|
||||
{LNP_PATH}/models/final_backup
|
||||
print("备份完成 ✓")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 常见问题
|
||||
|
||||
| 报错 | 原因 | 解决方法 |
|
||||
|---|---|---|
|
||||
| `No module named 'rdkit'` | 依赖未安装 | 重新运行 Cell 2 |
|
||||
| `No module named 'lnp_ml'` | 路径未配置 | 重新运行 Cell 3 |
|
||||
| `Drive not mounted` | Drive 未挂载 | 重新运行 Cell 1 |
|
||||
| `Mountpoint must not already contain files` | Drive 已挂载但缓存混乱 | 菜单栏 → 运行时 → 重新启动运行时，再从 Cell 1 开始 |
|
||||
| `RuntimeError: size mismatch` | OGB 版本差异 | 已在代码里修复，正常现象 |
|
||||
| `KeyError: test_results.json` | 测试未完成 | 确认 Cell 8 的第一条命令输出了 `Saved test results` |
|
||||
| 重新连接后变量丢失 | Colab 断开重连 | 从 Cell 1 重跑，Cell 6/7 已有权重可跳过 |
|
||||
|
||||
---
|
||||
|
||||
## 如果已有训练好的模型，直接从 Cell 8 开始
|
||||
|
||||
只需运行 Cell 1 → Cell 2 → Cell 3 → Cell 8，跳过 Cell 4-7。
|
||||
|
||||
---
|
||||
|
||||
## 【从零开始】如果 lnp_project_moleculestm 不存在
|
||||
|
||||
### Step A：从 BioT5 版本复制代码
|
||||
|
||||
```python
|
||||
import os
|
||||
os.makedirs("/content/drive/MyDrive/lnp_project_moleculestm", exist_ok=True)
|
||||
|
||||
!cp -r "/content/drive/MyDrive/lnp_project_biot5-plus/lnp_ml" \
|
||||
"/content/drive/MyDrive/lnp_project_moleculestm/lnp_ml"
|
||||
print("代码复制完成 ✓")
|
||||
```
|
||||
|
||||
### Step B：下载 MoleculeSTM GNN 权重（约 3GB，需 1-2 分钟）
|
||||
|
||||
```python
|
||||
from huggingface_hub import snapshot_download
|
||||
import os
|
||||
|
||||
os.makedirs("/content/drive/MyDrive/lnp_project_moleculestm/moleculestm_weights", exist_ok=True)
|
||||
|
||||
snapshot_download(
|
||||
repo_id="chao1224/MoleculeSTM",
|
||||
repo_type="model",
|
||||
local_dir="/content/drive/MyDrive/lnp_project_moleculestm/moleculestm_weights",
|
||||
allow_patterns="*Graph*"
|
||||
)
|
||||
print("权重下载完成 ✓")
|
||||
```
|
||||
|
||||
### Step C：修改代码文件
|
||||
|
||||
运行以下四个 cell 完成代码适配：
|
||||
|
||||
```python
|
||||
# Step C-1：更新 __init__.py
|
||||
path = "/content/drive/MyDrive/lnp_project_moleculestm/lnp_ml/lnp_ml/modeling/encoders/__init__.py"
|
||||
new_content = """from lnp_ml.modeling.encoders.rdkit_encoder import CachedRDKitEncoder
|
||||
from lnp_ml.modeling.encoders.mpnn_encoder import CachedMPNNEncoder
|
||||
from .llm_encoder import LLMEncoder
|
||||
from .moleculestm_encoder import MoleculeSTMEncoder
|
||||
|
||||
__all__ = ["CachedRDKitEncoder", "CachedMPNNEncoder", "LLMEncoder", "MoleculeSTMEncoder"]
|
||||
"""
|
||||
with open(path, 'w') as f:
|
||||
f.write(new_content)
|
||||
print("__init__.py ✓")
|
||||
```
|
||||
|
||||
```python
|
||||
# Step C-2：更新 models.py
|
||||
path = "/content/drive/MyDrive/lnp_project_moleculestm/lnp_ml/lnp_ml/modeling/models.py"
|
||||
with open(path, 'r') as f:
|
||||
content = f.read()
|
||||
content = content.replace(
|
||||
"from lnp_ml.modeling.encoders import CachedRDKitEncoder, CachedMPNNEncoder, LLMEncoder",
|
||||
"from lnp_ml.modeling.encoders import CachedRDKitEncoder, CachedMPNNEncoder, LLMEncoder, MoleculeSTMEncoder"
|
||||
)
|
||||
content = content.replace(
|
||||
"self.llm_encoder = LLMEncoder(",
|
||||
"self.llm_encoder = MoleculeSTMEncoder("
|
||||
)
|
||||
content = content.replace(
|
||||
"model_path=llm_model_path,",
|
||||
"checkpoint_path=llm_model_path,"
|
||||
)
|
||||
content = content.replace(
|
||||
"llm_checkpoint_path=llm_model_path,",
|
||||
"llm_model_path=llm_model_path,"
|
||||
)
|
||||
with open(path, 'w') as f:
|
||||
f.write(content)
|
||||
print("models.py ✓")
|
||||
```
|
||||
|
||||
```python
|
||||
# Step C-3：修复 verify 脚本
|
||||
path = "/content/drive/MyDrive/lnp_project_moleculestm/lnp_ml/verify_llm_encoder.py"
|
||||
with open(path, 'r') as f:
|
||||
content = f.read()
|
||||
content = content.replace(
|
||||
"from lnp_ml.modeling.encoders.llm_encoder import LLMEncoder",
|
||||
"from lnp_ml.modeling.encoders.moleculestm_encoder import MoleculeSTMEncoder as LLMEncoder"
|
||||
)
|
||||
content = content.replace(
|
||||
"encoder = LLMEncoder(\n model_path=model_path,",
|
||||
"encoder = LLMEncoder(\n checkpoint_path=model_path,"
|
||||
)
|
||||
content = content.replace(
|
||||
"encoder._llm.named_parameters()",
|
||||
"encoder._gnn.named_parameters()"
|
||||
)
|
||||
content = content.replace(
|
||||
"model.llm_encoder._llm.parameters()",
|
||||
"model.llm_encoder._gnn.parameters()"
|
||||
)
|
||||
with open(path, 'w') as f:
|
||||
f.write(content)
|
||||
print("verify_llm_encoder.py ✓")
|
||||
```
|
||||
|
||||
```python
|
||||
# Step C-4：修复 pretrain.py
|
||||
path = "/content/drive/MyDrive/lnp_project_moleculestm/lnp_ml/lnp_ml/modeling/pretrain.py"
|
||||
with open(path, 'r') as f:
|
||||
content = f.read()
|
||||
content = content.replace(
|
||||
"llm_checkpoint_path=llm_model_path,",
|
||||
"llm_model_path=llm_model_path,"
|
||||
)
|
||||
with open(path, 'w') as f:
|
||||
f.write(content)
|
||||
print("pretrain.py ✓")
|
||||
```
|
||||
|
||||
完成后从 Cell 5 开始正常运行。
|
||||
|
||||
304
docs/colab_molt5_quickstart.md
Normal file
304
docs/colab_molt5_quickstart.md
Normal file
@ -0,0 +1,304 @@
|
||||
# LNP-ML + MolT5 快速上手指南
|
||||
|
||||
适用于：拿到代码包、从未运行过任何代码的新组员。
|
||||
|
||||
论文：Edwards et al., "Translation between Molecules and Natural Language", EMNLP 2022
|
||||
模型：laituan245/molt5-base（109M 参数，无需手动下载，训练时自动从 HuggingFace 拉取）
|
||||
|
||||
---
|
||||
|
||||
## 前提条件（开始前确认）
|
||||
|
||||
- Google Drive 的 `MyDrive/lnp_project_molt5/` 下已有以下内容：
|
||||
- `lnp_ml/`：项目代码（从 BioT5 版本复制来的，已修改）
|
||||
- MolT5 权重无需提前下载，训练时自动拉取
|
||||
|
||||
如果以上文件夹不存在，先看文末的【从零开始】章节。
|
||||
|
||||
---
|
||||
|
||||
## 每次重新打开 Colab 都需要运行
|
||||
|
||||
### Cell 1：挂载 Drive
|
||||
|
||||
```python
|
||||
from google.colab import drive
|
||||
drive.mount('/content/drive')
|
||||
```
|
||||
|
||||
> 弹出授权页面时，确认选择有 `lnp_project_molt5` 的 Google 账号。
|
||||
|
||||
---
|
||||
|
||||
### Cell 2：安装依赖
|
||||
|
||||
```python
|
||||
!pip install rdkit loguru transformers sentencepiece optuna -q
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Cell 3：配置路径
|
||||
|
||||
```python
|
||||
import os, sys
|
||||
|
||||
LNP_PATH = "/content/drive/MyDrive/lnp_project_molt5/lnp_ml"
|
||||
|
||||
sys.path.insert(0, LNP_PATH)
|
||||
os.chdir(LNP_PATH)
|
||||
print("路径配置完成 ✓")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Cell 4：确认文件完整
|
||||
|
||||
```python
|
||||
checks = [
|
||||
f"{LNP_PATH}/lnp_ml/modeling/models.py",
|
||||
f"{LNP_PATH}/lnp_ml/modeling/encoders/molt5_encoder.py",
|
||||
f"{LNP_PATH}/lnp_ml/modeling/encoders/__init__.py",
|
||||
f"{LNP_PATH}/data/processed/train.parquet",
|
||||
f"{LNP_PATH}/data/processed/val.parquet",
|
||||
f"{LNP_PATH}/data/processed/test.parquet",
|
||||
f"{LNP_PATH}/data/processed/train_pretrain.parquet",
|
||||
f"{LNP_PATH}/data/processed/val_pretrain.parquet",
|
||||
]
|
||||
for f in checks:
|
||||
print(f"{'✓' if os.path.exists(f) else '✗ 缺失'} {f}")
|
||||
```
|
||||
|
||||
**所有文件都显示 ✓ 再继续。**
|
||||
|
||||
---
|
||||
|
||||
## 正式流程
|
||||
|
||||
### Cell 5：验证 MolT5 模块
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python {LNP_PATH}/verify_llm_encoder.py \
|
||||
--model_path "laituan245/molt5-base" \
|
||||
--lnp_repo_path "{LNP_PATH}"
|
||||
```
|
||||
|
||||
**正常输出：**
|
||||
```
|
||||
TEST 1 PASSED ✓
|
||||
TEST 2 PASSED ✓
|
||||
TEST 3 PASSED ✓
|
||||
TEST 4 PASSED ✓
|
||||
ALL TESTS PASSED ✓
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Cell 6：预训练（约 10 分钟）
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python -m lnp_ml.modeling.pretrain main \
|
||||
--train-path data/processed/train_pretrain.parquet \
|
||||
--val-path data/processed/val_pretrain.parquet \
|
||||
--epochs 50 \
|
||||
--lr 1e-4 \
|
||||
--device cuda
|
||||
```
|
||||
|
||||
完成后确认权重存在：
|
||||
```python
|
||||
print("✓ 预训练权重存在" if os.path.exists(f"{LNP_PATH}/models/pretrain_delivery.pt")
|
||||
else "✗ 预训练权重未生成")
|
||||
```
|
||||
|
||||
参考结果：`Best val_loss: 0.5867`
|
||||
|
||||
---
|
||||
|
||||
### Cell 7：正式训练（约 20 分钟）
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python lnp_ml/modeling/final_train_optuna_cv.py \
|
||||
--init-from-pretrain models/pretrain_delivery.pt \
|
||||
--use-llm \
|
||||
--llm-model-path "laituan245/molt5-base" \
|
||||
--llm-device cuda \
|
||||
--n-trials 20 \
|
||||
--epochs-per-trial 30 \
|
||||
--seed 42 \
|
||||
--device cuda \
|
||||
--output-dir models/final
|
||||
```
|
||||
|
||||
完成后确认：
|
||||
```python
|
||||
print("✓ 模型权重存在" if os.path.exists(f"{LNP_PATH}/models/final/model.pt")
|
||||
else "✗ 模型未生成")
|
||||
```
|
||||
|
||||
参考结果：`Best val_loss: 2.3131`
|
||||
|
||||
---
|
||||
|
||||
### Cell 8：测试评估
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python lnp_ml/modeling/predict.py test \
|
||||
--test-path data/processed/test.parquet \
|
||||
--model-path models/final/model.pt \
|
||||
--output-path models/final/test_results.json
|
||||
|
||||
import json
|
||||
with open(f'{LNP_PATH}/models/final/test_results.json') as f:
|
||||
results = json.load(f)
|
||||
|
||||
print("=== 分类任务 ===")
|
||||
for task in ['pdi', 'ee', 'toxic']:
|
||||
m = results['detailed_metrics'][task]
|
||||
print(f" {task}: acc={m['accuracy']:.4f}, f1={m['f1']:.4f}")
|
||||
|
||||
print("\n=== 回归任务 ===")
|
||||
for task in ['size', 'delivery']:
|
||||
m = results['detailed_metrics'][task]
|
||||
print(f" {task}: R²={m['r2']:.4f}, RMSE={m['rmse']:.4f}")
|
||||
|
||||
print("\n=== 分布任务 ===")
|
||||
m = results['detailed_metrics']['biodist']
|
||||
print(f" biodist: KL={m['kl_divergence']:.4f}, JS={m['js_divergence']:.4f}")
|
||||
```
|
||||
|
||||
**参考指标：**
|
||||
|
||||
| 任务 | 指标 | 参考值 |
|
||||
|---|---|---|
|
||||
| delivery | R² | ~0.60 |
|
||||
| size | R² | ~0.44 |
|
||||
| pdi | acc | ~0.75 |
|
||||
| ee | acc | ~0.70 |
|
||||
| toxic | acc | ~0.98 |
|
||||
| biodist | KL | ~0.73 |
|
||||
|
||||
---
|
||||
|
||||
### Cell 9：备份模型
|
||||
|
||||
```python
|
||||
!cp -r {LNP_PATH}/models/final \
|
||||
{LNP_PATH}/models/final_backup
|
||||
print("备份完成 ✓")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 如果已有训练好的模型，直接从 Cell 8 开始
|
||||
|
||||
只需运行 Cell 1 → Cell 2 → Cell 3 → Cell 8，跳过 Cell 4-7。
|
||||
|
||||
---
|
||||
|
||||
## 常见问题
|
||||
|
||||
| 报错 | 原因 | 解决方法 |
|
||||
|---|---|---|
|
||||
| `No module named 'rdkit'` | 依赖未安装 | 重新运行 Cell 2 |
|
||||
| `No module named 'lnp_ml'` | 路径未配置 | 重新运行 Cell 3 |
|
||||
| `'NoneType' has no attribute 'Study'` | optuna 未安装 | 重新运行 Cell 2 |
|
||||
| `Drive not mounted` | Drive 未挂载 | 重新运行 Cell 1 |
|
||||
| `Mountpoint must not already contain files` | Drive 已挂载但缓存混乱 | 菜单栏 → 运行时 → 重新启动运行时，再从 Cell 1 开始 |
|
||||
| Colab 断开后重连 | 环境变量丢失 | 从 Cell 1 重跑，Cell 6/7 已有权重可跳过 |
|
||||
|
||||
---
|
||||
|
||||
## 【从零开始】如果 lnp_project_molt5 不存在
|
||||
|
||||
### Step A：从 BioT5 版本复制代码
|
||||
|
||||
```python
|
||||
import os
|
||||
os.makedirs("/content/drive/MyDrive/lnp_project_molt5", exist_ok=True)
|
||||
|
||||
!cp -r "/content/drive/MyDrive/lnp_project_biot5-plus/lnp_ml" \
|
||||
"/content/drive/MyDrive/lnp_project_molt5/lnp_ml"
|
||||
print("代码复制完成 ✓")
|
||||
```
|
||||
|
||||
### Step B：上传 molt5_encoder.py
|
||||
|
||||
```python
|
||||
from google.colab import files
|
||||
import shutil
|
||||
|
||||
uploaded = files.upload() # 选择 molt5_encoder.py
|
||||
|
||||
shutil.copy(
|
||||
'molt5_encoder.py',
|
||||
'/content/drive/MyDrive/lnp_project_molt5/lnp_ml/lnp_ml/modeling/encoders/molt5_encoder.py'
|
||||
)
|
||||
print("上传完成 ✓")
|
||||
```
|
||||
|
||||
### Step C：修改代码文件
|
||||
|
||||
```python
|
||||
# Step C-1：更新 __init__.py
|
||||
path = "/content/drive/MyDrive/lnp_project_molt5/lnp_ml/lnp_ml/modeling/encoders/__init__.py"
|
||||
new_content = """from lnp_ml.modeling.encoders.rdkit_encoder import CachedRDKitEncoder
|
||||
from lnp_ml.modeling.encoders.mpnn_encoder import CachedMPNNEncoder
|
||||
from .llm_encoder import LLMEncoder
|
||||
from .molt5_encoder import MolT5Encoder
|
||||
|
||||
__all__ = ["CachedRDKitEncoder", "CachedMPNNEncoder", "LLMEncoder", "MolT5Encoder"]
|
||||
"""
|
||||
with open(path, 'w') as f:
|
||||
f.write(new_content)
|
||||
print("__init__.py ✓")
|
||||
```
|
||||
|
||||
```python
|
||||
# Step C-2：更新 models.py
|
||||
path = "/content/drive/MyDrive/lnp_project_molt5/lnp_ml/lnp_ml/modeling/models.py"
|
||||
with open(path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
content = content.replace(
|
||||
"from lnp_ml.modeling.encoders import CachedRDKitEncoder, CachedMPNNEncoder, LLMEncoder",
|
||||
"from lnp_ml.modeling.encoders import CachedRDKitEncoder, CachedMPNNEncoder, LLMEncoder, MolT5Encoder"
|
||||
)
|
||||
content = content.replace(
|
||||
"self.llm_encoder = LLMEncoder(",
|
||||
"self.llm_encoder = MolT5Encoder("
|
||||
)
|
||||
with open(path, 'w') as f:
|
||||
f.write(content)
|
||||
print("models.py ✓")
|
||||
```
|
||||
|
||||
```python
|
||||
# Step C-3：修复 verify 脚本
|
||||
path = "/content/drive/MyDrive/lnp_project_molt5/lnp_ml/verify_llm_encoder.py"
|
||||
with open(path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
content = content.replace(
|
||||
"from lnp_ml.modeling.encoders.llm_encoder import LLMEncoder",
|
||||
"from lnp_ml.modeling.encoders.molt5_encoder import MolT5Encoder as LLMEncoder"
|
||||
)
|
||||
content = content.replace(
|
||||
"encoder._llm.named_parameters()",
|
||||
"encoder._encoder.named_parameters()"
|
||||
)
|
||||
content = content.replace(
|
||||
"model.llm_encoder._llm.parameters()",
|
||||
"model.llm_encoder._encoder.parameters()"
|
||||
)
|
||||
with open(path, 'w') as f:
|
||||
f.write(content)
|
||||
print("verify_llm_encoder.py ✓")
|
||||
```
|
||||
|
||||
完成后从 Cell 5 开始正常运行。
|
||||
|
||||
323
docs/colab_selfiested_quickstart.md
Normal file
323
docs/colab_selfiested_quickstart.md
Normal file
@ -0,0 +1,323 @@
|
||||
# LNP-ML + SELFIES-TED 快速上手指南
|
||||
|
||||
适用于：拿到代码包、从未运行过任何代码的新组员。
|
||||
|
||||
论文：IBM, "A Large Encoder-Decoder Family of Foundation Models for Chemical Language" (2024-2025)
|
||||
模型：ibm/materials.selfies-ted（358M 参数，BART 架构，无需手动下载，训练时自动从 HuggingFace 拉取）
|
||||
|
||||
---
|
||||
|
||||
## ⚠️ 与前几个模型的三个关键区别
|
||||
|
||||
1. **必须用 GPU 运行时**（模型 358M，CPU 会非常慢）。开始前先确认硬件加速器选了 GPU。
|
||||
2. **需要额外安装 selfies 库**（用于 SMILES → SELFIES 转换）。
|
||||
3. **hidden size 是 1024**（前几个模型是 768 或 300），这已经在 encoder 代码里处理好了。
|
||||
|
||||
---
|
||||
|
||||
## 前提条件（开始前确认）
|
||||
|
||||
- Google Drive 的 `MyDrive/lnp_project_selfiested/` 下已有：
|
||||
- `lnp_ml/`：项目代码（从 BioT5 版本复制来的，已修改）
|
||||
- SELFIES-TED 权重无需提前下载，训练时自动拉取
|
||||
|
||||
如果文件夹不存在，先看文末的【从零开始】章节。
|
||||
|
||||
---
|
||||
|
||||
## 第一步：确认 GPU
|
||||
|
||||
菜单栏 → 代码执行程序 → 更改运行时类型 → 硬件加速器选 **GPU** → 保存。
|
||||
|
||||
然后运行：
|
||||
```python
|
||||
import torch
|
||||
print(f"GPU 可用: {torch.cuda.is_available()}")
|
||||
```
|
||||
显示 `True` 再继续。如果是 `False`，按上面的步骤切换 GPU。
|
||||
|
||||
---
|
||||
|
||||
## 每次重新打开 Colab 都需要运行
|
||||
|
||||
### Cell 1：挂载 Drive
|
||||
|
||||
```python
|
||||
from google.colab import drive
|
||||
drive.mount('/content/drive')
|
||||
```
|
||||
|
||||
> 弹出授权页面时，确认选择有 `lnp_project_selfiested` 的 Google 账号。
|
||||
|
||||
---
|
||||
|
||||
### Cell 2：安装依赖（注意多了 selfies）
|
||||
|
||||
```python
|
||||
!pip install rdkit loguru transformers sentencepiece optuna selfies -q
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Cell 3：配置路径
|
||||
|
||||
```python
|
||||
import os, sys
|
||||
|
||||
LNP_PATH = "/content/drive/MyDrive/lnp_project_selfiested/lnp_ml"
|
||||
|
||||
sys.path.insert(0, LNP_PATH)
|
||||
os.chdir(LNP_PATH)
|
||||
|
||||
import torch
|
||||
print(f"路径配置完成 ✓ | GPU 可用: {torch.cuda.is_available()}")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Cell 4：确认文件完整
|
||||
|
||||
```python
|
||||
checks = [
|
||||
f"{LNP_PATH}/lnp_ml/modeling/models.py",
|
||||
f"{LNP_PATH}/lnp_ml/modeling/encoders/selfiested_encoder.py",
|
||||
f"{LNP_PATH}/lnp_ml/modeling/encoders/__init__.py",
|
||||
f"{LNP_PATH}/data/processed/train.parquet",
|
||||
f"{LNP_PATH}/data/processed/val.parquet",
|
||||
f"{LNP_PATH}/data/processed/test.parquet",
|
||||
f"{LNP_PATH}/data/processed/train_pretrain.parquet",
|
||||
f"{LNP_PATH}/data/processed/val_pretrain.parquet",
|
||||
]
|
||||
for f in checks:
|
||||
print(f"{'✓' if os.path.exists(f) else '✗ 缺失'} {f}")
|
||||
```
|
||||
|
||||
**所有文件都显示 ✓ 再继续。**
|
||||
|
||||
---
|
||||
|
||||
## 正式流程
|
||||
|
||||
### Cell 5：验证 SELFIES-TED 模块
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python {LNP_PATH}/verify_llm_encoder.py \
|
||||
--model_path "ibm/materials.selfies-ted" \
|
||||
--lnp_repo_path "{LNP_PATH}"
|
||||
```
|
||||
|
||||
**正常输出：**
|
||||
```
|
||||
TEST 1 PASSED ✓
|
||||
TEST 2 PASSED ✓
|
||||
TEST 3 PASSED ✓
|
||||
TEST 4 PASSED ✓
|
||||
ALL TESTS PASSED ✓
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Cell 6：预训练（约 10 分钟）
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python -m lnp_ml.modeling.pretrain main \
|
||||
--train-path data/processed/train_pretrain.parquet \
|
||||
--val-path data/processed/val_pretrain.parquet \
|
||||
--epochs 50 \
|
||||
--lr 1e-4 \
|
||||
--device cuda
|
||||
```
|
||||
|
||||
完成后确认权重存在：
|
||||
```python
|
||||
print("✓ 预训练权重存在" if os.path.exists(f"{LNP_PATH}/models/pretrain_delivery.pt")
|
||||
else "✗ 预训练权重未生成")
|
||||
```
|
||||
|
||||
参考结果：`Best val_loss: 0.5659`
|
||||
|
||||
---
|
||||
|
||||
### Cell 7：正式训练（约 25-30 分钟，比前几个模型略慢）
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python lnp_ml/modeling/final_train_optuna_cv.py \
|
||||
--init-from-pretrain models/pretrain_delivery.pt \
|
||||
--use-llm \
|
||||
--llm-model-path "ibm/materials.selfies-ted" \
|
||||
--llm-device cuda \
|
||||
--n-trials 20 \
|
||||
--epochs-per-trial 30 \
|
||||
--seed 42 \
|
||||
--device cuda \
|
||||
--output-dir models/final
|
||||
```
|
||||
|
||||
完成后确认：
|
||||
```python
|
||||
print("✓ 模型权重存在" if os.path.exists(f"{LNP_PATH}/models/final/model.pt")
|
||||
else "✗ 模型未生成")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Cell 8：测试评估
|
||||
|
||||
```python
|
||||
!PYTHONPATH={LNP_PATH} \
|
||||
python lnp_ml/modeling/predict.py test \
|
||||
--test-path data/processed/test.parquet \
|
||||
--model-path models/final/model.pt \
|
||||
--output-path models/final/test_results.json
|
||||
|
||||
import json
|
||||
with open(f'{LNP_PATH}/models/final/test_results.json') as f:
|
||||
results = json.load(f)
|
||||
|
||||
print("=== 分类任务 ===")
|
||||
for task in ['pdi', 'ee', 'toxic']:
|
||||
m = results['detailed_metrics'][task]
|
||||
print(f" {task}: acc={m['accuracy']:.4f}, f1={m['f1']:.4f}")
|
||||
|
||||
print("\n=== 回归任务 ===")
|
||||
for task in ['size', 'delivery']:
|
||||
m = results['detailed_metrics'][task]
|
||||
print(f" {task}: R²={m['r2']:.4f}, RMSE={m['rmse']:.4f}")
|
||||
|
||||
print("\n=== 分布任务 ===")
|
||||
m = results['detailed_metrics']['biodist']
|
||||
print(f" biodist: KL={m['kl_divergence']:.4f}, JS={m['js_divergence']:.4f}")
|
||||
```
|
||||
|
||||
**参考指标：**
|
||||
|
||||
| 任务 | 指标 | 参考值 |
|
||||
|---|---|---|
|
||||
| delivery | R² | ~0.62 |
|
||||
| size | R² | ~0.38 |
|
||||
| pdi | acc | ~0.66 |
|
||||
| ee | acc | ~0.68 |
|
||||
| toxic | acc | ~0.94 |
|
||||
| biodist | KL | ~0.64 |
|
||||
|
||||
---
|
||||
|
||||
### Cell 9：备份模型
|
||||
|
||||
```python
|
||||
!cp -r {LNP_PATH}/models/final \
|
||||
{LNP_PATH}/models/final_backup
|
||||
print("备份完成 ✓")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 如果已有训练好的模型，直接从 Cell 8 开始
|
||||
|
||||
只需运行 Cell 1 → Cell 2 → Cell 3 → Cell 8，跳过 Cell 4-7。
|
||||
|
||||
---
|
||||
|
||||
## 常见问题
|
||||
|
||||
| 报错 | 原因 | 解决方法 |
|
||||
|---|---|---|
|
||||
| `Device: cpu` 或训练极慢 | 没选 GPU 运行时 | 菜单 → 更改运行时类型 → 选 GPU |
|
||||
| `No module named 'selfies'` | selfies 库未装 | 重新运行 Cell 2 |
|
||||
| `No module named 'rdkit'` | 依赖未安装 | 重新运行 Cell 2 |
|
||||
| `No module named 'lnp_ml'` | 路径未配置 | 重新运行 Cell 3 |
|
||||
| `'NoneType' has no attribute 'Study'` | optuna 未安装 | 重新运行 Cell 2 |
|
||||
| `Mountpoint must not already contain files` | Drive 缓存混乱 | 菜单 → 运行时 → 重新启动运行时，再从 Cell 1 开始 |
|
||||
| Colab 断开后重连 | 环境变量丢失 | 从 Cell 1 重跑，Cell 6/7 已有权重可跳过 |
|
||||
|
||||
---
|
||||
|
||||
## 【从零开始】如果 lnp_project_selfiested 不存在
|
||||
|
||||
### Step A：从 BioT5 版本复制代码
|
||||
|
||||
```python
|
||||
import os
|
||||
os.makedirs("/content/drive/MyDrive/lnp_project_selfiested", exist_ok=True)
|
||||
|
||||
!cp -r "/content/drive/MyDrive/lnp_project_biot5-plus/lnp_ml" \
|
||||
"/content/drive/MyDrive/lnp_project_selfiested/lnp_ml"
|
||||
print("代码复制完成 ✓")
|
||||
```
|
||||
|
||||
### Step B：上传 selfiested_encoder.py
|
||||
|
||||
```python
|
||||
from google.colab import files
|
||||
import shutil
|
||||
|
||||
uploaded = files.upload() # 选择 selfiested_encoder.py
|
||||
|
||||
shutil.copy(
|
||||
'selfiested_encoder.py',
|
||||
'/content/drive/MyDrive/lnp_project_selfiested/lnp_ml/lnp_ml/modeling/encoders/selfiested_encoder.py'
|
||||
)
|
||||
print("上传完成 ✓")
|
||||
```
|
||||
|
||||
### Step C：修改代码文件
|
||||
|
||||
```python
|
||||
# Step C-1：更新 __init__.py
|
||||
path = "/content/drive/MyDrive/lnp_project_selfiested/lnp_ml/lnp_ml/modeling/encoders/__init__.py"
|
||||
with open(path, 'w') as f:
|
||||
f.write("""from lnp_ml.modeling.encoders.rdkit_encoder import CachedRDKitEncoder
|
||||
from lnp_ml.modeling.encoders.mpnn_encoder import CachedMPNNEncoder
|
||||
from .llm_encoder import LLMEncoder
|
||||
from .selfiested_encoder import SELFIESTEDEncoder
|
||||
|
||||
__all__ = ["CachedRDKitEncoder", "CachedMPNNEncoder", "LLMEncoder", "SELFIESTEDEncoder"]
|
||||
""")
|
||||
print("__init__.py ✓")
|
||||
```
|
||||
|
||||
```python
|
||||
# Step C-2：更新 models.py
|
||||
path = "/content/drive/MyDrive/lnp_project_selfiested/lnp_ml/lnp_ml/modeling/models.py"
|
||||
with open(path, 'r') as f:
|
||||
content = f.read()
|
||||
content = content.replace(
|
||||
"from lnp_ml.modeling.encoders import CachedRDKitEncoder, CachedMPNNEncoder, LLMEncoder",
|
||||
"from lnp_ml.modeling.encoders import CachedRDKitEncoder, CachedMPNNEncoder, LLMEncoder, SELFIESTEDEncoder"
|
||||
)
|
||||
content = content.replace(
|
||||
"self.llm_encoder = LLMEncoder(",
|
||||
"self.llm_encoder = SELFIESTEDEncoder("
|
||||
)
|
||||
with open(path, 'w') as f:
|
||||
f.write(content)
|
||||
print("models.py ✓")
|
||||
```
|
||||
|
||||
```python
|
||||
# Step C-3：修复 verify 脚本（属性名是 _model）
|
||||
path = "/content/drive/MyDrive/lnp_project_selfiested/lnp_ml/verify_llm_encoder.py"
|
||||
with open(path, 'r') as f:
|
||||
content = f.read()
|
||||
content = content.replace(
|
||||
"from lnp_ml.modeling.encoders.llm_encoder import LLMEncoder",
|
||||
"from lnp_ml.modeling.encoders.selfiested_encoder import SELFIESTEDEncoder as LLMEncoder"
|
||||
)
|
||||
content = content.replace(
|
||||
"encoder._llm.named_parameters()",
|
||||
"encoder._model.named_parameters()"
|
||||
)
|
||||
content = content.replace(
|
||||
"model.llm_encoder._llm.parameters()",
|
||||
"model.llm_encoder._model.parameters()"
|
||||
)
|
||||
with open(path, 'w') as f:
|
||||
f.write(content)
|
||||
print("verify_llm_encoder.py ✓")
|
||||
```
|
||||
|
||||
完成后从【第一步：确认 GPU】开始正常运行。
|
||||
|
||||
@ -1,5 +1,34 @@
|
||||
from lnp_ml.modeling.encoders.rdkit_encoder import CachedRDKitEncoder
|
||||
from lnp_ml.modeling.encoders.mpnn_encoder import CachedMPNNEncoder
|
||||
|
||||
__all__ = ["CachedRDKitEncoder", "CachedMPNNEncoder"]
|
||||
# ============ LLM / 分子编码器（可选） ============
|
||||
# 这些 encoder 在使用 use_llm=True 时才会被实例化。
|
||||
# 为避免在未安装对应依赖（如 torch_geometric、selfies）时导入即报错，
|
||||
# 采用惰性导入：实际类在 build_llm_encoder() 内按需 import。
|
||||
|
||||
__all__ = [
|
||||
"CachedRDKitEncoder",
|
||||
"CachedMPNNEncoder",
|
||||
# LLM encoders（按需惰性导入，见 models.build_llm_encoder）
|
||||
"LLMEncoder",
|
||||
"MolT5Encoder",
|
||||
"MoleculeSTMEncoder",
|
||||
"SELFIESTEDEncoder",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
"""惰性导入 LLM encoder 类，避免未装可选依赖时 import 失败。"""
|
||||
if name == "LLMEncoder":
|
||||
from lnp_ml.modeling.encoders.llm_encoder import LLMEncoder
|
||||
return LLMEncoder
|
||||
if name == "MolT5Encoder":
|
||||
from lnp_ml.modeling.encoders.molt5_encoder import MolT5Encoder
|
||||
return MolT5Encoder
|
||||
if name == "MoleculeSTMEncoder":
|
||||
from lnp_ml.modeling.encoders.moleculestm_encoder import MoleculeSTMEncoder
|
||||
return MoleculeSTMEncoder
|
||||
if name == "SELFIESTEDEncoder":
|
||||
from lnp_ml.modeling.encoders.selfiested_encoder import SELFIESTEDEncoder
|
||||
return SELFIESTEDEncoder
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
226
lnp_ml/modeling/encoders/llm_encoder.py
Normal file
226
lnp_ml/modeling/encoders/llm_encoder.py
Normal file
@ -0,0 +1,226 @@
|
||||
"""
|
||||
BioT5-plus LLM 编码器
|
||||
|
||||
功能：SMILES -> BioT5-plus encoder (frozen) -> mean pooling -> Linear -> [B, d_model]
|
||||
产出 V5 token，追加到 LNPModel 的 attention pooling 序列中。
|
||||
|
||||
放置位置：lnp_ml/modeling/encoders/llm_encoder.py
|
||||
|
||||
依赖：
|
||||
pip install transformers sentencepiece
|
||||
|
||||
加载方式（两种都支持）：
|
||||
# 方式 A：本地路径（推荐，Colab 挂载 Drive 后使用）
|
||||
LLMEncoder(model_path="/content/drive/MyDrive/biot5-plus-base")
|
||||
|
||||
# 方式 B：HuggingFace 在线（需要网络）
|
||||
LLMEncoder(model_path="QizhiPei/biot5-plus-base")
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# BioT5-plus 分子输入格式：<bom> SMILES/SELFIES <eom>
|
||||
BOM_TOKEN = "<bom>"
|
||||
EOM_TOKEN = "<eom>"
|
||||
|
||||
|
||||
class LLMEncoder(nn.Module):
|
||||
"""
|
||||
基于 BioT5-plus 的冻结分子编码器。
|
||||
|
||||
流程：
|
||||
SMILES
|
||||
-> "<bom>" + SMILES + "<eom>"
|
||||
-> T5Tokenizer（BioT5-plus 自带完整词表，直接 from_pretrained）
|
||||
-> T5EncoderModel (frozen, no_grad)
|
||||
-> last_hidden_state [B, seq_len, hidden_size]
|
||||
-> attention-mask mean pooling -> [B, hidden_size]
|
||||
-> LayerNorm + Linear + ReLU + Dropout
|
||||
-> [B, d_model] ← V5 token
|
||||
|
||||
Args:
|
||||
model_path : BioT5-plus 本地路径或 HuggingFace ID
|
||||
本地示例："content/drive/MyDrive/biot5-plus-base"
|
||||
在线示例："QizhiPei/biot5-plus-base"
|
||||
d_model : 输出维度，与 LNPModel.d_model 一致，默认 256
|
||||
dropout : projection 层 dropout，默认 0.1
|
||||
device : 推理设备，Colab 上用 "cuda"
|
||||
max_length : tokenizer 截断长度，默认 512
|
||||
use_cache : 是否缓存已编码 SMILES 的向量，避免重复推理
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str,
|
||||
d_model: int = 256,
|
||||
dropout: float = 0.1,
|
||||
device: str = "cuda",
|
||||
max_length: int = 512,
|
||||
use_cache: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.model_path = model_path
|
||||
self.d_model = d_model
|
||||
self.dropout_rate = dropout
|
||||
self.max_length = max_length
|
||||
self.use_cache = use_cache
|
||||
self._device_str = device
|
||||
|
||||
# LLM（延迟加载，首次 forward 时才初始化）
|
||||
self._tokenizer = None
|
||||
self._llm = None
|
||||
self._hidden_size: Optional[int] = None
|
||||
self._initialized: bool = False
|
||||
|
||||
# SMILES -> pooled numpy 向量的内存缓存
|
||||
self._cache: Dict[str, np.ndarray] = {}
|
||||
|
||||
# projection 层在 _lazy_init 里直接赋值注册
|
||||
# 注意：不能在这里 self.projection = None，否则 add_module 会冲突
|
||||
# 用 _projection_ready 标记是否已初始化
|
||||
self._projection_ready: bool = False
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 延迟初始化 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _lazy_init(self) -> None:
|
||||
"""首次 forward 时加载模型，避免导入时就占满显存。"""
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
logger.info(f"[LLMEncoder] Loading BioT5-plus from: {self.model_path}")
|
||||
from transformers import T5Tokenizer, T5EncoderModel
|
||||
|
||||
# BioT5-plus 的 tokenizer checkpoint 里已内置完整词表（含 SELFIES token）
|
||||
# 直接 from_pretrained 即可，不需要手动 add_tokens
|
||||
self._tokenizer = T5Tokenizer.from_pretrained(
|
||||
self.model_path,
|
||||
model_max_length=self.max_length,
|
||||
)
|
||||
|
||||
# encoder-only 模型（比 T5ForConditionalGeneration 省约一半显存）
|
||||
self._llm = T5EncoderModel.from_pretrained(self.model_path)
|
||||
self._llm = self._llm.to(self._device_str)
|
||||
self._llm.eval()
|
||||
|
||||
# 完全冻结 LLM，不参与梯度更新
|
||||
for param in self._llm.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
self._hidden_size = self._llm.config.d_model
|
||||
logger.info(
|
||||
f"[LLMEncoder] hidden_size={self._hidden_size}, "
|
||||
f"vocab_size={len(self._tokenizer)}, frozen=True"
|
||||
)
|
||||
|
||||
# projection 层（唯一可训练的部分）
|
||||
# 直接赋值给 self.projection，PyTorch 会自动注册为子模块
|
||||
self.projection = nn.Sequential(
|
||||
nn.LayerNorm(self._hidden_size),
|
||||
nn.Linear(self._hidden_size, self.d_model),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(self.dropout_rate),
|
||||
).to(self._device_str)
|
||||
|
||||
self._projection_ready = True
|
||||
self._initialized = True
|
||||
logger.info("[LLMEncoder] Ready.")
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 内部编码 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
@torch.no_grad()
|
||||
def _encode_batch_raw(self, smiles_list: List[str]) -> np.ndarray:
|
||||
"""
|
||||
对一批 SMILES 过 BioT5-plus encoder，返回 mean-pooled numpy 向量。
|
||||
|
||||
Returns:
|
||||
np.ndarray shape [N, hidden_size] dtype float32
|
||||
"""
|
||||
# 加分子起止标记，与 BioT5-plus 预训练格式一致
|
||||
formatted = [f"{BOM_TOKEN}{s}{EOM_TOKEN}" for s in smiles_list]
|
||||
|
||||
encoded = self._tokenizer(
|
||||
formatted,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.max_length,
|
||||
return_tensors="pt",
|
||||
)
|
||||
input_ids = encoded["input_ids"].to(self._device_str) # [B, L]
|
||||
attention_mask = encoded["attention_mask"].to(self._device_str) # [B, L]
|
||||
|
||||
# LLM 推理（@no_grad 已在装饰器处理）
|
||||
outputs = self._llm(input_ids=input_ids, attention_mask=attention_mask)
|
||||
hidden = outputs.last_hidden_state # [B, L, H]
|
||||
|
||||
# attention-mask mean pooling（忽略 padding token）
|
||||
mask = attention_mask.unsqueeze(-1).float() # [B, L, 1]
|
||||
pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9) # [B, H]
|
||||
|
||||
return pooled.cpu().numpy().astype(np.float32)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# forward #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def forward(self, smiles_list: List[str]) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
smiles_list: 长度为 B 的 SMILES 字符串列表
|
||||
|
||||
Returns:
|
||||
torch.Tensor shape [B, d_model]，在模型所在设备上
|
||||
"""
|
||||
self._lazy_init()
|
||||
|
||||
# 找出未缓存的 SMILES
|
||||
to_compute = (
|
||||
[s for s in smiles_list if s not in self._cache]
|
||||
if self.use_cache
|
||||
else list(smiles_list)
|
||||
)
|
||||
|
||||
# 批量计算并写入缓存
|
||||
if to_compute:
|
||||
new_vecs = self._encode_batch_raw(to_compute) # [N, H]
|
||||
if self.use_cache:
|
||||
for smi, vec in zip(to_compute, new_vecs):
|
||||
self._cache[smi] = vec
|
||||
|
||||
# 按原始顺序组装结果
|
||||
pooled_np = (
|
||||
np.stack([self._cache[s] for s in smiles_list])
|
||||
if self.use_cache
|
||||
else new_vecs
|
||||
)
|
||||
|
||||
pooled = torch.from_numpy(pooled_np).to(self._device_str) # [B, H]
|
||||
return self.projection(pooled) # [B, d_model]
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 工具属性 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
@property
|
||||
def hidden_size(self) -> int:
|
||||
self._lazy_init()
|
||||
return self._hidden_size
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
"""清空 SMILES 向量缓存。"""
|
||||
self._cache.clear()
|
||||
|
||||
@property
|
||||
def cache_size(self) -> int:
|
||||
return len(self._cache)
|
||||
76
lnp_ml/modeling/encoders/molecule_gnn_model.py
Normal file
76
lnp_ml/modeling/encoders/molecule_gnn_model.py
Normal file
@ -0,0 +1,76 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch_geometric.nn import (MessagePassing, global_add_pool,
|
||||
global_max_pool, global_mean_pool)
|
||||
from torch_geometric.utils import degree
|
||||
from torch_scatter import scatter_add
|
||||
from ogb.graphproppred.mol_encoder import AtomEncoder, BondEncoder
|
||||
|
||||
|
||||
class GINConv(MessagePassing):
|
||||
def __init__(self, emb_dim, aggr="add"):
|
||||
super(GINConv, self).__init__(aggr=aggr)
|
||||
self.mlp = torch.nn.Sequential(
|
||||
torch.nn.Linear(emb_dim, 2*emb_dim),
|
||||
torch.nn.BatchNorm1d(2*emb_dim),
|
||||
torch.nn.ReLU(),
|
||||
torch.nn.Linear(2*emb_dim, emb_dim)
|
||||
)
|
||||
self.eps = torch.nn.Parameter(torch.Tensor([0]))
|
||||
self.bond_encoder = BondEncoder(emb_dim=emb_dim)
|
||||
|
||||
def forward(self, x, edge_index, edge_attr):
|
||||
edge_embedding = self.bond_encoder(edge_attr)
|
||||
out = self.mlp((1 + self.eps) * x + self.propagate(edge_index, x=x, edge_attr=edge_embedding))
|
||||
return out
|
||||
|
||||
def message(self, x_j, edge_attr):
|
||||
return F.relu(x_j + edge_attr)
|
||||
|
||||
def update(self, aggr_out):
|
||||
return aggr_out
|
||||
|
||||
|
||||
class GNN(nn.Module):
|
||||
def __init__(self, num_layer, emb_dim, JK="last", drop_ratio=0., gnn_type="gin"):
|
||||
if num_layer < 2:
|
||||
raise ValueError("Number of GNN layers must be greater than 1.")
|
||||
super(GNN, self).__init__()
|
||||
self.drop_ratio = drop_ratio
|
||||
self.num_layer = num_layer
|
||||
self.JK = JK
|
||||
self.atom_encoder = AtomEncoder(emb_dim)
|
||||
self.gnns = nn.ModuleList()
|
||||
for layer in range(num_layer):
|
||||
if gnn_type == "gin":
|
||||
self.gnns.append(GINConv(emb_dim, aggr="add"))
|
||||
self.batch_norms = nn.ModuleList()
|
||||
for layer in range(num_layer):
|
||||
self.batch_norms.append(nn.BatchNorm1d(emb_dim))
|
||||
|
||||
def forward(self, *argv):
|
||||
if len(argv) == 3:
|
||||
x, edge_index, edge_attr = argv[0], argv[1], argv[2]
|
||||
elif len(argv) == 1:
|
||||
data = argv[0]
|
||||
x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
|
||||
else:
|
||||
raise ValueError("unmatched number of arguments.")
|
||||
x = self.atom_encoder(x)
|
||||
h_list = [x]
|
||||
for layer in range(self.num_layer):
|
||||
h = self.gnns[layer](h_list[layer], edge_index, edge_attr)
|
||||
h = self.batch_norms[layer](h)
|
||||
if layer == self.num_layer - 1:
|
||||
h = F.dropout(h, self.drop_ratio, training=self.training)
|
||||
else:
|
||||
h = F.dropout(F.relu(h), self.drop_ratio, training=self.training)
|
||||
h_list.append(h)
|
||||
if self.JK == "last":
|
||||
return h_list[-1]
|
||||
elif self.JK == "sum":
|
||||
h_list = [h.unsqueeze_(0) for h in h_list]
|
||||
return torch.sum(torch.cat(h_list, dim=0), dim=0)[0]
|
||||
else:
|
||||
return h_list[-1]
|
||||
261
lnp_ml/modeling/encoders/moleculestm_encoder.py
Normal file
261
lnp_ml/modeling/encoders/moleculestm_encoder.py
Normal file
@ -0,0 +1,261 @@
|
||||
"""
|
||||
MoleculeSTM GNN 编码器
|
||||
|
||||
功能：SMILES -> 分子图 -> GNN (frozen) -> global mean pooling -> Linear -> [B, d_model]
|
||||
产出 V5 token，追加到 LNPModel 的 attention pooling 序列中。
|
||||
|
||||
放置位置：lnp_ml/modeling/encoders/moleculestm_encoder.py
|
||||
|
||||
依赖：
|
||||
pip install torch_geometric torch_scatter torch_sparse ogb
|
||||
|
||||
使用方式（替换 llm_encoder.py，接口完全一致）：
|
||||
encoder = MoleculeSTMEncoder(
|
||||
checkpoint_path="/content/drive/MyDrive/lnp_project_moleculestm/moleculestm_weights/pretrained_MoleculeSTM/SciBERT-Graph-3e-5-1-1e-4-1-InfoNCE-0.1-32-32/molecule_model.pth",
|
||||
d_model=256,
|
||||
device="cuda",
|
||||
)
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch_geometric.data import Data, Batch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# GNN 默认参数（与 MoleculeSTM 预训练时一致）
|
||||
GNN_NUM_LAYER = 5
|
||||
GNN_EMB_DIM = 300
|
||||
GNN_JK = "last"
|
||||
GNN_DROP_RATIO = 0.0
|
||||
GNN_TYPE = "gin"
|
||||
|
||||
|
||||
def smiles_to_graph(smiles: str) -> Optional[Data]:
|
||||
"""
|
||||
SMILES -> PyG Data 对象（原子特征 + 键特征）。
|
||||
使用 OGB 的原子/键编码方式，与 MoleculeSTM 预训练时一致。
|
||||
"""
|
||||
from rdkit import Chem
|
||||
from ogb.utils.mol import smiles2graph
|
||||
|
||||
try:
|
||||
graph = smiles2graph(smiles)
|
||||
if graph is None:
|
||||
return None
|
||||
|
||||
x = torch.tensor(graph['node_feat'], dtype=torch.long)
|
||||
edge_index = torch.tensor(graph['edge_index'], dtype=torch.long)
|
||||
edge_attr = torch.tensor(graph['edge_feat'], dtype=torch.long)
|
||||
|
||||
return Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to convert SMILES to graph: {smiles}, error: {e}")
|
||||
return None
|
||||
|
||||
|
||||
class MoleculeSTMEncoder(nn.Module):
|
||||
"""
|
||||
基于 MoleculeSTM GNN 的冻结分子编码器。
|
||||
|
||||
流程：
|
||||
SMILES
|
||||
-> RDKit + OGB 转分子图 Data
|
||||
-> Batch 打包
|
||||
-> GNN (frozen, 5层 GIN, emb_dim=300)
|
||||
-> global mean pooling -> [B, 300]
|
||||
-> LayerNorm + Linear(300→d_model) + ReLU + Dropout
|
||||
-> [B, d_model] ← V5 token
|
||||
|
||||
Args:
|
||||
checkpoint_path : molecule_model.pth 的路径
|
||||
d_model : 输出维度，与 LNPModel.d_model 一致，默认 256
|
||||
dropout : projection 层 dropout，默认 0.1
|
||||
device : 推理设备，Colab 上用 "cuda"
|
||||
use_cache : 是否缓存已编码 SMILES 的向量
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
checkpoint_path: str,
|
||||
d_model: int = 256,
|
||||
dropout: float = 0.1,
|
||||
device: str = "cuda",
|
||||
use_cache: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.checkpoint_path = checkpoint_path
|
||||
self.d_model = d_model
|
||||
self.dropout_rate = dropout
|
||||
self._device_str = device
|
||||
self.use_cache = use_cache
|
||||
|
||||
# GNN（延迟加载）
|
||||
self._gnn = None
|
||||
self._initialized: bool = False
|
||||
|
||||
# SMILES -> pooled numpy 向量的内存缓存
|
||||
self._cache: Dict[str, np.ndarray] = {}
|
||||
|
||||
# projection 层在 _lazy_init 后赋值
|
||||
self.projection: Optional[nn.Sequential] = None
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 延迟初始化 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _lazy_init(self) -> None:
|
||||
"""首次 forward 时加载 GNN 权重。"""
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
logger.info(f"[MoleculeSTMEncoder] Loading GNN from: {self.checkpoint_path}")
|
||||
|
||||
from lnp_ml.modeling.encoders.molecule_gnn_model import GNN
|
||||
|
||||
# 初始化 GNN
|
||||
self._gnn = GNN(
|
||||
num_layer=GNN_NUM_LAYER,
|
||||
emb_dim=GNN_EMB_DIM,
|
||||
JK=GNN_JK,
|
||||
drop_ratio=GNN_DROP_RATIO,
|
||||
gnn_type=GNN_TYPE,
|
||||
)
|
||||
|
||||
# 加载预训练权重
|
||||
# 权重保存在 GNN_graphpred 里，key 有 molecule_node_model. 前缀，需要去掉
|
||||
state_dict = torch.load(self.checkpoint_path, map_location="cpu")
|
||||
new_state_dict = {}
|
||||
for k, v in state_dict.items():
|
||||
if k.startswith("molecule_node_model."):
|
||||
new_state_dict[k[len("molecule_node_model."):]] = v
|
||||
# 跳过 size 不匹配的 key（OGB 版本差异导致原子特征维度不同）
|
||||
model_state = self._gnn.state_dict()
|
||||
filtered = {
|
||||
k: v for k, v in new_state_dict.items()
|
||||
if k in model_state and v.shape == model_state[k].shape
|
||||
}
|
||||
model_state.update(filtered)
|
||||
self._gnn.load_state_dict(model_state)
|
||||
logger.info(f"[MoleculeSTMEncoder] Loaded {len(filtered)}/{len(new_state_dict)} keys from checkpoint")
|
||||
self._gnn = self._gnn.to(self._device_str)
|
||||
self._gnn.eval()
|
||||
|
||||
# 完全冻结 GNN
|
||||
for param in self._gnn.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
logger.info(f"[MoleculeSTMEncoder] GNN loaded, emb_dim={GNN_EMB_DIM}, frozen=True")
|
||||
|
||||
# projection 层（唯一可训练的部分）
|
||||
self.projection = nn.Sequential(
|
||||
nn.LayerNorm(GNN_EMB_DIM),
|
||||
nn.Linear(GNN_EMB_DIM, self.d_model),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(self.dropout_rate),
|
||||
).to(self._device_str)
|
||||
|
||||
self._initialized = True
|
||||
logger.info("[MoleculeSTMEncoder] Ready.")
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 内部编码 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
@torch.no_grad()
|
||||
def _encode_batch_raw(self, smiles_list: List[str]) -> np.ndarray:
|
||||
"""
|
||||
对一批 SMILES 过 GNN，返回 mean-pooled numpy 向量。
|
||||
|
||||
Returns:
|
||||
np.ndarray shape [N, GNN_EMB_DIM] dtype float32
|
||||
"""
|
||||
from torch_geometric.nn import global_mean_pool
|
||||
|
||||
# SMILES -> 分子图
|
||||
graphs = []
|
||||
valid_indices = []
|
||||
for i, smi in enumerate(smiles_list):
|
||||
g = smiles_to_graph(smi)
|
||||
if g is not None:
|
||||
graphs.append(g)
|
||||
valid_indices.append(i)
|
||||
else:
|
||||
logger.warning(f"Invalid SMILES skipped: {smi}")
|
||||
|
||||
if not graphs:
|
||||
raise ValueError("All SMILES in batch are invalid")
|
||||
|
||||
# 打包成 batch
|
||||
batch = Batch.from_data_list(graphs).to(self._device_str)
|
||||
|
||||
# GNN forward
|
||||
node_repr = self._gnn(batch.x, batch.edge_index, batch.edge_attr)
|
||||
|
||||
# Global mean pooling -> [N_valid, emb_dim]
|
||||
graph_repr = global_mean_pool(node_repr, batch.batch)
|
||||
|
||||
result = graph_repr.cpu().numpy().astype(np.float32)
|
||||
|
||||
# 如果有 invalid SMILES，用零向量填充
|
||||
if len(valid_indices) < len(smiles_list):
|
||||
full_result = np.zeros((len(smiles_list), GNN_EMB_DIM), dtype=np.float32)
|
||||
for out_idx, in_idx in enumerate(valid_indices):
|
||||
full_result[in_idx] = result[out_idx]
|
||||
return full_result
|
||||
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# forward #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def forward(self, smiles_list: List[str]) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
smiles_list: 长度为 B 的 SMILES 字符串列表
|
||||
|
||||
Returns:
|
||||
torch.Tensor shape [B, d_model]
|
||||
"""
|
||||
self._lazy_init()
|
||||
|
||||
# 找出未缓存的 SMILES
|
||||
to_compute = (
|
||||
[s for s in smiles_list if s not in self._cache]
|
||||
if self.use_cache
|
||||
else list(smiles_list)
|
||||
)
|
||||
|
||||
# 批量计算并写入缓存
|
||||
if to_compute:
|
||||
new_vecs = self._encode_batch_raw(to_compute)
|
||||
if self.use_cache:
|
||||
for smi, vec in zip(to_compute, new_vecs):
|
||||
self._cache[smi] = vec
|
||||
|
||||
# 按原始顺序组装
|
||||
pooled_np = (
|
||||
np.stack([self._cache[s] for s in smiles_list])
|
||||
if self.use_cache
|
||||
else new_vecs
|
||||
)
|
||||
|
||||
pooled = torch.from_numpy(pooled_np).to(self._device_str)
|
||||
return self.projection(pooled)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 工具属性 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
self._cache.clear()
|
||||
|
||||
@property
|
||||
def cache_size(self) -> int:
|
||||
return len(self._cache)
|
||||
203
lnp_ml/modeling/encoders/molt5_encoder.py
Normal file
203
lnp_ml/modeling/encoders/molt5_encoder.py
Normal file
@ -0,0 +1,203 @@
|
||||
"""
|
||||
MolT5 编码器
|
||||
|
||||
论文: "Translation between Molecules and Natural Language" (Edwards et al., EMNLP 2022)
|
||||
模型: laituan245/molt5-base (109M 参数, T5 架构)
|
||||
|
||||
功能：SMILES -> T5Encoder (frozen) -> mean pooling -> Linear -> [B, d_model]
|
||||
产出 V5 token，追加到 LNPModel 的 attention pooling 序列中。
|
||||
|
||||
放置位置：lnp_ml/modeling/encoders/molt5_encoder.py
|
||||
|
||||
接口与 llm_encoder.py 完全一致，models.py 只需修改初始化时的类名。
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MOLT5_MODEL_NAME = "laituan245/molt5-base"
|
||||
MOLT5_HIDDEN_SIZE = 768 # molt5-base 的 hidden size
|
||||
|
||||
|
||||
class MolT5Encoder(nn.Module):
|
||||
"""
|
||||
基于 MolT5-base 的冻结分子编码器。
|
||||
|
||||
流程：
|
||||
SMILES 字符串列表 [B]
|
||||
-> T5Tokenizer (molt5-base)
|
||||
-> T5EncoderModel (frozen, 109M)
|
||||
-> attention-mask mean pooling -> [B, 768]
|
||||
-> LayerNorm + Linear(768->d_model) + ReLU + Dropout
|
||||
-> [B, d_model] ← V5 token
|
||||
|
||||
Args:
|
||||
model_path : HuggingFace model name 或本地路径，默认 laituan245/molt5-base
|
||||
d_model : 输出维度，与 LNPModel.d_model 一致，默认 256
|
||||
dropout : projection 层 dropout，默认 0.1
|
||||
device : 推理设备
|
||||
use_cache : 是否缓存已编码 SMILES 的向量
|
||||
max_length : tokenizer 最大长度，默认 512
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str = MOLT5_MODEL_NAME,
|
||||
d_model: int = 256,
|
||||
dropout: float = 0.1,
|
||||
device: str = "cuda",
|
||||
use_cache: bool = True,
|
||||
max_length: int = 512,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.model_path = model_path
|
||||
self.d_model = d_model
|
||||
self.dropout_rate = dropout
|
||||
self._device_str = device
|
||||
self.use_cache = use_cache
|
||||
self.max_length = max_length
|
||||
|
||||
# 延迟加载
|
||||
self._tokenizer = None
|
||||
self._encoder = None
|
||||
self._initialized: bool = False
|
||||
|
||||
# SMILES -> pooled numpy 向量的内存缓存
|
||||
self._cache: Dict[str, np.ndarray] = {}
|
||||
|
||||
# projection 层在 _lazy_init 后赋值
|
||||
self.projection: Optional[nn.Sequential] = None
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 延迟初始化 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _lazy_init(self) -> None:
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
logger.info(f"[MolT5Encoder] Loading MolT5 from: {self.model_path}")
|
||||
|
||||
from transformers import T5Tokenizer, T5EncoderModel
|
||||
|
||||
self._tokenizer = T5Tokenizer.from_pretrained(
|
||||
self.model_path,
|
||||
model_max_length=self.max_length,
|
||||
)
|
||||
self._encoder = T5EncoderModel.from_pretrained(self.model_path)
|
||||
self._encoder = self._encoder.to(self._device_str)
|
||||
self._encoder.eval()
|
||||
|
||||
# 完全冻结 T5 encoder
|
||||
for param in self._encoder.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
n_params = sum(p.numel() for p in self._encoder.parameters())
|
||||
logger.info(
|
||||
f"[MolT5Encoder] Loaded, params={n_params:,}, frozen=True"
|
||||
)
|
||||
|
||||
# projection 层（唯一可训练部分）
|
||||
self.projection = nn.Sequential(
|
||||
nn.LayerNorm(MOLT5_HIDDEN_SIZE),
|
||||
nn.Linear(MOLT5_HIDDEN_SIZE, self.d_model),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(self.dropout_rate),
|
||||
).to(self._device_str)
|
||||
|
||||
self._initialized = True
|
||||
logger.info("[MolT5Encoder] Ready.")
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 内部编码 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
@torch.no_grad()
|
||||
def _encode_batch_raw(self, smiles_list: List[str]) -> np.ndarray:
|
||||
"""
|
||||
对一批 SMILES 过 T5Encoder，返回 mean-pooled numpy 向量。
|
||||
|
||||
Returns:
|
||||
np.ndarray shape [N, MOLT5_HIDDEN_SIZE] dtype float32
|
||||
"""
|
||||
# MolT5 在 SMILES 前加 "SMILES: " 前缀效果更好（原论文做法）
|
||||
inputs_text = [f"SMILES: {s}" for s in smiles_list]
|
||||
|
||||
encoding = self._tokenizer(
|
||||
inputs_text,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.max_length,
|
||||
)
|
||||
input_ids = encoding["input_ids"].to(self._device_str)
|
||||
attention_mask = encoding["attention_mask"].to(self._device_str)
|
||||
|
||||
outputs = self._encoder(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
# last_hidden_state: [B, seq_len, hidden]
|
||||
hidden = outputs.last_hidden_state
|
||||
|
||||
# attention-mask mean pooling
|
||||
mask_expanded = attention_mask.unsqueeze(-1).float()
|
||||
pooled = (hidden * mask_expanded).sum(dim=1) / mask_expanded.sum(dim=1).clamp(min=1e-9)
|
||||
|
||||
return pooled.cpu().numpy().astype(np.float32)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# forward #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def forward(self, smiles_list: List[str]) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
smiles_list: 长度为 B 的 SMILES 字符串列表
|
||||
|
||||
Returns:
|
||||
torch.Tensor shape [B, d_model]
|
||||
"""
|
||||
self._lazy_init()
|
||||
|
||||
# 找出未缓存的 SMILES
|
||||
to_compute = (
|
||||
[s for s in smiles_list if s not in self._cache]
|
||||
if self.use_cache
|
||||
else list(smiles_list)
|
||||
)
|
||||
|
||||
# 批量计算并写入缓存
|
||||
if to_compute:
|
||||
new_vecs = self._encode_batch_raw(to_compute)
|
||||
if self.use_cache:
|
||||
for smi, vec in zip(to_compute, new_vecs):
|
||||
self._cache[smi] = vec
|
||||
|
||||
# 按原始顺序组装
|
||||
pooled_np = (
|
||||
np.stack([self._cache[s] for s in smiles_list])
|
||||
if self.use_cache
|
||||
else new_vecs
|
||||
)
|
||||
|
||||
pooled = torch.from_numpy(pooled_np).to(self._device_str)
|
||||
return self.projection(pooled)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 工具属性 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
self._cache.clear()
|
||||
|
||||
@property
|
||||
def cache_size(self) -> int:
|
||||
return len(self._cache)
|
||||
246
lnp_ml/modeling/encoders/selfiested_encoder.py
Normal file
246
lnp_ml/modeling/encoders/selfiested_encoder.py
Normal file
@ -0,0 +1,246 @@
|
||||
"""
|
||||
SELFIES-TED 编码器
|
||||
|
||||
论文: "A Large Encoder-Decoder Family of Foundation Models for Chemical Language" (IBM, 2024-2025)
|
||||
模型: ibm/materials.selfies-ted (358M 参数, BART 架构, SELFIES 输入)
|
||||
|
||||
功能：SMILES -> SELFIES -> BART encoder (frozen) -> mean pooling -> Linear -> [B, d_model]
|
||||
产出 V5 token，追加到 LNPModel 的 attention pooling 序列中。
|
||||
|
||||
放置位置：lnp_ml/modeling/encoders/selfiested_encoder.py
|
||||
|
||||
依赖：
|
||||
pip install transformers selfies
|
||||
|
||||
注意：
|
||||
- SELFIES-TED 的 hidden size 是 1024（不是 BioT5/MolT5 的 768）
|
||||
- 输入需要先把 SMILES 转成 SELFIES，转换失败时回退为零向量
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SELFIESTED_MODEL_NAME = "ibm/materials.selfies-ted"
|
||||
SELFIESTED_HIDDEN_SIZE = 1024 # selfies-ted 的 BART hidden size
|
||||
|
||||
|
||||
def smiles_to_selfies(smiles: str) -> Optional[str]:
|
||||
"""
|
||||
SMILES -> SELFIES 字符串，并按官方做法在 token 之间加空格。
|
||||
转换失败返回 None。
|
||||
"""
|
||||
import selfies as sf
|
||||
|
||||
try:
|
||||
selfies = sf.encoder(smiles)
|
||||
if selfies is None:
|
||||
return None
|
||||
# 官方做法：把 "][" 替换成 "] ["，让 tokenizer 正确分词
|
||||
selfies = selfies.replace("][", "] [")
|
||||
return selfies
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to convert SMILES to SELFIES: {smiles}, error: {e}")
|
||||
return None
|
||||
|
||||
|
||||
class SELFIESTEDEncoder(nn.Module):
|
||||
"""
|
||||
基于 SELFIES-TED 的冻结分子编码器。
|
||||
|
||||
流程：
|
||||
SMILES 字符串列表 [B]
|
||||
-> SMILES -> SELFIES 转换
|
||||
-> SELFIES-TED tokenizer
|
||||
-> BART encoder (frozen, 358M)
|
||||
-> attention-mask mean pooling -> [B, 1024]
|
||||
-> LayerNorm + Linear(1024->d_model) + ReLU + Dropout
|
||||
-> [B, d_model] ← V5 token
|
||||
|
||||
Args:
|
||||
model_path : HuggingFace model name 或本地路径，默认 ibm/materials.selfies-ted
|
||||
d_model : 输出维度，与 LNPModel.d_model 一致，默认 256
|
||||
dropout : projection 层 dropout，默认 0.1
|
||||
device : 推理设备
|
||||
use_cache : 是否缓存已编码 SMILES 的向量
|
||||
max_length : tokenizer 最大长度，默认 128（官方示例用 128）
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: str = SELFIESTED_MODEL_NAME,
|
||||
d_model: int = 256,
|
||||
dropout: float = 0.1,
|
||||
device: str = "cuda",
|
||||
use_cache: bool = True,
|
||||
max_length: int = 128,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.model_path = model_path
|
||||
self.d_model = d_model
|
||||
self.dropout_rate = dropout
|
||||
self._device_str = device
|
||||
self.use_cache = use_cache
|
||||
self.max_length = max_length
|
||||
|
||||
# 延迟加载
|
||||
self._tokenizer = None
|
||||
self._model = None
|
||||
self._initialized: bool = False
|
||||
|
||||
# SMILES -> pooled numpy 向量的内存缓存
|
||||
self._cache: Dict[str, np.ndarray] = {}
|
||||
|
||||
# projection 层在 _lazy_init 后赋值
|
||||
self.projection: Optional[nn.Sequential] = None
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 延迟初始化 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _lazy_init(self) -> None:
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
logger.info(f"[SELFIESTEDEncoder] Loading SELFIES-TED from: {self.model_path}")
|
||||
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
|
||||
self._tokenizer = AutoTokenizer.from_pretrained(self.model_path)
|
||||
self._model = AutoModel.from_pretrained(self.model_path)
|
||||
self._model = self._model.to(self._device_str)
|
||||
self._model.eval()
|
||||
|
||||
# 完全冻结 BART 模型
|
||||
for param in self._model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
n_params = sum(p.numel() for p in self._model.parameters())
|
||||
logger.info(
|
||||
f"[SELFIESTEDEncoder] Loaded, params={n_params:,}, frozen=True"
|
||||
)
|
||||
|
||||
# projection 层（唯一可训练部分），注意输入维度是 1024
|
||||
self.projection = nn.Sequential(
|
||||
nn.LayerNorm(SELFIESTED_HIDDEN_SIZE),
|
||||
nn.Linear(SELFIESTED_HIDDEN_SIZE, self.d_model),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(self.dropout_rate),
|
||||
).to(self._device_str)
|
||||
|
||||
self._initialized = True
|
||||
logger.info("[SELFIESTEDEncoder] Ready.")
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 内部编码 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
@torch.no_grad()
|
||||
def _encode_batch_raw(self, smiles_list: List[str]) -> np.ndarray:
|
||||
"""
|
||||
对一批 SMILES 过 SELFIES-TED encoder，返回 mean-pooled numpy 向量。
|
||||
|
||||
Returns:
|
||||
np.ndarray shape [N, SELFIESTED_HIDDEN_SIZE] dtype float32
|
||||
"""
|
||||
# SMILES -> SELFIES，记录哪些转换成功
|
||||
selfies_list = []
|
||||
valid_indices = []
|
||||
for i, smi in enumerate(smiles_list):
|
||||
s = smiles_to_selfies(smi)
|
||||
if s is not None:
|
||||
selfies_list.append(s)
|
||||
valid_indices.append(i)
|
||||
else:
|
||||
logger.warning(f"SMILES->SELFIES failed, will use zero vector: {smi}")
|
||||
|
||||
if not selfies_list:
|
||||
# 全部转换失败，返回全零
|
||||
return np.zeros((len(smiles_list), SELFIESTED_HIDDEN_SIZE), dtype=np.float32)
|
||||
|
||||
token = self._tokenizer(
|
||||
selfies_list,
|
||||
return_tensors='pt',
|
||||
max_length=self.max_length,
|
||||
truncation=True,
|
||||
padding='max_length',
|
||||
)
|
||||
input_ids = token['input_ids'].to(self._device_str)
|
||||
attention_mask = token['attention_mask'].to(self._device_str)
|
||||
|
||||
# 只用 encoder
|
||||
outputs = self._model.encoder(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
hidden = outputs.last_hidden_state # [N_valid, seq_len, 1024]
|
||||
|
||||
# attention-mask mean pooling
|
||||
mask_expanded = attention_mask.unsqueeze(-1).float()
|
||||
pooled = (hidden * mask_expanded).sum(dim=1) / mask_expanded.sum(dim=1).clamp(min=1e-9)
|
||||
|
||||
result = pooled.cpu().numpy().astype(np.float32)
|
||||
|
||||
# 如果有转换失败的，用零向量填充回原始位置
|
||||
if len(valid_indices) < len(smiles_list):
|
||||
full_result = np.zeros((len(smiles_list), SELFIESTED_HIDDEN_SIZE), dtype=np.float32)
|
||||
for out_idx, in_idx in enumerate(valid_indices):
|
||||
full_result[in_idx] = result[out_idx]
|
||||
return full_result
|
||||
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# forward #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def forward(self, smiles_list: List[str]) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
smiles_list: 长度为 B 的 SMILES 字符串列表
|
||||
|
||||
Returns:
|
||||
torch.Tensor shape [B, d_model]
|
||||
"""
|
||||
self._lazy_init()
|
||||
|
||||
# 找出未缓存的 SMILES
|
||||
to_compute = (
|
||||
[s for s in smiles_list if s not in self._cache]
|
||||
if self.use_cache
|
||||
else list(smiles_list)
|
||||
)
|
||||
|
||||
# 批量计算并写入缓存
|
||||
if to_compute:
|
||||
new_vecs = self._encode_batch_raw(to_compute)
|
||||
if self.use_cache:
|
||||
for smi, vec in zip(to_compute, new_vecs):
|
||||
self._cache[smi] = vec
|
||||
|
||||
# 按原始顺序组装
|
||||
pooled_np = (
|
||||
np.stack([self._cache[s] for s in smiles_list])
|
||||
if self.use_cache
|
||||
else new_vecs
|
||||
)
|
||||
|
||||
pooled = torch.from_numpy(pooled_np).to(self._device_str)
|
||||
return self.projection(pooled)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 工具属性 #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
self._cache.clear()
|
||||
|
||||
@property
|
||||
def cache_size(self) -> int:
|
||||
return len(self._cache)
|
||||
@ -176,8 +176,20 @@ def create_model(
|
||||
dropout: float = 0.1,
|
||||
use_mpnn: bool = False,
|
||||
mpnn_device: str = "cpu",
|
||||
# ============ LLM 相关 ============
|
||||
use_llm: bool = False,
|
||||
llm_type: str = "biot5",
|
||||
llm_model_path: str = "",
|
||||
llm_device: str = "cuda",
|
||||
) -> Union[LNPModel, LNPModelWithoutMPNN]:
|
||||
"""创建模型"""
|
||||
llm_kwargs = dict(
|
||||
use_llm=use_llm,
|
||||
llm_type=llm_type,
|
||||
llm_model_path=llm_model_path,
|
||||
llm_device=llm_device,
|
||||
)
|
||||
|
||||
if use_mpnn:
|
||||
ensemble_paths = find_mpnn_ensemble_paths()
|
||||
return LNPModel(
|
||||
@ -189,6 +201,7 @@ def create_model(
|
||||
dropout=dropout,
|
||||
mpnn_ensemble_paths=ensemble_paths,
|
||||
mpnn_device=mpnn_device,
|
||||
**llm_kwargs,
|
||||
)
|
||||
else:
|
||||
return LNPModelWithoutMPNN(
|
||||
@ -198,6 +211,7 @@ def create_model(
|
||||
fusion_strategy=fusion_strategy,
|
||||
head_hidden_dim=head_hidden_dim,
|
||||
dropout=dropout,
|
||||
**llm_kwargs,
|
||||
)
|
||||
|
||||
|
||||
@ -232,7 +246,6 @@ def load_pretrain_weights_to_model(
|
||||
|
||||
|
||||
# ============ 3-fold Optuna 调参 ============
|
||||
|
||||
def run_optuna_cv(
|
||||
full_dataset: LNPDataset,
|
||||
strata: np.ndarray,
|
||||
@ -249,6 +262,11 @@ def run_optuna_cv(
|
||||
pretrain_config: Optional[Dict] = None,
|
||||
load_delivery_head: bool = True,
|
||||
rdkit_cache: Optional[Dict] = None,
|
||||
# ============ LLM 相关 ============
|
||||
use_llm: bool = False,
|
||||
llm_type: str = "biot5",
|
||||
llm_model_path: str = "",
|
||||
llm_device: str = "cuda",
|
||||
) -> Tuple[Dict, int, optuna.Study]:
|
||||
"""
|
||||
使用全量数据做 3-fold CV Optuna 超参搜索。
|
||||
@ -335,6 +353,10 @@ def run_optuna_cv(
|
||||
dropout=dropout,
|
||||
use_mpnn=use_mpnn,
|
||||
mpnn_device=device.type,
|
||||
use_llm=use_llm,
|
||||
llm_type=llm_type,
|
||||
llm_model_path=llm_model_path,
|
||||
llm_device=llm_device,
|
||||
)
|
||||
if rdkit_cache is not None:
|
||||
model.rdkit_encoder._cache = rdkit_cache
|
||||
@ -404,7 +426,6 @@ def run_optuna_cv(
|
||||
|
||||
|
||||
# ============ 主流程 ============
|
||||
|
||||
@app.command()
|
||||
def main(
|
||||
input_path: Path = INTERIM_DATA_DIR / "internal.csv",
|
||||
@ -427,6 +448,11 @@ def main(
|
||||
load_delivery_head: bool = False,
|
||||
# MPNN
|
||||
use_mpnn: bool = False,
|
||||
# LLM（新增）
|
||||
use_llm: bool = False,
|
||||
llm_type: str = "biot5",
|
||||
llm_model_path: str = "",
|
||||
llm_device: str = "cuda",
|
||||
# 设备
|
||||
device: str = "cuda" if torch.cuda.is_available() else "cpu",
|
||||
):
|
||||
@ -507,6 +533,10 @@ def main(
|
||||
pretrain_config=pretrain_config,
|
||||
load_delivery_head=load_delivery_head,
|
||||
rdkit_cache=rdkit_cache,
|
||||
use_llm=use_llm,
|
||||
llm_type=llm_type,
|
||||
llm_model_path=llm_model_path,
|
||||
llm_device=llm_device,
|
||||
)
|
||||
|
||||
# 保存最佳参数
|
||||
@ -564,6 +594,10 @@ def main(
|
||||
dropout=best_params["dropout"],
|
||||
use_mpnn=use_mpnn,
|
||||
mpnn_device=device.type,
|
||||
use_llm=use_llm,
|
||||
llm_type=llm_type,
|
||||
llm_model_path=llm_model_path,
|
||||
llm_device=llm_device,
|
||||
)
|
||||
model.rdkit_encoder._cache = rdkit_cache
|
||||
|
||||
@ -611,8 +645,12 @@ def main(
|
||||
"head_hidden_dim": best_params["head_hidden_dim"],
|
||||
"dropout": best_params["dropout"],
|
||||
"use_mpnn": use_mpnn,
|
||||
"use_llm": use_llm,
|
||||
"llm_type": llm_type,
|
||||
"llm_model_path": llm_model_path,
|
||||
"llm_device": llm_device,
|
||||
}
|
||||
|
||||
|
||||
torch.save({
|
||||
"model_state_dict": train_result["final_state"],
|
||||
"config": config,
|
||||
@ -655,5 +693,4 @@ def main(
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
|
||||
app()
|
||||
@ -1,4 +1,18 @@
|
||||
"""LNP 多任务预测模型"""
|
||||
"""LNP 多任务预测模型
|
||||
|
||||
本文件在原始 LNPModel 基础上，新增了可选的 LLM 分子编码器分支（V5 token）。
|
||||
|
||||
LLM 分支说明：
|
||||
- 默认 use_llm=False，模型行为与原版完全一致，不影响既有功能。
|
||||
- use_llm=True 时，通过 llm_type 选择一个分子编码器，把 SMILES 编码成
|
||||
一个额外的 V5 token，追加到 fusion 序列末尾（n_tokens + 1）。
|
||||
- 支持的 llm_type:
|
||||
"biot5" -> BioT5-plus (encoders/llm_encoder.py)
|
||||
"molt5" -> MolT5 (encoders/molt5_encoder.py)
|
||||
"moleculestm" -> MoleculeSTM GNN (encoders/moleculestm_encoder.py)
|
||||
"selfiested" -> SELFIES-TED (encoders/selfiested_encoder.py)
|
||||
- LLM 主体全程冻结，仅其内部 projection 层参与训练。
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@ -10,6 +24,7 @@ from lnp_ml.modeling.heads import MultiTaskHead
|
||||
|
||||
|
||||
PoolingStrategy = Literal["concat", "avg", "max", "attention"]
|
||||
LLMType = Literal["biot5", "molt5", "moleculestm", "selfiested"]
|
||||
|
||||
|
||||
# Token 维度配置（根据 ARCHITECTURE.md）
|
||||
@ -18,7 +33,7 @@ DEFAULT_INPUT_DIMS = {
|
||||
"mpnn": 600, # D-MPNN embedding
|
||||
"morgan": 1024, # Morgan fingerprint
|
||||
"maccs": 167, # MACCS keys
|
||||
"desc": 210, # RDKit descriptors
|
||||
"desc": 217, # RDKit descriptors（新版 RDKit 为 217 维）
|
||||
# Channel B: 配方/实验条件
|
||||
"comp": 5, # 配方比例
|
||||
"phys": 12, # 物理参数 one-hot
|
||||
@ -30,6 +45,64 @@ DEFAULT_INPUT_DIMS = {
|
||||
TOKEN_ORDER = ["mpnn", "morgan", "maccs", "desc", "comp", "phys", "help", "exp"]
|
||||
|
||||
|
||||
def build_llm_encoder(
|
||||
llm_type: str,
|
||||
llm_model_path: str,
|
||||
d_model: int,
|
||||
dropout: float,
|
||||
llm_device: str,
|
||||
) -> nn.Module:
|
||||
"""
|
||||
LLM encoder 工厂：根据 llm_type 实例化对应的分子编码器。
|
||||
|
||||
四个 encoder 接口一致：forward(List[str]) -> [B, d_model]，
|
||||
内部 LLM 冻结，仅 projection 可训练。
|
||||
|
||||
注意不同 encoder 的路径参数名不同：
|
||||
biot5 / molt5 / selfiested -> model_path
|
||||
moleculestm -> checkpoint_path
|
||||
"""
|
||||
llm_type = llm_type.lower()
|
||||
|
||||
if llm_type == "biot5":
|
||||
from lnp_ml.modeling.encoders import LLMEncoder
|
||||
return LLMEncoder(
|
||||
model_path=llm_model_path,
|
||||
d_model=d_model,
|
||||
dropout=dropout,
|
||||
device=llm_device,
|
||||
)
|
||||
elif llm_type == "molt5":
|
||||
from lnp_ml.modeling.encoders import MolT5Encoder
|
||||
return MolT5Encoder(
|
||||
model_path=llm_model_path,
|
||||
d_model=d_model,
|
||||
dropout=dropout,
|
||||
device=llm_device,
|
||||
)
|
||||
elif llm_type == "moleculestm":
|
||||
from lnp_ml.modeling.encoders import MoleculeSTMEncoder
|
||||
return MoleculeSTMEncoder(
|
||||
checkpoint_path=llm_model_path,
|
||||
d_model=d_model,
|
||||
dropout=dropout,
|
||||
device=llm_device,
|
||||
)
|
||||
elif llm_type == "selfiested":
|
||||
from lnp_ml.modeling.encoders import SELFIESTEDEncoder
|
||||
return SELFIESTEDEncoder(
|
||||
model_path=llm_model_path,
|
||||
d_model=d_model,
|
||||
dropout=dropout,
|
||||
device=llm_device,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown llm_type: {llm_type!r}. "
|
||||
f"Choose from: biot5, molt5, moleculestm, selfiested"
|
||||
)
|
||||
|
||||
|
||||
class LNPModel(nn.Module):
|
||||
"""
|
||||
LNP 药物递送性能预测模型。
|
||||
@ -37,10 +110,11 @@ class LNPModel(nn.Module):
|
||||
架构流程:
|
||||
1. Encoders: SMILES -> 化学特征; tabular -> 配方/实验特征
|
||||
2. TokenProjector: 统一到 d_model
|
||||
3. Stack: [B, 8, d_model]
|
||||
3. Stack: [B, n_tokens, d_model]
|
||||
4. CrossModalAttention: Channel A (化学) <-> Channel B (配方/实验)
|
||||
5. FusionLayer: [B, 8, d_model] -> [B, fusion_dim]
|
||||
6. MultiTaskHead: 多任务预测
|
||||
5. (可选) LLM 分支: SMILES -> V5 token，追加到序列末尾
|
||||
6. FusionLayer: [B, n_tokens(+1), d_model] -> [B, fusion_dim]
|
||||
7. MultiTaskHead: 多任务预测
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@ -62,12 +136,19 @@ class LNPModel(nn.Module):
|
||||
mpnn_device: str = "cpu",
|
||||
# 输入维度配置
|
||||
input_dims: Optional[Dict[str, int]] = None,
|
||||
# ============ LLM 分支（可选） ============
|
||||
use_llm: bool = False,
|
||||
llm_type: Optional[str] = None,
|
||||
llm_model_path: str = "",
|
||||
llm_device: str = "cuda",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.input_dims = input_dims or DEFAULT_INPUT_DIMS
|
||||
self.d_model = d_model
|
||||
self.use_mpnn = mpnn_checkpoint is not None or mpnn_ensemble_paths is not None
|
||||
self.use_llm = use_llm
|
||||
self.llm_type = llm_type
|
||||
|
||||
# ============ Encoders ============
|
||||
# RDKit encoder (always used)
|
||||
@ -83,6 +164,20 @@ class LNPModel(nn.Module):
|
||||
else:
|
||||
self.mpnn_encoder = None
|
||||
|
||||
# ============ LLM encoder (optional) ============
|
||||
if self.use_llm:
|
||||
if not llm_type:
|
||||
raise ValueError("use_llm=True 时必须指定 llm_type")
|
||||
self.llm_encoder = build_llm_encoder(
|
||||
llm_type=llm_type,
|
||||
llm_model_path=llm_model_path,
|
||||
d_model=d_model,
|
||||
dropout=dropout,
|
||||
llm_device=llm_device,
|
||||
)
|
||||
else:
|
||||
self.llm_encoder = None
|
||||
|
||||
# ============ Token Projector ============
|
||||
# 根据是否使用 MPNN 调整输入维度
|
||||
proj_input_dims = {k: v for k, v in self.input_dims.items()}
|
||||
@ -108,9 +203,13 @@ class LNPModel(nn.Module):
|
||||
)
|
||||
|
||||
# ============ Fusion Layer ============
|
||||
# 启用 LLM 时，fusion 序列多一个 V5 token
|
||||
n_fusion_tokens = n_tokens + 1 if self.use_llm else n_tokens
|
||||
self.n_fusion_tokens = n_fusion_tokens
|
||||
|
||||
self.fusion = FusionLayer(
|
||||
d_model=d_model,
|
||||
n_tokens=n_tokens,
|
||||
n_tokens=n_fusion_tokens,
|
||||
strategy=fusion_strategy,
|
||||
)
|
||||
|
||||
@ -128,39 +227,35 @@ class LNPModel(nn.Module):
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
内部方法：编码 SMILES 和 tabular，返回 stacked tokens。
|
||||
|
||||
|
||||
Returns:
|
||||
stacked: [B, n_tokens, d_model]
|
||||
"""
|
||||
# 获取目标设备（从 tabular 数据推断）
|
||||
device = tabular["comp"].device
|
||||
|
||||
|
||||
# 1. Encode SMILES
|
||||
rdkit_features = self.rdkit_encoder(smiles) # {"morgan", "maccs", "desc"}
|
||||
rdkit_features = self.rdkit_encoder(smiles)
|
||||
|
||||
# 2. 合并所有特征
|
||||
all_features: Dict[str, torch.Tensor] = {}
|
||||
|
||||
# MPNN 特征（如果启用）
|
||||
if self.use_mpnn:
|
||||
mpnn_features = self.mpnn_encoder(smiles)
|
||||
all_features["mpnn"] = mpnn_features["mpnn"].to(device)
|
||||
|
||||
# RDKit 特征（移到正确设备）
|
||||
all_features["morgan"] = rdkit_features["morgan"].to(device)
|
||||
all_features["maccs"] = rdkit_features["maccs"].to(device)
|
||||
all_features["desc"] = rdkit_features["desc"].to(device)
|
||||
|
||||
# Tabular 特征（已在正确设备上）
|
||||
all_features["comp"] = tabular["comp"]
|
||||
all_features["phys"] = tabular["phys"]
|
||||
all_features["help"] = tabular["help"]
|
||||
all_features["exp"] = tabular["exp"]
|
||||
|
||||
# 3. Token Projector: 统一维度
|
||||
projected = self.token_projector(all_features) # Dict[str, [B, d_model]]
|
||||
# 3. Token Projector
|
||||
projected = self.token_projector(all_features)
|
||||
|
||||
# 4. Stack tokens: [B, n_tokens, d_model]
|
||||
# 4. Stack tokens
|
||||
if self.use_mpnn:
|
||||
token_order = ["mpnn", "morgan", "maccs", "desc", "comp", "phys", "help", "exp"]
|
||||
else:
|
||||
@ -169,6 +264,28 @@ class LNPModel(nn.Module):
|
||||
stacked = torch.stack([projected[k] for k in token_order], dim=1)
|
||||
return stacked
|
||||
|
||||
def _attended_with_llm(
|
||||
self,
|
||||
smiles: List[str],
|
||||
tabular: Dict[str, torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
编码 -> 投影 -> cross-attention -> (可选) 追加 V5 token。
|
||||
|
||||
Returns:
|
||||
attended: [B, n_tokens, d] (use_llm=False)
|
||||
或 [B, n_tokens + 1, d] (use_llm=True)
|
||||
"""
|
||||
stacked = self._encode_and_project(smiles, tabular)
|
||||
attended = self.cross_attention(stacked)
|
||||
|
||||
if self.use_llm:
|
||||
device = attended.device
|
||||
v5 = self.llm_encoder(smiles).to(device) # [B, d]
|
||||
attended = torch.cat([attended, v5.unsqueeze(1)], dim=1) # [B, n+1, d]
|
||||
|
||||
return attended
|
||||
|
||||
def forward_from_projected(
|
||||
self,
|
||||
stacked: torch.Tensor,
|
||||
@ -176,14 +293,7 @@ class LNPModel(nn.Module):
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
从已投影的 stacked tokens 开始 forward，用于 Captum 归因。
|
||||
|
||||
Args:
|
||||
stacked: [B, n_tokens, d_model] TokenProjector 输出后 stack 的张量
|
||||
task: 指定单任务名 ("size", "pdi", "ee", "delivery", "biodist", "toxic")。
|
||||
若为 None，返回 delivery head 的标量输出。
|
||||
|
||||
Returns:
|
||||
[B, 1] 或 [B, num_classes] 对应任务的预测输出
|
||||
注意：该路径不经过 LLM 分支（归因针对主干 token）。
|
||||
"""
|
||||
attended = self.cross_attention(stacked)
|
||||
fused = self.fusion(attended)
|
||||
@ -208,23 +318,10 @@ class LNPModel(nn.Module):
|
||||
base_projected: torch.Tensor,
|
||||
task: Optional[str] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
用原始特征替换 base_projected 中指定 token 的投影，然后 forward。
|
||||
|
||||
用于对单个 token 内部特征做 Captum 归因（如 desc 的 210 维）。
|
||||
|
||||
Args:
|
||||
raw_feature: [B, input_dim] 某个 token 的原始特征
|
||||
feature_key: token 名称，如 "desc"
|
||||
base_projected: [B, n_tokens, d_model] 其他 token 已投影好的张量
|
||||
task: 任务名
|
||||
|
||||
Returns:
|
||||
对应任务的预测输出
|
||||
"""
|
||||
"""用原始特征替换指定 token 的投影后 forward（Captum 归因用）。"""
|
||||
projected = self.token_projector.projectors[feature_key](raw_feature)
|
||||
gate = torch.sigmoid(self.token_projector.weights[feature_key])
|
||||
projected = projected * gate # [B, d_model]
|
||||
projected = projected * gate
|
||||
|
||||
token_order = list(self.token_projector.keys)
|
||||
token_idx = token_order.index(feature_key)
|
||||
@ -240,26 +337,10 @@ class LNPModel(nn.Module):
|
||||
tabular: Dict[str, torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Backbone forward：编码 -> 投影 -> 注意力 -> 融合，不经过任务头。
|
||||
|
||||
用于 pretrain 阶段或需要提取特征的场景。
|
||||
|
||||
Args:
|
||||
smiles: SMILES 字符串列表，长度为 B
|
||||
tabular: Dict[str, Tensor]
|
||||
|
||||
Returns:
|
||||
fused: [B, fusion_dim] 融合后的特征向量
|
||||
Backbone forward：编码 -> 投影 -> 注意力 ->（可选 V5）-> 融合，不经过任务头。
|
||||
"""
|
||||
# 编码 + 投影 + stack
|
||||
stacked = self._encode_and_project(smiles, tabular)
|
||||
|
||||
# Cross Modal Attention
|
||||
attended = self.cross_attention(stacked)
|
||||
|
||||
# Fusion
|
||||
attended = self._attended_with_llm(smiles, tabular)
|
||||
fused = self.fusion(attended)
|
||||
|
||||
return fused
|
||||
|
||||
def forward_delivery(
|
||||
@ -267,16 +348,7 @@ class LNPModel(nn.Module):
|
||||
smiles: List[str],
|
||||
tabular: Dict[str, torch.Tensor],
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
仅预测 delivery（用于 pretrain）。
|
||||
|
||||
Args:
|
||||
smiles: SMILES 字符串列表，长度为 B
|
||||
tabular: Dict[str, Tensor]
|
||||
|
||||
Returns:
|
||||
delivery: [B, 1] 预测的 delivery 值
|
||||
"""
|
||||
"""仅预测 delivery（用于 pretrain）。"""
|
||||
fused = self.forward_backbone(smiles, tabular)
|
||||
return self.head.delivery_head(fused)
|
||||
|
||||
@ -285,26 +357,7 @@ class LNPModel(nn.Module):
|
||||
smiles: List[str],
|
||||
tabular: Dict[str, torch.Tensor],
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""
|
||||
完整的多任务 forward。
|
||||
|
||||
Args:
|
||||
smiles: SMILES 字符串列表，长度为 B
|
||||
tabular: Dict[str, Tensor]，包含:
|
||||
- "comp": [B, 5] 配方比例
|
||||
- "phys": [B, 12] 物理参数
|
||||
- "help": [B, 4] Helper lipid
|
||||
- "exp": [B, 32] 实验条件
|
||||
|
||||
Returns:
|
||||
Dict[str, Tensor]:
|
||||
- "size": [B, 1]
|
||||
- "pdi": [B, 4]
|
||||
- "ee": [B, 3]
|
||||
- "delivery": [B, 1]
|
||||
- "biodist": [B, 7]
|
||||
- "toxic": [B, 2]
|
||||
"""
|
||||
"""完整的多任务 forward。"""
|
||||
fused = self.forward_backbone(smiles, tabular)
|
||||
outputs = self.head(fused)
|
||||
return outputs
|
||||
@ -314,22 +367,18 @@ class LNPModel(nn.Module):
|
||||
self.rdkit_encoder.clear_cache()
|
||||
if self.mpnn_encoder is not None:
|
||||
self.mpnn_encoder.clear_cache()
|
||||
if self.llm_encoder is not None and hasattr(self.llm_encoder, "clear_cache"):
|
||||
self.llm_encoder.clear_cache()
|
||||
|
||||
def get_backbone_state_dict(self) -> Dict[str, torch.Tensor]:
|
||||
"""
|
||||
获取 backbone 部分的 state_dict（不含任务头）。
|
||||
|
||||
包含: token_projector, cross_attention, fusion
|
||||
"""
|
||||
"""获取 backbone 部分的 state_dict（不含任务头）。"""
|
||||
backbone_keys = []
|
||||
for name in self.state_dict().keys():
|
||||
if name.startswith(("token_projector.", "cross_attention.", "fusion.")):
|
||||
if name.startswith(("token_projector.", "cross_attention.", "fusion.", "llm_encoder.")):
|
||||
backbone_keys.append(name)
|
||||
|
||||
return {k: v for k, v in self.state_dict().items() if k in backbone_keys}
|
||||
|
||||
def get_delivery_head_state_dict(self) -> Dict[str, torch.Tensor]:
|
||||
"""获取 delivery head 的 state_dict"""
|
||||
return {
|
||||
k: v for k, v in self.state_dict().items()
|
||||
if k.startswith("head.delivery_head.")
|
||||
@ -341,28 +390,17 @@ class LNPModel(nn.Module):
|
||||
load_delivery_head: bool = True,
|
||||
strict: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
从预训练 checkpoint 加载 backbone 和（可选）delivery head 权重。
|
||||
|
||||
Args:
|
||||
pretrain_state_dict: 预训练模型的 state_dict
|
||||
load_delivery_head: 是否加载 delivery head 权重
|
||||
strict: 是否严格匹配（默认 False，允许缺失/多余的键）
|
||||
"""
|
||||
# 筛选要加载的参数
|
||||
"""从预训练 checkpoint 加载 backbone 和（可选）delivery head 权重。"""
|
||||
keys_to_load = []
|
||||
for name in pretrain_state_dict.keys():
|
||||
# Backbone 部分
|
||||
if name.startswith(("token_projector.", "cross_attention.", "fusion.")):
|
||||
if name.startswith(("token_projector.", "cross_attention.", "fusion.", "llm_encoder.")):
|
||||
keys_to_load.append(name)
|
||||
# Delivery head（可选）
|
||||
elif load_delivery_head and name.startswith("head.delivery_head."):
|
||||
keys_to_load.append(name)
|
||||
|
||||
|
||||
filtered_state_dict = {k: v for k, v in pretrain_state_dict.items() if k in keys_to_load}
|
||||
|
||||
# 加载权重
|
||||
missing, unexpected = [], []
|
||||
|
||||
unexpected = []
|
||||
model_state = self.state_dict()
|
||||
for k, v in filtered_state_dict.items():
|
||||
if k in model_state:
|
||||
@ -372,15 +410,15 @@ class LNPModel(nn.Module):
|
||||
unexpected.append(f"{k} (shape mismatch: {model_state[k].shape} vs {v.shape})")
|
||||
else:
|
||||
unexpected.append(k)
|
||||
|
||||
|
||||
self.load_state_dict(model_state, strict=False)
|
||||
|
||||
if strict and (missing or unexpected):
|
||||
raise RuntimeError(f"Missing keys: {missing}, Unexpected keys: {unexpected}")
|
||||
|
||||
if strict and unexpected:
|
||||
raise RuntimeError(f"Unexpected keys: {unexpected}")
|
||||
|
||||
|
||||
class LNPModelWithoutMPNN(LNPModel):
|
||||
"""不使用 MPNN 的简化版本"""
|
||||
"""不使用 MPNN 的简化版本（实验默认使用此类）。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@ -391,6 +429,11 @@ class LNPModelWithoutMPNN(LNPModel):
|
||||
head_hidden_dim: int = 128,
|
||||
dropout: float = 0.1,
|
||||
input_dims: Optional[Dict[str, int]] = None,
|
||||
# ============ LLM 分支（可选） ============
|
||||
use_llm: bool = False,
|
||||
llm_type: Optional[str] = None,
|
||||
llm_model_path: str = "",
|
||||
llm_device: str = "cuda",
|
||||
) -> None:
|
||||
# 移除 mpnn 维度
|
||||
dims = input_dims or DEFAULT_INPUT_DIMS.copy()
|
||||
@ -406,5 +449,8 @@ class LNPModelWithoutMPNN(LNPModel):
|
||||
mpnn_checkpoint=None,
|
||||
mpnn_ensemble_paths=None,
|
||||
input_dims=dims,
|
||||
use_llm=use_llm,
|
||||
llm_type=llm_type,
|
||||
llm_model_path=llm_model_path,
|
||||
llm_device=llm_device,
|
||||
)
|
||||
|
||||
|
||||
@ -225,6 +225,11 @@ def main(
|
||||
mpnn_checkpoint: Optional[str] = None,
|
||||
mpnn_ensemble_paths: Optional[str] = None, # 逗号分隔的路径列表
|
||||
mpnn_device: str = "cpu",
|
||||
# ============ LLM 参数（可选） ============
|
||||
use_llm: bool = False,
|
||||
llm_type: str = "biot5",
|
||||
llm_model_path: str = "",
|
||||
llm_device: str = "cuda",
|
||||
# 训练参数
|
||||
batch_size: int = 64,
|
||||
lr: float = 1e-4,
|
||||
@ -292,6 +297,10 @@ def main(
|
||||
mpnn_checkpoint=mpnn_checkpoint,
|
||||
mpnn_ensemble_paths=ensemble_paths_list,
|
||||
mpnn_device=mpnn_device,
|
||||
use_llm=use_llm,
|
||||
llm_type=llm_type,
|
||||
llm_model_path=llm_model_path,
|
||||
llm_device=llm_device,
|
||||
)
|
||||
else:
|
||||
model = LNPModelWithoutMPNN(
|
||||
@ -301,6 +310,10 @@ def main(
|
||||
fusion_strategy=fusion_strategy,
|
||||
head_hidden_dim=head_hidden_dim,
|
||||
dropout=dropout,
|
||||
use_llm=use_llm,
|
||||
llm_type=llm_type,
|
||||
llm_model_path=llm_model_path,
|
||||
llm_device=llm_device,
|
||||
)
|
||||
|
||||
n_params_total = sum(p.numel() for p in model.parameters())
|
||||
@ -340,6 +353,10 @@ def main(
|
||||
"head_hidden_dim": head_hidden_dim,
|
||||
"dropout": dropout,
|
||||
"use_mpnn": enable_mpnn,
|
||||
"use_llm": use_llm,
|
||||
"llm_type": llm_type,
|
||||
"llm_model_path": llm_model_path,
|
||||
"llm_device": llm_device,
|
||||
},
|
||||
"best_val_loss": result["best_val_loss"],
|
||||
},
|
||||
@ -412,6 +429,10 @@ def test(
|
||||
dropout=config["dropout"],
|
||||
mpnn_ensemble_paths=ensemble_paths,
|
||||
mpnn_device=mpnn_device,
|
||||
use_llm=config.get("use_llm", False),
|
||||
llm_type=config.get("llm_type", "biot5"),
|
||||
llm_model_path=config.get("llm_model_path", ""),
|
||||
llm_device=config.get("llm_device", "cuda"),
|
||||
)
|
||||
else:
|
||||
model = LNPModelWithoutMPNN(
|
||||
@ -421,6 +442,10 @@ def test(
|
||||
fusion_strategy=config["fusion_strategy"],
|
||||
head_hidden_dim=config["head_hidden_dim"],
|
||||
dropout=config["dropout"],
|
||||
use_llm=config.get("use_llm", False),
|
||||
llm_type=config.get("llm_type", "biot5"),
|
||||
llm_model_path=config.get("llm_model_path", ""),
|
||||
llm_device=config.get("llm_device", "cuda"),
|
||||
)
|
||||
|
||||
model.load_state_dict(checkpoint["model_state_dict"])
|
||||
|
||||
280
verify_llm_encoder.py
Normal file
280
verify_llm_encoder.py
Normal file
@ -0,0 +1,280 @@
|
||||
"""
|
||||
LLM 编码器验证脚本（统一版）
|
||||
|
||||
对四种 LLM encoder 中的任意一种做四项检查：
|
||||
TEST 1: encoder 独立输出（shape / NaN / 梯度 / 缓存 / batch 一致性）
|
||||
TEST 2: 与 LNPModel 集成（任务头维度 / biodist 和为1 / 反向传播 / fusion token 数）
|
||||
TEST 3: 可复现性（同种子相同输出，不同种子不同输出）
|
||||
TEST 4: 过拟合 sanity（小数据上 loss 下降、R² > 0）
|
||||
|
||||
用法:
|
||||
python verify_llm_encoder.py \
|
||||
--model_path <llm_model_path> \
|
||||
--llm_type <biot5|molt5|moleculestm|selfiested> \
|
||||
--lnp_repo_path <仓库根目录>
|
||||
|
||||
不同 encoder 的可训练主体属性名不同，本脚本会按 llm_type 自动选择：
|
||||
biot5 -> encoder._llm
|
||||
molt5 -> encoder._encoder
|
||||
moleculestm -> encoder._gnn
|
||||
selfiested -> encoder._model
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
# llm_type -> (encoder 模块, 类名, 冻结主体属性名)
|
||||
LLM_REGISTRY = {
|
||||
"biot5": ("lnp_ml.modeling.encoders.llm_encoder", "LLMEncoder", "_llm"),
|
||||
"molt5": ("lnp_ml.modeling.encoders.molt5_encoder", "MolT5Encoder", "_encoder"),
|
||||
"moleculestm": ("lnp_ml.modeling.encoders.moleculestm_encoder", "MoleculeSTMEncoder", "_gnn"),
|
||||
"selfiested": ("lnp_ml.modeling.encoders.selfiested_encoder", "SELFIESTEDEncoder", "_model"),
|
||||
}
|
||||
|
||||
|
||||
def set_seed(seed: int = 42) -> None:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
print(f"[Seed] All seeds fixed to {seed}")
|
||||
|
||||
|
||||
def build_encoder(llm_type, model_path, d_model, device):
|
||||
"""按 llm_type 构造 encoder，处理 model_path / checkpoint_path 参数名差异。"""
|
||||
import importlib
|
||||
module_name, class_name, _ = LLM_REGISTRY[llm_type]
|
||||
module = importlib.import_module(module_name)
|
||||
EncoderCls = getattr(module, class_name)
|
||||
|
||||
if llm_type == "moleculestm":
|
||||
return EncoderCls(checkpoint_path=model_path, d_model=d_model, device=device)
|
||||
else:
|
||||
return EncoderCls(model_path=model_path, d_model=d_model, device=device)
|
||||
|
||||
|
||||
def get_frozen_body(encoder, llm_type):
|
||||
"""取出该 encoder 的冻结主体（用于检查参数是否冻结）。"""
|
||||
_, _, attr = LLM_REGISTRY[llm_type]
|
||||
return getattr(encoder, attr)
|
||||
|
||||
|
||||
SAMPLE_SMILES = [
|
||||
"CCCCCCCCCCCCCCCC(=O)OCC(COC(=O)CCCCCCCCCCCCCCC)OC(=O)CCCCCCC",
|
||||
"CCN(CC)CCOC(=O)C1=CC=CC=C1",
|
||||
"CC(C)(C)OC(=O)N1CCC(CC1)C(=O)O",
|
||||
"O=C(O)c1ccccc1O",
|
||||
]
|
||||
|
||||
|
||||
def make_tabular(batch_size, device):
|
||||
return {
|
||||
"comp": torch.randn(batch_size, 5, device=device),
|
||||
"phys": torch.randn(batch_size, 12, device=device),
|
||||
"help": torch.randn(batch_size, 4, device=device),
|
||||
"exp": torch.randn(batch_size, 32, device=device),
|
||||
}
|
||||
|
||||
|
||||
def test_1_standalone(llm_type, model_path, d_model, device):
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 1: encoder standalone")
|
||||
print("=" * 60)
|
||||
|
||||
print("\n[1a] Testing output shape...")
|
||||
encoder = build_encoder(llm_type, model_path, d_model, device)
|
||||
out = encoder(SAMPLE_SMILES)
|
||||
assert out.shape == (len(SAMPLE_SMILES), d_model), f"bad shape {out.shape}"
|
||||
print(f" Output shape: {out.shape} OK")
|
||||
|
||||
print("[1b] Checking for NaN/Inf...")
|
||||
assert not torch.isnan(out).any() and not torch.isinf(out).any()
|
||||
print(f" No NaN/Inf OK")
|
||||
print(f" Output stats: mean={out.mean():.4f}, std={out.std():.4f}, "
|
||||
f"min={out.min():.4f}, max={out.max():.4f}")
|
||||
|
||||
print("[1c] Checking gradients...")
|
||||
body = get_frozen_body(encoder, llm_type)
|
||||
n_trainable_body = sum(p.numel() for p in body.parameters() if p.requires_grad)
|
||||
assert n_trainable_body == 0, f"frozen body has {n_trainable_body} trainable params"
|
||||
print(f" Frozen body ({0} trainable params) OK")
|
||||
n_proj = sum(p.numel() for p in encoder.projection.parameters() if p.requires_grad)
|
||||
assert n_proj > 0
|
||||
print(f" Projection trainable ({n_proj:,} params) OK")
|
||||
loss = encoder(SAMPLE_SMILES).sum()
|
||||
loss.backward()
|
||||
grad_ok = any(p.grad is not None for p in encoder.projection.parameters())
|
||||
assert grad_ok
|
||||
print(f" Gradient flows to projection OK")
|
||||
|
||||
print("[1e] Single vs batch consistency...")
|
||||
encoder.eval()
|
||||
with torch.no_grad():
|
||||
batch_out = encoder(SAMPLE_SMILES)
|
||||
single_out = torch.cat([encoder([s]) for s in SAMPLE_SMILES], dim=0)
|
||||
assert torch.allclose(batch_out, single_out, atol=1e-4)
|
||||
print(f" Batch == single OK")
|
||||
print("\nTEST 1 PASSED")
|
||||
|
||||
|
||||
def test_2_integration(llm_type, model_path, d_model, device):
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 2: Integration with LNPModel")
|
||||
print("=" * 60)
|
||||
|
||||
from lnp_ml.modeling.models import LNPModelWithoutMPNN
|
||||
model = LNPModelWithoutMPNN(
|
||||
d_model=d_model, use_llm=True, llm_type=llm_type,
|
||||
llm_model_path=model_path, llm_device=device,
|
||||
).to(device)
|
||||
|
||||
print("\n[2a] Testing full forward pass shapes...")
|
||||
tab = make_tabular(len(SAMPLE_SMILES), device)
|
||||
outputs = model(SAMPLE_SMILES, tab)
|
||||
expected = {"size": 1, "pdi": 4, "ee": 3, "delivery": 1, "biodist": 7, "toxic": 2}
|
||||
for k, dim in expected.items():
|
||||
assert outputs[k].shape == (len(SAMPLE_SMILES), dim), f"{k}: {outputs[k].shape}"
|
||||
print(f" outputs['{k}']: {tuple(outputs[k].shape)} OK")
|
||||
|
||||
print("[2b] Checking biodist sums to 1...")
|
||||
sums = outputs["biodist"].sum(dim=-1)
|
||||
assert torch.allclose(sums, torch.ones_like(sums), atol=1e-4)
|
||||
print(f" biodist sums: {[round(s.item(),4) for s in sums]} OK")
|
||||
|
||||
print("[2c] Testing backward pass...")
|
||||
loss = sum(o.sum() for o in outputs.values())
|
||||
loss.backward()
|
||||
proj_grad = any(p.grad is not None for p in model.llm_encoder.projection.parameters())
|
||||
assert proj_grad
|
||||
print(f" llm_encoder.projection gets gradients OK")
|
||||
body = get_frozen_body(model.llm_encoder, llm_type)
|
||||
body_grad = any(p.grad is not None and p.grad.abs().sum() > 0 for p in body.parameters())
|
||||
assert not body_grad
|
||||
print(f" LLM body has no gradients (frozen) OK")
|
||||
|
||||
print("[2d] Checking fusion token count...")
|
||||
# 无 MPNN 时主干 7 个 token + 1 个 V5 = 8
|
||||
assert model.n_fusion_tokens == 8, f"n_fusion_tokens={model.n_fusion_tokens}"
|
||||
print(f" fusion.n_tokens={model.n_fusion_tokens} OK")
|
||||
print("\nTEST 2 PASSED")
|
||||
|
||||
|
||||
def test_3_reproducibility(llm_type, model_path, d_model, device):
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 3: Reproducibility")
|
||||
print("=" * 60)
|
||||
|
||||
print("\n[3a] Running twice with same seed=42...")
|
||||
set_seed(42)
|
||||
enc1 = build_encoder(llm_type, model_path, d_model, device)
|
||||
enc1.eval()
|
||||
with torch.no_grad():
|
||||
o1 = enc1(SAMPLE_SMILES)
|
||||
set_seed(42)
|
||||
enc2 = build_encoder(llm_type, model_path, d_model, device)
|
||||
enc2.eval()
|
||||
with torch.no_grad():
|
||||
o2 = enc2(SAMPLE_SMILES)
|
||||
assert torch.allclose(o1, o2, atol=1e-5)
|
||||
print(f" Same seed -> identical outputs OK")
|
||||
|
||||
print("[3b] Running with different seeds...")
|
||||
set_seed(123)
|
||||
enc3 = build_encoder(llm_type, model_path, d_model, device)
|
||||
enc3.eval()
|
||||
with torch.no_grad():
|
||||
o3 = enc3(SAMPLE_SMILES)
|
||||
assert not torch.allclose(o1, o3, atol=1e-5)
|
||||
print(f" Different seeds -> different outputs OK")
|
||||
print("\nTEST 3 PASSED")
|
||||
|
||||
|
||||
def test_4_overfit(llm_type, model_path, d_model, device):
|
||||
print("\n" + "=" * 60)
|
||||
print("TEST 4: Overfit sanity (gradient + metric check)")
|
||||
print("=" * 60)
|
||||
set_seed(42)
|
||||
|
||||
from lnp_ml.modeling.models import LNPModelWithoutMPNN
|
||||
model = LNPModelWithoutMPNN(
|
||||
d_model=d_model, use_llm=True, llm_type=llm_type,
|
||||
llm_model_path=model_path, llm_device=device,
|
||||
).to(device)
|
||||
|
||||
n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
print(f"\n Trainable params: {n_trainable:,}")
|
||||
|
||||
smiles = SAMPLE_SMILES * 2 # 8 samples
|
||||
tab = make_tabular(len(smiles), device)
|
||||
target = torch.randn(len(smiles), 1, device=device)
|
||||
|
||||
opt = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=1e-3)
|
||||
lossfn = nn.MSELoss()
|
||||
|
||||
print(f" Training 50 steps on {len(smiles)} samples...")
|
||||
first_loss = None
|
||||
model.train()
|
||||
for step in range(1, 51):
|
||||
opt.zero_grad()
|
||||
pred = model.forward_delivery(smiles, tab)
|
||||
loss = lossfn(pred, target)
|
||||
loss.backward()
|
||||
opt.step()
|
||||
if first_loss is None:
|
||||
first_loss = loss.item()
|
||||
if step % 10 == 0:
|
||||
print(f" step {step:3d}: loss={loss.item():.4f}")
|
||||
|
||||
last_loss = loss.item()
|
||||
assert last_loss < first_loss
|
||||
print(f" Loss decreased: {first_loss:.4f} -> {last_loss:.4f} OK")
|
||||
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
pred = model.forward_delivery(smiles, tab)
|
||||
ss_res = ((pred - target) ** 2).sum().item()
|
||||
ss_tot = ((target - target.mean()) ** 2).sum().item()
|
||||
r2 = 1 - ss_res / (ss_tot + 1e-9)
|
||||
rmse = (ss_res / len(smiles)) ** 0.5
|
||||
print(f"\n After 50 steps on {len(smiles)} samples (overfit test):")
|
||||
print(f" R2 = {r2:.4f} (should be > 0 after overfitting)")
|
||||
print(f" RMSE = {rmse:.4f}")
|
||||
assert r2 > 0
|
||||
print(f" R2 > 0 OK")
|
||||
print("\nTEST 4 PASSED")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model_path", required=True, help="LLM 权重路径或 HuggingFace ID")
|
||||
parser.add_argument("--llm_type", required=True,
|
||||
choices=["biot5", "molt5", "moleculestm", "selfiested"])
|
||||
parser.add_argument("--lnp_repo_path", default=".", help="仓库根目录")
|
||||
parser.add_argument("--d_model", type=int, default=256)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.lnp_repo_path not in sys.path:
|
||||
sys.path.insert(0, args.lnp_repo_path)
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
print(f"Device: {device}")
|
||||
set_seed(42)
|
||||
|
||||
test_1_standalone(args.llm_type, args.model_path, args.d_model, device)
|
||||
test_2_integration(args.llm_type, args.model_path, args.d_model, device)
|
||||
test_3_reproducibility(args.llm_type, args.model_path, args.d_model, device)
|
||||
test_4_overfit(args.llm_type, args.model_path, args.d_model, device)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("ALL TESTS PASSED")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
x
Reference in New Issue
Block a user