Add privacy-clean core modeling components

.gitignore

@@ -0,0 +1,95 @@
+# Data and datasets
+data/
+dataset/
+datasets/
+raw/
+interim/
+processed/
+external/
+
+# Results and experiment artifacts
+reports/
+results/
+outputs/
+output/
+artifacts/
+figures/
+predictions/
+failure_case_export*/
+wandb/
+runs/
+logs/
+
+# Models and checkpoints
+models/
+checkpoints/
+MolE_ckpt/
+*.pt
+*.pth
+*.ckpt
+*.bin
+*.safetensors
+*.onnx
+
+# Structured and serialized data
+*.csv
+*.tsv
+*.xlsx
+*.xls
+*.parquet
+*.feather
+*.arrow
+*.json
+*.jsonl
+*.pkl
+*.pickle
+*.joblib
+*.npy
+*.npz
+*.h5
+*.hdf5
+*.db
+*.sqlite
+*.sqlite3
+
+# Images and documents
+*.png
+*.jpg
+*.jpeg
+*.webp
+*.svg
+*.pdf
+
+# Secrets
+.env
+.env.*
+*.pem
+*.key
+*.crt
+
+# Environments and caches
+.venv/
+venv/
+env/
+.pixi/
+.molformer_env/
+.molformer_libs/
+.molformer_transformers_v5/
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.ipynb_checkpoints/
+
+# Logs and temporary files
+*.log
+*.out
+*.bak
+*.bak_*
+*.before_*
+*.tmp
+*.swp
+*.swo
+*~
+.DS_Store

LICENSE

@@ -0,0 +1,10 @@
+
+The MIT License (MIT)
+Copyright (c) 2026, Your name (or your organization/company/team)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+

README.md

@@ -0,0 +1,35 @@
+# LNP Modeling Components
+
+This repository contains reusable neural network architecture components for molecular and formulation modeling.
+
+## Included
+
+- Molecular feature encoders
+- Token projection layers
+- Set-based attention modules
+- Mixture-of-experts layers
+- Feature fusion modules
+- Prediction heads
+
+## Excluded
+
+This repository does not include:
+
+- Training or evaluation datasets
+- Processed data
+- Model checkpoints
+- Precomputed embeddings
+- Experimental results
+- Prediction outputs
+- Server-specific configuration
+- Private file paths
+
+Users must provide their own authorized data and runtime configuration.
+
+## Installation
+
+    pip install -e .
+
+## License
+
+See `LICENSE`.

lnp_ml/__init__.py

@@ -0,0 +1,3 @@
+"""Core neural network components for LNP modeling."""
+
+__version__ = "0.1.0"

lnp_ml/config.py

@@ -0,0 +1,5 @@
+"""Runtime-independent package configuration."""
+
+from pathlib import Path
+
+PACKAGE_ROOT = Path(__file__).resolve().parent

lnp_ml/featurization/__init__.py

@@ -0,0 +1,4 @@
+from lnp_ml.featurization.smiles import RDKitFeaturizer
+
+__all__ = ["RDKitFeaturizer"]
+

lnp_ml/featurization/smiles.py

@@ -0,0 +1,167 @@
+"""SMILES 分子特征提取器"""
+
+import logging
+import warnings
+from dataclasses import dataclass, field
+from typing import List, Optional, Dict
+import numpy as np
+
+# Suppress RDKit deprecation warnings
+from rdkit import RDLogger
+RDLogger.DisableLog("rdApp.*")
+
+from rdkit import Chem
+from rdkit.Chem import AllChem, MACCSkeys, Descriptors
+
+import torch
+from chemprop.utils import load_checkpoint
+from chemprop.features import mol2graph
+
+_chemprop_logger = logging.getLogger("chemprop.load_checkpoint")
+_chemprop_logger.setLevel(logging.WARNING)
+
+
+@dataclass
+class RDKitFeaturizer:
+    """
+    SMILES -> RDKit 特征向量，返回 Dict[str, np.ndarray]。
+    """
+
+    morgan_radius: int = 2
+    morgan_nbits: int = 1024
+
+    def transform(self, smiles_list: List[str]) -> Dict[str, np.ndarray]:
+        """SMILES 特征字典 -> value: (N, D_i) arrays"""
+        encoded = [self._encode_one(s) for s in smiles_list]
+        return {
+            "morgan": np.vstack([enc["morgan"] for enc in encoded]),
+            "maccs": np.vstack([enc["maccs"] for enc in encoded]),
+            "desc": np.vstack([enc["desc"] for enc in encoded])
+        }
+
+    def _encode_morgan(self, mol: Chem.Mol) -> np.ndarray:
+        return np.array(AllChem.GetMorganFingerprintAsBitVect(
+            mol, radius=self.morgan_radius, nBits=self.morgan_nbits
+        ).ToList(), dtype=np.float32)
+
+    def _encode_maccs(self, mol: Chem.Mol) -> np.ndarray:
+        return np.array(MACCSkeys.GenMACCSKeys(mol).ToList(), dtype=np.float32)
+    
+    def _encode_desc(self, mol: Chem.Mol) -> np.ndarray:
+        # 使用 float64 计算，然后 clip 到 float32 范围，避免 overflow
+        desc_values = list(Descriptors.CalcMolDescriptors(mol).values())
+        arr = np.array(desc_values, dtype=np.float64)
+        # 替换 inf/nan，clip 到 float32 范围
+        arr = np.nan_to_num(arr, nan=0.0, posinf=1e10, neginf=-1e10)
+        arr = np.clip(arr, -1e10, 1e10)
+        return arr.astype(np.float32)
+
+    def _encode_one(self, smiles: str) -> Dict[str, np.ndarray]:
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            raise ValueError(f"Invalid SMILES: {smiles!r}")
+        
+        return {
+            "morgan": self._encode_morgan(mol),
+            "maccs": self._encode_maccs(mol),
+            "desc": self._encode_desc(mol)
+        }
+
+
+@dataclass
+class MPNNFeaturizer:
+    """
+    SMILES -> D-MPNN 预训练特征向量 (N, hidden_size=600)。
+    
+    从训练好的 chemprop 模型中提取 D-MPNN 编码器的输出作为分子特征。
+    
+    Args:
+        checkpoint_path: 模型检查点路径（.pt文件）
+        device: 计算设备 ("cpu" 或 "cuda")
+        ensemble_paths: 可选，多个模型路径列表用于集成（取平均）
+    """
+    checkpoint_path: Optional[str] = None
+    device: str = "cpu"
+    ensemble_paths: Optional[List[str]] = None
+    
+    # 内部状态（不由用户设置）
+    _encoders: List = field(default_factory=list, init=False, repr=False)
+    _hidden_size: int = field(default=0, init=False, repr=False)
+    _initialized: bool = field(default=False, init=False, repr=False)
+    
+    def __post_init__(self) -> None:
+        """延迟初始化，在首次调用 transform 时加载模型"""
+        if self.checkpoint_path is None and self.ensemble_paths is None:
+            raise ValueError("必须提供 checkpoint_path 或 ensemble_paths")
+        
+    def _lazy_init(self) -> None:
+        """延迟加载模型，避免在创建对象时就加载大模型"""
+        if self._initialized:
+            return
+            
+        device = torch.device(self.device)
+        paths = self.ensemble_paths or [self.checkpoint_path]
+        
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=FutureWarning)
+            for path in paths:
+                model = load_checkpoint(path, device=device, logger=_chemprop_logger)
+                model.eval()
+                encoder = model.encoder.encoder[0]
+                for param in encoder.parameters():
+                    param.requires_grad = False
+                self._encoders.append(encoder)
+                self._hidden_size = encoder.hidden_size
+        
+        self._initialized = True
+    
+    def transform(self, smiles_list: List[str]) -> Dict[str, np.ndarray]:
+        """
+        SMILES 列表 -> tuple of (N, hidden_size) array
+        
+        Args:
+            smiles_list: SMILES 字符串列表
+            
+        Returns:
+            tuple 包含一个形状为 (N, hidden_size) 的 numpy 数组
+            如果使用集成模型，返回所有模型输出的平均值
+        """
+        self._lazy_init()
+        
+        # 验证 SMILES 有效性
+        for smi in smiles_list:
+            mol = Chem.MolFromSmiles(smi)
+            if mol is None:
+                raise ValueError(f"Invalid SMILES: {smi!r}")
+        
+        # 构建分子图（批量处理）
+        batch_mol_graph = mol2graph(smiles_list)
+        
+        # 从所有编码器提取特征
+        all_features = []
+        with torch.no_grad():
+            for encoder in self._encoders:
+                features = encoder(batch_mol_graph)
+                all_features.append(features.cpu().numpy())
+        
+        # 如果是集成模型，取平均
+        if len(all_features) > 1:
+            features_array = np.mean(all_features, axis=0).astype(np.float32)
+        else:
+            features_array = all_features[0].astype(np.float32)
+        
+        return {
+            "mpnn": features_array
+        }
+    
+    @property
+    def hidden_size(self) -> int:
+        """返回特征维度"""
+        self._lazy_init()
+        return self._hidden_size
+    
+    @property
+    def n_models(self) -> int:
+        """返回集成模型数量"""
+        self._lazy_init()
+        return len(self._encoders)

lnp_ml/modeling/__init__.py

@@ -0,0 +1 @@
+"""Model architecture components."""

lnp_ml/modeling/encoders/__init__.py

@@ -0,0 +1,7 @@
+from lnp_ml.modeling.encoders.rdkit_encoder import CachedRDKitEncoder
+from lnp_ml.modeling.encoders.mpnn_encoder import CachedMPNNEncoder
+
+__all__ = ["CachedRDKitEncoder", "CachedMPNNEncoder"]
+
+from lnp_ml.modeling.encoders.chemeleon_encoder import CheMeleonEmbeddingEncoder
+from lnp_ml.modeling.encoders.unimol_encoder import UniMolEmbeddingEncoder

lnp_ml/modeling/encoders/chemeleon_encoder.py

@@ -0,0 +1,48 @@
+"""CheMeleon 预计算指纹的查表编码器"""
+from pathlib import Path
+from typing import Dict, List
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class CheMeleonEmbeddingEncoder(nn.Module):
+    """从预计算 .npz 缓存按 SMILES 查 CheMeleon 指纹，返回 {"chemeleon": [B, D]}。"""
+
+    def __init__(self, cache_path: str) -> None:
+        super().__init__()
+        self.cache_path = str(cache_path)
+        self._table: Dict[str, np.ndarray] = {}
+        self._embed_dim: int = 0
+        self._load()
+
+    def _load(self) -> None:
+        path = Path(self.cache_path)
+        if not path.exists():
+            raise FileNotFoundError(
+                f"CheMeleon 缓存不存在: {path}。请先运行 scripts/precompute_chemeleon.py。"
+            )
+        data = np.load(path, allow_pickle=True)
+        embeddings = np.asarray(data["embeddings"], dtype=np.float32)
+        if embeddings.ndim != 2:
+            raise ValueError(f"embeddings 应为 2D，实际 {embeddings.shape}")
+        self._embed_dim = int(embeddings.shape[1])
+        self._table = {str(s): embeddings[i] for i, s in enumerate(data["smiles"])}
+
+    def forward(self, smiles_list: List[str]) -> Dict[str, torch.Tensor]:
+        missing = [s for s in smiles_list if s not in self._table]
+        if missing:
+            raise KeyError(
+                f"{len(missing)} 个 SMILES 不在 CheMeleon 缓存中，请重跑 precompute_chemeleon.py。"
+                f"示例: {missing[:3]}"
+            )
+        mat = np.stack([self._table[s] for s in smiles_list])
+        return {"chemeleon": torch.from_numpy(mat)}
+
+    def clear_cache(self) -> None:
+        """查表器无临时缓存，仅为接口对齐。"""
+
+    @property
+    def embed_dim(self) -> int:
+        return self._embed_dim

