# LLM 分子编码器分支（LLM Encoder Branch）

本模块在 LNP 多任务预测模型上新增了一个**可选的 LLM 分子编码器分支**。它把可电离脂质的 SMILES 送入一个冻结的预训练分子模型，提取语义向量，作为一个额外的 token（V5）追加到模型的 fusion 序列中参与多任务预测。

支持四种分子编码器，可通过参数自由切换，用于横向对比：

| llm_type | 模型 | 年份 / 发表 | 参数量 | hidden | 输入 |
|---|---|---|---|---|---|
| `biot5` | BioT5-plus | 2024, arXiv | ~250M | 768 | SMILES |
| `molt5` | MolT5-base | 2022, EMNLP | 109M | 768 | SMILES |
| `moleculestm` | MoleculeSTM GNN | 2023, Nature MI | ~2M (GNN) | 300 | 分子图 |
| `selfiested` | SELFIES-TED | 2025, IBM | 358M | 1024 | SELFIES |

---

## 1. 设计说明

- **非侵入式**：默认 `use_llm=False`，此时模型行为与原版完全一致，不影响既有功能。只有显式开启才会启用 LLM 分支。
- **统一接口**：四个 encoder 的 `forward(List[str]) -> [B, d_model]` 接口完全一致，通过 `models.build_llm_encoder()` 工厂函数按 `llm_type` 实例化。
- **冻结主体**：LLM / GNN 主体全程冻结，只有内部一层 `projection`（LayerNorm + Linear + ReLU + Dropout）参与训练，参数量增加极小。
- **惰性导入**：`encoders/__init__.py` 对四个 LLM encoder 采用惰性导入，未安装 `torch_geometric` / `selfies` 等可选依赖时，只要不使用对应 encoder，导入不会报错。

---

## 2. 文件清单

新增文件（均在 `lnp_ml/modeling/encoders/` 下）：

```
encoders/
├── llm_encoder.py            # BioT5-plus  (类名 LLMEncoder)
├── molt5_encoder.py          # MolT5       (类名 MolT5Encoder)
├── moleculestm_encoder.py    # MoleculeSTM (类名 MoleculeSTMEncoder)
├── molecule_gnn_model.py     # MoleculeSTM 的 GNN 网络定义（被上一个文件 import）
└── selfiested_encoder.py     # SELFIES-TED (类名 SELFIESTEDEncoder)
```

修改文件：

```
lnp_ml/modeling/models.py            # 新增 use_llm / llm_type 等参数 + build_llm_encoder 工厂 + V5 追加逻辑；desc 维度 210 -> 217
lnp_ml/modeling/encoders/__init__.py # 惰性导出四个 LLM encoder
verify_llm_encoder.py                # LLM 模块验证脚本（四项测试）
```

> 注：`models.py` 中 `DEFAULT_INPUT_DIMS["desc"]` 已从 210 更新为 **217**，以匹配新版 RDKit 的描述符维度。

---

## 3. 依赖安装

**基础依赖（所有模型都需要）：**
```bash
pip install rdkit loguru transformers sentencepiece optuna
```

**按模型追加：**

| 模型 | 额外依赖 |
|---|---|
| biot5 | 无（基础依赖即可） |
| molt5 | 无 |
| moleculestm | `pip install torch_geometric torch_scatter torch_sparse ogb` |
| selfiested | `pip install selfies` |

---

## 4. 预训练权重

| 模型 | 是否需手动下载 | 说明 |
|---|---|---|
| **molt5** | 否 | 训练时自动从 HuggingFace 拉取 `laituan245/molt5-base` |
| **selfiested** | 否 | 训练时自动从 HuggingFace 拉取 `ibm/materials.selfies-ted` |
| **biot5** | 是 | 需下载 BioT5-plus 权重文件夹，见下方 |
| **moleculestm** | 是 | 需下载 MoleculeSTM GNN 权重，见下方 |

**BioT5-plus 权重下载：**
从 HuggingFace `QizhiPei/biot5-plus-base` 下载整个权重文件夹（含 `config.json`、`tokenizer_config.json` 等），放到本地任意路径，训练时用 `--llm-model-path <该路径>` 指定。也可直接传 HuggingFace ID `QizhiPei/biot5-plus-base` 在线加载。

**MoleculeSTM GNN 权重下载：**
```python
from huggingface_hub import snapshot_download
snapshot_download(
    repo_id="chao1224/MoleculeSTM",
    repo_type="model",
    local_dir="<本地目录>/moleculestm_weights",
    allow_patterns="*Graph*",
)
```
下载后使用的 GNN 权重路径为：
```
<本地目录>/moleculestm_weights/pretrained_MoleculeSTM/SciBERT-Graph-3e-5-1-1e-4-1-InfoNCE-0.1-32-32/molecule_model.pth
```
训练时用 `--llm-model-path <该 .pth 路径>` 指定。

---

## 5. 用法

