File size: 3,850 Bytes
5c5099a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | #!/usr/bin/env python3
"""
SMI-TED BBBP 分类验证 — 启用 FlagGems 算子。
使用 encode() 提取嵌入 + XGBoost 分类，对标官方 notebook ROC-AUC = 0.9194。
用法:
python scripts/inference/validate_bbbp_flaggems.py
# 排除特定算子
python scripts/inference/validate_bbbp_flaggems.py --flaggems-exclude mm gelu
"""
import sys
import os
import time
import argparse
import warnings
warnings.filterwarnings("ignore")
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
_PROJECT_ROOT = os.path.dirname(os.path.dirname(_SCRIPT_DIR))
sys.path.insert(0, _PROJECT_ROOT)
import torch
import pandas as pd
import numpy as np
from xgboost import XGBClassifier
from sklearn.metrics import roc_auc_score
import flag_gems
from models.smi_ted.smi_ted_light.load import load_smi_ted, normalize_smiles
def main():
parser = argparse.ArgumentParser(description="SMI-TED BBBP 分类验证 (FlagGems)")
parser.add_argument("--flaggems-exclude", type=str, nargs="*", default=["mm", "gelu"],
help="排除导致海光崩溃的算子 (默认: mm gelu)")
args = parser.parse_args()
model_dir = os.path.join(_PROJECT_ROOT, "models", "smi_ted", "smi_ted_light")
train_path = os.path.join(_PROJECT_ROOT, "models", "smi_ted", "finetune", "moleculenet", "bbbp", "train.csv")
test_path = os.path.join(_PROJECT_ROOT, "models", "smi_ted", "finetune", "moleculenet", "bbbp", "test.csv")
device = "cuda" if torch.cuda.is_available() else "cpu"
# ── 启用 FlagGems ──
exclude = set(args.flaggems_exclude)
log_path = os.path.join(_PROJECT_ROOT, "flag_gems_enable.txt")
flag_gems.enable(unused=exclude if exclude else None, record=True, once=True, path=log_path)
print(f"FlagGems 已启用 (version: {flag_gems.__version__})")
if exclude:
print(f" 排除: {exclude}")
print()
# ── 加载数据 ──
train_df = pd.read_csv(train_path)
test_df = pd.read_csv(test_path)
train_df["norm_smiles"] = train_df["smiles"].apply(normalize_smiles)
test_df["norm_smiles"] = test_df["smiles"].apply(normalize_smiles)
train_df = train_df.dropna()
test_df = test_df.dropna()
print(f"数据集: BBBP (血脑屏障渗透)")
print(f" 训练集: {len(train_df)} 分子")
print(f" 测试集: {len(test_df)} 分子")
# ── 加载模型 ──
print("加载模型...")
model = load_smi_ted(folder=model_dir, ckpt_filename="smi-ted-Light_40.pt")
model.eval()
if device == "cuda":
model = model.cuda()
# ── 提取嵌入 ──
print("提取训练集嵌入...")
t0 = time.time()
with torch.no_grad():
train_emb = model.encode(train_df["norm_smiles"].tolist())
print(f" 训练集: {train_emb.shape} 耗时: {time.time() - t0:.1f}s")
print("提取测试集嵌入...")
t0 = time.time()
with torch.no_grad():
test_emb = model.encode(test_df["norm_smiles"].tolist())
print(f" 测试集: {test_emb.shape} 耗时: {time.time() - t0:.1f}s")
# ── XGBoost 分类 ──
print("训练 XGBoost 分类器...")
xgb = XGBClassifier(n_estimators=2000, learning_rate=0.04, max_depth=8)
xgb.fit(train_emb, train_df["p_np"])
y_prob = xgb.predict_proba(test_emb)[:, 1]
roc_auc = roc_auc_score(test_df["p_np"], y_prob)
# ── 输出结果 ──
print(f"\n{'=' * 50}")
print(f"BBBP 分类验证结果 (FlagGems: 已启用)")
print(f"{'=' * 50}")
print(f" ROC-AUC: {roc_auc:.4f}")
print(f" 官方参考值: 0.9194")
print(f" 差异: {abs(roc_auc - 0.9194):.4f}")
if abs(roc_auc - 0.9194) < 0.02:
print(f"\n ✓ 验证通过 (差异 < 0.02)")
else:
print(f"\n ✗ 验证未通过!")
print(f"{'=' * 50}")
if __name__ == "__main__":
main()
|