scripts/inference/validate_bbbp_native.py

#!/usr/bin/env python3
"""
SMI-TED BBBP 分类验证 — 不启用 FlagGems。
使用 encode() 提取嵌入 + XGBoost 分类，对标官方 notebook ROC-AUC = 0.9194。

用法:
  python scripts/inference/validate_bbbp_native.py
"""

import sys
import os
import time
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

from models.smi_ted.smi_ted_light.load import load_smi_ted, normalize_smiles


def main():
    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"

    # ── 加载数据 ──
    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()