lnp_ml/modeling/encoders/mpnn_encoder.py

@@ -0,0 +1,67 @@
+from typing import List, Optional, Dict
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+from lnp_ml.featurization.smiles import MPNNFeaturizer
+
+
+class CachedMPNNEncoder(nn.Module):
+    """
+    带内存缓存的 D-MPNN 特征提取模块。
+
+    - 使用预训练 chemprop 模型的 encoder 提取特征
+    - 不可训练，不参与反向传播
+    - 缓存已计算的 SMILES 特征，避免重复计算
+    - forward 返回 Dict[str, Tensor]，key: "mpnn"
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: Optional[str] = None,
+        ensemble_paths: Optional[List[str]] = None,
+        device: str = "cpu",
+    ) -> None:
+        super().__init__()
+        self._featurizer = MPNNFeaturizer(
+            checkpoint_path=checkpoint_path,
+            ensemble_paths=ensemble_paths,
+            device=device,
+        )
+        self._cache: Dict[str, np.ndarray] = {}
+
+    def forward(self, smiles_list: List[str]) -> Dict[str, torch.Tensor]:
+        """
+        SMILES 列表 -> Dict[str, Tensor]
+
+        Returns:
+            {"mpnn": (N, hidden_size)}
+        """
+        # 分离：已缓存 vs 需计算
+        to_compute = [s for s in smiles_list if s not in self._cache]
+
+        # 批量计算未缓存的
+        if to_compute:
+            new_features = self._featurizer.transform(to_compute)
+            for idx, smiles in enumerate(to_compute):
+                self._cache[smiles] = new_features["mpnn"][idx]
+
+        # 按原顺序组装结果
+        return {
+            "mpnn": torch.from_numpy(np.stack([self._cache[s] for s in smiles_list]))
+        }
+
+    def clear_cache(self) -> None:
+        """清空缓存"""
+        self._cache.clear()
+
+    @property
+    def cache_size(self) -> int:
+        """当前缓存的 SMILES 数量"""
+        return len(self._cache)
+
+    @property
+    def hidden_size(self) -> int:
+        """特征维度"""
+        return self._featurizer.hidden_size

lnp_ml/modeling/encoders/rdkit_encoder.py

@@ -0,0 +1,58 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from typing import List, Dict
+
+from lnp_ml.featurization.smiles import RDKitFeaturizer
+
+
+class CachedRDKitEncoder(nn.Module):
+    """
+    带内存缓存的 RDKit 特征提取模块。
+
+    - 不可训练，不参与反向传播
+    - 缓存已计算的 SMILES 特征，避免重复计算
+    - forward 返回 Dict[str, Tensor]，keys: "morgan", "maccs", "desc"
+    """
+
+    def __init__(self, morgan_radius: int = 2, morgan_nbits: int = 1024) -> None:
+        super().__init__()
+        self._featurizer = RDKitFeaturizer(
+            morgan_radius=morgan_radius,
+            morgan_nbits=morgan_nbits,
+        )
+        self._cache: Dict[str, Dict[str, np.ndarray]] = {}
+
+    def forward(self, smiles_list: List[str]) -> Dict[str, torch.Tensor]:
+        """
+        SMILES 列表 -> Dict[str, Tensor]
+
+        Returns:
+            {"morgan": (N, 1024), "maccs": (N, 167), "desc": (N, 210)}
+        """
+        # 分离：已缓存 vs 需计算
+        to_compute = [s for s in smiles_list if s not in self._cache]
+
+        # 批量计算未缓存的
+        if to_compute:
+            new_features = self._featurizer.transform(to_compute)
+            for idx, smiles in enumerate(to_compute):
+                self._cache[smiles] = {
+                    k: new_features[k][idx] for k in new_features
+                }
+
+        # 按原顺序组装结果
+        keys = ["morgan", "maccs", "desc"]
+        return {
+            k: torch.from_numpy(np.stack([self._cache[s][k] for s in smiles_list]))
+            for k in keys
+        }
+
+    def clear_cache(self) -> None:
+        """清空缓存"""
+        self._cache.clear()
+
+    @property
+    def cache_size(self) -> int:
+        """当前缓存的 SMILES 数量"""
+        return len(self._cache)

lnp_ml/modeling/encoders/tabular_encoder.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+from typing import Dict
+
+
+class TabularEncoder(nn.Module):
+    
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, tabular_data: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        # Input: Dict with keys 'comp', 'phys', 'help', 'exp'
+        #   Each value is a tensor [B, D_i] where D_i is the feature dimension
+        # Output: Same dict (pass-through, features already grouped by DataLoader)
+        
+        # The DataLoader (trainer.py) already groups features correctly:
+        # - 'comp': [B, 9] - composition features
+        # - 'phys': [B, 9] - physical features (including processed PDI)
+        # - 'help': [B, 4] - helper lipid one-hot features
+        # - 'exp': [B, 20] - experimental condition one-hot features (including processed Purity)
+        
+        # Simply return the dict as-is
+        # If we wanted to add learned transformations, we could add linear layers here
+        return tabular_data

lnp_ml/modeling/encoders/unimol_encoder.py

@@ -0,0 +1,48 @@
+"""UniMol 预计算表征的查表编码器"""
+from pathlib import Path
+from typing import Dict, List
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class UniMolEmbeddingEncoder(nn.Module):
+    """从预计算 .npz 缓存按 SMILES 查 UniMol 表征，返回 {"unimol": [B, D]}。"""
+
+    def __init__(self, cache_path: str) -> None:
+        super().__init__()
+        self.cache_path = str(cache_path)
+        self._table: Dict[str, np.ndarray] = {}
+        self._embed_dim: int = 0
+        self._load()
+
+    def _load(self) -> None:
+        path = Path(self.cache_path)
+        if not path.exists():
+            raise FileNotFoundError(
+                f"UniMol 缓存不存在: {path}。请先运行 scripts/precompute_unimol.py。"
+            )
+        data = np.load(path, allow_pickle=True)
+        embeddings = np.asarray(data["embeddings"], dtype=np.float32)
+        if embeddings.ndim != 2:
+            raise ValueError(f"embeddings 应为 2D，实际 {embeddings.shape}")
+        self._embed_dim = int(embeddings.shape[1])
+        self._table = {str(s): embeddings[i] for i, s in enumerate(data["smiles"])}
+
+    def forward(self, smiles_list: List[str]) -> Dict[str, torch.Tensor]:
+        missing = [s for s in smiles_list if s not in self._table]
+        if missing:
+            raise KeyError(
+                f"{len(missing)} 个 SMILES 不在 UniMol 缓存中，请重跑 precompute_unimol.py。"
+                f"示例: {missing[:3]}"
+            )
+        mat = np.stack([self._table[s] for s in smiles_list])
+        return {"unimol": torch.from_numpy(mat)}
+
+    def clear_cache(self) -> None:
+        """查表器无临时缓存，仅为接口对齐。"""
+
+    @property
+    def embed_dim(self) -> int:
+        return self._embed_dim

lnp_ml/modeling/heads.py

@@ -0,0 +1,126 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict
+
+
+class RegressionHead(nn.Module):
+    """回归任务头：输出单个 float 值"""
+
+    def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, 1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """[B, in_dim] -> [B, 1]"""
+        return self.net(x)
+
+
+class ClassificationHead(nn.Module):
+    """分类任务头：输出 logits"""
+
+    def __init__(
+        self, in_dim: int, num_classes: int, hidden_dim: int = 128, dropout: float = 0.1
+    ) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, num_classes),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """[B, in_dim] -> [B, num_classes] (logits)"""
+        return self.net(x)
+
+
+class DistributionHead(nn.Module):
+    """分布任务头：输出和为 1 的概率分布（用于 Biodistribution）"""
+
+    def __init__(
+        self, in_dim: int, num_outputs: int, hidden_dim: int = 128, dropout: float = 0.1
+    ) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, num_outputs),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """[B, in_dim] -> [B, num_outputs] (softmax, sum=1)"""
+        logits = self.net(x)
+        return F.softmax(logits, dim=-1)
+
+
+class MultiTaskHead(nn.Module):
+    """
+    多任务预测头，根据任务配置自动创建对应的子头。
+
+    输出:
+        - size: [B, 1] 回归
+        - pdi: [B, 2] 分类 logits
+        - ee: [B, 3] 分类 logits
+        - delivery: [B, 1] 回归
+        - biodist: [B, 7] softmax 分布
+        - toxic: [B, 2] 二分类 logits
+    """
+
+    def __init__(self, in_dim: int, hidden_dim: int = 128, dropout: float = 0.1) -> None:
+        super().__init__()
+
+        # size: 回归 (log-transformed)
+        size_dropout = min(0.5, dropout + 0.2)
+        self.size_head = RegressionHead(in_dim, hidden_dim, size_dropout)
+
+        # PDI: 2 分类
+        self.pdi_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
+
+        # Encapsulation Efficiency: 3 分类
+        self.ee_head = ClassificationHead(in_dim, num_classes=3, hidden_dim=hidden_dim, dropout=dropout)
+
+        # Continuous target: regression
+        self.delivery_head = RegressionHead(in_dim, hidden_dim, dropout)
+
+        # Biodistribution: 7 输出，softmax (sum=1)
+        self.biodist_head = DistributionHead(in_dim, num_outputs=7, hidden_dim=hidden_dim, dropout=dropout)
+
+        # toxic: 二分类
+        self.toxic_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
+
+        # 不确定性加权（Kendall 2018）：每个任务一个可学习 log σ²，初始 0
+        self.log_vars = nn.ParameterDict({
+            t: nn.Parameter(torch.zeros(()))
+            for t in ["size", "delivery", "pdi", "ee", "toxic", "biodist"]
+        })
+
+    def forward(self, x: torch.Tensor, x_numeric: torch.Tensor = None) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            x: [B, in_dim] fusion 层输出
+
+        Returns:
+            Dict with keys:
+                - "size": [B, 1]
+                - "pdi": [B, 2] logits
+                - "ee": [B, 3] logits
+                - "delivery": [B, 1]
+                - "biodist": [B, 7] probabilities (sum=1)
+                - "toxic": [B, 2] logits
+        """
+        xn = x if x_numeric is None else x_numeric
+        return {
+            "size": self.size_head(xn),
+            "pdi": self.pdi_head(x),
+            "ee": self.ee_head(x),
+            "delivery": self.delivery_head(xn),
+            "biodist": self.biodist_head(x),
+            "toxic": self.toxic_head(x),
+        }

lnp_ml/modeling/layers/__init__.py

@@ -0,0 +1,16 @@
+from lnp_ml.modeling.layers.token_projector import TokenProjector
+from lnp_ml.modeling.layers.bidirectional_cross_attention import CrossModalAttention
+from lnp_ml.modeling.layers.set_transformer import SetTransformer
+from lnp_ml.modeling.layers.fusion import FusionLayer, ResidualConcatFusion
+from lnp_ml.modeling.layers.moe import MoEBlock
+from lnp_ml.modeling.layers.llm_prompt import LLMPromptEncoder
+
+__all__ = [
+    "TokenProjector",
+    "CrossModalAttention",
+    "SetTransformer",
+    "FusionLayer",
+    "ResidualConcatFusion",
+    "MoEBlock",
+    "LLMPromptEncoder",
+]

