# 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】开始正常运行。