RYDE-WORK/lnp_ml
1
0
Fork 0
mirror of https://github.com/RYDE-WORK/lnp_ml.git synced 2026-07-25 15:35:25 +08:00
Code Issues Packages Projects Releases Wiki Activity
lnp_ml/models/abl_full/s1_moe/seed42/local_uncommitted.patch

648 lines
26 KiB
Diff
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

diff --git a/lnp_ml/interpretability/token_importance.py b/lnp_ml/interpretability/token_importance.py
index ee0d1d8..605a19b 100644
--- a/lnp_ml/interpretability/token_importance.py
+++ b/lnp_ml/interpretability/token_importance.py
@@ -40,10 +40,7 @@ BIODIST_ORGANS = ["lymph_nodes", "heart", "liver", "spleen", "lung", "kidney", "
def get_token_names(model: Union[LNPModel, LNPModelWithoutMPNN]) -> List[str]:
- if model.use_mpnn:
- names = ["mpnn", "morgan", "maccs", "desc", "comp", "phys", "help", "exp"]
- else:
- names = ["morgan", "maccs", "desc", "comp", "phys", "help", "exp"]
+ names = list(model.token_order)
if getattr(model, "moe", None) is not None:
names.append("moe")
return names
@@ -300,7 +297,8 @@ def plot_token_importance(
vals_sorted = normed[order]
n_tokens = len(token_names)
- split_idx = 4 if "mpnn" in token_names else 3
+ _mol_tokens = {"mpnn", "chemeleon", "unimol", "morgan", "maccs", "desc"}
+ split_idx = sum(1 for n in token_names if n in _mol_tokens)
channel_a_set = set(token_names[:split_idx])
colors = [color_a if n in channel_a_set else color_b for n in names_sorted]
diff --git a/lnp_ml/modeling/benchmark.py b/lnp_ml/modeling/benchmark.py
index 6ac6b20..af9cba0 100644
--- a/lnp_ml/modeling/benchmark.py
+++ b/lnp_ml/modeling/benchmark.py
@@ -34,7 +34,7 @@ def find_mpnn_ensemble_paths(base_dir: Path = DEFAULT_MPNN_ENSEMBLE_DIR) -> List
from lnp_ml.modeling.models import LNPModel, LNPModelWithoutMPNN
-
+from lnp_ml.utils.seed import set_global_seed
app = typer.Typer()
@@ -172,6 +172,8 @@ def train_fold(
early_stopping = EarlyStopping(patience=patience)
best_val_loss = float("inf")
+ best_val_rmse = 0.0
+ best_val_r2 = 0.0
best_state = None
history = []
@@ -202,6 +204,8 @@ def train_fold(
if val_metrics["loss"] < best_val_loss:
best_val_loss = val_metrics["loss"]
+ best_val_rmse = val_metrics.get("rmse", 0)
+ best_val_r2 = val_metrics.get("r2", 0)
best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
logger.info(f" -> New best val_loss: {best_val_loss:.4f}")
@@ -239,8 +243,8 @@ def train_fold(
return {
"fold_idx": fold_idx,
"best_val_loss": best_val_loss,
- "best_val_rmse": history[-1]["val_rmse"] if history else 0,
- "best_val_r2": history[-1]["val_r2"] if history else 0,
+ "best_val_rmse": best_val_rmse,
+ "best_val_r2": best_val_r2,
"epochs_trained": len(history),
}
@@ -255,6 +259,8 @@ def create_model(
use_mpnn: bool = False,
mpnn_ensemble_paths: Optional[List[str]] = None,
mpnn_device: str = "cpu",
+ chemeleon_cache: Optional[str] = None,
+ unimol_cache: Optional[str] = None,
) -> nn.Module:
"""创建模型实例"""
if use_mpnn:
@@ -267,6 +273,8 @@ def create_model(
dropout=dropout,
mpnn_ensemble_paths=mpnn_ensemble_paths,
mpnn_device=mpnn_device,
+ chemeleon_cache_path=chemeleon_cache,
+ unimol_cache_path=unimol_cache,
)
else:
return LNPModelWithoutMPNN(
@@ -276,6 +284,8 @@ def create_model(
fusion_strategy=fusion_strategy,
head_hidden_dim=head_hidden_dim,
dropout=dropout,
+ chemeleon_cache_path=chemeleon_cache,
+ unimol_cache_path=unimol_cache,
)
@@ -295,12 +305,18 @@ def main(
mpnn_checkpoint: Optional[str] = None,
mpnn_ensemble_paths: Optional[str] = None,
mpnn_device: str = "cpu",
+ use_chemeleon: bool = False,
+ chemeleon_cache: str = "data/processed/chemeleon_embeddings.npz",
+ use_unimol: bool = False,
+ unimol_cache: str = "data/processed/unimol_embeddings.npz",
# 训练参数
batch_size: int = 64,
lr: float = 1e-4,
weight_decay: float = 1e-5,
epochs: int = 50,
patience: int = 10,
+ # 随机种子
+ seed: int = 42,
# 设备
device: str = "cuda" if torch.cuda.is_available() else "cpu",
):
@@ -311,6 +327,8 @@ def main(
使用 --use-mpnn 启用 MPNN encoder。
"""
logger.info(f"Using device: {device}")
+ set_global_seed(seed)
+ logger.info(f"Global seed set to {seed}")
device = torch.device(device)
# 解析 MPNN 参数
@@ -349,6 +367,11 @@ def main(
"dropout": dropout,
"use_mpnn": use_mpnn,
"mpnn_ensemble_paths": mpnn_paths,
+ "use_chemeleon": use_chemeleon,
+ "chemeleon_cache": chemeleon_cache if use_chemeleon else None,
+ "use_unimol": use_unimol,
+ "unimol_cache": unimol_cache if use_unimol else None,
+ "seed": seed,
"lr": lr,
"weight_decay": weight_decay,
"batch_size": batch_size,
@@ -405,6 +428,8 @@ def main(
use_mpnn=use_mpnn,
mpnn_ensemble_paths=mpnn_paths,
mpnn_device=device.type,
+ chemeleon_cache=(chemeleon_cache if use_chemeleon else None),
+ unimol_cache=(unimol_cache if use_unimol else None),
)
model = model.to(device)
diff --git a/lnp_ml/modeling/encoders/__init__.py b/lnp_ml/modeling/encoders/__init__.py
index 2ab762a..eb78d5c 100644
--- a/lnp_ml/modeling/encoders/__init__.py
+++ b/lnp_ml/modeling/encoders/__init__.py
@@ -1,5 +1,11 @@
from lnp_ml.modeling.encoders.rdkit_encoder import CachedRDKitEncoder
from lnp_ml.modeling.encoders.mpnn_encoder import CachedMPNNEncoder
+from lnp_ml.modeling.encoders.chemeleon_encoder import CheMeleonEmbeddingEncoder
+from lnp_ml.modeling.encoders.unimol_encoder import UniMolEmbeddingEncoder
-__all__ = ["CachedRDKitEncoder", "CachedMPNNEncoder"]
-
+__all__ = [
+ "CachedRDKitEncoder",
+ "CachedMPNNEncoder",
+ "CheMeleonEmbeddingEncoder",
+ "UniMolEmbeddingEncoder",
+]
\ No newline at end of file
diff --git a/lnp_ml/modeling/final_train_optuna_cv.py b/lnp_ml/modeling/final_train_optuna_cv.py
index d3c7a97..4b0a8b8 100644
--- a/lnp_ml/modeling/final_train_optuna_cv.py
+++ b/lnp_ml/modeling/final_train_optuna_cv.py
@@ -176,6 +176,8 @@ def create_model(
dropout: float = 0.1,
use_mpnn: bool = False,
mpnn_device: str = "cpu",
+ chemeleon_cache: Optional[str] = None,
+ unimol_cache: Optional[str] = None,
# ============ MoE 相关（新增） ============
use_moe: bool = False,
moe_n_experts: int = 4,
@@ -203,6 +205,8 @@ def create_model(
dropout=dropout,
mpnn_ensemble_paths=ensemble_paths,
mpnn_device=mpnn_device,
+ chemeleon_cache_path=chemeleon_cache,
+ unimol_cache_path=unimol_cache,
**moe_kwargs,
)
else:
@@ -213,6 +217,8 @@ def create_model(
fusion_strategy=fusion_strategy,
head_hidden_dim=head_hidden_dim,
dropout=dropout,
+ chemeleon_cache_path=chemeleon_cache,
+ unimol_cache_path=unimol_cache,
**moe_kwargs,
)
@@ -258,6 +264,8 @@ def run_optuna_cv(
batch_size: int = 32,
n_folds: int = 3,
use_mpnn: bool = False,
+ chemeleon_cache: Optional[str] = None,
+ unimol_cache: Optional[str] = None,
seed: int = 42,
study_path: Optional[Path] = None,
pretrain_state_dict: Optional[Dict] = None,
@@ -355,6 +363,8 @@ def run_optuna_cv(
dropout=dropout,
use_mpnn=use_mpnn,
mpnn_device=device.type,
+ chemeleon_cache=chemeleon_cache,
+ unimol_cache=unimol_cache,
use_moe=use_moe,
moe_n_experts=moe_n_experts,
moe_top_k=moe_top_k,
@@ -451,7 +461,12 @@ def main(
load_delivery_head: bool = False,
# MPNN
use_mpnn: bool = False,
- # MoE（新增）
+ # CheMeleon
+ use_chemeleon: bool = False,
+ chemeleon_cache: str = "data/processed/chemeleon_embeddings.npz",
+ use_unimol: bool = False,
+ unimol_cache: str = "data/processed/unimol_embeddings.npz",
+ # MoE
use_moe: bool = False,
moe_n_experts: int = 4,
moe_top_k: int = 2,
@@ -531,6 +546,8 @@ def main(
batch_size=batch_size,
n_folds=n_folds,
use_mpnn=use_mpnn,
+ chemeleon_cache=(chemeleon_cache if use_chemeleon else None),
+ unimol_cache=(unimol_cache if use_unimol else None),
seed=seed,
study_path=study_path,
pretrain_state_dict=pretrain_state_dict,
@@ -599,6 +616,8 @@ def main(
dropout=best_params["dropout"],
use_mpnn=use_mpnn,
mpnn_device=device.type,
+ chemeleon_cache=(chemeleon_cache if use_chemeleon else None),
+ unimol_cache=(unimol_cache if use_unimol else None),
use_moe=use_moe,
moe_n_experts=moe_n_experts,
moe_top_k=moe_top_k,
@@ -651,6 +670,10 @@ def main(
"head_hidden_dim": best_params["head_hidden_dim"],
"dropout": best_params["dropout"],
"use_mpnn": use_mpnn,
+ "use_chemeleon": use_chemeleon,
+ "chemeleon_cache": chemeleon_cache if use_chemeleon else None,
+ "use_unimol": use_unimol,
+ "unimol_cache": unimol_cache if use_unimol else None,
"use_moe": use_moe,
"moe_n_experts": moe_n_experts,
"moe_top_k": moe_top_k,
diff --git a/lnp_ml/modeling/layers/llm_prompt.py b/lnp_ml/modeling/layers/llm_prompt.py
index 21b7e65..b5c6987 100644
--- a/lnp_ml/modeling/layers/llm_prompt.py
+++ b/lnp_ml/modeling/layers/llm_prompt.py
@@ -54,9 +54,14 @@ class LLMPromptEncoder(nn.Module):
_is_t5 = "t5" in _name_l
_is_qwen = "qwen" in _name_l
self._is_qwen = _is_qwen
+ _is_biot5 = "biot5" in _name_l
+ self._is_biot5 = _is_biot5
self.tokenizer = AutoTokenizer.from_pretrained(
- model_name_or_path, trust_remote_code=_is_qwen)
+ model_name_or_path,
+ trust_remote_code=_is_qwen,
+ use_fast=not _is_biot5,
+ )
if _is_qwen and self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
@@ -208,6 +213,19 @@ class LLMPromptEncoder(nn.Module):
return "[" + ", ".join(f"{x:.3f}" for x in v) + "]"
return f"{float(v):.3f}"
+ def _fmt_mol(self, smiles: str) -> str:
+ """按 backbone 期望格式化分子。
+ BioT5：SMILES -> SELFIES，用 <bom>...<eom> 紧贴包裹（官方格式，token 间无空格）；
+ 其他 backbone：原样返回 SMILES。"""
+ if not getattr(self, "_is_biot5", False):
+ return smiles
+ try:
+ import selfies as sf
+ sfs = sf.encoder(smiles) # CCO -> [C][C][O]
+ except Exception:
+ return smiles # 转换失败退回 SMILES，避免整批中断
+ return f"<bom>{sfs}<eom>"
+
def _build_rag_prompt(self, target_smiles: str, neighbors) -> str:
"""构造 RAG prompt：原始 SMILES + 邻居多任务结果（numeric）。"""
blocks = []
@@ -215,7 +233,7 @@ class LLMPromptEncoder(nn.Module):
ex = nb["extra"]
blocks.append(
f"Retrieved sample {rank}:\n"
- f"SMILES: {nb['smiles']}\n"
+ f"Molecule: {self._fmt_mol(nb['smiles'])}\n"
f"Similarity score: {nb['sim']:.3f}\n"
f"delivery_log: {self._fmt(nb['delivery'])}\n"
f"size_z: {self._fmt(ex.get('size'))}\n"
@@ -228,7 +246,7 @@ class LLMPromptEncoder(nn.Module):
return (
"Task: Encode the target LNP molecule into a retrieval-aware representation "
"for downstream multi-task property prediction. Do not output predictions.\n\n"
- f"[Target Molecule]\nSMILES: {target_smiles}\n\n"
+ f"[Target Molecule]\nMolecule: {self._fmt_mol(target_smiles)}\n\n"
"[Retrieved Similar LNP Samples]\n"
"Retrieved from the training set by fingerprint similarity, with their known "
"multi-task outcomes (delivery_log, size_z, pdi_class, ee_class, toxic, biodist; "
@@ -241,7 +259,7 @@ class LLMPromptEncoder(nn.Module):
def _get_prompt(self, s: str) -> str:
if not self.use_rag:
- return s
+ return self._fmt_mol(s)
key = f"{self._rag_pool_id}::{s}"
if key not in self._prompt_cache:
self._prompt_cache[key] = self._build_rag_prompt(s, self._retrieve_topk(s))
@@ -326,7 +344,8 @@ class LLMPromptEncoder(nn.Module):
uniq = list(dict.fromkeys(missing))
for i in range(0, len(uniq), 256):
chunk = uniq[i:i + 256]
- enc = self.tokenizer(chunk, padding=True, truncation=True,
+ mols = [self._fmt_mol(s) for s in chunk]
+ enc = self.tokenizer(mols, padding=True, truncation=True,
max_length=self.max_length, return_tensors="pt").to(device)
out = self.encoder(**enc).last_hidden_state
pooled = self._mean_pool(out, enc["attention_mask"])
@@ -335,7 +354,8 @@ class LLMPromptEncoder(nn.Module):
return torch.stack([self._cache[s] for s in smiles]).to(device)
def _encode_trainable(self, smiles, device):
- enc = self.tokenizer(list(smiles), padding=True, truncation=True,
+ mols = [self._fmt_mol(s) for s in smiles]
+ enc = self.tokenizer(mols, padding=True, truncation=True,
max_length=self.max_length, return_tensors="pt").to(device)
out = self.encoder(**enc).last_hidden_state
return self._mean_pool(out, enc["attention_mask"])
diff --git a/lnp_ml/modeling/models.py b/lnp_ml/modeling/models.py
index 1428075..87e6643 100644
--- a/lnp_ml/modeling/models.py
+++ b/lnp_ml/modeling/models.py
@@ -4,7 +4,12 @@ import torch
import torch.nn as nn
from typing import Dict, List, Optional, Literal
-from lnp_ml.modeling.encoders import CachedRDKitEncoder, CachedMPNNEncoder
+from lnp_ml.modeling.encoders import (
+ CachedRDKitEncoder,
+ CachedMPNNEncoder,
+ CheMeleonEmbeddingEncoder,
+ UniMolEmbeddingEncoder,
+)
from lnp_ml.modeling.layers import (
TokenProjector,
SetTransformer,
@@ -91,6 +96,10 @@ class LNPModel(nn.Module):
mpnn_checkpoint: Optional[str] = None,
mpnn_ensemble_paths: Optional[List[str]] = None,
mpnn_device: str = "cpu",
+ # CheMeleon encoder
+ chemeleon_cache_path: Optional[str] = None,
+ # UniMol encoder
+ unimol_cache_path: Optional[str] = None,
# 输入维度配置
input_dims: Optional[Dict[str, int]] = None,
# ============ MoE 相关 ============
@@ -121,6 +130,8 @@ class LNPModel(nn.Module):
self.input_dims = input_dims or DEFAULT_INPUT_DIMS
self.d_model = d_model
self.use_mpnn = mpnn_checkpoint is not None or mpnn_ensemble_paths is not None
+ self.use_chemeleon = chemeleon_cache_path is not None
+ self.use_unimol = unimol_cache_path is not None
# ============ Encoders ============
self.rdkit_encoder = CachedRDKitEncoder()
@@ -133,6 +144,18 @@ class LNPModel(nn.Module):
else:
self.mpnn_encoder = None
+ if self.use_chemeleon:
+ self.chemeleon_encoder = CheMeleonEmbeddingEncoder(cache_path=chemeleon_cache_path)
+ self.input_dims = {**self.input_dims, "chemeleon": self.chemeleon_encoder.embed_dim}
+ else:
+ self.chemeleon_encoder = None
+
+ if self.use_unimol:
+ self.unimol_encoder = UniMolEmbeddingEncoder(cache_path=unimol_cache_path)
+ self.input_dims = {**self.input_dims, "unimol": self.unimol_encoder.embed_dim}
+ else:
+ self.unimol_encoder = None
+
# ============ Token Projector ============
proj_input_dims = {k: v for k, v in self.input_dims.items()}
if not self.use_mpnn:
@@ -143,8 +166,16 @@ class LNPModel(nn.Module):
dropout=dropout,
)
- # token 顺序与化学侧 token 数
- self.chem_keys = CHEM_KEYS_WITH_MPNN if self.use_mpnn else CHEM_KEYS_NO_MPNN
+ # token 顺序：可选 embedding（mpnn/chemeleon）排在指纹类 token 之前
+ chem_keys: List[str] = []
+ if self.use_mpnn:
+ chem_keys.append("mpnn")
+ if self.use_chemeleon:
+ chem_keys.append("chemeleon")
+ if self.use_unimol:
+ chem_keys.append("unimol")
+ chem_keys += ["morgan", "maccs", "desc"]
+ self.chem_keys = chem_keys
self.tab_keys = TAB_KEYS
self.token_order = self.chem_keys + self.tab_keys
self.split_idx = len(self.chem_keys)
@@ -233,6 +264,10 @@ class LNPModel(nn.Module):
if self.use_mpnn:
mpnn_features = self.mpnn_encoder(smiles)
all_features["mpnn"] = mpnn_features["mpnn"].to(device)
+ if self.use_chemeleon:
+ all_features["chemeleon"] = self.chemeleon_encoder(smiles)["chemeleon"].to(device)
+ if self.use_unimol:
+ all_features["unimol"] = self.unimol_encoder(smiles)["unimol"].to(device)
all_features["morgan"] = rdkit_features["morgan"].to(device)
all_features["maccs"] = rdkit_features["maccs"].to(device)
all_features["desc"] = rdkit_features["desc"].to(device)
@@ -297,6 +332,15 @@ class LNPModel(nn.Module):
if task is None:
task = "delivery"
+ x_reg = getattr(self, "_last_pooled", None) if getattr(self, "reg_bypass", True) else None
+ x_for = {
+ "size": x_reg if x_reg is not None else fused,
+ "pdi": fused,
+ "ee": fused,
+ "delivery": x_reg if x_reg is not None else fused,
+ "biodist": fused,
+ "toxic": fused,
+ }
task_heads = {
"size": self.head.size_head,
"pdi": self.head.pdi_head,
@@ -305,7 +349,7 @@ class LNPModel(nn.Module):
"biodist": self.head.biodist_head,
"toxic": self.head.toxic_head,
}
- return task_heads[task](fused)
+ return task_heads[task](x_for[task])
def forward_replacing_token(
self,
@@ -347,7 +391,8 @@ class LNPModel(nn.Module):
) -> torch.Tensor:
"""仅预测 delivery（用于 pretrain）。返回 [B, 1]。"""
fused = self.forward_backbone(smiles, tabular)
- return self.head.delivery_head(fused)
+ x_reg = getattr(self, "_last_pooled", None) if getattr(self, "reg_bypass", True) else None
+ return self.head.delivery_head(x_reg if x_reg is not None else fused)
def forward(
self,
@@ -369,6 +414,10 @@ class LNPModel(nn.Module):
self.rdkit_encoder.clear_cache()
if self.mpnn_encoder is not None:
self.mpnn_encoder.clear_cache()
+ if self.chemeleon_encoder is not None:
+ self.chemeleon_encoder.clear_cache()
+ if self.unimol_encoder is not None:
+ self.unimol_encoder.clear_cache()
if self.llm_prompt is not None and hasattr(self.llm_prompt, "clear_cache"):
self.llm_prompt.clear_cache()
@@ -442,6 +491,8 @@ class LNPModelWithoutMPNN(LNPModel):
head_hidden_dim: int = 128,
dropout: float = 0.1,
input_dims: Optional[Dict[str, int]] = None,
+ chemeleon_cache_path: Optional[str] = None,
+ unimol_cache_path: Optional[str] = None,
# ============ MoE 相关 ============
use_moe: bool = False,
moe_n_experts: int = 4,
@@ -478,6 +529,8 @@ class LNPModelWithoutMPNN(LNPModel):
dropout=dropout,
mpnn_checkpoint=None,
mpnn_ensemble_paths=None,
+ chemeleon_cache_path=chemeleon_cache_path,
+ unimol_cache_path=unimol_cache_path,
input_dims=dims,
reg_bypass=reg_bypass,
use_moe=use_moe,
diff --git a/lnp_ml/modeling/nested_cv_optuna.py b/lnp_ml/modeling/nested_cv_optuna.py
index 7c34e63..6b0a72c 100644
--- a/lnp_ml/modeling/nested_cv_optuna.py
+++ b/lnp_ml/modeling/nested_cv_optuna.py
@@ -195,6 +195,8 @@ def create_model(
dropout: float = 0.1,
use_mpnn: bool = False,
mpnn_device: str = "cpu",
+ chemeleon_cache: Optional[str] = None,
+ unimol_cache: Optional[str] = None,
set_transformer_block: str = "sab",
# MoE
use_moe: bool = False,
@@ -218,6 +220,8 @@ def create_model(
moe_jitter_noise=moe_jitter_noise,
use_retrieval=use_retrieval,
retr_feature_dim=retr_feature_dim,
+ chemeleon_cache_path=chemeleon_cache,
+ unimol_cache_path=unimol_cache,
**(llm_kwargs or {}),
)
@@ -418,6 +422,8 @@ def run_inner_optuna(
batch_size: int = 32,
n_inner_folds: int = 3,
use_mpnn: bool = False,
+ chemeleon_cache: Optional[str] = None,
+ unimol_cache: Optional[str] = None,
seed: int = 42,
study_path: Optional[Path] = None,
pretrain_state_dict: Optional[Dict] = None,
@@ -536,6 +542,8 @@ def run_inner_optuna(
dropout=dropout,
use_mpnn=use_mpnn,
mpnn_device=device.type,
+ chemeleon_cache=chemeleon_cache,
+ unimol_cache=unimol_cache,
use_moe=use_moe,
moe_n_experts=moe_ne_t,
moe_top_k=moe_tk_t,
@@ -631,6 +639,8 @@ def _run_single_outer_fold(
batch_size: int,
n_inner_folds: int,
use_mpnn: bool,
+ chemeleon_cache: Optional[str],
+ unimol_cache: Optional[str],
seed: int,
pretrain_state_dict: Optional[Dict],
pretrain_config: Optional[Dict],
@@ -732,6 +742,8 @@ def _run_single_outer_fold(
batch_size=batch_size,
n_inner_folds=n_inner_folds,
use_mpnn=use_mpnn,
+ chemeleon_cache=chemeleon_cache,
+ unimol_cache=unimol_cache,
seed=seed + outer_fold,
study_path=study_path,
pretrain_state_dict=pretrain_state_dict,
@@ -789,6 +801,8 @@ def _run_single_outer_fold(
dropout=best_params["dropout"],
use_mpnn=use_mpnn,
mpnn_device=device.type,
+ chemeleon_cache=chemeleon_cache,
+ unimol_cache=unimol_cache,
use_moe=use_moe,
moe_n_experts=best_params.get("moe_n_experts", moe_n_experts),
moe_top_k=best_params.get("moe_top_k", moe_top_k),
@@ -853,6 +867,10 @@ def _run_single_outer_fold(
"set_transformer_block": best_params.get("set_transformer_block", "sab"),
"dropout": best_params["dropout"],
"use_mpnn": use_mpnn,
+ "use_chemeleon": chemeleon_cache is not None,
+ "chemeleon_cache": chemeleon_cache,
+ "use_unimol": unimol_cache is not None,
+ "unimol_cache": unimol_cache,
"use_moe": use_moe,
"moe_n_experts": moe_n_experts,
"moe_top_k": moe_top_k,
@@ -920,6 +938,11 @@ def main(
load_delivery_head: bool = False,
# MPNN
use_mpnn: bool = False,
+ # CheMeleon
+ use_chemeleon: bool = False,
+ chemeleon_cache: str = "data/processed/chemeleon_embeddings.npz",
+ use_unimol: bool = False,
+ unimol_cache: str = "data/processed/unimol_embeddings.npz",
n_repeats: int = 1,
repeat_seed_step: int = 1000,
# MoE（消融开关）
@@ -1050,6 +1073,8 @@ def main(
batch_size=batch_size,
n_inner_folds=n_inner_folds,
use_mpnn=use_mpnn,
+ chemeleon_cache=(chemeleon_cache if use_chemeleon else None),
+ unimol_cache=(unimol_cache if use_unimol else None),
seed=seed,
pretrain_state_dict=pretrain_state_dict,
pretrain_config=pretrain_config,
diff --git a/lnp_ml/modeling/predict.py b/lnp_ml/modeling/predict.py
index a5cf836..4b02f44 100644
--- a/lnp_ml/modeling/predict.py
+++ b/lnp_ml/modeling/predict.py
@@ -64,6 +64,8 @@ def load_model(
llm_lora_dropout=config.get("llm_lora_dropout", 0.05),
mpnn_ensemble_paths=ensemble_paths,
mpnn_device=mpnn_device,
+ chemeleon_cache_path=config.get("chemeleon_cache"),
+ unimol_cache_path=config.get("unimol_cache"),
)
else:
model = LNPModelWithoutMPNN(
@@ -80,6 +82,8 @@ def load_model(
llm_lora_r=config.get("llm_lora_r", 8),
llm_lora_alpha=config.get("llm_lora_alpha", 16),
llm_lora_dropout=config.get("llm_lora_dropout", 0.05),
+ chemeleon_cache_path=config.get("chemeleon_cache"),
+ unimol_cache_path=config.get("unimol_cache"),
)
model.load_state_dict(checkpoint["model_state_dict"], strict=False)
diff --git a/lnp_ml/modeling/pretrain.py b/lnp_ml/modeling/pretrain.py
index 54a39ea..48154ad 100644
--- a/lnp_ml/modeling/pretrain.py
+++ b/lnp_ml/modeling/pretrain.py
@@ -243,6 +243,10 @@ def main(
mpnn_checkpoint: Optional[str] = None,
mpnn_ensemble_paths: Optional[str] = None,
mpnn_device: str = "cpu",
+ use_chemeleon: bool = False,
+ chemeleon_cache: str = "data/processed/chemeleon_embeddings.npz",
+ use_unimol: bool = False,
+ unimol_cache: str = "data/processed/unimol_embeddings.npz",
# 训练参数
batch_size: int = 64,
lr: float = 1e-4,
@@ -324,6 +328,8 @@ def main(
llm_lora_r=llm_lora_r,
llm_lora_alpha=llm_lora_alpha,
llm_lora_dropout=llm_lora_dropout,
+ chemeleon_cache_path=(chemeleon_cache if use_chemeleon else None),
+ unimol_cache_path=(unimol_cache if use_unimol else None),
)
if enable_mpnn:
model = LNPModel(
@@ -373,6 +379,8 @@ def main(
"head_hidden_dim": head_hidden_dim,
"dropout": dropout,
"use_mpnn": enable_mpnn,
+ "use_chemeleon": use_chemeleon,
+ "use_unimol": use_unimol,
"use_moe": use_moe,
"moe_n_experts": moe_n_experts,
"moe_top_k": moe_top_k,
Reference in New Issue View Git Blame Copy Permalink