diff --git a/.gitignore b/.gitignore index d0f020b..2b64c96 100644 --- a/.gitignore +++ b/.gitignore @@ -188,4 +188,14 @@ cython_debug/ .ruff_cache/ # PyPI configuration file -.pypirc \ No newline at end of file +.pypirc + +# ============ LLM encoder branch additions ============ +# 大文件 / 模型权重（不进仓库） +*.pt +*.pth +*.ckpt +*.safetensors + +# 数据 +*.parquet diff --git a/README_llm_encoders.md b/README_llm_encoders.md new file mode 100644 index 0000000..98b6593 --- /dev/null +++ b/README_llm_encoders.md @@ -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 \ + --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 节对照表）。 diff --git a/RESULTS.md b/RESULTS.md new file mode 100644 index 0000000..8800f77 --- /dev/null +++ b/RESULTS.md @@ -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`。 diff --git a/docs/colab_biot5_quickstart.md b/docs/colab_biot5_quickstart.md new file mode 100644 index 0000000..912d3e2 --- /dev/null +++ b/docs/colab_biot5_quickstart.md @@ -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。 diff --git a/docs/colab_moleculestm_quickstart.md b/docs/colab_moleculestm_quickstart.md new file mode 100644 index 0000000..58311b7 --- /dev/null +++ b/docs/colab_moleculestm_quickstart.md @@ -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 开始正常运行。 + diff --git a/docs/colab_molt5_quickstart.md b/docs/colab_molt5_quickstart.md new file mode 100644 index 0000000..4cff66c --- /dev/null +++ b/docs/colab_molt5_quickstart.md @@ -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 开始正常运行。 + diff --git a/docs/colab_selfiested_quickstart.md b/docs/colab_selfiested_quickstart.md new file mode 100644 index 0000000..657b546 --- /dev/null +++ b/docs/colab_selfiested_quickstart.md @@ -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】开始正常运行。 + diff --git a/lnp_ml/modeling/encoders/__init__.py b/lnp_ml/modeling/encoders/__init__.py index 2ab762a..2b540d7 100644 --- a/lnp_ml/modeling/encoders/__init__.py +++ b/lnp_ml/modeling/encoders/__init__.py @@ -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}") diff --git a/lnp_ml/modeling/encoders/llm_encoder.py b/lnp_ml/modeling/encoders/llm_encoder.py new file mode 100644 index 0000000..9147fe2 --- /dev/null +++ b/lnp_ml/modeling/encoders/llm_encoder.py @@ -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 分子输入格式： SMILES/SELFIES +BOM_TOKEN = "" +EOM_TOKEN = "" + + +class LLMEncoder(nn.Module): + """ + 基于 BioT5-plus 的冻结分子编码器。 + + 流程： + SMILES + -> "" + SMILES + "" + -> 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) \ No newline at end of file diff --git a/lnp_ml/modeling/encoders/molecule_gnn_model.py b/lnp_ml/modeling/encoders/molecule_gnn_model.py new file mode 100644 index 0000000..cefabba --- /dev/null +++ b/lnp_ml/modeling/encoders/molecule_gnn_model.py @@ -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] diff --git a/lnp_ml/modeling/encoders/moleculestm_encoder.py b/lnp_ml/modeling/encoders/moleculestm_encoder.py new file mode 100644 index 0000000..b8b6fe5 --- /dev/null +++ b/lnp_ml/modeling/encoders/moleculestm_encoder.py @@ -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) diff --git a/lnp_ml/modeling/encoders/molt5_encoder.py b/lnp_ml/modeling/encoders/molt5_encoder.py new file mode 100644 index 0000000..0e91fbc --- /dev/null +++ b/lnp_ml/modeling/encoders/molt5_encoder.py @@ -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) diff --git a/lnp_ml/modeling/encoders/selfiested_encoder.py b/lnp_ml/modeling/encoders/selfiested_encoder.py new file mode 100644 index 0000000..cf8347a --- /dev/null +++ b/lnp_ml/modeling/encoders/selfiested_encoder.py @@ -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) diff --git a/lnp_ml/modeling/final_train_optuna_cv.py b/lnp_ml/modeling/final_train_optuna_cv.py index 9e2173f..3a93e49 100644 --- a/lnp_ml/modeling/final_train_optuna_cv.py +++ b/lnp_ml/modeling/final_train_optuna_cv.py @@ -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() \ No newline at end of file diff --git a/lnp_ml/modeling/models.py b/lnp_ml/modeling/models.py index c85a609..201c248 100644 --- a/lnp_ml/modeling/models.py +++ b/lnp_ml/modeling/models.py @@ -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, ) - diff --git a/lnp_ml/modeling/pretrain.py b/lnp_ml/modeling/pretrain.py index b7d9d66..85d3358 100644 --- a/lnp_ml/modeling/pretrain.py +++ b/lnp_ml/modeling/pretrain.py @@ -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"]) diff --git a/verify_llm_encoder.py b/verify_llm_encoder.py new file mode 100644 index 0000000..fd28d4b --- /dev/null +++ b/verify_llm_encoder.py @@ -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_type \ + --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()