lnp_ml/modeling/layers/bidirectional_cross_attention.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn as nn
+from typing import Tuple
+
+
+class CrossAttentionLayer(nn.Module):
+    """单层双向交叉注意力"""
+
+    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1) -> None:
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+
+        # A -> B: A as Q, B as K/V
+        self.cross_attn_a2b = nn.MultiheadAttention(
+            embed_dim=d_model,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+        # B -> A: B as Q, A as K/V
+        self.cross_attn_b2a = nn.MultiheadAttention(
+            embed_dim=d_model,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+
+        # LayerNorm + FFN for channel A
+        self.norm_a1 = nn.LayerNorm(d_model)
+        self.norm_a2 = nn.LayerNorm(d_model)
+        self.ffn_a = nn.Sequential(
+            nn.Linear(d_model, d_model * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_model * 4, d_model),
+            nn.Dropout(dropout),
+        )
+
+        # LayerNorm + FFN for channel B
+        self.norm_b1 = nn.LayerNorm(d_model)
+        self.norm_b2 = nn.LayerNorm(d_model)
+        self.ffn_b = nn.Sequential(
+            nn.Linear(d_model, d_model * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_model * 4, d_model),
+            nn.Dropout(dropout),
+        )
+
+    def forward(
+        self, a: torch.Tensor, b: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            a: [B, seq_len, d_model]
+            b: [B, seq_len, d_model]
+
+        Returns:
+            (a_out, b_out): 更新后的两个 channel
+        """
+        # Cross attention: A attends to B
+        a_attn, _ = self.cross_attn_a2b(query=a, key=b, value=b)
+        a = self.norm_a1(a + a_attn)
+        a = self.norm_a2(a + self.ffn_a(a))
+
+        # Cross attention: B attends to A
+        b_attn, _ = self.cross_attn_b2a(query=b, key=a, value=a)
+        b = self.norm_b1(b + b_attn)
+        b = self.norm_b2(b + self.ffn_b(b))
+
+        return a, b
+
+
+class CrossModalAttention(nn.Module):
+    """
+    双向交叉注意力模块。
+
+    输入 stacked tokens [B, 8, d_model]，split 成两个 channel 后执行
+    n_layers 层双向交叉注意力。
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        n_layers: int,
+        split_idx: int = 4,
+        dropout: float = 0.1,
+    ) -> None:
+        """
+        Args:
+            d_model: 特征维度
+            num_heads: 注意力头数，d_head = d_model / num_heads
+            n_layers: 交叉注意力层数
+            split_idx: channel split 的位置，默认 4 (0:4, 4:)
+            dropout: dropout 比例
+        """
+        super().__init__()
+        self.split_idx = split_idx
+
+        self.layers = nn.ModuleList([
+            CrossAttentionLayer(d_model, num_heads, dropout)
+            for _ in range(n_layers)
+        ])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, 8, d_model] stacked tokens
+
+        Returns:
+            [B, 8, d_model] 融合后的 tokens
+        """
+        # Split: [B, 8, d_model] -> [B, 4, d_model], [B, 4, d_model]
+        a = x[:, : self.split_idx, :]
+        b = x[:, self.split_idx :, :]
+
+        # N layers of bidirectional cross attention
+        for layer in self.layers:
+            a, b = layer(a, b)
+
+        # Concat back: [B, 8, d_model]
+        return torch.cat([a, b], dim=1)
+

lnp_ml/modeling/layers/fusion.py

@@ -0,0 +1,159 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, List, Literal, Optional, Tuple, Union
+
+
+PoolingStrategy = Literal["concat", "avg", "max", "attention"]
+
+
+class FusionLayer(nn.Module):
+    """
+    将多个 token 融合成单个向量。
+
+    输入: Dict[str, Tensor] 或 [B, n_tokens, d_model]
+    输出: [B, fusion_dim]
+
+    策略:
+        - concat: [B, n_tokens, d_model] -> [B, n_tokens * d_model]
+        - avg: [B, n_tokens, d_model] -> [B, d_model]
+        - max: [B, n_tokens, d_model] -> [B, d_model]
+        - attention: [B, n_tokens, d_model] -> [B, d_model] (learnable attention pooling)
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_tokens: int,
+        strategy: PoolingStrategy = "attention",
+    ) -> None:
+        """
+        Args:
+            d_model: 每个 token 的维度
+            n_tokens: token 数量（如 8）
+            strategy: 融合策略
+        """
+        super().__init__()
+        self.d_model = d_model
+        self.n_tokens = n_tokens
+        self.strategy = strategy
+
+        if strategy == "concat":
+            self.fusion_dim = n_tokens * d_model
+        else:
+            self.fusion_dim = d_model
+
+        # Attention pooling: learnable query
+        if strategy == "attention":
+            self.attn_query = nn.Parameter(torch.randn(1, 1, d_model))
+            self.attn_proj = nn.Linear(d_model, d_model)
+
+    def forward(
+        self,
+        x: Union[Dict[str, torch.Tensor], torch.Tensor],
+        return_attn_weights: bool = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Args:
+            x: Dict[str, Tensor] 每个 [B, d_model]，或已 stack 的 [B, n_tokens, d_model]
+            return_attn_weights: 若为 True 且策略为 attention，额外返回 attn_weights [B, n_tokens]
+
+        Returns:
+            return_attn_weights=False: [B, fusion_dim]
+            return_attn_weights=True:  ([B, fusion_dim], [B, n_tokens])
+        """
+        if isinstance(x, dict):
+            x = torch.stack(list(x.values()), dim=1)
+
+        if self.strategy == "concat":
+            out = x.flatten(start_dim=1)
+            return (out, None) if return_attn_weights else out
+
+        elif self.strategy == "avg":
+            out = x.mean(dim=1)
+            return (out, None) if return_attn_weights else out
+
+        elif self.strategy == "max":
+            out = x.max(dim=1).values
+            return (out, None) if return_attn_weights else out
+
+        elif self.strategy == "attention":
+            return self._attention_pooling(x, return_attn_weights)
+
+        else:
+            raise ValueError(f"Unknown strategy: {self.strategy}")
+
+    def _attention_pooling(
+        self, x: torch.Tensor, return_attn_weights: bool = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Attention pooling: 用可学习 query 对 tokens 做加权求和
+
+        Args:
+            x: [B, n_tokens, d_model]
+            return_attn_weights: 是否返回权重
+
+        Returns:
+            return_attn_weights=False: [B, d_model]
+            return_attn_weights=True:  ([B, d_model], [B, n_tokens])
+        """
+        B = x.size(0)
+        query = self.attn_query.expand(B, -1, -1)
+
+        keys = self.attn_proj(x)
+        scores = torch.bmm(query, keys.transpose(1, 2)) / (self.d_model ** 0.5)
+        attn_weights = F.softmax(scores, dim=-1)  # [B, 1, n_tokens]
+
+        out = torch.bmm(attn_weights, x).squeeze(1)  # [B, d_model]
+
+        if return_attn_weights:
+            return out, attn_weights.squeeze(1)  # [B, n_tokens]
+        return out
+
+
+class ResidualConcatFusion(nn.Module):
+    """对真实 token 做 attention pooling，再用零初始化门把 MoE/LLM 旁路以残差方式加入。
+
+    g_moe / g_llm 初始为 0 → +moe/+llm 起点严格等于 baseline；
+    旁路只有确实有用时才会被训练打开，从机制上保证“加了不会更差”。
+    """
+
+    def __init__(self, d_model: int, strategy: PoolingStrategy = "attention") -> None:
+        super().__init__()
+        if strategy == "concat":
+            raise ValueError("ResidualConcatFusion 不支持 concat（token 数随开关变化）")
+        self.d_model = d_model
+        self.pool = FusionLayer(d_model=d_model, n_tokens=1, strategy=strategy)
+        self.fusion_dim = self.pool.fusion_dim
+        # 零初始化门控（可学习标量），旁路初始不参与
+        self.g_moe = nn.Parameter(torch.zeros(()))
+        self.g_llm = nn.Parameter(torch.zeros(()))
+        self.g_retr = nn.Parameter(torch.zeros(()))  # 检索旁路零初始化门控
+
+    def forward(
+        self,
+        chem: torch.Tensor,
+        tab: torch.Tensor,
+        f_moe: Optional[torch.Tensor] = None,
+        f_llm: Optional[torch.Tensor] = None,
+        f_retr: Optional[torch.Tensor] = None,
+        return_attn_weights: bool = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        # 只对真实 token（chem + tab）做注意力池化，旁路不参与 softmax 竞争
+        seq = torch.cat([chem, tab], dim=1)  # [B, n_chem + n_cond, d_model]
+        pooled = self.pool(seq, return_attn_weights=return_attn_weights)
+        if return_attn_weights:
+            pooled, attn = pooled
+
+        out = pooled
+        if f_moe is not None:
+            out = out + self.g_moe * f_moe  # 残差 + 零初始化门
+        if f_llm is not None:
+            out = out + self.g_llm * f_llm
+        if f_retr is not None:
+            out = out + self.g_retr * f_retr  # 检索旁路，零初始化门保证起点=不开
+
+        # 回归旁路：额外返回 pooled（纯 chem+tab，不含 f_llm/f_moe/f_retr）
+        if return_attn_weights:
+            return out, pooled, attn
+        return out, pooled

lnp_ml/modeling/layers/llm_prompt.py

@@ -0,0 +1,356 @@
+"""LLM 分子特征分支：SMILES 文本 + chem'/tab soft token 的混合 prompt 编码。"""
+
+import os
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+
+
+DEFAULT_MOLT5_PATH = os.environ.get("MOLT5_PATH", "models/molt5-base")
+
+# 检索池支持的额外多任务（除 delivery 外）
+_EXTRA_TASKS = ["size", "pdi", "ee", "toxic", "biodist"]
+
+
+class LLMPromptEncoder(nn.Module):
+    """用 LLM 编码分子，输出 F_llm [B, d_model]。
+
+    三种模式：
+      - use_soft_prompt=True：原始 SMILES 文本(+RAG 邻居) + chem'/tab soft token，
+        inputs_embeds 注入，全程带梯度（配合 LoRA/QLoRA 真微调），不缓存特征。
+      - use_rag=True 且非 soft：旧的冻结+缓存 RAG 路径。
+      - 其他：冻结缓存 / 可训练 mean-pool。
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_chem_tokens: int = 4,
+        n_cond_tokens: int = 4,
+        model_name_or_path: str = DEFAULT_MOLT5_PATH,
+        freeze: bool = True,
+        use_lora: bool = False,
+        use_qlora: bool = False,
+        lora_r: int = 8,
+        lora_alpha: int = 16,
+        lora_dropout: float = 0.05,
+        max_length: int = 256,
+        use_rag: bool = False,
+        rag_top_k: int = 4,
+        use_soft_prompt: bool = False,
+    ) -> None:
+        super().__init__()
+        from transformers import AutoTokenizer, T5EncoderModel, AutoModel
+
+        self.use_lora = use_lora or use_qlora
+        self.use_qlora = use_qlora
+        self.max_length = max_length
+        self.use_rag = use_rag
+        self.rag_top_k = rag_top_k
+        self.use_soft_prompt = use_soft_prompt
+
+        _name_l = model_name_or_path.lower()
+        _is_t5 = "t5" in _name_l
+        _is_qwen = "qwen" in _name_l
+        self._is_qwen = _is_qwen
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path, trust_remote_code=_is_qwen)
+        if _is_qwen and self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        if _is_t5:
+            self.encoder = T5EncoderModel.from_pretrained(model_name_or_path)
+            self.hidden_size = self.encoder.config.d_model
+        elif _is_qwen and use_qlora:
+            from transformers import BitsAndBytesConfig
+            bnb = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_compute_dtype=torch.bfloat16,
+            )
+            self.encoder = AutoModel.from_pretrained(
+                model_name_or_path, trust_remote_code=True,
+                quantization_config=bnb, torch_dtype=torch.bfloat16)
+            self.hidden_size = self.encoder.config.hidden_size
+        elif _is_qwen:
+            # 微调(use_lora)时用 4-bit 量化(QLoRA)省显存；纯冻结时用 fp16
+            if use_lora:
+                from transformers import BitsAndBytesConfig
+                _bnb = BitsAndBytesConfig(
+                    load_in_4bit=True, bnb_4bit_quant_type="nf4",
+                    bnb_4bit_compute_dtype=torch.float16,
+                    bnb_4bit_use_double_quant=True)
+                self.encoder = AutoModel.from_pretrained(
+                    model_name_or_path, trust_remote_code=True,
+                    quantization_config=_bnb, device_map={"": 0})
+            else:
+                self.encoder = AutoModel.from_pretrained(
+                    model_name_or_path, trust_remote_code=True, torch_dtype=torch.float16)
+            self.hidden_size = self.encoder.config.hidden_size
+        else:
+            self.encoder = AutoModel.from_pretrained(model_name_or_path)
+            self.hidden_size = self.encoder.config.hidden_size
+
+        if self.use_lora:
+            self._apply_lora(lora_r, lora_alpha, lora_dropout, prepare_kbit=use_qlora)
+        elif freeze:
+            for p in self.encoder.parameters():
+                p.requires_grad = False
+        self._frozen = freeze and not self.use_lora
+
+        # soft token 投影：chem'(d_model) / tab(d_model) -> hidden_size
+        if use_soft_prompt:
+            self.chem_to_llm = nn.Linear(d_model, self.hidden_size)
+            self.tab_to_llm = nn.Linear(d_model, self.hidden_size)
+        else:
+            self.chem_to_llm = None
+            self.tab_to_llm = None
+
+        self.proj_down = nn.Sequential(
+            nn.Linear(self.hidden_size, d_model), nn.LayerNorm(d_model)
+        )
+
+        # 缓存：冻结句向量缓存 / 文本 prompt 缓存
+        self._cache: Dict[str, torch.Tensor] = {}
+        self._prompt_cache: Dict[str, str] = {}
+
+        # RAG 检索池
+        self._rag_pool_smiles: List[str] = []
+        self._rag_pool_labels = None          # delivery [N]
+        self._rag_pool_extra = None           # dict: task -> (values, valid)
+        self._rag_pool_fps = None
+        self._rag_pool_id: str = "none"
+
+    def _apply_lora(self, r, alpha, dropout, prepare_kbit=False):
+        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+
+        if prepare_kbit:
+            self.encoder = prepare_model_for_kbit_training(self.encoder)
+        else:
+            for p in self.encoder.parameters():
+                p.requires_grad = False
+
+        # 4-bit 量化模型(QLoRA)需先 prepare，才能正确接收梯度
+        if getattr(self.encoder, "is_loaded_in_4bit", False) or getattr(self.encoder, "is_loaded_in_8bit", False):
+            from peft import prepare_model_for_kbit_training
+            self.encoder = prepare_model_for_kbit_training(self.encoder)
+        for p in self.encoder.parameters():
+            p.requires_grad = False
+        # 不同架构注意力层命名不同：
+        # T5: q/k/v/o；Roberta/ChemBERTa: query/key/value；Qwen/Llama: q_proj/k_proj/v_proj/o_proj
+        _enc_name = type(self.encoder).__name__.lower()
+        if getattr(self, "_is_qwen", False) or "qwen" in _enc_name or "llama" in _enc_name:
+            _targets = ["q_proj", "k_proj", "v_proj", "o_proj"]
+        elif "t5" in _enc_name or hasattr(self.encoder.config, "d_model"):
+            _targets = ["q", "k", "v", "o"]
+        else:
+            _targets = ["query", "key", "value"]
+        cfg = LoraConfig(r=r, lora_alpha=alpha, lora_dropout=dropout,
+                         target_modules=_targets, bias="none")
+        self.encoder = get_peft_model(self.encoder, cfg)
+
+    # ---------- 检索池 ----------
+    def set_retrieval_pool(self, smiles_list, labels, pool_id="train", extra_labels=None):
+        """设置 RAG 检索池（防泄漏：只传训练集分子与标签）。
+        labels: delivery [N]; extra_labels: dict task -> (values, valid_bool)。"""
+        import numpy as np
+        from lnp_ml.modeling.retrieval import _smiles_to_fp
+        self._rag_pool_smiles = list(smiles_list)
+        self._rag_pool_labels = np.asarray(labels, dtype=np.float32).reshape(-1)
+        self._rag_pool_extra = extra_labels
+        self._rag_pool_fps = [_smiles_to_fp(s) for s in self._rag_pool_smiles]
+        if pool_id != self._rag_pool_id:
+            self._cache = {k: v for k, v in self._cache.items() if not k.startswith("RAG::")}
+            self._prompt_cache.clear()
+            self._rag_pool_id = pool_id
+
+    def _retrieve_topk(self, query_smiles: str):
+        """检索 Top-K 邻居，返回 [dict(smiles, sim, delivery, extra)]；排除查询分子自己。"""
+        from rdkit import DataStructs
+        from lnp_ml.modeling.retrieval import _smiles_to_fp
+        qfp = _smiles_to_fp(query_smiles)
+        if qfp is None or self._rag_pool_fps is None:
+            return []
+        sims = []
+        for j, (smi, fp) in enumerate(zip(self._rag_pool_smiles, self._rag_pool_fps)):
+            if fp is None or smi == query_smiles:
+                continue
+            sims.append((j, DataStructs.TanimotoSimilarity(qfp, fp)))
+        sims.sort(key=lambda t: t[1], reverse=True)
+        out = []
+        for j, sim in sims[: self.rag_top_k]:
+            extra = {}
+            if self._rag_pool_extra is not None:
+                for task in _EXTRA_TASKS:
+                    if task in self._rag_pool_extra:
+                        vals, valid = self._rag_pool_extra[task]
+                        extra[task] = vals[j] if bool(valid[j]) else None
+                    else:
+                        extra[task] = None
+            out.append({
+                "smiles": self._rag_pool_smiles[j],
+                "sim": sim,
+                "delivery": float(self._rag_pool_labels[j]),
+                "extra": extra,
+            })
+        return out
+
+    @staticmethod
+    def _fmt(v, kind="float"):
+        if v is None:
+            return "unknown"
+        if kind == "int":
+            return str(int(v))
+        if kind == "vec":
+            return "[" + ", ".join(f"{x:.3f}" for x in v) + "]"
+        return f"{float(v):.3f}"
+
+    def _build_rag_prompt(self, target_smiles: str, neighbors) -> str:
+        """构造 RAG prompt：原始 SMILES + 邻居多任务结果（numeric）。"""
+        blocks = []
+        for rank, nb in enumerate(neighbors, 1):
+            ex = nb["extra"]
+            blocks.append(
+                f"Retrieved sample {rank}:\n"
+                f"SMILES: {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"
+                f"pdi_class: {self._fmt(ex.get('pdi'), 'int')}\n"
+                f"ee_class: {self._fmt(ex.get('ee'), 'int')}\n"
+                f"toxic: {self._fmt(ex.get('toxic'), 'int')}\n"
+                f"biodist: {self._fmt(ex.get('biodist'), 'vec')}"
+            )
+        retrieved_block = "\n\n".join(blocks) if blocks else "(no retrieved samples)"
+        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"
+            "[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; "
+            "'unknown' means the measurement is missing):\n\n"
+            f"{retrieved_block}\n\n"
+            "[Encoding Instructions]\n"
+            "Capture the target structure and the retrieval evidence (structural similarity "
+            "and consistency of retrieved outcomes) into your internal representation."
+        )
+
+    def _get_prompt(self, s: str) -> str:
+        if not self.use_rag:
+            return 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))
+        return self._prompt_cache[key]
+    def _rag_encode_batch(self, prompts, device):
+        """编码一批 RAG prompt，取最后有效 token。grad 由调用方上下文决定。"""
+        outs = []
+        for i in range(0, len(prompts), 4):
+            bt = prompts[i:i+4]
+            enc = self.tokenizer(
+                bt, padding=True, truncation=True,
+                max_length=self.max_length, return_tensors="pt",
+            ).to(device)
+            out = self.encoder(**enc).last_hidden_state          # [B,L,H]
+            lengths = enc["attention_mask"].sum(1) - 1
+            b = torch.arange(out.size(0), device=device)
+            outs.append(out[b, lengths.long(), :].float())       # [B,H]
+        return torch.cat(outs, 0)
+
+    def _encode_rag(self, smiles, device):
+        """RAG 编码。冻结时：no_grad + 缓存（快）。微调时：算梯度 + 不缓存。"""
+        if self._frozen:
+            # 冻结路径：缓存复用，no_grad
+            keys = [f"RAG::{self._rag_pool_id}::{s}" for s in smiles]
+            to_compute = [s for s, k in zip(smiles, keys) if k not in self._cache]
+            if to_compute:
+                uniq = list(dict.fromkeys(to_compute))
+                prompts = [self._build_rag_prompt(s, self._retrieve_topk(s)) for s in uniq]
+                with torch.no_grad():
+                    feat = self._rag_encode_batch(prompts, device)
+                for s, v in zip(uniq, feat):
+                    self._cache[f"RAG::{self._rag_pool_id}::{s}"] = v.float().cpu()
+            return torch.stack([self._cache[k] for k in keys]).to(device)
+        else:
+            # 微调路径：每次重新编码，保留计算图（算梯度），不缓存
+            prompts = [self._build_rag_prompt(s, self._retrieve_topk(s)) for s in smiles]
+            return self._rag_encode_batch(prompts, device)
+
+    # ---------- soft-prompt 编码（带梯度，不缓存特征）----------
+    def _encode_softrag(self, smiles, chem, tab, device) -> torch.Tensor:
+        prompts = [self._get_prompt(s) for s in smiles]
+        enc = self.tokenizer(
+            prompts, padding=True, truncation=True,
+            max_length=self.max_length, return_tensors="pt",
+        ).to(device)
+
+        embed_layer = self.encoder.get_input_embeddings()
+        text_embeds = embed_layer(enc["input_ids"])     # [B, L, H]
+        text_mask = enc["attention_mask"]               # [B, L]
+
+        soft_list = []
+        if self.chem_to_llm is not None and chem is not None:
+            soft_list.append(self.chem_to_llm(chem))    # [B, n_chem, H]
+        if self.tab_to_llm is not None and tab is not None:
+            soft_list.append(self.tab_to_llm(tab))      # [B, n_tab, H]
+
+        if soft_list:
+            soft = torch.cat(soft_list, dim=1).to(text_embeds.dtype)
+            inputs_embeds = torch.cat([soft, text_embeds], dim=1)
+            soft_mask = torch.ones(soft.size(0), soft.size(1),
+                                   device=device, dtype=text_mask.dtype)
+            attn_mask = torch.cat([soft_mask, text_mask], dim=1)
+        else:
+            inputs_embeds = text_embeds
+            attn_mask = text_mask
+
+        out = self.encoder(inputs_embeds=inputs_embeds, attention_mask=attn_mask).last_hidden_state
+        lengths = attn_mask.sum(1) - 1
+        b = torch.arange(out.size(0), device=device)
+        feat = out[b, lengths.long(), :]                # [B, H] 最后有效 token
+        return self.proj_down(feat.float())
+
+    # ---------- 旧路径（保留，向后兼容）----------
+    def _mean_pool(self, last_hidden, mask):
+        m = mask.unsqueeze(-1).float()
+        return (last_hidden * m).sum(1) / m.sum(1).clamp(min=1e-6)
+
+    @torch.no_grad()
+    def _encode_frozen(self, smiles, device):
+        missing = [s for s in smiles if s not in self._cache]
+        if missing:
+            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,
+                                     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"])
+                for s, v in zip(chunk, pooled):
+                    self._cache[s] = v.cpu()
+        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,
+                             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"])
+
+    def forward(self, smiles: List[str], chem: Optional[torch.Tensor] = None,
+                tab: Optional[torch.Tensor] = None) -> torch.Tensor:
+        device = self.proj_down[0].weight.device
+        if self.use_soft_prompt:
+            return self._encode_softrag(smiles, chem, tab, device)
+        if self.use_rag:
+            return self._encode_softrag(smiles, None, None, device)
+        feat = self._encode_frozen(smiles, device) if self._frozen \
+            else self._encode_trainable(smiles, device)
+        return self.proj_down(feat)
+
+    def clear_cache(self) -> None:
+        self._cache.clear()
+        self._prompt_cache.clear()

lnp_ml/modeling/layers/moe.py

@@ -0,0 +1,229 @@
+"""
+MoE (Mixture-of-Experts) 模块：sample-level + 跨模态路由。
+
+设计要点（与现有 8-token 架构对齐）：
+    - Router 输入：tab token 池化后的 [B, d]（即配方/实验条件向量）。
+    - Expert 输入：chem token flatten 后的 [B, T_chem * d]（化学侧 token）。
+    - 路由粒度：每个样本只算一次 router → gates [B, K]。
+    - Top-k 稀疏激活 + 可选训练态 jitter 噪声。
+    - 返回 load-balancing aux loss，由 trainer 端按权重加进总 loss。
+
+输出:
+    F_moe: [B, d_model]   与单个 token 同维度，便于追加到 fusion 序列中。
+    extras: dict           诊断与监控信息（aux loss、gates 等）。
+"""
+
+from typing import Dict, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MoEAttentionPool(nn.Module):
+    """与 FusionLayer 同款的 attention pooling，参数独立。
+
+    将一组 token [B, T, d] 池化成单个查询向量 [B, d]，用作 router 的输入。
+    """
+
+    def __init__(self, d_model: int) -> None:
+        super().__init__()
+        self.d_model = d_model
+        self.query = nn.Parameter(torch.randn(1, 1, d_model))
+        self.proj = nn.Linear(d_model, d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T, d_model]
+
+        Returns:
+            [B, d_model]
+        """
+        B = x.size(0)
+        q = self.query.expand(B, -1, -1)                              # [B, 1, d]
+        k = self.proj(x)                                              # [B, T, d]
+        scores = torch.bmm(q, k.transpose(1, 2)) / (self.d_model ** 0.5)
+        weights = F.softmax(scores, dim=-1)                           # [B, 1, T]
+        return torch.bmm(weights, x).squeeze(1)                       # [B, d]
+
+
+class MoERouter(nn.Module):
+    """单层 softmax router + Top-k 稀疏激活。
+
+    Args:
+        d_model: router 输入维度。
+        n_experts: 专家数量 K。
+        top_k: 每个样本激活的专家数。
+        jitter_noise: 训练态加在 logits 上的均匀噪声幅度，用于探索；评估态自动关闭。
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_experts: int,
+        top_k: int = 2,
+        jitter_noise: float = 0.0,
+    ) -> None:
+        super().__init__()
+        if not 1 <= top_k <= n_experts:
+            raise ValueError(f"top_k 必须在 [1, n_experts={n_experts}]，收到 {top_k}")
+        self.n_experts = n_experts
+        self.top_k = top_k
+        self.jitter_noise = jitter_noise
+        self.linear = nn.Linear(d_model, n_experts)
+
+    def forward(
+        self, q: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            q: [B, d_model] 路由查询向量。
+
+        Returns:
+            gates:       [B, n_experts] top-k 后再归一化的稀疏概率。
+            probs_full:  [B, n_experts] 未掩码的 softmax 概率（用于 aux loss）。
+            expert_mask: [B, n_experts] 0/1 掩码，标记被激活的专家。
+        """
+        logits = self.linear(q)                                       # [B, K]
+        if self.training and self.jitter_noise > 0.0:
+            noise = (torch.rand_like(logits) - 0.5) * self.jitter_noise
+            logits = logits + noise
+
+        probs_full = F.softmax(logits, dim=-1)                        # [B, K]
+
+        topk_vals, topk_idx = probs_full.topk(self.top_k, dim=-1)     # [B, k]
+        expert_mask = torch.zeros_like(probs_full)
+        expert_mask.scatter_(1, topk_idx, 1.0)
+
+        gates = probs_full * expert_mask
+        gates = gates / gates.sum(dim=-1, keepdim=True).clamp(min=1e-9)
+        return gates, probs_full, expert_mask
+
+
+class MoEExpert(nn.Module):
+    """单个专家 MLP：in_dim → hidden_dim → out_dim。"""
+
+    def __init__(
+        self, in_dim: int, hidden_dim: int, out_dim: int, dropout: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, out_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class MoEBlock(nn.Module):
+    """
+    Sample-level 跨模态 MoE。
+
+    流程:
+        1. tab tokens [B, T_tab, d] -> attention pool -> q_tab [B, d]
+        2. q_tab -> Router -> gates [B, K] (+ probs_full / expert_mask)
+        3. chem tokens [B, T_chem, d] -> flatten -> [B, T_chem * d]
+        4. K 个 Expert MLP 并行处理 -> stack [B, K, d]
+        5. 加权求和 -> F_moe [B, d]
+        6. 计算 load-balancing aux loss
+
+    Args:
+        d_model: token 维度。
+        n_chem_tokens: chem 侧 token 数（决定 expert 输入维度）。
+        n_experts: 专家数量 K。
+        top_k: 每个样本激活的专家数。
+        expert_hidden_mult: expert 中间层维度 = expert_hidden_mult * d_model。
+        dropout: expert 内部 dropout。
+        jitter_noise: router 训练态噪声幅度，0.0 表示关闭。
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_chem_tokens: int = 4,
+        n_experts: int = 4,
+        top_k: int = 2,
+        expert_hidden_mult: int = 2,
+        dropout: float = 0.1,
+        jitter_noise: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.d_model = d_model
+        self.n_chem_tokens = n_chem_tokens
+        self.n_experts = n_experts
+        self.top_k = top_k
+
+        self.tab_pool = MoEAttentionPool(d_model)
+        self.router = MoERouter(
+            d_model, n_experts, top_k=top_k, jitter_noise=jitter_noise,
+        )
+
+        expert_in = d_model * n_chem_tokens
+        expert_hidden = d_model * expert_hidden_mult
+        self.experts = nn.ModuleList([
+            MoEExpert(expert_in, expert_hidden, d_model, dropout=dropout)
+            for _ in range(n_experts)
+        ])
+
+    def forward(
+        self,
+        chem: torch.Tensor,
+        tab: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        """
+        Args:
+            chem: [B, T_chem, d_model] 化学侧 token（router 不看）。
+            tab:  [B, T_tab,  d_model] 配方/实验侧 token（router 看）。
+
+        Returns:
+            F_moe: [B, d_model]
+            extras: {
+                "lb_loss": 标量 load-balancing aux loss（带梯度），
+                "gates":   [B, K] 稀疏归一后的门控（detached），
+                "probs":   [B, K] 原始 softmax 概率（detached），
+            }
+        """
+        if chem.size(1) != self.n_chem_tokens:
+            raise ValueError(
+                f"chem token 数不匹配：期望 {self.n_chem_tokens}, "
+                f"实际 {chem.size(1)}"
+            )
+
+        q_tab = self.tab_pool(tab)                                     # [B, d]
+        gates, probs_full, expert_mask = self.router(q_tab)            # 各 [B, K]
+
+        flat = chem.flatten(start_dim=1)                               # [B, T_chem * d]
+        expert_outs = torch.stack(
+            [expert(flat) for expert in self.experts], dim=1,
+        )                                                              # [B, K, d]
+
+        F_moe = (gates.unsqueeze(-1) * expert_outs).sum(dim=1)         # [B, d]
+
+        lb_loss = self._load_balancing_loss(probs_full, expert_mask)
+
+        return F_moe, {
+            "lb_loss": lb_loss,
+            "gates": gates.detach(),
+            "probs": probs_full.detach(),
+        }
+
+    def _load_balancing_loss(
+        self,
+        probs_full: torch.Tensor,
+        expert_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Switch Transformer 风格的 load-balancing loss。
+
+            f_i = 该 batch 中 expert_i 被激活的样本占比（top-k mask 的均值）
+            p_i = 该 batch 中 expert_i 的 softmax 概率均值
+            loss = K * Σ_i f_i * p_i
+
+        理想情况下 f_i 与 p_i 都接近 1/K，loss ≈ 1。
+        """
+        f = expert_mask.mean(dim=0)                                   # [K]
+        p = probs_full.mean(dim=0)                                    # [K]
+        return self.n_experts * (f * p).sum()

lnp_ml/modeling/layers/set_transformer.py

@@ -0,0 +1,123 @@
+"""Set Transformer 集合编码器。
+输入/输出形状均为 [B, n_tokens, d_model]，
+n_tokens 可变（支持 3 或 4）。
+"""
+
+import torch
+import torch.nn as nn
+
+
+class MAB(nn.Module):
+    """多头注意力块：MAB(Q, K) = LN(H + FFN(H))，H = LN(Q + MultiHeadAttn(Q, K, K))。"""
+
+    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1, ln: bool = True) -> None:
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"
+        self.attn = nn.MultiheadAttention(
+            d_model, num_heads, dropout=dropout, batch_first=True
+        )
+        self.norm1 = nn.LayerNorm(d_model) if ln else nn.Identity()
+        self.norm2 = nn.LayerNorm(d_model) if ln else nn.Identity()
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, d_model * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_model * 4, d_model),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            q: [B, n_q, d_model]
+            k: [B, n_k, d_model]
+
+        Returns:
+            [B, n_q, d_model]
+        """
+        attn_out, _ = self.attn(q, k, k)
+        h = self.norm1(q + attn_out)
+        return self.norm2(h + self.ffn(h))
+
+
+class SAB(nn.Module):
+    """集合自注意力块：SAB(X) = MAB(X, X)。"""
+
+    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1, ln: bool = True) -> None:
+        super().__init__()
+        self.mab = MAB(d_model, num_heads, dropout, ln)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """[B, n, d_model] -> [B, n, d_model]"""
+        return self.mab(x, x)
+
+
+class ISAB(nn.Module):
+    """诱导点集合注意力块：用 m 个可学习诱导点降低大集合的注意力复杂度。"""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        num_inducing: int = 16,
+        dropout: float = 0.1,
+        ln: bool = True,
+    ) -> None:
+        super().__init__()
+        # 可学习诱导点
+        self.inducing = nn.Parameter(torch.empty(1, num_inducing, d_model))
+        nn.init.xavier_uniform_(self.inducing)
+        self.mab_in = MAB(d_model, num_heads, dropout, ln)
+        self.mab_out = MAB(d_model, num_heads, dropout, ln)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """[B, n, d_model] -> [B, n, d_model]"""
+        inducing = self.inducing.expand(x.size(0), -1, -1)  # [B, m, d_model]
+        h = self.mab_in(inducing, x)  # [B, m, d_model]
+        return self.mab_out(x, h)     # [B, n, d_model]
+
+
+class SetTransformer(nn.Module):
+    """对化学 token 集合做 N 层集合自注意力，形状保持 [B, n_tokens, d_model]。"""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int = 8,
+        n_layers: int = 4,
+        dropout: float = 0.1,
+        block: str = "sab",
+        num_inducing: int = 16,
+        ln: bool = True,
+    ) -> None:
+        """
+        Args:
+            d_model: token 维度
+            num_heads: 注意力头数，d_head = d_model / num_heads
+            n_layers: 集合注意力层数
+            dropout: dropout 比例
+            block: 注意力块类型，"sab"（全自注意力）或 "isab"（诱导点）
+            num_inducing: ISAB 的诱导点数量（block="isab" 时生效）
+            ln: 是否使用 LayerNorm
+        """
+        super().__init__()
+        self.block_type = block
+
+        self.layers = nn.ModuleList()
+        for _ in range(n_layers):
+            if block == "isab":
+                self.layers.append(ISAB(d_model, num_heads, num_inducing, dropout, ln))
+            else:
+                self.layers.append(SAB(d_model, num_heads, dropout, ln))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, n_tokens, d_model] 化学 token 集合
+
+        Returns:
+            [B, n_tokens, d_model] 集合编码后的 token
+        """
+        for layer in self.layers:
+            x = layer(x)
+        return x

lnp_ml/modeling/layers/token_projector.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+from typing import Dict
+
+
+class TokenProjector(nn.Module):
+    """
+    将不同维度的特征投影到统一的 d_model 维度。
+
+    每个特征分支的流程：
+    [B, input_dim_i] -> BN -> Linear -> [B, d_model] -> ReLU -> BN -> Dropout -> * sigmoid(weight_i)
+    """
+
+    def __init__(
+        self,
+        input_dims: Dict[str, int],
+        d_model: int,
+        dropout: float = 0.1,
+    ) -> None:
+        """
+        Args:
+            input_dims: 各特征的输入维度，如 {"morgan": 1024, "maccs": 167, "desc": 210}
+            d_model: 统一的输出维度
+            dropout: dropout 比例
+        """
+        super().__init__()
+        self.keys = list(input_dims.keys())
+
+        # 为每个特征分支创建投影层
+        self.projectors = nn.ModuleDict()
+        for key, in_dim in input_dims.items():
+            self.projectors[key] = nn.Sequential(
+                nn.BatchNorm1d(in_dim),
+                nn.Linear(in_dim, d_model),
+                nn.ReLU(),
+                nn.BatchNorm1d(d_model),
+                nn.Dropout(dropout),
+            )
+
+        # 每个分支的可学习权重（初始化为 0，sigmoid 后为 0.5）
+        self.weights = nn.ParameterDict({
+            key: nn.Parameter(torch.zeros(1)) for key in self.keys
+        })
+
+    def forward(self, features: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            features: Dict[str, Tensor]，每个 tensor 形状为 (B, input_dim_i)
+
+        Returns:
+            Dict[str, Tensor]，每个 tensor 形状为 (B, d_model)
+        """
+        out = {}
+        for key in self.keys:
+            x = self.projectors[key](features[key])
+            w = torch.sigmoid(self.weights[key])
+            out[key] = x * w
+        return out
+

lnp_ml/modeling/models.py

@@ -0,0 +1,603 @@
+"""LNP 多任务预测模型"""
+
+import torch
+import torch.nn as nn
+from typing import Dict, List, Optional, Literal
+
+from lnp_ml.modeling.encoders import CachedRDKitEncoder, CachedMPNNEncoder, CheMeleonEmbeddingEncoder, UniMolEmbeddingEncoder
+from lnp_ml.modeling.layers import (
+    TokenProjector,
+    SetTransformer,
+    ResidualConcatFusion,
+    MoEBlock,
+    LLMPromptEncoder,
+)
+from lnp_ml.modeling.layers.llm_prompt import DEFAULT_MOLT5_PATH
+from lnp_ml.modeling.heads import MultiTaskHead
+
+
+PoolingStrategy = Literal["attention", "avg", "max"]
+
+
+# Token 维度配置
+def _infer_desc_dim() -> int:
+    """RDKit 描述符数量随版本变化，运行时实测，避免写死导致 BatchNorm 维度不匹配。"""
+    from rdkit import Chem
+    from rdkit.Chem import Descriptors
+    return len(Descriptors.CalcMolDescriptors(Chem.MolFromSmiles("CCO")))
+
+
+_DESC_DIM = _infer_desc_dim()  # 本机 = 210
+
+DEFAULT_INPUT_DIMS = {
+    # Channel A: 化学特征
+    "mpnn": 600,        # D-MPNN embedding
+    "molformer": 768,   # MolFormer-XL 离线 embedding
+    "graphmvp": 300,    # GraphMVP GIN 离线 embedding
+    "grover": 3200,     # GROVER dualtrans fingerprint (atom+bond)
+    "morgan": 1024,     # Morgan fingerprint
+    "maccs": 167,       # MACCS keys
+    "desc": _DESC_DIM,  # RDKit descriptors（随版本动态）
+    # Channel B: 配方/实验条件
+    "comp": 5,          # 配方比例
+    "phys": 12,         # 物理参数 one-hot
+    "help": 4,          # Helper lipid one-hot
+    "exp": 32,          # 实验条件 one-hot
+}
+
+# 化学 / 配方 token 的键顺序
+CHEM_KEYS_WITH_MPNN = ["mpnn", "morgan", "maccs", "desc"]
+CHEM_KEYS_NO_MPNN = ["morgan", "maccs", "desc"]
+CHEM_KEYS_WITH_MOLFORMER = ["molformer", "morgan", "maccs", "desc"]
+CHEM_KEYS_WITH_GRAPHMVP = ["graphmvp", "morgan", "maccs", "desc"]
+CHEM_KEYS_WITH_GROVER = ["grover", "morgan", "maccs", "desc"]
+CHEM_KEYS_WITH_CHEMELEON = ["chemeleon", "morgan", "maccs", "desc"]
+CHEM_KEYS_WITH_UNIMOL = ["unimol", "morgan", "maccs", "desc"]
+TAB_KEYS = ["comp", "phys", "help", "exp"]
+
+# backbone 权重前缀（用于预训练加载与导出）
+BACKBONE_PREFIXES = (
+    "token_projector.",
+    "set_transformer.",
+    "fusion.",
+    "moe.",
+    "llm_prompt.",
+)
+# 冻结的 MolT5 encoder 权重前缀，不纳入 backbone（由本地权重加载，不进 checkpoint）
+LLM_FROZEN_PREFIX = "llm_prompt.encoder."
+
+
+class LNPModel(nn.Module):
+    """
+    LNP 药物递送性能预测模型。
+
+    架构流程:
+        1. Encoders: SMILES -> 化学特征; tabular -> 配方/实验特征
+        2. TokenProjector: 统一到 d_model
+        3. SetTransformer: 对化学 token 集合做置换等变编码 -> chem'
+        4. MoE (可选): router 看 tab，expert 吃 chem' -> F_moe
+        5. LLM (可选): chem'/tab 注入 MolT5 prompt -> F_llm
+        6. ResidualConcatFusion: 拼接 chem'/tab/F_moe/F_llm -> attention pooling
+        7. MultiTaskHead: 多任务预测
+    """
+
+    def __init__(
+        self,
+        # 模型维度
+        d_model: int = 256,
+        # Set Transformer
+        num_heads: int = 8,
+        n_attn_layers: int = 4,
+        set_transformer_block: str = "sab",
+        # Fusion
+        fusion_strategy: PoolingStrategy = "attention",
+        # Head
+        head_hidden_dim: int = 128,
+        # Dropout
+        dropout: float = 0.1,
+        # MPNN encoder
+        mpnn_checkpoint: Optional[str] = None,
+        mpnn_ensemble_paths: Optional[List[str]] = None,
+        mpnn_device: str = "cpu",
+        # 输入维度配置
+        input_dims: Optional[Dict[str, int]] = None,
+        # ============ MoE 相关 ============
+        use_moe: bool = False,
+        moe_n_experts: int = 4,
+        moe_top_k: int = 2,
+        moe_expert_hidden_mult: int = 2,
+        moe_jitter_noise: float = 0.0,
+        # ============ 回归旁路开关 ============
+        reg_bypass: str = "on",
+        # ============ LLM 相关 ============
+        use_llm: bool = False,
+        llm_model_path: str = DEFAULT_MOLT5_PATH,
+        llm_freeze: bool = True,
+        llm_use_lora: bool = False,
+        llm_use_qlora: bool = False,
+        use_soft_prompt: bool = False,
+        llm_lora_r: int = 8,
+        llm_lora_alpha: int = 16,
+        llm_lora_dropout: float = 0.05,
+        use_rag: bool = False,
+        rag_top_k: int = 4,
+        use_retrieval: bool = False,
+        retr_feature_dim: int = 0,
+        # ===== 离线预训练分子 embedding（三选一）=====
+        molformer_emb_path: Optional[str] = None,
+        molformer_smiles_csv: Optional[str] = None,
+        graphmvp_emb_path: Optional[str] = None,
+        graphmvp_smiles_csv: Optional[str] = None,
+        grover_emb_path: Optional[str] = None,
+        grover_smiles_csv: Optional[str] = None,
+        chemeleon_cache_path: Optional[str] = None,
+        unimol_cache_path: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+
+        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
+
+        # 离线预训练 embedding：{smiles: vector} 查表
+        def _load_offline_emb(emb_path, csv_path):
+            if emb_path is None:
+                return False, {}
+            import numpy as _np, pandas as _pd
+            _E = _np.load(emb_path)
+            _sm = _pd.read_csv(csv_path, low_memory=False)["smiles"].astype(str).tolist()
+            assert len(_sm) == _E.shape[0], f"SMILES({len(_sm)}) vs emb({_E.shape[0]}) 不匹配"
+            return True, {sm: _E[i] for i, sm in enumerate(_sm)}
+
+        self.use_molformer, self._molformer_lut = _load_offline_emb(molformer_emb_path, molformer_smiles_csv)
+        self.use_graphmvp, self._graphmvp_lut = _load_offline_emb(graphmvp_emb_path, graphmvp_smiles_csv)
+        self.use_grover, self._grover_lut = _load_offline_emb(grover_emb_path, grover_smiles_csv)
+        self.use_chemeleon = chemeleon_cache_path is not 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
+        self.use_unimol = unimol_cache_path is not 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
+        assert sum([self.use_molformer, self.use_graphmvp, self.use_grover, self.use_chemeleon, self.use_unimol]) <= 1, \
+            "molformer/graphmvp/grover/chemeleon/unimol 只能启用其中一个"
+
+        # ============ Encoders ============
+        self.rdkit_encoder = CachedRDKitEncoder()
+        if self.use_mpnn:
+            self.mpnn_encoder = CachedMPNNEncoder(
+                checkpoint_path=mpnn_checkpoint,
+                ensemble_paths=mpnn_ensemble_paths,
+                device=mpnn_device,
+            )
+        else:
+            self.mpnn_encoder = None
+
+        # ============ Token Projector ============
+        proj_input_dims = {k: v for k, v in self.input_dims.items()}
+        if not self.use_mpnn:
+            proj_input_dims.pop("mpnn", None)
+        for _k, _flag in (("molformer", "use_molformer"), ("graphmvp", "use_graphmvp"), ("grover", "use_grover"), ("chemeleon", "use_chemeleon"), ("unimol", "use_unimol")):
+            if not getattr(self, _flag, False):
+                proj_input_dims.pop(_k, None)
+        self.token_projector = TokenProjector(
+            input_dims=proj_input_dims,
+            d_model=d_model,
+            dropout=dropout,
+        )
+
+        # token 顺序与化学侧 token 数
+        if self.use_mpnn:
+            self.chem_keys = CHEM_KEYS_WITH_MPNN
+        elif self.use_molformer:
+            self.chem_keys = CHEM_KEYS_WITH_MOLFORMER
+        elif self.use_graphmvp:
+            self.chem_keys = CHEM_KEYS_WITH_GRAPHMVP
+        elif self.use_grover:
+            self.chem_keys = CHEM_KEYS_WITH_GROVER
+        elif self.use_chemeleon:
+            self.chem_keys = CHEM_KEYS_WITH_CHEMELEON
+        elif self.use_unimol:
+            self.chem_keys = CHEM_KEYS_WITH_UNIMOL
+        else:
+            self.chem_keys = CHEM_KEYS_NO_MPNN
+        self.tab_keys = TAB_KEYS
+        self.token_order = self.chem_keys + self.tab_keys
+        self.split_idx = len(self.chem_keys)
+
+        # ============ Set Transformer ============
+        self.set_transformer = SetTransformer(
+            d_model=d_model,
+            num_heads=num_heads,
+            n_layers=n_attn_layers,
+            dropout=dropout,
+            block=set_transformer_block,
+        )
+
+        # ============ MoE Block (可选) ============
+        self.use_moe = use_moe
+        self._last_moe_extras: Optional[Dict[str, torch.Tensor]] = None
+        if use_moe:
+            self.moe = MoEBlock(
+                d_model=d_model,
+                n_chem_tokens=self.split_idx,
+                n_experts=moe_n_experts,
+                top_k=moe_top_k,
+                expert_hidden_mult=moe_expert_hidden_mult,
+                dropout=dropout,
+                jitter_noise=moe_jitter_noise,
+            )
+        else:
+            self.moe = None
+
+        # ============ LLM Prompt (可选) ============
+        self.use_llm = use_llm
+        self.reg_bypass = (str(reg_bypass).lower() == "on")
+        if use_llm:
+            self.llm_prompt = LLMPromptEncoder(
+                d_model=d_model,
+                n_chem_tokens=self.split_idx,
+                n_cond_tokens=len(self.tab_keys),
+                model_name_or_path=llm_model_path,
+                freeze=llm_freeze,
+                use_lora=llm_use_lora,
+                use_qlora=llm_use_qlora,
+                lora_r=llm_lora_r,
+                lora_alpha=llm_lora_alpha,
+                lora_dropout=llm_lora_dropout,
+                use_rag=use_rag,
+                rag_top_k=rag_top_k,
+                use_soft_prompt=use_soft_prompt,
+            )
+        else:
+            self.llm_prompt = None
+
+        # ============ 检索增强 (可选) ============
+        self.use_retrieval = use_retrieval
+        self.retriever = None  # 由 nested_cv 每折用训练集构建后赋值
+        if use_retrieval:
+            self.retr_proj = nn.Linear(retr_feature_dim, d_model)
+        else:
+            self.retr_proj = None
+
+        # ============ Residual Concat + Fusion ============
+        self.fusion = ResidualConcatFusion(d_model=d_model, strategy=fusion_strategy)
+
+        # ============ Multi-Task Head ============
+        self.head = MultiTaskHead(
+            in_dim=self.fusion.fusion_dim,
+            hidden_dim=head_hidden_dim,
+            dropout=dropout,
+        )
+
+    def _encode_and_project(
+        self,
+        smiles: List[str],
+        tabular: Dict[str, torch.Tensor],
+    ) -> torch.Tensor:
+        """
+        编码 SMILES 和 tabular，返回 stacked tokens。
+
+        Returns:
+            stacked: [B, n_tokens, d_model]，顺序为 chem 在前、tab 在后
+        """
+        device = tabular["comp"].device
+
+        rdkit_features = self.rdkit_encoder(smiles)
+
+        all_features: Dict[str, torch.Tensor] = {}
+        if self.use_mpnn:
+            mpnn_features = self.mpnn_encoder(smiles)
+            all_features["mpnn"] = mpnn_features["mpnn"].to(device)
+        import numpy as _np
+        for _key, _flag, _lut in (
+            ("molformer", "use_molformer", "_molformer_lut"),
+            ("graphmvp", "use_graphmvp", "_graphmvp_lut"),
+            ("grover", "use_grover", "_grover_lut"),
+        ):
+            if getattr(self, _flag, False):
+                _lookup = getattr(self, _lut)
+                _v = _np.stack([_lookup[sm] for sm in smiles])
+                all_features[_key] = torch.from_numpy(_v).float().to(device)
+        if self.use_chemeleon:
+            all_features["chemeleon"] = self.chemeleon_encoder(smiles)["chemeleon"].float().to(device)
+        if self.use_unimol:
+            all_features["unimol"] = self.unimol_encoder(smiles)["unimol"].float().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)
+        all_features["comp"] = tabular["comp"]
+        all_features["phys"] = tabular["phys"]
+        all_features["help"] = tabular["help"]
+        all_features["exp"] = tabular["exp"]
+
+        projected = self.token_projector(all_features)
+        stacked = torch.stack([projected[k] for k in self.token_order], dim=1)
+        return stacked
+
+    def _backbone_from_stacked(
+        self, stacked: torch.Tensor, smiles: Optional[List[str]] = None
+    ) -> torch.Tensor:
+        chem = stacked[:, : self.split_idx, :]
+        tab = stacked[:, self.split_idx :, :]
+
+        chem = self.set_transformer(chem)  # chem'
+
+        f_moe = None
+        if self.moe is not None:
+            f_moe, extras = self.moe(chem, tab)
+            self._last_moe_extras = extras
+        else:
+            self._last_moe_extras = None
+
+        f_llm = None
+        if self.llm_prompt is not None and smiles is not None:
+            if getattr(self.llm_prompt, "use_soft_prompt", False):
+                f_llm = self.llm_prompt(smiles, chem=chem, tab=tab)
+            else:
+                f_llm = self.llm_prompt(smiles)
+
+        f_retr = None
+        if self.retr_proj is not None and self.retriever is not None and smiles is not None:
+            _feats = self.retriever.query_batch(smiles, exclude_self="auto")
+            _feats_t = torch.as_tensor(_feats, dtype=chem.dtype, device=chem.device)
+            f_retr = self.retr_proj(_feats_t)
+
+        fused, pooled = self.fusion(chem, tab, f_moe=f_moe, f_llm=f_llm, f_retr=f_retr)
+        self._last_pooled = pooled  # 纯数值向量，供回归 head 使用
+        return fused
+
+    def forward_from_projected(
+        self,
+        stacked: torch.Tensor,
+        task: Optional[str] = None,
+    ) -> torch.Tensor:
+        """
+        从已投影的 stacked tokens 开始 forward，用于 Captum 归因。
+
+        Args:
+            stacked: [B, n_tokens, d_model]
+            task: 单任务名；None 时返回 delivery head 输出。
+
+        Returns:
+            对应任务的预测输出。
+        """
+        fused = self._backbone_from_stacked(stacked)
+
+        if task is None:
+            task = "delivery"
+
+        task_heads = {
+            "size": self.head.size_head,
+            "pdi": self.head.pdi_head,
+            "ee": self.head.ee_head,
+            "delivery": self.head.delivery_head,
+            "biodist": self.head.biodist_head,
+            "toxic": self.head.toxic_head,
+        }
+        return task_heads[task](fused)
+
+    def forward_replacing_token(
+        self,
+        raw_feature: torch.Tensor,
+        feature_key: str,
+        base_projected: torch.Tensor,
+        task: Optional[str] = None,
+    ) -> torch.Tensor:
+        """
+        用原始特征替换 base_projected 中指定 token 的投影，然后 forward。
+
+        用于对单个 token 内部特征做 Captum 归因（如 desc 的 210 维）。
+        """
+        projected = self.token_projector.projectors[feature_key](raw_feature)
+        gate = torch.sigmoid(self.token_projector.weights[feature_key])
+        projected = projected * gate
+
+        token_order = list(self.token_projector.keys)
+        token_idx = token_order.index(feature_key)
+
+        stacked = base_projected.clone()
+        stacked[:, token_idx, :] = projected
+
+        return self.forward_from_projected(stacked, task=task)
+
+    def forward_backbone(
+        self,
+        smiles: List[str],
+        tabular: Dict[str, torch.Tensor],
+    ) -> torch.Tensor:
+        """Backbone forward：编码 -> 投影 -> set transformer -> (MoE/LLM) -> 融合。"""
+        stacked = self._encode_and_project(smiles, tabular)
+        return self._backbone_from_stacked(stacked, smiles=smiles)
+
+    def forward_delivery(
+        self,
+        smiles: List[str],
+        tabular: Dict[str, torch.Tensor],
+    ) -> torch.Tensor:
+        """仅预测 delivery（用于 pretrain）。返回 [B, 1]。"""
+        fused = self.forward_backbone(smiles, tabular)
+        x_numeric = (
+            getattr(self, "_last_pooled", None)
+            if getattr(self, "reg_bypass", True)
+            else None
+        )
+        return self.head.delivery_head(
+            fused if x_numeric is None else x_numeric
+        )
+
+    def forward(
+        self,
+        smiles: List[str],
+        tabular: Dict[str, torch.Tensor],
+    ) -> Dict[str, torch.Tensor]:
+        """
+        完整的多任务 forward。
+
+        Returns:
+            Dict[str, Tensor]: size [B,1], pdi [B,4], ee [B,3],
+            delivery [B,1], biodist [B,7], toxic [B,2]
+        """
+        fused = self.forward_backbone(smiles, tabular)
+        return self.head(fused, getattr(self, "_last_pooled", None) if getattr(self, "reg_bypass", True) else None)
+
+    def clear_cache(self) -> None:
+        """清空所有 encoder 的缓存"""
+        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()
+
+    def get_last_moe_extras(self) -> Optional[Dict[str, torch.Tensor]]:
+        """返回最近一次 forward 中 MoE 模块的副产物（aux loss、gates、probs）。"""
+        return self._last_moe_extras
+
+    def get_backbone_state_dict(self) -> Dict[str, torch.Tensor]:
+        """
+        获取 backbone 部分的 state_dict（不含任务头，且排除冻结的 MolT5 encoder）。
+        """
+        return {
+            k: v for k, v in self.state_dict().items()
+            if k.startswith(BACKBONE_PREFIXES) and not k.startswith(LLM_FROZEN_PREFIX)
+        }
+
+    def get_delivery_head_state_dict(self) -> Dict[str, torch.Tensor]:
+        """获取 delivery head 的 state_dict"""
+        return {
+            k: v for k, v in self.state_dict().items()
+            if k.startswith("head.delivery_head.")
+        }
+
+    def load_pretrain_weights(
+        self,
+        pretrain_state_dict: Dict[str, torch.Tensor],
+        load_delivery_head: bool = True,
+        strict: bool = False,
+    ) -> None:
+        """
+        从预训练 checkpoint 加载 backbone 和（可选）delivery head 权重。
+
+        冻结的 MolT5 encoder 权重不在加载范围内。
+        """
+        keys_to_load = []
+        for name in pretrain_state_dict.keys():
+            if name.startswith(BACKBONE_PREFIXES) and not name.startswith(LLM_FROZEN_PREFIX):
+                keys_to_load.append(name)
+            elif load_delivery_head and name.startswith("head.delivery_head."):
+                keys_to_load.append(name)
+
+        filtered_state_dict = {
+            k: v for k, v in pretrain_state_dict.items() if k in keys_to_load
+        }
+
+        unexpected = []
+        model_state = self.state_dict()
+        for k, v in filtered_state_dict.items():
+            if k in model_state and model_state[k].shape == v.shape:
+                model_state[k] = v
+            else:
+                unexpected.append(k)
+
+
+        self.load_state_dict(model_state, strict=False)
+
+        if strict and unexpected:
+            raise RuntimeError(f"Unexpected keys: {unexpected}")
+
+
+class LNPModelWithoutMPNN(LNPModel):
+    """不使用 MPNN 的简化版本（化学 token 为 3 个）"""
+
+    def __init__(
+        self,
+        d_model: int = 256,
+        num_heads: int = 8,
+        n_attn_layers: int = 4,
+        set_transformer_block: str = "sab",
+        fusion_strategy: PoolingStrategy = "attention",
+        head_hidden_dim: int = 128,
+        dropout: float = 0.1,
+        input_dims: Optional[Dict[str, int]] = None,
+        # ============ MoE 相关 ============
+        use_moe: bool = False,
+        moe_n_experts: int = 4,
+        moe_top_k: int = 2,
+        moe_expert_hidden_mult: int = 2,
+        moe_jitter_noise: float = 0.0,
+        # ============ 回归旁路开关 ============
+        reg_bypass: str = "on",
+        # ============ LLM 相关 ============
+        use_llm: bool = False,
+        llm_model_path: str = DEFAULT_MOLT5_PATH,
+        llm_freeze: bool = True,
+        llm_use_lora: bool = False,
+        llm_use_qlora: bool = False,
+        use_soft_prompt: bool = False,
+        llm_lora_r: int = 8,
+        llm_lora_alpha: int = 16,
+        llm_lora_dropout: float = 0.05,
+        use_rag: bool = False,
+        rag_top_k: int = 4,
+        use_retrieval: bool = False,
+        retr_feature_dim: int = 0,
+        molformer_emb_path: Optional[str] = None,
+        molformer_smiles_csv: Optional[str] = None,
+        graphmvp_emb_path: Optional[str] = None,
+        graphmvp_smiles_csv: Optional[str] = None,
+        grover_emb_path: Optional[str] = None,
+        grover_smiles_csv: Optional[str] = None,
+        chemeleon_cache_path: Optional[str] = None,
+        unimol_cache_path: Optional[str] = None,
+    ) -> None:
+        dims = input_dims or DEFAULT_INPUT_DIMS.copy()
+        dims.pop("mpnn", None)
+
+        super().__init__(
+            d_model=d_model,
+            num_heads=num_heads,
+            n_attn_layers=n_attn_layers,
+            set_transformer_block=set_transformer_block,
+            fusion_strategy=fusion_strategy,
+            head_hidden_dim=head_hidden_dim,
+            dropout=dropout,
+            mpnn_checkpoint=None,
+            mpnn_ensemble_paths=None,
+            input_dims=dims,
+            reg_bypass=reg_bypass,
+            use_moe=use_moe,
+            moe_n_experts=moe_n_experts,
+            moe_top_k=moe_top_k,
+            moe_expert_hidden_mult=moe_expert_hidden_mult,
+            moe_jitter_noise=moe_jitter_noise,
+            use_llm=use_llm,
+            use_rag=use_rag,
+            rag_top_k=rag_top_k,
+            use_retrieval=use_retrieval,
+            retr_feature_dim=retr_feature_dim,
+            llm_model_path=llm_model_path,
+            llm_freeze=llm_freeze,
+            llm_use_lora=llm_use_lora,
+            llm_use_qlora=llm_use_qlora,
+            use_soft_prompt=use_soft_prompt,
+            llm_lora_r=llm_lora_r,
+            llm_lora_alpha=llm_lora_alpha,
+            llm_lora_dropout=llm_lora_dropout,
+            molformer_emb_path=molformer_emb_path,
+            molformer_smiles_csv=molformer_smiles_csv,
+            graphmvp_emb_path=graphmvp_emb_path,
+            graphmvp_smiles_csv=graphmvp_smiles_csv,
+            grover_emb_path=grover_emb_path,
+            grover_smiles_csv=grover_smiles_csv,
+            chemeleon_cache_path=chemeleon_cache_path,
+            unimol_cache_path=unimol_cache_path,
+        )

lnp_ml/modeling/retrieval.py

@@ -0,0 +1,85 @@
+"""Morgan 指纹检索：RAG prompt 构造与检索旁路共用。
+
+- _smiles_to_fp: SMILES -> RDKit ExplicitBitVect（供 Tanimoto 相似度）。
+- MorganRetriever: 训练集指纹检索器，query_batch 返回每个查询的 3 维检索特征。
+"""
+
+from typing import List, Optional
+
+import numpy as np
+from rdkit import Chem, DataStructs
+from rdkit.Chem import AllChem
+
+# 与 RDKitFeaturizer 的 morgan token 保持一致
+MORGAN_RADIUS = 2
+MORGAN_NBITS = 1024
+
+
+def _smiles_to_fp(smiles: Optional[str], radius: int = MORGAN_RADIUS, n_bits: int = MORGAN_NBITS):
+    """SMILES -> Morgan ExplicitBitVect；非法/空返回 None。"""
+    if not smiles:
+        return None
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return None
+    return AllChem.GetMorganFingerprintAsBitVect(mol, radius=radius, nBits=n_bits)
+
+
+class MorganRetriever:
+    """基于 Morgan + Tanimoto 的训练集检索器（用于 use_retrieval 旁路）。
+
+    query_batch 为每个查询返回 3 维特征：
+        [相似度加权邻居标签均值, top1 相似度, top-k 平均相似度]
+    """
+
+    def __init__(self, smiles_list: List[str], labels, k: int = 5) -> None:
+        self.k = int(k)
+        self.pool_smiles = list(smiles_list)
+
+        labels = np.asarray(labels, dtype=np.float32)
+        if labels.ndim == 1:
+            labels = labels.reshape(-1, 1)
+        self.pool_labels = labels  # [N, D]，旁路只用第 0 列
+
+        fps = [_smiles_to_fp(s) for s in self.pool_smiles]
+        # 只保留可解析的池分子
+        self._valid_pos = [i for i, fp in enumerate(fps) if fp is not None]
+        self._valid_fps = [fps[i] for i in self._valid_pos]
+        self._valid_labels = self.pool_labels[self._valid_pos, 0] if self._valid_pos else np.zeros(0, np.float32)
+        self._valid_smiles = [self.pool_smiles[i] for i in self._valid_pos]
+        self._fp_cache = {}  # query 指纹缓存
+
+    def _query_fp(self, smiles: str):
+        if smiles not in self._fp_cache:
+            self._fp_cache[smiles] = _smiles_to_fp(smiles)
+        return self._fp_cache[smiles]
+
+    def query_batch(self, smiles_list: List[str], exclude_self="auto") -> np.ndarray:
+        return np.stack([self._query_one(s, exclude_self) for s in smiles_list]).astype(np.float32)
+
+    def _query_one(self, smiles: str, exclude_self) -> np.ndarray:
+        qfp = self._query_fp(smiles)
+        if qfp is None or not self._valid_fps:
+            return np.zeros(3, dtype=np.float32)
+
+        # 向量化 Tanimoto（C 层批量），替代 Python 逐条
+        sims = np.asarray(DataStructs.BulkTanimotoSimilarity(qfp, self._valid_fps), dtype=np.float32)
+
+        if exclude_self:
+            self_mask = np.fromiter((s == smiles for s in self._valid_smiles), dtype=bool, count=len(sims))
+            sims = np.where(self_mask, -1.0, sims)
+
+        k = min(self.k, int((sims >= 0).sum()))
+        if k <= 0:
+            return np.zeros(3, dtype=np.float32)
+
+        # argpartition 取 top-k（O(N)），再对这 k 个排序
+        top = np.argpartition(-sims, k - 1)[:k]
+        top = top[np.argsort(-sims[top])]
+        top_sims = sims[top]
+        top_labels = self._valid_labels[top]
+
+        w = np.clip(top_sims, 0.0, None)
+        wsum = float(w.sum())
+        weighted_label = float((w * top_labels).sum() / wsum) if wsum > 0 else float(top_labels.mean())
+        return np.array([weighted_label, float(top_sims[0]), float(top_sims.mean())], dtype=np.float32)

lnp_ml/utils/seed.py

@@ -0,0 +1,19 @@
+"""全局随机种子工具。"""
+
+import random
+
+import numpy as np
+import torch
+
+
+def set_global_seed(seed: int) -> None:
+    """固定 random / numpy / torch 种子，提升可复现性。
+
+    不开启 cudnn deterministic，以免显著拖慢训练；如需严格复现可自行追加：
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)

pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["flit_core >=3.2,<4"]
+build-backend = "flit_core.buildapi"
+
+[project]
+name = "lnp_ml"
+version = "0.0.1"
+description = "Reusable neural network components for molecular and formulation modeling."
+license = { file = "LICENSE" }
+readme = "README.md"
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "License :: OSI Approved :: MIT License"
+]
+requires-python = ">=3.8"
+
+[tool.ruff]
+line-length = 99
+src = ["lnp_ml"]
+include = ["pyproject.toml", "lnp_ml/**/*.py"]
+
+[tool.ruff.lint]
+extend-select = ["I"]
+
+[tool.ruff.lint.isort]
+known-first-party = ["lnp_ml"]
+force-sort-within-sections = true

requirements.txt

@@ -0,0 +1,25 @@
+# 严格遵循 pixi.toml 的依赖版本
+# 注意: lnp_ml 本地包在 Dockerfile 中单独安装
+
+# conda dependencies (in pixi.toml [dependencies])
+loguru
+tqdm
+typer
+
+# pypi dependencies (in pixi.toml [pypi-dependencies])
+chemprop==1.7.0
+setuptools
+pandas>=2.0.3,<3
+openpyxl>=3.1.5,<4
+python-dotenv>=1.0.1,<2
+pyarrow>=17.0.0,<18
+fastparquet>=2024.2.0,<2025
+fastapi>=0.124.4,<0.125
+streamlit>=1.40.1,<2
+httpx>=0.28.1,<0.29
+uvicorn>=0.33.0,<0.34
+optuna>=4.5.0,<5
+captum>=0.7.0
+transformers>=4.30,<4.46
+sentencepiece
+protobuf