mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 16:23:20 +08:00
7eff2b3d02
三个预训练分子编码器作为额外的 chemical token 接入（三选一，互斥）：
- MolFormer-XL 768d（SMILES 序列，11亿分子预训练）
- GraphMVP 300d（2D 图自监督，ICLR'22）
- GROVER 3200d（图 Transformer，NeurIPS'20）
实现：embedding 离线编码存 npy，模型内按 {smiles: vector} 查表注入，
不改动前向逻辑，默认全关、不影响既有实验。
- models.py: 新增三个 CHEM_KEYS_WITH_*、统一的 _load_offline_emb 查表、
proj_input_dims 按开关裁剪、forward 注入；子类同步转发参数
- nested_cv_optuna.py / pretrain.py: 新增 --{molformer,graphmvp,grover}-emb/-csv
- scripts_run/encode_embeddings/: 三个离线编码脚本 + README（含权重来源与
新旧环境兼容补丁说明）
- results/pretrained_encoders/: 五组 15 折结果 summary
(trials15/inner3/repeats3/seed42, 均带 external 预训练)
注：npy embedding 与预训练权重未入库（可由编码脚本复现）
52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
"""MolFormer-XL 离线 embedding (768d)
|
||
|
||
依赖：transformers>=4.57（需 transformers.masking_utils）。若与项目 pin 的
|
||
4.45 冲突，装到独立目录后用 sys.path 隔离：
|
||
pip install --target=/tmp/mf_libs "transformers>=4.57" --no-deps
|
||
pip install --target=/tmp/mf_libs tokenizers>=0.22,<=0.23 --no-deps
|
||
# 本脚本开头 sys.path.insert(0, '/tmp/mf_libs')
|
||
|
||
用法：
|
||
python encode_molformer.py --csv data/interim/internal.csv \
|
||
--out data/interim/molformer_emb.npy
|
||
"""
|
||
import argparse
|
||
import numpy as np
|
||
import pandas as pd
|
||
import torch
|
||
|
||
MODEL = "ibm/MoLFormer-XL-both-10pct"
|
||
|
||
|
||
def main(csv_path: str, out_npy: str, col: str = "smiles", batch: int = 32) -> None:
|
||
from transformers import AutoModel, AutoTokenizer
|
||
|
||
smiles = pd.read_csv(csv_path, low_memory=False)[col].astype(str).tolist()
|
||
print(f"{len(smiles)} 行, {len(set(smiles))} 唯一分子")
|
||
|
||
tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
|
||
model = AutoModel.from_pretrained(
|
||
MODEL, trust_remote_code=True, deterministic_eval=True
|
||
).cuda().eval()
|
||
|
||
embs = []
|
||
with torch.no_grad():
|
||
for i in range(0, len(smiles), batch):
|
||
b = tok(smiles[i:i + batch], padding=True, truncation=True,
|
||
max_length=512, return_tensors="pt")
|
||
b = {k: v.cuda() for k, v in b.items()}
|
||
embs.append(model(**b).pooler_output.cpu().numpy())
|
||
|
||
E = np.vstack(embs)
|
||
np.save(out_npy, E)
|
||
print(f"{out_npy} {E.shape}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
p = argparse.ArgumentParser()
|
||
p.add_argument("--csv", required=True)
|
||
p.add_argument("--out", required=True)
|
||
p.add_argument("--col", default="smiles")
|
||
a = p.parse_args()
|
||
main(a.csv, a.out, a.col)
|