lnp_ml/docs/colab_moleculestm_quickstart.md

LNP-ML + MoleculeSTM GNN 快速上手指南

适用于：拿到代码包、从未运行过任何代码的新组员。

---

前提条件（开始前确认）

• Google Drive 的 MyDrive/lnp_project_moleculestm/ 下已有以下内容：
  • lnp_ml/：项目代码（从 BioT5 版本复制来的，已修改）
  • moleculestm_weights/：GNN 预训练权重

如果以上文件夹不存在，先看文末的【从零开始】章节。

---

每次重新打开 Colab 都需要运行

Cell 1：挂载 Drive

from google.colab import drive
drive.mount('/content/drive')

弹出授权页面时，确认选择有 lnp_project_moleculestm 的 Google 账号。

---

Cell 2：安装依赖

!pip install rdkit loguru transformers sentencepiece optuna \
             torch_geometric torch_scatter torch_sparse ogb -q

---

Cell 3：配置路径

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：确认文件完整

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 模块

!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 分钟）

!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

完成后确认权重存在：

print("✓ 预训练权重存在" if os.path.exists(f"{LNP_PATH}/models/pretrain_delivery.pt")
      else "✗ 预训练权重未生成")

参考结果：Best val_loss: 0.5667

---

Cell 7：正式训练（约 20 分钟）

!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

完成后确认：

print("✓ 模型权重存在" if os.path.exists(f"{LNP_PATH}/models/final/model.pt")
      else "✗ 模型未生成")

参考结果：Best val_loss: 2.4433

---

Cell 8：测试评估

!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：备份模型

!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 版本复制代码

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 分钟）

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 完成代码适配：

# 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 ✓")

# 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 ✓")

# 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 ✓")

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