mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 19:13:21 +08:00
118 lines
4.3 KiB
Python
118 lines
4.3 KiB
Python
"""在 moleculestm 环境中运行：把数据集里所有唯一 SMILES 编码成 MoleculeSTM(Graph 版) 表征缓存。
|
||
|
||
用法（工作目录 = lnp_ml 根）:
|
||
conda activate moleculestm
|
||
export HF_HUB_OFFLINE=1 # 只用本地权重
|
||
python scripts/precompute_moleculestm.py \
|
||
--stm-repo /path/to/MoleculeSTM \
|
||
--ckpt MoleculeSTM_ckpt/pretrained_MoleculeSTM/<某个Graph变体>/molecule_model.pth \
|
||
--out data/processed/moleculestm_embeddings.npz
|
||
|
||
产物 .npz 含 "smiles" (N,) 与 "embeddings" (N, D=300)，供主项目查表使用。
|
||
"""
|
||
import argparse
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import pandas as pd
|
||
import torch
|
||
|
||
DATA_FILES = [
|
||
"data/interim/internal.csv",
|
||
"data/external/all_data_LiON.csv",
|
||
]
|
||
GLOB_FILES = [
|
||
"data/external/all_amine_split_for_LiON/cv_*/train.csv",
|
||
"data/external/all_amine_split_for_LiON/cv_*/test.csv",
|
||
]
|
||
|
||
|
||
def collect_smiles(root: Path) -> list:
|
||
seen: set = set()
|
||
files = [root / f for f in DATA_FILES]
|
||
for pattern in GLOB_FILES:
|
||
files += sorted(root.glob(pattern))
|
||
for f in files:
|
||
if not f.exists():
|
||
print(f"跳过不存在的文件: {f}")
|
||
continue
|
||
df = pd.read_csv(f, low_memory=False)
|
||
if "smiles" in df.columns:
|
||
seen.update(df["smiles"].dropna().astype(str).tolist())
|
||
return sorted(seen)
|
||
|
||
|
||
def build_encoder(stm_repo: str, ckpt: str, device: str):
|
||
"""加载 MoleculeSTM Graph 分子编码器（GNN_graphpred 包装的 GIN）。"""
|
||
sys.path.insert(0, stm_repo)
|
||
from MoleculeSTM.models import GNN, GNN_graphpred # 官方模型定义
|
||
|
||
node = GNN(num_layer=5, emb_dim=300, JK="last", drop_ratio=0.0, gnn_type="gin")
|
||
model = GNN_graphpred(
|
||
num_layer=5, emb_dim=300, num_tasks=1, JK="last",
|
||
graph_pooling="mean", molecule_node_model=node,
|
||
)
|
||
state = torch.load(ckpt, map_location="cpu")
|
||
model.load_state_dict(state) # 与官方一致，key 完全匹配
|
||
model.eval().to(device)
|
||
return model
|
||
|
||
|
||
def smiles_to_data(smiles: str, stm_repo: str):
|
||
"""SMILES -> torch_geometric Data（用官方 mol_to_graph_data_obj_simple）。"""
|
||
from rdkit import Chem
|
||
# 官方图特征化函数（不同版本路径可能是 datasets.utils 或 datasets.molecule_datasets）
|
||
from MoleculeSTM.datasets.utils import mol_to_graph_data_obj_simple
|
||
|
||
mol = Chem.MolFromSmiles(smiles)
|
||
if mol is None:
|
||
return None
|
||
return mol_to_graph_data_obj_simple(mol)
|
||
|
||
|
||
@torch.no_grad()
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--stm-repo", required=True, help="官方 MoleculeSTM 仓库路径")
|
||
ap.add_argument("--ckpt", required=True, help="molecule_model.pth 路径")
|
||
ap.add_argument("--out", default="data/processed/moleculestm_embeddings.npz")
|
||
ap.add_argument("--root", default=".", help="lnp_ml 仓库根目录")
|
||
ap.add_argument("--batch-size", type=int, default=256)
|
||
ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
||
args = ap.parse_args()
|
||
|
||
from torch_geometric.data import DataLoader
|
||
from torch_geometric.nn import global_mean_pool
|
||
|
||
model = build_encoder(args.stm_repo, args.ckpt, args.device)
|
||
smiles = collect_smiles(Path(args.root))
|
||
print(f"收集到 {len(smiles)} 个唯一 SMILES，开始编码...")
|
||
|
||
# 图特征化（记录成功的 SMILES，跳过 RDKit 解析失败的）
|
||
data_list, kept = [], []
|
||
for s in smiles:
|
||
d = smiles_to_data(s, args.stm_repo)
|
||
if d is not None:
|
||
data_list.append(d)
|
||
kept.append(s)
|
||
if len(kept) < len(smiles):
|
||
print(f"警告: {len(smiles) - len(kept)} 个 SMILES 无法图特征化，已跳过")
|
||
|
||
loader = DataLoader(data_list, batch_size=args.batch_size, shuffle=False)
|
||
chunks = []
|
||
for batch in loader:
|
||
batch = batch.to(args.device)
|
||
graph_repr, _ = model(batch.x, batch.edge_index, batch.edge_attr, batch.batch) # [B, 300]
|
||
chunks.append(graph_repr.cpu().numpy().astype(np.float32))
|
||
print(f" 已编码 {sum(c.shape[0] for c in chunks)}/{len(kept)}")
|
||
embeddings = np.vstack(chunks)
|
||
|
||
out = Path(args.out)
|
||
out.parent.mkdir(parents=True, exist_ok=True)
|
||
np.savez_compressed(out, smiles=np.array(kept), embeddings=embeddings)
|
||
print(f"已保存 {embeddings.shape} 到 {out}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |