mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 19:13:21 +08:00
59 lines
2.2 KiB
Python
59 lines
2.2 KiB
Python
"""在 mole 环境中运行：把数据集所有唯一 SMILES 编码成 MolE 表征缓存。
|
||
|
||
用法（工作目录 = lnp_ml 根，且已 pip install -e mole_public）:
|
||
pyenv activate mole
|
||
python scripts/precompute_mole.py --ckpt MolE_ckpt/model.pth \
|
||
--out data/processed/mole_embeddings.npz
|
||
|
||
产物 .npz 含 "smiles" (N,) 与 "embeddings" (N, D=8000)。
|
||
"""
|
||
import argparse
|
||
from pathlib import Path
|
||
import numpy as np
|
||
import pandas as pd
|
||
|
||
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()
|
||
files = [root / f for f in DATA_FILES]
|
||
for pat in GLOB_FILES:
|
||
files += sorted(root.glob(pat))
|
||
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 main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--ckpt", default="MolE_ckpt/model.pth", help="MolE 权重(或权重目录)")
|
||
ap.add_argument("--out", default="data/processed/mole_embeddings.npz")
|
||
ap.add_argument("--root", default=".")
|
||
ap.add_argument("--batch-size", type=int, default=32)
|
||
ap.add_argument("--num-workers", type=int, default=4)
|
||
args = ap.parse_args()
|
||
|
||
from mole import mole_predict
|
||
|
||
smiles = collect_smiles(Path(args.root))
|
||
print(f"收集到 {len(smiles)} 个唯一 SMILES，开始 MolE 编码...")
|
||
|
||
emb = mole_predict.encode(
|
||
smiles=smiles, pretrained_model=args.ckpt,
|
||
batch_size=args.batch_size, num_workers=args.num_workers,
|
||
)
|
||
emb = np.asarray(emb, dtype=np.float32)
|
||
print(f"embeddings.shape = {emb.shape}") # 期望 (N, 8000)
|
||
assert emb.shape[0] == len(smiles), "encode 丢弃了部分 SMILES，需对齐后再存"
|
||
|
||
out = Path(args.out); out.parent.mkdir(parents=True, exist_ok=True)
|
||
np.savez_compressed(out, smiles=np.array(smiles), embeddings=emb)
|
||
print(f"已保存 {emb.shape} 到 {out}")
|
||
|
||
if __name__ == "__main__":
|
||
main() |