### 5.1 代码中直接构造模型

```python
from lnp_ml.modeling.models import LNPModelWithoutMPNN

# 不启用 LLM（与原版行为一致）
model = LNPModelWithoutMPNN(d_model=256)

# 启用 LLM 分支（以 MolT5 为例）
model = LNPModelWithoutMPNN(
    d_model=256,
    use_llm=True,
    llm_type="molt5",                       # biot5 / molt5 / moleculestm / selfiested
    llm_model_path="laituan245/molt5-base", # 见上表，不同模型路径不同
    llm_device="cuda",
)
```

各模型对应的 `llm_type` 和 `llm_model_path`：

| llm_type | llm_model_path 示例 |
|---|---|
| `biot5` | `QizhiPei/biot5-plus-base` 或本地 BioT5 权重文件夹路径 |
| `molt5` | `laituan245/molt5-base` |
| `moleculestm` | `.../moleculestm_weights/.../molecule_model.pth` |
| `selfiested` | `ibm/materials.selfies-ted` |

### 5.2 验证 LLM 模块（训练前建议先跑）

```bash
python verify_llm_encoder.py \
    --model_path <llm_model_path> \
    --lnp_repo_path <仓库根目录>
```
四项测试（输出维度、维度对齐、可复现性、过拟合 sanity）全部 PASS 再训练。

### 5.3 训练

> 训练脚本（`final_train_optuna_cv.py` / `pretrain.py`）需要支持 `--use-llm` / `--llm-type` / `--llm-model-path` / `--llm-device` 参数。各模型的完整 Colab 操作步骤见 `docs/` 下四份 quickstart 指南。

---

## 6. 实验结果

四个模型在 LNP 测试集上的对比（核心指标 delivery R²，baseline 为不含 LLM 的版本）：

| 任务 | Baseline | MolT5 | MoleculeSTM | BioT5-plus | SELFIES-TED |
|---|---|---|---|---|---|
| delivery R² | 0.373 | 0.602 | 0.552 | **0.633** | 0.618 |
| size R² | 0.215 | 0.440 | 0.408 | 0.423 | 0.381 |
| pdi acc | 0.786 | 0.748 | 0.710 | 0.695 | 0.664 |
| ee acc | 0.679 | 0.702 | 0.679 | 0.679 | 0.679 |
| toxic acc | 0.980 | 0.980 | 0.951 | 0.960 | 0.941 |
| biodist KL | 0.293 | 0.727 | 0.646 | 0.703 | 0.643 |

**主要结论：**
1. 四个 LLM 编码器都显著提升了核心回归任务，delivery R² 从 baseline 的 0.37 最高提升到 0.63。
2. 领域专用预训练比堆参数更重要：BioT5-plus（针对生物分子预训练）效果最好，而参数最大、最新的 SELFIES-TED 仅排第二。
3. 序列模型整体优于图模型：三个基于 SMILES/SELFIES 的模型均优于基于图的 MoleculeSTM。
4. MolT5 性价比最高：仅 109M 参数，delivery R² 排第三且与第一仅差 0.03，分类任务表现最均衡。
5. 共同局限：biodist（分布预测）四个 LLM 均略低于 baseline，属多任务学习中的 trade-off。

> 各模型每个任务的详细指标见 `results/<模型名>/` 目录。

---

## 附：训练脚本改动

本 PR 已经更新了两个训练脚本以支持 LLM 分支（MoE 相关代码未引入）：

- `lnp_ml/modeling/final_train_optuna_cv.py`
- `lnp_ml/modeling/pretrain.py`

两者都新增了四个命令行参数，并透传给模型：

```
--use-llm                                          # 是否启用 LLM 分支
--llm-type {biot5,molt5,moleculestm,selfiested}    # 选择哪个 encoder
--llm-model-path <路径或HF_ID>                     # 见第 4/5 节
--llm-device cuda                                  # LLM 推理设备
```

预训练保存的 checkpoint 的 config 中也会记录 LLM 设置，恢复模型时自动读取。

### 训练示例（以 MolT5 为例）

```bash
# 预训练
python -m lnp_ml.modeling.pretrain main \
    --train-path data/processed/train_pretrain.parquet \
    --val-path   data/processed/val_pretrain.parquet \
    --use-llm --llm-type molt5 \
    --llm-model-path laituan245/molt5-base \
    --llm-device cuda --epochs 50 --lr 1e-4 --device cuda

# 正式训练
python lnp_ml/modeling/final_train_optuna_cv.py \
    --init-from-pretrain models/pretrain_delivery.pt \
    --use-llm --llm-type molt5 \
    --llm-model-path laituan245/molt5-base \
    --llm-device cuda \
    --n-trials 20 --epochs-per-trial 30 --seed 42 \
    --device cuda --output-dir models/final
```

切换其他模型只需改 `--llm-type` 和 `--llm-model-path`（见第 5.1 节对照表）。