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，用 ... 紧贴包裹（官方格式，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"{sfs}" + 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,