Add full ablation results and CheMeleon/UniMol encoders

.gitignore

@@ -199,3 +199,6 @@ models/abl/
 models/qwen2.5-7b-instruct/
 models/**/*.sqlite3tests/
 models/**/*.sqlite3
+
+models/biot5-plus-base/
+data/processed/*.npz

chemeleon_fingerprint.py

@@ -0,0 +1,63 @@
+# chemeleon_fingerprint.py
+#
+# this file contains the class CheMeleonFingerprint which can be instantiated
+# and called to generate the CheMeleon learned embeddings for a list of SMILES
+# strings and/or RDKit Mols. you may wish to simply copy or download this file directly for use,
+# or adapt the code for your own purposes. No other files are required for it
+# to work, though you must `pip install 'chemprop>=2.2.0'` for this to run.
+#
+# run `python chemeleon_fingerprint.py` for a quick usage demo, otherwise you
+# should `import` the CheMeleonFingerprint class into your other code and use
+# it there (following the example at the bottom of this file) to generate
+# your learned fingerprints
+from pathlib import Path
+from urllib.request import urlretrieve
+
+import numpy as np
+import torch
+from chemprop import featurizers, nn
+from chemprop.data import BatchMolGraph
+from chemprop.models import MPNN
+from chemprop.nn import RegressionFFN
+from rdkit.Chem import Mol, MolFromSmiles
+
+
+class CheMeleonFingerprint:
+    def __init__(self, device: str | torch.device | None = None):
+        self.featurizer = featurizers.SimpleMoleculeMolGraphFeaturizer()
+        agg = nn.MeanAggregation()
+        ckpt_dir = Path().home() / ".chemprop"
+        ckpt_dir.mkdir(exist_ok=True)
+        mp_path = ckpt_dir / "chemeleon_mp.pt"
+        if not mp_path.exists():
+            urlretrieve(
+                r"https://zenodo.org/records/15460715/files/chemeleon_mp.pt",
+                mp_path,
+            )
+        chemeleon_mp = torch.load(mp_path, weights_only=True)
+        mp = nn.BondMessagePassing(**chemeleon_mp["hyper_parameters"])
+        mp.load_state_dict(chemeleon_mp["state_dict"])
+        self.model = MPNN(
+            message_passing=mp,
+            agg=agg,
+            predictor=RegressionFFN(input_dim=mp.output_dim),  # not actually used
+        )
+        self.model.eval()
+        if device is not None:
+            self.model.to(device=device)
+
+    def __call__(self, molecules: list[str | Mol]) -> np.ndarray:
+        bmg = BatchMolGraph(
+            [
+                self.featurizer(MolFromSmiles(m) if isinstance(m, str) else m)
+                for m in molecules
+            ]
+        )
+        bmg.to(device=self.model.device)
+        with torch.no_grad():
+            return self.model.fingerprint(bmg).numpy(force=True)
+
+
+if __name__ == "__main__":
+    chemeleon_fingerprint = CheMeleonFingerprint()
+    chemeleon_fingerprint(["C", "CC", MolFromSmiles("CCC")])

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 = list(model.token_order)
-        names = ["mpnn", "morgan", "maccs", "desc", "comp", "phys", "help", "exp"]
-    else:
-        names = ["morgan", "maccs", "desc", "comp", "phys", "help", "exp"]
     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]
 

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_rmse": best_val_rmse,
-        "best_val_r2": history[-1]["val_r2"] if history else 0,
+        "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)
         

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",
+]

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/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/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,

lnp_ml/modeling/layers/llm_prompt.py

@@ -54,9 +54,14 @@ class LLMPromptEncoder(nn.Module):
         _is_t5 = "t5" in _name_l
         _is_qwen = "qwen" in _name_l
         self._is_qwen = _is_qwen
+        _is_biot5 = "biot5" in _name_l
+        self._is_biot5 = _is_biot5
 
         self.tokenizer = AutoTokenizer.from_pretrained(
-            model_name_or_path, trust_remote_code=_is_qwen)
+            model_name_or_path,
+            trust_remote_code=_is_qwen,
+            use_fast=not _is_biot5, 
+        )
         if _is_qwen and self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
@@ -208,6 +213,19 @@ class LLMPromptEncoder(nn.Module):
             return "[" + ", ".join(f"{x:.3f}" for x in v) + "]"
         return f"{float(v):.3f}"
 
+    def _fmt_mol(self, smiles: str) -> str:
+        """按 backbone 期望格式化分子。
+        BioT5：SMILES -> SELFIES，用 <bom>...<eom> 紧贴包裹（官方格式，token 间无空格）；
+        其他 backbone：原样返回 SMILES。"""
+        if not getattr(self, "_is_biot5", False):
+            return smiles
+        try:
+            import selfies as sf
+            sfs = sf.encoder(smiles)          # CCO -> [C][C][O]
+        except Exception:
+            return smiles                      # 转换失败退回 SMILES，避免整批中断
+        return f"<bom>{sfs}<eom>"
+        
     def _build_rag_prompt(self, target_smiles: str, neighbors) -> str:
         """构造 RAG prompt：原始 SMILES + 邻居多任务结果（numeric）。"""
         blocks = []
@@ -215,7 +233,7 @@ class LLMPromptEncoder(nn.Module):
             ex = nb["extra"]
             blocks.append(
                 f"Retrieved sample {rank}:\n"
-                f"SMILES: {nb['smiles']}\n"
+                f"Molecule: {self._fmt_mol(nb['smiles'])}\n"
                 f"Similarity score: {nb['sim']:.3f}\n"
                 f"delivery_log: {self._fmt(nb['delivery'])}\n"
                 f"size_z: {self._fmt(ex.get('size'))}\n"
@@ -228,7 +246,7 @@ class LLMPromptEncoder(nn.Module):
         return (
             "Task: Encode the target LNP molecule into a retrieval-aware representation "
             "for downstream multi-task property prediction. Do not output predictions.\n\n"
-            f"[Target Molecule]\nSMILES: {target_smiles}\n\n"
+            f"[Target Molecule]\nMolecule: {self._fmt_mol(target_smiles)}\n\n"
             "[Retrieved Similar LNP Samples]\n"
             "Retrieved from the training set by fingerprint similarity, with their known "
             "multi-task outcomes (delivery_log, size_z, pdi_class, ee_class, toxic, biodist; "
@@ -241,7 +259,7 @@ class LLMPromptEncoder(nn.Module):
 
     def _get_prompt(self, s: str) -> str:
         if not self.use_rag:
-            return s
+            return self._fmt_mol(s)
         key = f"{self._rag_pool_id}::{s}"
         if key not in self._prompt_cache:
             self._prompt_cache[key] = self._build_rag_prompt(s, self._retrieve_topk(s))
@@ -326,7 +344,8 @@ class LLMPromptEncoder(nn.Module):
             uniq = list(dict.fromkeys(missing))
             for i in range(0, len(uniq), 256):
                 chunk = uniq[i:i + 256]
-                enc = self.tokenizer(chunk, padding=True, truncation=True,
+                mols = [self._fmt_mol(s) for s in chunk]
+                enc = self.tokenizer(mols, padding=True, truncation=True,
                                      max_length=self.max_length, return_tensors="pt").to(device)
                 out = self.encoder(**enc).last_hidden_state
                 pooled = self._mean_pool(out, enc["attention_mask"])
@@ -335,7 +354,8 @@ class LLMPromptEncoder(nn.Module):
         return torch.stack([self._cache[s] for s in smiles]).to(device)
 
     def _encode_trainable(self, smiles, device):
-        enc = self.tokenizer(list(smiles), padding=True, truncation=True,
+        mols = [self._fmt_mol(s) for s in smiles]
+        enc = self.tokenizer(mols, padding=True, truncation=True,
                              max_length=self.max_length, return_tensors="pt").to(device)
         out = self.encoder(**enc).last_hidden_state
         return self._mean_pool(out, enc["attention_mask"])

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 数
+        # token 顺序：可选 embedding（mpnn/chemeleon）排在指纹类 token 之前
-        self.chem_keys = CHEM_KEYS_WITH_MPNN if self.use_mpnn else CHEM_KEYS_NO_MPNN
+        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,

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,

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)

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,

models/abl_full/s1_baseline/seed42/local_uncommitted.patch

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

models/abl_full/s1_baseline/seed42/outer_fold_0/best_params.json

@@ -0,0 +1,12 @@
+{
+  "dropout": 0.2708449222533328,
+  "lr": 0.000991430410947854,
+  "weight_decay": 0.017147551517191443,
+  "backbone_lr_ratio": 0.43561339540605226,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_baseline/seed42/outer_fold_0/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 13}

models/abl_full/s1_baseline/seed42/outer_fold_0/history.json

@@ -0,0 +1,279 @@
+{
+  "train": [
+    {
+      "loss": 5.444678425788879,
+      "loss_size": 0.9013985113965141,
+      "loss_pdi": 0.6729189786646101,
+      "loss_ee": 1.054741895861096,
+      "loss_delivery": 1.0977894796265497,
+      "loss_biodist": 1.1086501479148865,
+      "loss_toxic": 0.6209922830263773
+    },
+    {
+      "loss": 4.5500569343566895,
+      "loss_size": 0.85751240948836,
+      "loss_pdi": 0.6458425356282128,
+      "loss_ee": 0.9811898039446937,
+      "loss_delivery": 0.9201048314571381,
+      "loss_biodist": 0.7063412467638651,
+      "loss_toxic": 0.46355681783623165
+    },
+    {
+      "loss": 3.628201378716363,
+      "loss_size": 0.4766200818121433,
+      "loss_pdi": 0.6188660065333048,
+      "loss_ee": 0.9291005432605743,
+      "loss_delivery": 0.8234397636519538,
+      "loss_biodist": 0.4827902648184035,
+      "loss_toxic": 0.322152310775386
+    },
+    {
+      "loss": 4.01716587278578,
+      "loss_size": 0.7821791080964936,
+      "loss_pdi": 0.6054948419332504,
+      "loss_ee": 0.9422276218732198,
+      "loss_delivery": 1.0460369272364511,
+      "loss_biodist": 0.4228086951706145,
+      "loss_toxic": 0.2583594545722008
+    },
+    {
+      "loss": 3.469016088379754,
+      "loss_size": 0.7554286056094699,
+      "loss_pdi": 0.5811512834495969,
+      "loss_ee": 0.8725419342517853,
+      "loss_delivery": 0.886976587275664,
+      "loss_biodist": 0.28543684093488586,
+      "loss_toxic": 0.12317715233398809
+    },
+    {
+      "loss": 2.846436858177185,
+      "loss_size": 0.4568066878451241,
+      "loss_pdi": 0.5172890192932553,
+      "loss_ee": 0.823758559094535,
+      "loss_delivery": 0.7545244072874387,
+      "loss_biodist": 0.23187503384219277,
+      "loss_toxic": 0.09306768234819174
+    },
+    {
+      "loss": 2.863664746284485,
+      "loss_size": 0.6340856614212195,
+      "loss_pdi": 0.46964752508534324,
+      "loss_ee": 0.7400971684190962,
+      "loss_delivery": 0.7597490027546883,
+      "loss_biodist": 0.21114367453588379,
+      "loss_toxic": 0.08032437338907686
+    },
+    {
+      "loss": 2.74409736527337,
+      "loss_size": 0.6125703155994415,
+      "loss_pdi": 0.4939524581034978,
+      "loss_ee": 0.7552461607588662,
+      "loss_delivery": 0.5236966917291284,
+      "loss_biodist": 0.20348657543460527,
+      "loss_toxic": 0.18123929373299083
+    },
+    {
+      "loss": 2.559107575151655,
+      "loss_size": 0.513363784386052,
+      "loss_pdi": 0.4523950500620736,
+      "loss_ee": 0.6855895270903906,
+      "loss_delivery": 0.6234937575128343,
+      "loss_biodist": 0.18743807822465897,
+      "loss_toxic": 0.12504185586132938
+    },
+    {
+      "loss": 2.491521100203196,
+      "loss_size": 0.4445962525076336,
+      "loss_pdi": 0.41342465662293965,
+      "loss_ee": 0.6792096479071511,
+      "loss_delivery": 0.7282942334810892,
+      "loss_biodist": 0.14249309400717416,
+      "loss_toxic": 0.11571477616154072
+    },
+    {
+      "loss": 2.3987844718827143,
+      "loss_size": 0.4384675555759006,
+      "loss_pdi": 0.4503721247116725,
+      "loss_ee": 0.6800311356782913,
+      "loss_delivery": 0.6577720418572426,
+      "loss_biodist": 0.1248554570807351,
+      "loss_toxic": 0.08295682887546718
+    },
+    {
+      "loss": 2.2998381389511957,
+      "loss_size": 0.39875652516881627,
+      "loss_pdi": 0.4235676692591773,
+      "loss_ee": 0.6691084769037035,
+      "loss_delivery": 0.6309656215210756,
+      "loss_biodist": 0.13997879127661386,
+      "loss_toxic": 0.0722581786620948
+    },
+    {
+      "loss": 2.238205717669593,
+      "loss_size": 0.3525017342633671,
+      "loss_pdi": 0.42884764985905754,
+      "loss_ee": 0.6448585126135085,
+      "loss_delivery": 0.6537541647752126,
+      "loss_biodist": 0.1366871322194735,
+      "loss_toxic": 0.05711782972017924
+    }
+  ],
+  "val": [
+    {
+      "loss": 4.17715326944987,
+      "loss_size": 0.40348024666309357,
+      "loss_pdi": 0.6616458892822266,
+      "loss_ee": 0.910798708597819,
+      "loss_delivery": 0.6187337581068277,
+      "loss_biodist": 1.0140363375345867,
+      "loss_toxic": 0.5890255868434906,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.3939393939393939,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.9921347300211587,
+      "loss_size": 0.4160850743452708,
+      "loss_pdi": 0.8274338642756144,
+      "loss_ee": 0.7452529668807983,
+      "loss_delivery": 0.6703582654396693,
+      "loss_biodist": 0.5891921321551005,
+      "loss_toxic": 0.7717546025911967,
+      "acc_pdi": 0.48484848484848486,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.85
+    },
+    {
+      "loss": 4.064125935236613,
+      "loss_size": 0.4211361159880956,
+      "loss_pdi": 1.008450706799825,
+      "loss_ee": 0.7111543019612631,
+      "loss_delivery": 0.712218443552653,
+      "loss_biodist": 0.40626858671506244,
+      "loss_toxic": 0.8385392477114996,
+      "acc_pdi": 0.45454545454545453,
+      "acc_ee": 0.42424242424242425,
+      "acc_toxic": 0.85
+    },
+    {
+      "loss": 5.119507948557536,
+      "loss_size": 0.5155423680941263,
+      "loss_pdi": 1.521255115667979,
+      "loss_ee": 0.6390946346024672,
+      "loss_delivery": 0.6691141823927561,
+      "loss_biodist": 0.29909205188353855,
+      "loss_toxic": 1.52191378424565,
+      "acc_pdi": 0.36363636363636365,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.85
+    },
+    {
+      "loss": 4.037670930226644,
+      "loss_size": 0.3509005631009738,
+      "loss_pdi": 1.1224110921223958,
+      "loss_ee": 0.5911329852727553,
+      "loss_delivery": 0.6472536226113638,
+      "loss_biodist": 0.21039783706267676,
+      "loss_toxic": 1.1335941478610039,
+      "acc_pdi": 0.48484848484848486,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.4177165031433105,
+      "loss_size": 0.4018242011467616,
+      "loss_pdi": 1.030598243077596,
+      "loss_ee": 0.6572648485501608,
+      "loss_delivery": 0.4684615011016528,
+      "loss_biodist": 0.30621210982402164,
+      "loss_toxic": 0.5808065546055635,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 2.925480365753174,
+      "loss_size": 0.3535038710882266,
+      "loss_pdi": 0.9870580832163492,
+      "loss_ee": 0.6889502070844173,
+      "loss_delivery": 0.37784648065765697,
+      "loss_biodist": 0.2775272900859515,
+      "loss_toxic": 0.2776398677378893,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 4.1336948076883955,
+      "loss_size": 0.6345704197883606,
+      "loss_pdi": 1.4422211448351543,
+      "loss_ee": 0.6099853285898765,
+      "loss_delivery": 0.5296423882246017,
+      "loss_biodist": 0.2524216075738271,
+      "loss_toxic": 0.7040676549077034,
+      "acc_pdi": 0.5454545454545454,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.174628575642904,
+      "loss_size": 0.43584080537160236,
+      "loss_pdi": 1.2121825615564983,
+      "loss_ee": 0.7292982923487822,
+      "loss_delivery": 0.36530304451783496,
+      "loss_biodist": 0.22861340890328088,
+      "loss_toxic": 0.2580638462677598,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.330165465672811,
+      "loss_size": 0.3679695477088292,
+      "loss_pdi": 1.3173633416493733,
+      "loss_ee": 0.6759289329250654,
+      "loss_delivery": 0.41639943420886993,
+      "loss_biodist": 0.21120227749149004,
+      "loss_toxic": 0.3905988199015458,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.40001114209493,
+      "loss_size": 0.3954329341650009,
+      "loss_pdi": 1.376650333404541,
+      "loss_ee": 0.7021505000690619,
+      "loss_delivery": 0.40341157217820484,
+      "loss_biodist": 0.2231835462152958,
+      "loss_toxic": 0.35528251078600687,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.4064132372538247,
+      "loss_size": 0.41711575786272687,
+      "loss_pdi": 1.419565737247467,
+      "loss_ee": 0.673577698568503,
+      "loss_delivery": 0.4514770656824112,
+      "loss_biodist": 0.19664333636562029,
+      "loss_toxic": 0.3087075824538867,
+      "acc_pdi": 0.5454545454545454,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.477815628051758,
+      "loss_size": 0.4257845679918925,
+      "loss_pdi": 1.4289001822471619,
+      "loss_ee": 0.6537623330950737,
+      "loss_delivery": 0.4686317443847656,
+      "loss_biodist": 0.2124811721344789,
+      "loss_toxic": 0.34616246446967125,
+      "acc_pdi": 0.5454545454545454,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9
+    }
+  ]
+}

models/abl_full/s1_baseline/seed42/outer_fold_0/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15, 17, 18, 19, 21, 22, 23, 24, 25, 26, 27, 28, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 52, 53, 54, 55, 56, 57, 58, 61, 62, 63, 64, 65, 66, 67, 68, 71, 72, 73, 75, 76, 77, 80, 81, 85, 86, 88, 89, 90, 92, 94, 95, 96, 98, 99, 100, 101, 103, 104, 105, 106, 107, 108, 109, 111, 113, 114, 115, 116, 117, 118, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 133, 134, 135, 136, 138, 139, 141, 144, 145, 146, 147, 148, 149, 150, 151, 153, 154, 155, 156, 157, 158, 159, 160, 163, 164, 167, 168, 169, 170, 171, 172, 175, 176, 178, 179, 180, 182, 183, 184, 185, 188, 189, 190, 191, 192, 193, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 206, 208, 210, 211, 212, 213, 216, 217, 218, 219, 220, 222, 224, 225, 227, 228, 229, 230, 231, 232, 233, 234, 236, 237, 239, 241, 242, 244, 245, 247, 248, 249, 250, 252, 253, 255, 256, 257, 258, 259, 260, 261, 262, 265, 267, 268, 269, 270, 271, 272, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 285, 286, 287, 288, 291, 292, 294, 296, 297, 298, 299, 300, 301, 303, 304, 306, 307, 308, 309, 310, 311, 312, 313, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 332, 333, 335, 336, 337, 338, 339, 340, 342, 343, 344, 345, 346, 347, 348, 349, 351, 352, 354, 355, 356, 357, 358, 359, 360, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 374, 375, 377, 379, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389, 390, 391, 393, 395, 396, 397, 398, 399, 400, 401, 403, 404, 405, 406, 407, 408, 409, 411, 412, 413, 414, 415, 417, 418, 419], "outer_test_idx": [7, 16, 20, 29, 51, 59, 60, 69, 70, 74, 78, 79, 82, 83, 84, 87, 91, 93, 97, 102, 110, 112, 119, 120, 121, 132, 137, 140, 142, 143, 152, 161, 162, 165, 166, 173, 174, 177, 181, 186, 187, 194, 205, 207, 209, 214, 215, 221, 223, 226, 235, 238, 240, 243, 246, 251, 254, 263, 264, 266, 273, 284, 289, 290, 293, 295, 302, 305, 314, 315, 331, 334, 341, 350, 353, 361, 373, 376, 378, 392, 394, 402, 410, 416]}

models/abl_full/s1_baseline/seed42/outer_fold_0/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 83,
+    "mse": 1.566862974660202,
+    "rmse": 1.2517439732869506,
+    "mae": 0.5613356892872288,
+    "r2": 0.04104129132919676
+  },
+  "delivery": {
+    "n_samples": 58,
+    "mse": 0.47028553218171953,
+    "rmse": 0.6857736741678843,
+    "mae": 0.5210600533875925,
+    "r2": 0.40223870972577835
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7976190476190477,
+    "precision": 0.7458333333333333,
+    "recall": 0.7523164647184604,
+    "f1": 0.7489010022859153
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.7142857142857143,
+    "precision": 0.6966386554621847,
+    "recall": 0.7485714285714287,
+    "f1": 0.700510874324916
+  },
+  "toxic": {
+    "n_samples": 58,
+    "accuracy": 0.9655172413793104,
+    "precision": 0.75,
+    "recall": 0.9821428571428572,
+    "f1": 0.8242424242424242
+  },
+  "biodist": {
+    "n_samples": 58,
+    "kl_divergence": 0.20122327746488497,
+    "js_divergence": 0.04018146399397734
+  }
+}

models/abl_full/s1_baseline/seed42/outer_fold_1/best_params.json

@@ -0,0 +1,12 @@
+{
+  "dropout": 0.32248545844049914,
+  "lr": 0.0009826983395023184,
+  "weight_decay": 0.012778379350481496,
+  "backbone_lr_ratio": 0.06065542527628115,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_baseline/seed42/outer_fold_1/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 9}

models/abl_full/s1_baseline/seed42/outer_fold_1/history.json

@@ -0,0 +1,195 @@
+{
+  "train": [
+    {
+      "loss": 5.70767883459727,
+      "loss_size": 1.2613731573025386,
+      "loss_pdi": 0.6852067477173276,
+      "loss_ee": 1.058294650581148,
+      "loss_delivery": 0.9235104901923074,
+      "loss_biodist": 1.1679435107443068,
+      "loss_toxic": 0.6260289185576968
+    },
+    {
+      "loss": 4.797239462534587,
+      "loss_size": 1.140162585510148,
+      "loss_pdi": 0.6206662886672549,
+      "loss_ee": 0.9612810247474246,
+      "loss_delivery": 0.8655857410695817,
+      "loss_biodist": 0.776539769437578,
+      "loss_toxic": 0.4614737547106213
+    },
+    {
+      "loss": 4.001818444993761,
+      "loss_size": 0.875577992035283,
+      "loss_pdi": 0.5799383173386256,
+      "loss_ee": 0.8913398087024689,
+      "loss_delivery": 0.8193241639269723,
+      "loss_biodist": 0.52018670241038,
+      "loss_toxic": 0.342042191988892
+    },
+    {
+      "loss": 3.9948329793082342,
+      "loss_size": 1.1431896595491304,
+      "loss_pdi": 0.551549611820115,
+      "loss_ee": 0.8537947800424364,
+      "loss_delivery": 0.8136444911360741,
+      "loss_biodist": 0.39836518797609544,
+      "loss_toxic": 0.2655619970626301
+    },
+    {
+      "loss": 3.3755411439471774,
+      "loss_size": 0.7370148715045717,
+      "loss_pdi": 0.5458221783240637,
+      "loss_ee": 0.8433284295929803,
+      "loss_delivery": 0.818140533235338,
+      "loss_biodist": 0.2898569570647346,
+      "loss_toxic": 0.16872831413315403
+    },
+    {
+      "loss": 3.533170633845859,
+      "loss_size": 1.0178725578718715,
+      "loss_pdi": 0.5446212440729141,
+      "loss_ee": 0.8055897222624885,
+      "loss_delivery": 0.7446173818575011,
+      "loss_biodist": 0.27136093543635476,
+      "loss_toxic": 0.1793254835324155
+    },
+    {
+      "loss": 3.5833877987331815,
+      "loss_size": 1.083082848952876,
+      "loss_pdi": 0.5406820194588767,
+      "loss_ee": 0.8283255994319916,
+      "loss_delivery": 0.7340445145964622,
+      "loss_biodist": 0.2850555007656415,
+      "loss_toxic": 0.1476559311979347
+    },
+    {
+      "loss": 3.5580843024783664,
+      "loss_size": 1.059447083208296,
+      "loss_pdi": 0.5500213387939665,
+      "loss_ee": 0.7789086335235171,
+      "loss_delivery": 0.7687801553143395,
+      "loss_biodist": 0.2777090138859219,
+      "loss_toxic": 0.1581853452242083
+    },
+    {
+      "loss": 3.5390172799428306,
+      "loss_size": 1.0753879389829106,
+      "loss_pdi": 0.5188433958424462,
+      "loss_ee": 0.8118821216954125,
+      "loss_delivery": 0.7294342882103391,
+      "loss_biodist": 0.2630891485346688,
+      "loss_toxic": 0.17479956853720877
+    }
+  ],
+  "val": [
+    {
+      "loss": 3.5459298292795816,
+      "loss_size": 0.5026707717527946,
+      "loss_pdi": 0.6600943009058634,
+      "loss_ee": 0.9537328879038492,
+      "loss_delivery": 0.13466781392829338,
+      "loss_biodist": 0.8540196816126505,
+      "loss_toxic": 0.4517778853575389,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.696969696969697,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.0298144022623696,
+      "loss_size": 0.4697304304378728,
+      "loss_pdi": 0.6688773433367411,
+      "loss_ee": 0.8737930854161581,
+      "loss_delivery": 0.15286462505658469,
+      "loss_biodist": 0.5575888554255167,
+      "loss_toxic": 0.31474941968917847,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.783087968826294,
+      "loss_size": 0.4554748572409153,
+      "loss_pdi": 0.7089216311772665,
+      "loss_ee": 0.8464933832486471,
+      "loss_delivery": 0.1867634654045105,
+      "loss_biodist": 0.3984869917233785,
+      "loss_toxic": 0.19690123200416565,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.6162113348642984,
+      "loss_size": 0.5064889093240103,
+      "loss_pdi": 0.6168197989463806,
+      "loss_ee": 0.6544777154922485,
+      "loss_delivery": 0.4863942662874858,
+      "loss_biodist": 0.32840614517529804,
+      "loss_toxic": 0.03221984952688217,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.4476131598154702,
+      "loss_size": 0.4997301201025645,
+      "loss_pdi": 0.6036792099475861,
+      "loss_ee": 0.5978536307811737,
+      "loss_delivery": 0.3749092022577922,
+      "loss_biodist": 0.3613563229640325,
+      "loss_toxic": 0.0169179222236077,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.8181818181818182,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.435731569925944,
+      "loss_size": 0.47308515508969623,
+      "loss_pdi": 0.6390049755573273,
+      "loss_ee": 0.6266860961914062,
+      "loss_delivery": 0.37278643250465393,
+      "loss_biodist": 0.32002559800942737,
+      "loss_toxic": 0.015139088034629822,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.379436413447062,
+      "loss_size": 0.43454373876253766,
+      "loss_pdi": 0.6299058894316355,
+      "loss_ee": 0.6250458657741547,
+      "loss_delivery": 0.38385624686876935,
+      "loss_biodist": 0.3001476426919301,
+      "loss_toxic": 0.017312646533052128,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.347810427347819,
+      "loss_size": 0.432897428671519,
+      "loss_pdi": 0.6165774265925089,
+      "loss_ee": 0.6293469866116842,
+      "loss_delivery": 0.3661603331565857,
+      "loss_biodist": 0.29900068541367847,
+      "loss_toxic": 0.015287932939827442,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.361335039138794,
+      "loss_size": 0.43744757771492004,
+      "loss_pdi": 0.6209465364615122,
+      "loss_ee": 0.6247348288695017,
+      "loss_delivery": 0.37391554315884906,
+      "loss_biodist": 0.30248237649599713,
+      "loss_toxic": 0.013860152413447699,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/abl_full/s1_baseline/seed42/outer_fold_1/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 3, 5, 6, 7, 10, 11, 12, 13, 15, 16, 18, 19, 20, 21, 24, 26, 27, 29, 30, 31, 32, 33, 35, 36, 38, 39, 40, 42, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 65, 66, 67, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 97, 98, 99, 100, 101, 102, 103, 105, 106, 107, 108, 109, 110, 112, 114, 115, 118, 119, 120, 121, 123, 124, 125, 126, 128, 129, 130, 131, 132, 136, 137, 140, 141, 142, 143, 144, 145, 146, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 161, 162, 164, 165, 166, 167, 168, 169, 171, 172, 173, 174, 175, 176, 177, 178, 179, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 194, 195, 196, 197, 198, 199, 201, 202, 204, 205, 206, 207, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 225, 226, 228, 229, 231, 232, 233, 234, 235, 236, 237, 238, 240, 242, 243, 244, 246, 247, 248, 250, 251, 252, 254, 255, 256, 258, 259, 260, 261, 262, 263, 264, 266, 268, 269, 270, 271, 273, 274, 275, 276, 277, 278, 279, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 295, 296, 297, 300, 301, 302, 303, 304, 305, 306, 307, 308, 310, 311, 313, 314, 315, 316, 317, 318, 319, 321, 322, 323, 324, 325, 327, 328, 329, 330, 331, 333, 334, 335, 337, 338, 339, 340, 341, 342, 343, 344, 345, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 361, 362, 364, 365, 366, 367, 368, 369, 373, 374, 375, 376, 378, 379, 380, 381, 382, 384, 387, 388, 389, 391, 392, 393, 394, 396, 397, 398, 399, 402, 403, 404, 405, 406, 408, 409, 410, 411, 412, 413, 414, 416, 417, 419], "outer_test_idx": [1, 2, 4, 8, 9, 14, 17, 22, 23, 25, 28, 34, 37, 41, 43, 45, 63, 64, 68, 85, 96, 104, 111, 113, 116, 117, 122, 127, 133, 134, 135, 138, 139, 147, 160, 163, 170, 180, 191, 192, 193, 200, 203, 208, 224, 227, 230, 239, 241, 245, 249, 253, 257, 265, 267, 272, 280, 294, 298, 299, 309, 312, 320, 326, 332, 336, 346, 347, 360, 363, 370, 371, 372, 377, 383, 385, 386, 390, 395, 400, 401, 407, 415, 418]}

models/abl_full/s1_baseline/seed42/outer_fold_1/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.331702512441025,
+    "rmse": 0.5759362051833736,
+    "mae": 0.44149345545364277,
+    "r2": -0.01290154523011866
+  },
+  "delivery": {
+    "n_samples": 61,
+    "mse": 1.0846693440821236,
+    "rmse": 1.0414746007858875,
+    "mae": 0.6676197994011835,
+    "r2": 0.17898349918416012
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.5952380952380952,
+    "precision": 0.5436720142602496,
+    "recall": 0.5555555555555556,
+    "f1": 0.5360623781676414
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.5595238095238095,
+    "precision": 0.5063813813813813,
+    "recall": 0.5141049426763712,
+    "f1": 0.4997539067306509
+  },
+  "toxic": {
+    "n_samples": 61,
+    "accuracy": 0.9508196721311475,
+    "precision": 0.75,
+    "recall": 0.9741379310344828,
+    "f1": 0.8200589970501475
+  },
+  "biodist": {
+    "n_samples": 61,
+    "kl_divergence": 0.3161094471320723,
+    "js_divergence": 0.08013612372537363
+  }
+}

models/abl_full/s1_baseline/seed42/outer_fold_2/best_params.json

@@ -0,0 +1,12 @@
+{
+  "dropout": 0.10175834754304547,
+  "lr": 0.0006945620111173825,
+  "weight_decay": 0.01625672738871503,
+  "backbone_lr_ratio": 0.029086945171372888,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_baseline/seed42/outer_fold_2/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 15}

models/abl_full/s1_baseline/seed42/outer_fold_2/history.json

@@ -0,0 +1,321 @@
+{
+  "train": [
+    {
+      "loss": 5.202105376455519,
+      "loss_size": 0.848908067577415,
+      "loss_pdi": 0.6807463301552666,
+      "loss_ee": 1.0441955692238278,
+      "loss_delivery": 0.9348783335751958,
+      "loss_biodist": 1.1314214732911851,
+      "loss_toxic": 0.5710593147410287
+    },
+    {
+      "loss": 4.290187080701192,
+      "loss_size": 0.7784139240781466,
+      "loss_pdi": 0.6099397473865085,
+      "loss_ee": 0.958238508966234,
+      "loss_delivery": 0.8511772139204873,
+      "loss_biodist": 0.7236086759302351,
+      "loss_toxic": 0.3855020155509313
+    },
+    {
+      "loss": 3.6722990936703153,
+      "loss_size": 0.7381285068889459,
+      "loss_pdi": 0.5751698017120361,
+      "loss_ee": 0.8887821502155728,
+      "loss_delivery": 0.743710472351975,
+      "loss_biodist": 0.4884147254957093,
+      "loss_toxic": 0.255600235114495
+    },
+    {
+      "loss": 3.541553775469462,
+      "loss_size": 0.7249418058329158,
+      "loss_pdi": 0.5661465426286062,
+      "loss_ee": 0.8174009753598107,
+      "loss_delivery": 0.8127418044540617,
+      "loss_biodist": 0.40241866144869065,
+      "loss_toxic": 0.23975762041906515
+    },
+    {
+      "loss": 3.170949247148302,
+      "loss_size": 0.6672847900125716,
+      "loss_pdi": 0.5368095321787728,
+      "loss_ee": 0.7959727810488807,
+      "loss_delivery": 0.703557981385125,
+      "loss_biodist": 0.3367742705676291,
+      "loss_toxic": 0.15184742936657536
+    },
+    {
+      "loss": 2.9941213859452143,
+      "loss_size": 0.6181929343276553,
+      "loss_pdi": 0.5130738466978073,
+      "loss_ee": 0.7803180151515536,
+      "loss_delivery": 0.6597989669276608,
+      "loss_biodist": 0.3322647420896424,
+      "loss_toxic": 0.11287090265088612
+    },
+    {
+      "loss": 2.984050432840983,
+      "loss_size": 0.6329434472653601,
+      "loss_pdi": 0.5116611768802007,
+      "loss_ee": 0.7843350271383921,
+      "loss_delivery": 0.6358760777446959,
+      "loss_biodist": 0.3198155090212822,
+      "loss_toxic": 0.12358394089258379
+    },
+    {
+      "loss": 2.883889695008596,
+      "loss_size": 0.5605236656135983,
+      "loss_pdi": 0.5077957941426171,
+      "loss_ee": 0.7636184924178653,
+      "loss_delivery": 0.6490659183926053,
+      "loss_biodist": 0.32912411871883607,
+      "loss_toxic": 0.1002082825741834
+    },
+    {
+      "loss": 2.812773803869883,
+      "loss_size": 0.5342350660098923,
+      "loss_pdi": 0.4876205407910877,
+      "loss_ee": 0.7610062824355232,
+      "loss_delivery": 0.6225873910718493,
+      "loss_biodist": 0.3106722831726074,
+      "loss_toxic": 0.122382206738823
+    },
+    {
+      "loss": 2.7658647762404547,
+      "loss_size": 0.5677060633897781,
+      "loss_pdi": 0.48942267232471043,
+      "loss_ee": 0.7341696321964264,
+      "loss_delivery": 0.6172876771953371,
+      "loss_biodist": 0.2766479564209779,
+      "loss_toxic": 0.10702241336305936
+    },
+    {
+      "loss": 2.811511311266157,
+      "loss_size": 0.529818933043215,
+      "loss_pdi": 0.5082361698150635,
+      "loss_ee": 0.7545560250679652,
+      "loss_delivery": 0.6142487269308832,
+      "loss_biodist": 0.3137662361065547,
+      "loss_toxic": 0.11994944792240858
+    },
+    {
+      "loss": 2.7609912753105164,
+      "loss_size": 0.533905324836572,
+      "loss_pdi": 0.4897963073518541,
+      "loss_ee": 0.7325207938750585,
+      "loss_delivery": 0.6409336452682813,
+      "loss_biodist": 0.3105732450882594,
+      "loss_toxic": 0.08392345615559155
+    },
+    {
+      "loss": 2.5689589116308422,
+      "loss_size": 0.4789157559474309,
+      "loss_pdi": 0.47479958335558575,
+      "loss_ee": 0.7016067024734285,
+      "loss_delivery": 0.5635421830746863,
+      "loss_biodist": 0.27678533891836804,
+      "loss_toxic": 0.09953307381106748
+    },
+    {
+      "loss": 2.740408976872762,
+      "loss_size": 0.5218956611222691,
+      "loss_pdi": 0.4786536180310779,
+      "loss_ee": 0.7272893289724985,
+      "loss_delivery": 0.6642732901705636,
+      "loss_biodist": 0.2777698271804386,
+      "loss_toxic": 0.10078600504332119
+    },
+    {
+      "loss": 2.5524483058187695,
+      "loss_size": 0.3650745376944542,
+      "loss_pdi": 0.4811245517598258,
+      "loss_ee": 0.7493493192725711,
+      "loss_delivery": 0.5953619066211913,
+      "loss_biodist": 0.2748099168141683,
+      "loss_toxic": 0.11551791522651911
+    }
+  ],
+  "val": [
+    {
+      "loss": 7.533385117848714,
+      "loss_size": 3.32920112212499,
+      "loss_pdi": 0.6970377763112386,
+      "loss_ee": 1.1328806281089783,
+      "loss_delivery": 1.3549890021483104,
+      "loss_biodist": 0.6858217716217041,
+      "loss_toxic": 0.3650021553039551,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.42424242424242425,
+      "acc_toxic": 0.8
+    },
+    {
+      "loss": 7.252995093663533,
+      "loss_size": 3.4935885270436606,
+      "loss_pdi": 0.6707201798756918,
+      "loss_ee": 1.1492963035901387,
+      "loss_delivery": 1.3669734001159668,
+      "loss_biodist": 0.4113343358039856,
+      "loss_toxic": 0.210152397553126,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.900516748428345,
+      "loss_size": 3.424496034781138,
+      "loss_pdi": 0.6340553164482117,
+      "loss_ee": 1.0742388765017192,
+      "loss_delivery": 1.365605463584264,
+      "loss_biodist": 0.3165936271349589,
+      "loss_toxic": 0.14185781280199686,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.72283673286438,
+      "loss_size": 3.336984912554423,
+      "loss_pdi": 0.6350743969281515,
+      "loss_ee": 1.0475572347640991,
+      "loss_delivery": 1.363085389137268,
+      "loss_biodist": 0.28269702196121216,
+      "loss_toxic": 0.1187615195910136,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.573407570521037,
+      "loss_size": 3.2991819977760315,
+      "loss_pdi": 0.6176866690317789,
+      "loss_ee": 1.0387301047643025,
+      "loss_delivery": 1.3273769617080688,
+      "loss_biodist": 0.2601098616917928,
+      "loss_toxic": 0.09336130321025848,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.622515757878621,
+      "loss_size": 3.3836412330468497,
+      "loss_pdi": 0.5748510162035624,
+      "loss_ee": 1.0535104672114055,
+      "loss_delivery": 1.32452525695165,
+      "loss_biodist": 0.2572739322980245,
+      "loss_toxic": 0.09353671098748843,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.70836607615153,
+      "loss_size": 3.510503182808558,
+      "loss_pdi": 0.5809042056401571,
+      "loss_ee": 1.0374245643615723,
+      "loss_delivery": 1.2894654472668965,
+      "loss_biodist": 0.24987786014874777,
+      "loss_toxic": 0.10607268412907918,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.571467558542888,
+      "loss_size": 3.3893363177776337,
+      "loss_pdi": 0.5938868920008341,
+      "loss_ee": 1.0552860299746196,
+      "loss_delivery": 1.2768413424491882,
+      "loss_biodist": 0.23242007692654928,
+      "loss_toxic": 0.0916902410487334,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.619278033574422,
+      "loss_size": 3.401833007733027,
+      "loss_pdi": 0.6145266890525818,
+      "loss_ee": 1.0751490791638691,
+      "loss_delivery": 1.2833717465400696,
+      "loss_biodist": 0.2202401359875997,
+      "loss_toxic": 0.09399381031592687,
+      "acc_pdi": 0.8484848484848485,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.504025061925252,
+      "loss_size": 3.301579564809799,
+      "loss_pdi": 0.606462836265564,
+      "loss_ee": 1.086024324099223,
+      "loss_delivery": 1.2854044437408447,
+      "loss_biodist": 0.2099631130695343,
+      "loss_toxic": 0.08557485789060593,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.508427063624064,
+      "loss_size": 3.3275153736273446,
+      "loss_pdi": 0.6022421419620514,
+      "loss_ee": 1.096723775068919,
+      "loss_delivery": 1.2710585991541545,
+      "loss_biodist": 0.20578517516454062,
+      "loss_toxic": 0.07695226619640987,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.572851657867432,
+      "loss_size": 3.366280754407247,
+      "loss_pdi": 0.6043597459793091,
+      "loss_ee": 1.088720401128133,
+      "loss_delivery": 1.2916629314422607,
+      "loss_biodist": 0.21097198128700256,
+      "loss_toxic": 0.08279866352677345,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.4637378851572675,
+      "loss_size": 3.343488405148188,
+      "loss_pdi": 0.5962008237838745,
+      "loss_ee": 1.0601000785827637,
+      "loss_delivery": 1.2487306793530781,
+      "loss_biodist": 0.2008370558420817,
+      "loss_toxic": 0.08361160010099411,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.4439826011657715,
+      "loss_size": 3.307605812946955,
+      "loss_pdi": 0.6033458610375723,
+      "loss_ee": 1.064460555712382,
+      "loss_delivery": 1.2547081510225933,
+      "loss_biodist": 0.2014366090297699,
+      "loss_toxic": 0.08231960361202557,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 6.438955942789714,
+      "loss_size": 3.2952869534492493,
+      "loss_pdi": 0.605786015590032,
+      "loss_ee": 1.0861074129740398,
+      "loss_delivery": 1.2420463760693867,
+      "loss_biodist": 0.20020990570386252,
+      "loss_toxic": 0.0803364800910155,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.9
+    }
+  ]
+}

models/abl_full/s1_baseline/seed42/outer_fold_2/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 11, 13, 14, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 32, 33, 34, 36, 37, 38, 39, 40, 41, 43, 44, 45, 47, 48, 49, 50, 51, 55, 56, 57, 58, 59, 60, 61, 63, 64, 65, 68, 69, 70, 72, 74, 76, 78, 79, 80, 82, 83, 84, 85, 86, 87, 88, 91, 93, 94, 95, 96, 97, 100, 101, 102, 103, 104, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 142, 143, 145, 147, 148, 149, 150, 151, 152, 153, 156, 157, 158, 159, 160, 161, 162, 163, 165, 166, 167, 168, 170, 171, 172, 173, 174, 175, 177, 179, 180, 181, 183, 184, 185, 186, 187, 188, 191, 192, 193, 194, 196, 197, 199, 200, 202, 203, 204, 205, 207, 208, 209, 210, 211, 212, 213, 214, 215, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 232, 233, 235, 236, 237, 238, 239, 240, 241, 242, 243, 245, 246, 249, 250, 251, 252, 253, 254, 255, 256, 257, 259, 260, 261, 263, 264, 265, 266, 267, 268, 269, 271, 272, 273, 274, 275, 276, 280, 281, 282, 283, 284, 285, 287, 289, 290, 293, 294, 295, 296, 297, 298, 299, 300, 302, 303, 305, 307, 308, 309, 310, 311, 312, 313, 314, 315, 317, 318, 320, 321, 322, 323, 324, 326, 327, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 358, 360, 361, 363, 364, 365, 366, 368, 369, 370, 371, 372, 373, 376, 377, 378, 379, 380, 381, 382, 383, 385, 386, 387, 388, 389, 390, 392, 394, 395, 398, 399, 400, 401, 402, 404, 405, 407, 408, 409, 410, 412, 413, 414, 415, 416, 417, 418], "outer_test_idx": [10, 12, 15, 31, 35, 42, 46, 52, 53, 54, 62, 66, 67, 71, 73, 75, 77, 81, 89, 90, 92, 98, 99, 105, 123, 124, 141, 144, 146, 154, 155, 164, 169, 176, 178, 182, 189, 190, 195, 198, 201, 206, 216, 217, 218, 219, 231, 234, 244, 247, 248, 258, 262, 270, 277, 278, 279, 286, 288, 291, 292, 301, 304, 306, 316, 319, 325, 328, 343, 357, 359, 362, 367, 374, 375, 384, 391, 393, 396, 397, 403, 406, 411, 419]}

models/abl_full/s1_baseline/seed42/outer_fold_2/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.41810017905356084,
+    "rmse": 0.646606664869425,
+    "mae": 0.49784872390418533,
+    "r2": 0.17425961655133881
+  },
+  "delivery": {
+    "n_samples": 60,
+    "mse": 0.6790202850062702,
+    "rmse": 0.8240268715316693,
+    "mae": 0.571309769502841,
+    "r2": 0.09610631021195903
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7023809523809523,
+    "precision": 0.6423197492163009,
+    "recall": 0.6664222873900294,
+    "f1": 0.6480643539467068
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6428571428571429,
+    "precision": 0.5810820624546115,
+    "recall": 0.623015873015873,
+    "f1": 0.5905593448696896
+  },
+  "toxic": {
+    "n_samples": 61,
+    "accuracy": 0.9344262295081968,
+    "precision": 0.7142857142857143,
+    "recall": 0.9655172413793103,
+    "f1": 0.7821428571428571
+  },
+  "biodist": {
+    "n_samples": 60,
+    "kl_divergence": 0.28741923551931975,
+    "js_divergence": 0.07400667098115551
+  }
+}

models/abl_full/s1_baseline/seed42/outer_fold_3/best_params.json

@@ -0,0 +1,12 @@
+{
+  "dropout": 0.11798498673229349,
+  "lr": 0.0009464809649969671,
+  "weight_decay": 0.07051438027639334,
+  "backbone_lr_ratio": 0.8641000194900307,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_baseline/seed42/outer_fold_3/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 11}

models/abl_full/s1_baseline/seed42/outer_fold_3/history.json

@@ -0,0 +1,237 @@
+{
+  "train": [
+    {
+      "loss": 4.848361690839131,
+      "loss_size": 0.5232753165894084,
+      "loss_pdi": 0.6738916403717465,
+      "loss_ee": 1.0730043318536546,
+      "loss_delivery": 0.8802526924345229,
+      "loss_biodist": 1.132146613465415,
+      "loss_toxic": 0.5782980405622058
+    },
+    {
+      "loss": 3.917437288496229,
+      "loss_size": 0.4489060507880317,
+      "loss_pdi": 0.6260161598523458,
+      "loss_ee": 0.9411158363024393,
+      "loss_delivery": 0.8751034893923335,
+      "loss_biodist": 0.6718993153836992,
+      "loss_toxic": 0.3736401821176211
+    },
+    {
+      "loss": 3.2179792324701944,
+      "loss_size": 0.4107091782821549,
+      "loss_pdi": 0.5818840811649958,
+      "loss_ee": 0.864887641535865,
+      "loss_delivery": 0.7771476540300581,
+      "loss_biodist": 0.3617808139986462,
+      "loss_toxic": 0.23983760095304912
+    },
+    {
+      "loss": 3.418438990910848,
+      "loss_size": 0.5134996208879683,
+      "loss_pdi": 0.6955742753214307,
+      "loss_ee": 0.8767334156566196,
+      "loss_delivery": 0.779440215892262,
+      "loss_biodist": 0.3787732298175494,
+      "loss_toxic": 0.20739152230736282
+    },
+    {
+      "loss": 2.998665518230862,
+      "loss_size": 0.47621528017852044,
+      "loss_pdi": 0.5725314848952823,
+      "loss_ee": 0.865351163678699,
+      "loss_delivery": 0.6925184585981898,
+      "loss_biodist": 0.2738187892569436,
+      "loss_toxic": 0.14818634762842622
+    },
+    {
+      "loss": 3.0510948499043784,
+      "loss_size": 0.45409081793493694,
+      "loss_pdi": 0.5420019378264745,
+      "loss_ee": 0.8024691624773873,
+      "loss_delivery": 0.8633336284094386,
+      "loss_biodist": 0.23776587181621128,
+      "loss_toxic": 0.1845432592736971
+    },
+    {
+      "loss": 2.5850073099136353,
+      "loss_size": 0.3699224545723862,
+      "loss_pdi": 0.4901059567928314,
+      "loss_ee": 0.7250210146109263,
+      "loss_delivery": 0.7355813177095519,
+      "loss_biodist": 0.21237118128273222,
+      "loss_toxic": 0.08344074033407702
+    },
+    {
+      "loss": 2.3256689177619085,
+      "loss_size": 0.3629870389898618,
+      "loss_pdi": 0.4683772557311588,
+      "loss_ee": 0.6916663034094704,
+      "loss_delivery": 0.597500105905864,
+      "loss_biodist": 0.17209894148012003,
+      "loss_toxic": 0.06179508691679479
+    },
+    {
+      "loss": 2.231321738825904,
+      "loss_size": 0.3313147864407963,
+      "loss_pdi": 0.4590592268440459,
+      "loss_ee": 0.6846879339880414,
+      "loss_delivery": 0.5587979132930437,
+      "loss_biodist": 0.14582366517020595,
+      "loss_toxic": 0.08081445035188356
+    },
+    {
+      "loss": 2.142837623755137,
+      "loss_size": 0.25173141145043904,
+      "loss_pdi": 0.44459058509932625,
+      "loss_ee": 0.6690129935741425,
+      "loss_delivery": 0.5448936944206556,
+      "loss_biodist": 0.1397495832708147,
+      "loss_toxic": 0.12115205913626899
+    },
+    {
+      "loss": 2.030080166127947,
+      "loss_size": 0.3170225802395079,
+      "loss_pdi": 0.44736198335886,
+      "loss_ee": 0.6005619350406859,
+      "loss_delivery": 0.5095781040274434,
+      "loss_biodist": 0.14586213645007876,
+      "loss_toxic": 0.03706750710908738
+    }
+  ],
+  "val": [
+    {
+      "loss": 5.591053167978923,
+      "loss_size": 2.305419623851776,
+      "loss_pdi": 0.638962964216868,
+      "loss_ee": 0.996762752532959,
+      "loss_delivery": 0.7161052773396174,
+      "loss_biodist": 0.5862163702646891,
+      "loss_toxic": 0.3665180404980977,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 0.782608695652174
+    },
+    {
+      "loss": 4.8712999025980634,
+      "loss_size": 2.2451172123352685,
+      "loss_pdi": 0.6204447746276855,
+      "loss_ee": 0.759979248046875,
+      "loss_delivery": 0.6457759390274683,
+      "loss_biodist": 0.27716003855069477,
+      "loss_toxic": 0.32968173424402875,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.696969696969697,
+      "acc_toxic": 0.782608695652174
+    },
+    {
+      "loss": 4.636776884396871,
+      "loss_size": 2.234689603249232,
+      "loss_pdi": 0.5556902090708414,
+      "loss_ee": 0.6675189236799876,
+      "loss_delivery": 0.7458879252274832,
+      "loss_biodist": 0.16838649411996207,
+      "loss_toxic": 0.25733551383018494,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.396537661552429,
+      "loss_size": 2.1907527844111123,
+      "loss_pdi": 0.5864351193110148,
+      "loss_ee": 0.6704360147317251,
+      "loss_delivery": 0.7243756751219431,
+      "loss_biodist": 0.12566474080085754,
+      "loss_toxic": 0.09803215662638347,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.533458352088928,
+      "loss_size": 2.1551299691200256,
+      "loss_pdi": 0.47677676876386005,
+      "loss_ee": 0.7799407144387563,
+      "loss_delivery": 0.6638170480728149,
+      "loss_biodist": 0.349340816338857,
+      "loss_toxic": 0.11235619088013966,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.213073750336965,
+      "loss_size": 2.070941517750422,
+      "loss_pdi": 0.40856675306955975,
+      "loss_ee": 0.7304698725541433,
+      "loss_delivery": 0.7343880335489908,
+      "loss_biodist": 0.09980805218219757,
+      "loss_toxic": 0.16575131316979727,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.220274964968364,
+      "loss_size": 1.9418311367432277,
+      "loss_pdi": 0.5001972615718842,
+      "loss_ee": 0.8348930279413859,
+      "loss_delivery": 0.6054362754027048,
+      "loss_biodist": 0.13444699347019196,
+      "loss_toxic": 0.20829975605010986,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.8260869565217391
+    },
+    {
+      "loss": 4.313011169433594,
+      "loss_size": 2.2030676305294037,
+      "loss_pdi": 0.36614329119523364,
+      "loss_ee": 0.7657052477200826,
+      "loss_delivery": 0.7487889031569163,
+      "loss_biodist": 0.09725691378116608,
+      "loss_toxic": 0.12932762006918588,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 3.871891895929972,
+      "loss_size": 1.8283031185468037,
+      "loss_pdi": 0.4723495543003082,
+      "loss_ee": 0.7260879675547282,
+      "loss_delivery": 0.6667743722597758,
+      "loss_biodist": 0.10258982578913371,
+      "loss_toxic": 0.08184227099021275,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 3.892843246459961,
+      "loss_size": 1.8570303916931152,
+      "loss_pdi": 0.43714167674382526,
+      "loss_ee": 0.7338299651940664,
+      "loss_delivery": 0.6806402603785197,
+      "loss_biodist": 0.10608401397864024,
+      "loss_toxic": 0.08255672703186671,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 3.915009021759033,
+      "loss_size": 1.887580672899882,
+      "loss_pdi": 0.42275914549827576,
+      "loss_ee": 0.7281539936860403,
+      "loss_delivery": 0.6997166474660238,
+      "loss_biodist": 0.10753046969572704,
+      "loss_toxic": 0.07313075910011928,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8695652173913043
+    }
+  ]
+}

models/abl_full/s1_baseline/seed42/outer_fold_3/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 1, 2, 4, 5, 7, 8, 9, 10, 11, 12, 14, 15, 16, 17, 20, 21, 22, 23, 24, 25, 27, 28, 29, 30, 31, 32, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 45, 46, 47, 51, 52, 53, 54, 55, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 73, 74, 75, 76, 77, 78, 79, 81, 82, 83, 84, 85, 86, 87, 89, 90, 91, 92, 93, 95, 96, 97, 98, 99, 100, 102, 104, 105, 106, 109, 110, 111, 112, 113, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 127, 129, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 146, 147, 148, 150, 151, 152, 153, 154, 155, 156, 158, 159, 160, 161, 162, 163, 164, 165, 166, 168, 169, 170, 172, 173, 174, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 189, 190, 191, 192, 193, 194, 195, 198, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 212, 213, 214, 215, 216, 217, 218, 219, 221, 223, 224, 226, 227, 228, 230, 231, 233, 234, 235, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 251, 252, 253, 254, 255, 256, 257, 258, 259, 262, 263, 264, 265, 266, 267, 268, 269, 270, 272, 273, 274, 275, 276, 277, 278, 279, 280, 282, 284, 286, 288, 289, 290, 291, 292, 293, 294, 295, 297, 298, 299, 301, 302, 304, 305, 306, 307, 308, 309, 310, 312, 314, 315, 316, 317, 319, 320, 321, 325, 326, 327, 328, 331, 332, 334, 336, 340, 341, 343, 344, 346, 347, 348, 350, 352, 353, 357, 359, 360, 361, 362, 363, 364, 366, 367, 368, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 382, 383, 384, 385, 386, 387, 389, 390, 391, 392, 393, 394, 395, 396, 397, 398, 400, 401, 402, 403, 404, 406, 407, 410, 411, 412, 414, 415, 416, 418, 419], "outer_test_idx": [3, 6, 13, 18, 19, 26, 33, 44, 48, 49, 50, 56, 57, 72, 80, 88, 94, 101, 103, 107, 108, 114, 126, 128, 130, 145, 149, 157, 167, 171, 175, 188, 196, 197, 199, 211, 220, 222, 225, 229, 232, 236, 250, 260, 261, 271, 281, 283, 285, 287, 296, 300, 303, 311, 313, 318, 322, 323, 324, 329, 330, 333, 335, 337, 338, 339, 342, 345, 349, 351, 354, 355, 356, 358, 365, 369, 381, 388, 399, 405, 408, 409, 413, 417]}

models/abl_full/s1_baseline/seed42/outer_fold_3/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 83,
+    "mse": 1.7901166464922558,
+    "rmse": 1.3379524081566787,
+    "mae": 0.5590303131906562,
+    "r2": -0.012503619354828466
+  },
+  "delivery": {
+    "n_samples": 59,
+    "mse": 0.6595745693785202,
+    "rmse": 0.8121419638083728,
+    "mae": 0.5407938919172196,
+    "r2": 0.3588151539119675
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7261904761904762,
+    "precision": 0.661793372319688,
+    "recall": 0.6825513196480939,
+    "f1": 0.6686674669867947
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6428571428571429,
+    "precision": 0.5840653153153154,
+    "recall": 0.599270990447461,
+    "f1": 0.5746753246753248
+  },
+  "toxic": {
+    "n_samples": 60,
+    "accuracy": 0.95,
+    "precision": 0.75,
+    "recall": 0.9736842105263157,
+    "f1": 0.8198198198198199
+  },
+  "biodist": {
+    "n_samples": 60,
+    "kl_divergence": 0.20028836227278726,
+    "js_divergence": 0.049937479227155704
+  }
+}

models/abl_full/s1_baseline/seed42/outer_fold_4/best_params.json

@@ -0,0 +1,12 @@
+{
+  "dropout": 0.26919128361241457,
+  "lr": 0.0008909017692547573,
+  "weight_decay": 9.871919386761893e-05,
+  "backbone_lr_ratio": 0.10881403588853521,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_baseline/seed42/outer_fold_4/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 12}

models/abl_full/s1_baseline/seed42/outer_fold_4/history.json

@@ -0,0 +1,258 @@
+{
+  "train": [
+    {
+      "loss": 5.595231321122911,
+      "loss_size": 1.1839522479308977,
+      "loss_pdi": 0.6729257073667314,
+      "loss_ee": 1.075638880332311,
+      "loss_delivery": 0.9437730138500532,
+      "loss_biodist": 1.1260581413904827,
+      "loss_toxic": 0.6060592151350446
+    },
+    {
+      "loss": 4.639838006761339,
+      "loss_size": 1.0754313990473747,
+      "loss_pdi": 0.625369002421697,
+      "loss_ee": 0.9961114095316993,
+      "loss_delivery": 0.8382789691289266,
+      "loss_biodist": 0.721036563316981,
+      "loss_toxic": 0.4075711866219838
+    },
+    {
+      "loss": 4.094899270269606,
+      "loss_size": 1.004522333542506,
+      "loss_pdi": 0.5886125796371036,
+      "loss_ee": 0.9341508017645942,
+      "loss_delivery": 0.7943624585039086,
+      "loss_biodist": 0.5192612277136909,
+      "loss_toxic": 0.28273969143629074
+    },
+    {
+      "loss": 3.8222215043173895,
+      "loss_size": 1.0109948880142636,
+      "loss_pdi": 0.5654200729396608,
+      "loss_ee": 0.9026512834760878,
+      "loss_delivery": 0.8094180805815591,
+      "loss_biodist": 0.408863328397274,
+      "loss_toxic": 0.1601414912276798
+    },
+    {
+      "loss": 3.403739419248369,
+      "loss_size": 0.9653827233446969,
+      "loss_pdi": 0.5261630614598592,
+      "loss_ee": 0.8332475655608707,
+      "loss_delivery": 0.6998005616996024,
+      "loss_biodist": 0.31920282708273995,
+      "loss_toxic": 0.09139471620114313
+    },
+    {
+      "loss": 3.3004104693730674,
+      "loss_size": 0.9118132177326415,
+      "loss_pdi": 0.5160764538579516,
+      "loss_ee": 0.7787083884080251,
+      "loss_delivery": 0.6926687616440985,
+      "loss_biodist": 0.27825497256384957,
+      "loss_toxic": 0.15196955762803555
+    },
+    {
+      "loss": 3.1182049910227456,
+      "loss_size": 0.9116067811846733,
+      "loss_pdi": 0.5113644252220789,
+      "loss_ee": 0.7740151551034715,
+      "loss_delivery": 0.5904112358887991,
+      "loss_biodist": 0.2735174712207582,
+      "loss_toxic": 0.08817623591878349
+    },
+    {
+      "loss": 2.9697708752420215,
+      "loss_size": 0.7829584280649821,
+      "loss_pdi": 0.49709691603978473,
+      "loss_ee": 0.7715644571516249,
+      "loss_delivery": 0.650741753478845,
+      "loss_biodist": 0.23108594864606857,
+      "loss_toxic": 0.06917133280593488
+    },
+    {
+      "loss": 2.936929007371267,
+      "loss_size": 0.9003308109111257,
+      "loss_pdi": 0.4637180914481481,
+      "loss_ee": 0.7399231162336137,
+      "loss_delivery": 0.5586378359132342,
+      "loss_biodist": 0.22390511176652378,
+      "loss_toxic": 0.08055410016742018
+    },
+    {
+      "loss": 2.7768101361062794,
+      "loss_size": 0.6965880054566596,
+      "loss_pdi": 0.48824302355448407,
+      "loss_ee": 0.7363911502891116,
+      "loss_delivery": 0.5831497386097908,
+      "loss_biodist": 0.24472905115948784,
+      "loss_toxic": 0.059548021708097726
+    },
+    {
+      "loss": 2.8497932420836554,
+      "loss_size": 0.7591909679273764,
+      "loss_pdi": 0.48939237164126503,
+      "loss_ee": 0.7746303362978829,
+      "loss_delivery": 0.5287295778592428,
+      "loss_biodist": 0.2293989182346397,
+      "loss_toxic": 0.10032574077033335
+    },
+    {
+      "loss": 2.8877300487624273,
+      "loss_size": 0.7499017897579405,
+      "loss_pdi": 0.4924188272820579,
+      "loss_ee": 0.7267551057868533,
+      "loss_delivery": 0.5910006670488251,
+      "loss_biodist": 0.2306723656753699,
+      "loss_toxic": 0.12704977796723446
+    }
+  ],
+  "val": [
+    {
+      "loss": 4.951508442560832,
+      "loss_size": 0.7575126687685648,
+      "loss_pdi": 0.6334818204243978,
+      "loss_ee": 1.0574791034062703,
+      "loss_delivery": 0.7090578771506747,
+      "loss_biodist": 1.1586121320724487,
+      "loss_toxic": 0.6608855525652567,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 4.170979658762614,
+      "loss_size": 0.5657947659492493,
+      "loss_pdi": 0.5647561947504679,
+      "loss_ee": 0.93976891040802,
+      "loss_delivery": 0.7305525317788124,
+      "loss_biodist": 0.7686646183331808,
+      "loss_toxic": 0.6286381383736929,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 3.587951978047689,
+      "loss_size": 0.5118454098701477,
+      "loss_pdi": 0.49248048663139343,
+      "loss_ee": 0.8489340345064799,
+      "loss_delivery": 0.70010299111406,
+      "loss_biodist": 0.45328495899836224,
+      "loss_toxic": 0.5979513178269068,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 3.40149195988973,
+      "loss_size": 0.30642780661582947,
+      "loss_pdi": 0.39163154860337573,
+      "loss_ee": 0.6825251877307892,
+      "loss_delivery": 0.6704287938773632,
+      "loss_biodist": 0.19783767064412436,
+      "loss_toxic": 1.1384821409980457,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 3.612094004948934,
+      "loss_size": 0.4379969338575999,
+      "loss_pdi": 0.3864056666692098,
+      "loss_ee": 0.6370339145263036,
+      "loss_delivery": 0.7303227235873541,
+      "loss_biodist": 0.1583778622249762,
+      "loss_toxic": 1.2403487414121628,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.936262925465902,
+      "loss_size": 0.4848108838001887,
+      "loss_pdi": 0.4384245773156484,
+      "loss_ee": 0.6231235327819983,
+      "loss_delivery": 0.7736449837684631,
+      "loss_biodist": 0.15557912488778433,
+      "loss_toxic": 0.4637433315316836,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.9137487411499023,
+      "loss_size": 0.3753274182478587,
+      "loss_pdi": 0.36579150954882306,
+      "loss_ee": 0.6489318336049715,
+      "loss_delivery": 0.8045151382684708,
+      "loss_biodist": 0.13289600859085718,
+      "loss_toxic": 0.5830482244491577,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 3.071545680363973,
+      "loss_size": 0.40580152968565625,
+      "loss_pdi": 0.3794499784708023,
+      "loss_ee": 0.6523361951112747,
+      "loss_delivery": 0.8165061851032078,
+      "loss_biodist": 0.13434462249279022,
+      "loss_toxic": 0.6777304336428642,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.9691779613494873,
+      "loss_size": 0.4135140875975291,
+      "loss_pdi": 0.38700851301352185,
+      "loss_ee": 0.6639018009106318,
+      "loss_delivery": 0.8097411928077539,
+      "loss_biodist": 0.13824120288093886,
+      "loss_toxic": 0.5578225553035736,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.8584330876668296,
+      "loss_size": 0.38691479464372,
+      "loss_pdi": 0.3887040615081787,
+      "loss_ee": 0.6753135286271572,
+      "loss_delivery": 0.8105639110629758,
+      "loss_biodist": 0.13059528730809689,
+      "loss_toxic": 0.4719698553284009,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.887652635574341,
+      "loss_size": 0.3818443367878596,
+      "loss_pdi": 0.3719891682267189,
+      "loss_ee": 0.6681812778115273,
+      "loss_delivery": 0.8080568375686804,
+      "loss_biodist": 0.12677289421359697,
+      "loss_toxic": 0.5323583881060282,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.943228085835775,
+      "loss_size": 0.3836331069469452,
+      "loss_pdi": 0.37640736748774845,
+      "loss_ee": 0.6699009388685226,
+      "loss_delivery": 0.8047574162483215,
+      "loss_biodist": 0.12774107729395232,
+      "loss_toxic": 0.5802475971480211,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8888888888888888
+    }
+  ]
+}

models/abl_full/s1_baseline/seed42/outer_fold_4/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 23, 25, 26, 28, 29, 31, 33, 34, 35, 37, 41, 42, 43, 44, 45, 46, 48, 49, 50, 51, 52, 53, 54, 56, 57, 59, 60, 62, 63, 64, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 77, 78, 79, 80, 81, 82, 83, 84, 85, 87, 88, 89, 90, 91, 92, 93, 94, 96, 97, 98, 99, 101, 102, 103, 104, 105, 107, 108, 110, 111, 112, 113, 114, 116, 117, 119, 120, 121, 122, 123, 124, 126, 127, 128, 130, 132, 133, 134, 135, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 149, 152, 154, 155, 157, 160, 161, 162, 163, 164, 165, 166, 167, 169, 170, 171, 173, 174, 175, 176, 177, 178, 180, 181, 182, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 203, 205, 206, 207, 208, 209, 211, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 229, 230, 231, 232, 234, 235, 236, 238, 239, 240, 241, 243, 244, 245, 246, 247, 248, 249, 250, 251, 253, 254, 257, 258, 260, 261, 262, 263, 264, 265, 266, 267, 270, 271, 272, 273, 277, 278, 279, 280, 281, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 298, 299, 300, 301, 302, 303, 304, 305, 306, 309, 311, 312, 313, 314, 315, 316, 318, 319, 320, 322, 323, 324, 325, 326, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 341, 342, 343, 345, 346, 347, 349, 350, 351, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 365, 367, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 381, 383, 384, 385, 386, 388, 390, 391, 392, 393, 394, 395, 396, 397, 399, 400, 401, 402, 403, 405, 406, 407, 408, 409, 410, 411, 413, 415, 416, 417, 418, 419], "outer_test_idx": [0, 5, 11, 21, 24, 27, 30, 32, 36, 38, 39, 40, 47, 55, 58, 61, 65, 76, 86, 95, 100, 106, 109, 115, 118, 125, 129, 131, 136, 148, 150, 151, 153, 156, 158, 159, 168, 172, 179, 183, 184, 185, 202, 204, 210, 212, 213, 228, 233, 237, 242, 252, 255, 256, 259, 268, 269, 274, 275, 276, 282, 297, 307, 308, 310, 317, 321, 327, 340, 344, 348, 352, 364, 366, 368, 379, 380, 382, 387, 389, 398, 404, 412, 414]}

models/abl_full/s1_baseline/seed42/outer_fold_4/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.7761301185113868,
+    "rmse": 0.8809824734416609,
+    "mae": 0.4868293024677162,
+    "r2": -0.038558506277432514
+  },
+  "delivery": {
+    "n_samples": 58,
+    "mse": 0.8724045566943844,
+    "rmse": 0.9340259935860374,
+    "mae": 0.6374844717953739,
+    "r2": 0.15639919541989333
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.75,
+    "precision": 0.7014607425441266,
+    "recall": 0.7426686217008798,
+    "f1": 0.7105824446267432
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6666666666666666,
+    "precision": 0.6090232389251997,
+    "recall": 0.6632478632478632,
+    "f1": 0.6219283975033759
+  },
+  "toxic": {
+    "n_samples": 59,
+    "accuracy": 0.9661016949152542,
+    "precision": 0.8,
+    "recall": 0.9821428571428572,
+    "f1": 0.8659090909090909
+  },
+  "biodist": {
+    "n_samples": 58,
+    "kl_divergence": 0.23172865065491569,
+    "js_divergence": 0.06330909380413241
+  }
+}

models/abl_full/s1_baseline/seed42/pip_freeze.txt

@@ -0,0 +1,153 @@
+accelerate==1.0.1
+alabaster==0.7.13
+alembic==1.14.1
+altair==5.4.1
+annotated-doc==0.0.4
+annotated-types==0.7.0
+anyio==4.5.2
+attrs==25.3.0
+babel==2.18.0
+bitsandbytes==0.45.5
+blinker==1.8.2
+Brotli @ file:///croot/brotli-split_1714483155106/work
+cachetools==5.5.2
+captum==0.7.0
+certifi @ file:///croot/certifi_1725551672989/work/certifi
+charset-normalizer @ file:///croot/charset-normalizer_1721748349566/work
+chemprop==1.7.0
+click==8.1.8
+cloudpickle==3.1.2
+colorlog==6.10.1
+contourpy==1.1.1
+cramjam==2.11.0
+cycler==0.12.1
+descriptastorus==2.8.0
+docstring_parser==0.18.0
+docutils==0.20.1
+et_xmlfile==2.0.0
+exceptiongroup==1.3.1
+fastapi==0.124.4
+fastparquet==2024.2.0
+filelock @ file:///croot/filelock_1700591183607/work
+Flask==2.1.3
+fonttools==4.57.0
+fsspec==2025.3.0
+future==1.0.0
+gitdb==4.0.12
+GitPython==3.1.50
+gmpy2 @ file:///tmp/build/80754af9/gmpy2_1645455532332/work
+greenlet==3.1.1
+h11==0.16.0
+hf-xet==1.5.1
+hf_transfer==0.1.9
+httpcore==1.0.9
+httpx==0.28.1
+huggingface_hub==0.36.2
+hyperopt==0.2.7
+idna @ file:///croot/idna_1714398848350/work
+imagesize==1.5.0
+importlib_metadata==8.5.0
+importlib_resources==6.4.5
+iniconfig==2.1.0
+inquirerpy==0.3.4
+itsdangerous==2.2.0
+Jinja2 @ file:///croot/jinja2_1716993405101/work
+joblib==1.4.2
+jsonschema==4.23.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.7
+-e git+ssh://git@github.com/RYDE-WORK/lnp_ml.git@104dfef94c6d6f03eb63369b4c4edc2f6ed4437d#egg=lnp_ml
+loguru==0.7.3
+Mako==1.3.12
+markdown-it-py==3.0.0
+MarkupSafe @ file:///croot/markupsafe_1704205993651/work
+matplotlib==3.7.5
+mdurl==0.1.2
+mkl-fft==1.3.1
+mkl-random @ file:///tmp/build/80754af9/mkl_random_1626186064646/work
+mkl-service==2.4.0
+modelscope==1.31.0
+mpmath @ file:///croot/mpmath_1690848262763/work
+mypy_extensions==1.1.0
+narwhals==1.42.1
+networkx @ file:///croot/networkx_1690561992265/work
+numpy @ file:///croot/numpy_and_numpy_base_1682520569166/work
+openpyxl==3.1.5
+optuna==4.5.0
+packaging==24.2
+pandas==2.0.3
+pandas_flavor==0.7.0
+peft==0.13.2
+pfzy==0.3.4
+pillow @ file:///croot/pillow_1721059439630/work
+pkgutil_resolve_name==1.3.10
+pluggy==1.5.0
+prompt_toolkit==3.0.52
+protobuf==5.29.6
+psutil==7.2.2
+py4j==0.10.9.9
+pyarrow==17.0.0
+pydantic==2.10.6
+pydantic_core==2.27.2
+pydeck==0.9.2
+Pygments==2.19.2
+pyparsing==3.1.4
+PySocks @ file:///tmp/build/80754af9/pysocks_1605305779399/work
+pytest==8.3.5
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+pytz==2026.2
+PyYAML @ file:///croot/pyyaml_1728657952215/work
+rdkit==2024.3.5
+referencing==0.35.1
+regex==2024.11.6
+requests @ file:///croot/requests_1721410876868/work
+rich==13.9.4
+rpds-py==0.20.1
+safetensors==0.5.3
+scikit-learn==1.3.2
+scipy==1.10.1
+selfies==2.2.0
+sentencepiece==0.2.0
+shellingham==1.5.4
+six @ file:///tmp/build/80754af9/six_1644875935023/work
+smmap==5.0.3
+sniffio==1.3.1
+snowballstemmer==3.1.1
+Sphinx==7.1.2
+sphinx_rtd_theme==3.1.0
+sphinxcontrib-applehelp==1.0.4
+sphinxcontrib-devhelp==1.0.2
+sphinxcontrib-htmlhelp==2.0.1
+sphinxcontrib-jquery==4.1
+sphinxcontrib-jsmath==1.0.1
+sphinxcontrib-qthelp==1.0.3
+sphinxcontrib-serializinghtml==1.1.5
+SQLAlchemy==2.0.51
+starlette==0.44.0
+streamlit==1.40.1
+sympy @ file:///croot/sympy_1734622612703/work
+tenacity==9.0.0
+tensorboardX==2.6.2.2
+threadpoolctl==3.5.0
+tokenizers==0.20.3
+toml==0.10.2
+tomli==2.4.1
+torch==2.4.1
+torchvision==0.20.0
+tornado==6.4.2
+tqdm==4.68.2
+transformers==4.46.3
+triton==3.0.0
+typed-argument-parser==1.10.1
+typer==0.20.1
+typing-inspect==0.9.0
+typing_extensions==4.13.2
+tzdata==2026.2
+urllib3 @ file:///croot/urllib3_1727769808118/work
+uvicorn==0.33.0
+watchdog==4.0.2
+wcwidth==0.8.2
+Werkzeug==2.3.8
+xarray==2023.1.0
+zipp==3.20.2

models/abl_full/s1_baseline/seed42/repro.txt

@@ -0,0 +1,5 @@
+date: 2026-07-18T10:18:40+00:00
+git_commit: 104dfef94c6d6f03eb63369b4c4edc2f6ed4437d
+git_dirty_files: 23
+seed: 42  pythonhashseed: 42
+cmd: GPU=1 --batch-size 16 --use-mpnn

models/abl_full/s1_baseline/seed42/strata_info.json

@@ -0,0 +1,60 @@
+{
+  "original_strata_counts": {
+    "T0|P0|E0": "5",
+    "T0|P0|E1": "58",
+    "T0|P0|E2": "169",
+    "T0|P1|E0": "1",
+    "T0|P1|E1": "17",
+    "T0|P1|E2": "32",
+    "T0|P2|E2": "3",
+    "T1|P0|E2": "9",
+    "T1|P1|E2": "5",
+    "TNA|P0|E0": "28",
+    "TNA|P0|E1": "21",
+    "TNA|P0|E2": "20",
+    "TNA|P1|E0": "29",
+    "TNA|P1|E1": "7",
+    "TNA|P1|E2": "14",
+    "TNA|P2|E2": "1",
+    "TNA|P3|E0": "1"
+  },
+  "rare_strata": [
+    "T0|P1|E0",
+    "T0|P2|E2",
+    "TNA|P2|E2",
+    "TNA|P3|E0"
+  ],
+  "final_strata": [
+    "RARE",
+    "T0|P0|E0",
+    "T0|P0|E1",
+    "T0|P0|E2",
+    "T0|P1|E1",
+    "T0|P1|E2",
+    "T1|P0|E2",
+    "T1|P1|E2",
+    "TNA|P0|E0",
+    "TNA|P0|E1",
+    "TNA|P0|E2",
+    "TNA|P1|E0",
+    "TNA|P1|E1",
+    "TNA|P1|E2"
+  ],
+  "final_strata_counts": {
+    "RARE": "6",
+    "T0|P0|E0": "5",
+    "T0|P0|E1": "58",
+    "T0|P0|E2": "169",
+    "T0|P1|E1": "17",
+    "T0|P1|E2": "32",
+    "T1|P0|E2": "9",
+    "T1|P1|E2": "5",
+    "TNA|P0|E0": "28",
+    "TNA|P0|E1": "21",
+    "TNA|P0|E2": "20",
+    "TNA|P1|E0": "29",
+    "TNA|P1|E1": "7",
+    "TNA|P1|E2": "14"
+  },
+  "n_rare_merged": "6"
+}

models/abl_full/s1_baseline/seed42/summary.json

@@ -0,0 +1,352 @@
+{
+  "fold_results": [
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.2708449222533328,
+        "lr": 0.000991430410947854,
+        "weight_decay": 0.017147551517191443,
+        "backbone_lr_ratio": 0.43561339540605226,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 13,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.566862974660202,
+          "rmse": 1.2517439732869506,
+          "mae": 0.5613356892872288,
+          "r2": 0.04104129132919676
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.47028553218171953,
+          "rmse": 0.6857736741678843,
+          "mae": 0.5210600533875925,
+          "r2": 0.40223870972577835
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7976190476190477,
+          "precision": 0.7458333333333333,
+          "recall": 0.7523164647184604,
+          "f1": 0.7489010022859153
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6966386554621847,
+          "recall": 0.7485714285714287,
+          "f1": 0.700510874324916
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.20122327746488497,
+          "js_divergence": 0.04018146399397734
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.32248545844049914,
+        "lr": 0.0009826983395023184,
+        "weight_decay": 0.012778379350481496,
+        "backbone_lr_ratio": 0.06065542527628115,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 9,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.331702512441025,
+          "rmse": 0.5759362051833736,
+          "mae": 0.44149345545364277,
+          "r2": -0.01290154523011866
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.0846693440821236,
+          "rmse": 1.0414746007858875,
+          "mae": 0.6676197994011835,
+          "r2": 0.17898349918416012
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.5952380952380952,
+          "precision": 0.5436720142602496,
+          "recall": 0.5555555555555556,
+          "f1": 0.5360623781676414
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5595238095238095,
+          "precision": 0.5063813813813813,
+          "recall": 0.5141049426763712,
+          "f1": 0.4997539067306509
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.3161094471320723,
+          "js_divergence": 0.08013612372537363
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.10175834754304547,
+        "lr": 0.0006945620111173825,
+        "weight_decay": 0.01625672738871503,
+        "backbone_lr_ratio": 0.029086945171372888,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 15,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.41810017905356084,
+          "rmse": 0.646606664869425,
+          "mae": 0.49784872390418533,
+          "r2": 0.17425961655133881
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6790202850062702,
+          "rmse": 0.8240268715316693,
+          "mae": 0.571309769502841,
+          "r2": 0.09610631021195903
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6423197492163009,
+          "recall": 0.6664222873900294,
+          "f1": 0.6480643539467068
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5810820624546115,
+          "recall": 0.623015873015873,
+          "f1": 0.5905593448696896
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.28741923551931975,
+          "js_divergence": 0.07400667098115551
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.11798498673229349,
+        "lr": 0.0009464809649969671,
+        "weight_decay": 0.07051438027639334,
+        "backbone_lr_ratio": 0.8641000194900307,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.7901166464922558,
+          "rmse": 1.3379524081566787,
+          "mae": 0.5590303131906562,
+          "r2": -0.012503619354828466
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.6595745693785202,
+          "rmse": 0.8121419638083728,
+          "mae": 0.5407938919172196,
+          "r2": 0.3588151539119675
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.661793372319688,
+          "recall": 0.6825513196480939,
+          "f1": 0.6686674669867947
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5840653153153154,
+          "recall": 0.599270990447461,
+          "f1": 0.5746753246753248
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.95,
+          "precision": 0.75,
+          "recall": 0.9736842105263157,
+          "f1": 0.8198198198198199
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.20028836227278726,
+          "js_divergence": 0.049937479227155704
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.26919128361241457,
+        "lr": 0.0008909017692547573,
+        "weight_decay": 9.871919386761893e-05,
+        "backbone_lr_ratio": 0.10881403588853521,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 12,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.7761301185113868,
+          "rmse": 0.8809824734416609,
+          "mae": 0.4868293024677162,
+          "r2": -0.038558506277432514
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.8724045566943844,
+          "rmse": 0.9340259935860374,
+          "mae": 0.6374844717953739,
+          "r2": 0.15639919541989333
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.75,
+          "precision": 0.7014607425441266,
+          "recall": 0.7426686217008798,
+          "f1": 0.7105824446267432
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.6090232389251997,
+          "recall": 0.6632478632478632,
+          "f1": 0.6219283975033759
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.8659090909090909
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.23172865065491569,
+          "js_divergence": 0.06330909380413241
+        }
+      }
+    }
+  ],
+  "summary_stats": {
+    "size": {
+      "mse_mean": 0.976582486231686,
+      "mse_std": 0.5963594058421813,
+      "rmse_mean": 0.9386443449876178,
+      "rmse_std": 0.30907811286866,
+      "mae_mean": 0.5093074968606859,
+      "mae_std": 0.045639036325825846,
+      "r2_mean": 0.030267447403631187,
+      "r2_std": 0.07652419024801474
+    },
+    "delivery": {
+      "mse_mean": 0.7531908574686035,
+      "mse_std": 0.20899310967661966,
+      "rmse_mean": 0.8594886207759702,
+      "rmse_std": 0.12029201230848217,
+      "mae_mean": 0.5876535972008421,
+      "mae_std": 0.05617003172137069,
+      "r2_mean": 0.23850857369075168,
+      "r2_std": 0.11987061343757097
+    },
+    "pdi": {
+      "accuracy_mean": 0.7142857142857142,
+      "accuracy_std": 0.0673435029701474,
+      "precision_mean": 0.6590158423347396,
+      "precision_std": 0.06771930178345034,
+      "recall_mean": 0.6799028498026038,
+      "recall_std": 0.07047590125597848,
+      "f1_mean": 0.6624555292027603,
+      "f1_std": 0.0721248119992936
+    },
+    "ee": {
+      "accuracy_mean": 0.6452380952380952,
+      "accuracy_std": 0.05016977977548923,
+      "precision_mean": 0.5954381307077385,
+      "precision_std": 0.06114350922554513,
+      "recall_mean": 0.6296422195917994,
+      "recall_std": 0.07692682031615919,
+      "f1_mean": 0.5974855696207915,
+      "f1_std": 0.06531680188025496
+    },
+    "toxic": {
+      "accuracy_mean": 0.9533729675867818,
+      "accuracy_std": 0.01171647555533869,
+      "precision_mean": 0.7528571428571429,
+      "precision_std": 0.027330180671012855,
+      "recall_mean": 0.9755250194451648,
+      "recall_std": 0.006213851706247248,
+      "f1_mean": 0.822434637832868,
+      "f1_std": 0.026567667174673414
+    },
+    "biodist": {
+      "kl_divergence_mean": 0.24735379460879597,
+      "kl_divergence_std": 0.04673308028747034,
+      "js_divergence_mean": 0.06151416634635892,
+      "js_divergence_std": 0.014799767616865855
+    }
+  }
+}

models/abl_full/s1_both/seed42/local_uncommitted.patch

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

models/abl_full/s1_both/seed42/outer_fold_0/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.4108948503325973,
+  "lr": 0.0007340472227986773,
+  "weight_decay": 0.0007067135047679004,
+  "backbone_lr_ratio": 0.7321173243252591,
+  "moe_n_experts": 2,
+  "moe_top_k": 2,
+  "moe_expert_hidden_mult": 1,
+  "llm_lora_r": 32,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_both/seed42/outer_fold_0/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 14}

models/abl_full/s1_both/seed42/outer_fold_0/history.json

@@ -0,0 +1,328 @@
+{
+  "train": [
+    {
+      "loss": 5.400099200171393,
+      "loss_size": 0.8823040468064515,
+      "loss_pdi": 0.6812556498759502,
+      "loss_ee": 1.0784653425216675,
+      "loss_delivery": 1.0165383002242527,
+      "loss_biodist": 1.1444005257374532,
+      "loss_toxic": 0.5934148213347873,
+      "loss_moe_lb": 1.999999977446891
+    },
+    {
+      "loss": 4.204111298999271,
+      "loss_size": 0.8327704547224818,
+      "loss_pdi": 0.6399245173544497,
+      "loss_ee": 1.0011927485466003,
+      "loss_delivery": 0.845224013703095,
+      "loss_biodist": 0.5762553400284535,
+      "loss_toxic": 0.313652445537013,
+      "loss_moe_lb": 1.999999990334382
+    },
+    {
+      "loss": 3.768066009959659,
+      "loss_size": 0.8170974753998421,
+      "loss_pdi": 0.6213119199147096,
+      "loss_ee": 0.9148938607525181,
+      "loss_delivery": 0.8916980289929622,
+      "loss_biodist": 0.35265421726413676,
+      "loss_toxic": 0.18339438966752306,
+      "loss_moe_lb": 2.0
+    },
+    {
+      "loss": 3.7443795784099683,
+      "loss_size": 0.8354647985181293,
+      "loss_pdi": 0.6066541373729706,
+      "loss_ee": 0.8851948912079269,
+      "loss_delivery": 1.0085055058872379,
+      "loss_biodist": 0.29029415326344,
+      "loss_toxic": 0.1493166242498036,
+      "loss_moe_lb": 1.9999999871125091
+    },
+    {
+      "loss": 3.499276167637593,
+      "loss_size": 0.8638738892167002,
+      "loss_pdi": 0.5628914881396938,
+      "loss_ee": 0.9029648924196089,
+      "loss_delivery": 0.865693795117172,
+      "loss_biodist": 0.21755379871339411,
+      "loss_toxic": 0.12434523284507638,
+      "loss_moe_lb": 1.9999999935562547
+    },
+    {
+      "loss": 3.5330081018241675,
+      "loss_size": 0.7820458718248315,
+      "loss_pdi": 0.5551547770564621,
+      "loss_ee": 0.8933853394276386,
+      "loss_delivery": 0.989005317949262,
+      "loss_biodist": 0.20653205576377945,
+      "loss_toxic": 0.15761368533084524,
+      "loss_moe_lb": 1.999999977446891
+    },
+    {
+      "loss": 3.318131614375759,
+      "loss_size": 0.764020403696073,
+      "loss_pdi": 0.5151227450854069,
+      "loss_ee": 0.872619815774866,
+      "loss_delivery": 0.9253612577512458,
+      "loss_biodist": 0.1871767146764575,
+      "loss_toxic": 0.11265365003468154,
+      "loss_moe_lb": 1.9999999967781272
+    },
+    {
+      "loss": 3.1988877670184985,
+      "loss_size": 0.7646275287544405,
+      "loss_pdi": 0.5320992252311191,
+      "loss_ee": 0.8285162891890552,
+      "loss_delivery": 0.7588972472459883,
+      "loss_biodist": 0.20326888027626114,
+      "loss_toxic": 0.16523235014205598,
+      "loss_moe_lb": 1.9999999935562547
+    },
+    {
+      "loss": 2.9703359152819657,
+      "loss_size": 0.6830032755796974,
+      "loss_pdi": 0.501954268362071,
+      "loss_ee": 0.8167013405142604,
+      "loss_delivery": 0.7976664585640302,
+      "loss_biodist": 0.17228593788034208,
+      "loss_toxic": 0.06885367960343058,
+      "loss_moe_lb": 1.9999999871125091
+    },
+    {
+      "loss": 2.565713747127636,
+      "loss_size": 0.3230193958089158,
+      "loss_pdi": 0.5138459878193365,
+      "loss_ee": 0.7612508189033818,
+      "loss_delivery": 0.7678211121945768,
+      "loss_biodist": 0.16089659788318583,
+      "loss_toxic": 0.10352957973058487,
+      "loss_moe_lb": 1.9999999871125091
+    },
+    {
+      "loss": 2.964167301719253,
+      "loss_size": 0.6638369615617636,
+      "loss_pdi": 0.48985611103676463,
+      "loss_ee": 0.772022802282024,
+      "loss_delivery": 0.8712375029116064,
+      "loss_biodist": 0.15383520964029673,
+      "loss_toxic": 0.08995235672006635,
+      "loss_moe_lb": 1.9999999806687638
+    },
+    {
+      "loss": 2.6228387098054626,
+      "loss_size": 0.5342353371953642,
+      "loss_pdi": 0.463039289052422,
+      "loss_ee": 0.7535679332307867,
+      "loss_delivery": 0.7488482483738178,
+      "loss_biodist": 0.14714475761394244,
+      "loss_toxic": 0.051557072453834805,
+      "loss_moe_lb": 1.9999999967781272
+    },
+    {
+      "loss": 2.601656311267131,
+      "loss_size": 0.560483991042585,
+      "loss_pdi": 0.44463860022055135,
+      "loss_ee": 0.7165064908362724,
+      "loss_delivery": 0.769291910166676,
+      "loss_biodist": 0.14465474115835653,
+      "loss_toxic": 0.040929400137109634,
+      "loss_moe_lb": 1.9999999935562547
+    },
+    {
+      "loss": 2.4538051598780863,
+      "loss_size": 0.49438263657125264,
+      "loss_pdi": 0.4425109858448441,
+      "loss_ee": 0.7406848274372719,
+      "loss_delivery": 0.6696303382113173,
+      "loss_biodist": 0.15262043199225053,
+      "loss_toxic": 0.02877226605385265,
+      "loss_moe_lb": 1.999999990334382
+    }
+  ],
+  "val": [
+    {
+      "loss": 4.149928092956543,
+      "loss_size": 0.44090341343544426,
+      "loss_pdi": 0.6847947835922241,
+      "loss_ee": 1.0027188181877136,
+      "loss_delivery": 0.7304204341024161,
+      "loss_biodist": 0.8311094760894775,
+      "loss_toxic": 0.46597445011138916,
+      "loss_moe_lb": 2.0,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.800080585479736,
+      "loss_size": 0.49460249189287425,
+      "loss_pdi": 0.8068427205085754,
+      "loss_ee": 0.9725982129573822,
+      "loss_delivery": 0.6590877763926983,
+      "loss_biodist": 0.4985577404499054,
+      "loss_toxic": 0.3831261992454529,
+      "loss_moe_lb": 1.9999999523162841,
+      "acc_pdi": 0.5151515151515151,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.85
+    },
+    {
+      "loss": 3.882923984527588,
+      "loss_size": 0.48453460857272146,
+      "loss_pdi": 0.9176841259002686,
+      "loss_ee": 0.9536213561892509,
+      "loss_delivery": 0.7041484475135803,
+      "loss_biodist": 0.36803268045187,
+      "loss_toxic": 0.484267669566907,
+      "loss_moe_lb": 2.0,
+      "acc_pdi": 0.5151515151515151,
+      "acc_ee": 0.42424242424242425,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 4.024969005584717,
+      "loss_size": 0.5285034894943237,
+      "loss_pdi": 1.0518685221672057,
+      "loss_ee": 0.86401826441288,
+      "loss_delivery": 0.9207587957382202,
+      "loss_biodist": 0.18886734098196029,
+      "loss_toxic": 0.5104730214108713,
+      "loss_moe_lb": 1.9999999523162841,
+      "acc_pdi": 0.42424242424242425,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.8761045932769775,
+      "loss_size": 0.37897889912128446,
+      "loss_pdi": 1.301726371049881,
+      "loss_ee": 0.8393121339380741,
+      "loss_delivery": 0.8618740364909172,
+      "loss_biodist": 0.2399044305086136,
+      "loss_toxic": 0.32361946450546386,
+      "loss_moe_lb": 1.9999999761581422,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.42424242424242425,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.6892258644104006,
+      "loss_size": 0.44566770195961,
+      "loss_pdi": 1.0844948768615723,
+      "loss_ee": 0.7679000589996576,
+      "loss_delivery": 0.7768170319497585,
+      "loss_biodist": 0.20787838250398635,
+      "loss_toxic": 0.45133246323093773,
+      "loss_moe_lb": 2.0,
+      "acc_pdi": 0.5454545454545454,
+      "acc_ee": 0.3939393939393939,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.198654842376709,
+      "loss_size": 0.34323964416980746,
+      "loss_pdi": 1.1026686668395995,
+      "loss_ee": 0.7947472155094146,
+      "loss_delivery": 0.664758375287056,
+      "loss_biodist": 0.20282046347856522,
+      "loss_toxic": 0.1651931553366012,
+      "loss_moe_lb": 2.0,
+      "acc_pdi": 0.5151515151515151,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 2.689650774002075,
+      "loss_size": 0.28003521375358104,
+      "loss_pdi": 0.9093899607658387,
+      "loss_ee": 0.8276618905365467,
+      "loss_delivery": 0.5897062364965677,
+      "loss_biodist": 0.13434391915798188,
+      "loss_toxic": 0.033575344539713115,
+      "loss_moe_lb": 2.000000023841858,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.546375370025635,
+      "loss_size": 0.31427572667598724,
+      "loss_pdi": 1.1522727251052856,
+      "loss_ee": 0.8625796323642134,
+      "loss_delivery": 0.6758572697639466,
+      "loss_biodist": 0.17268332988023757,
+      "loss_toxic": 0.43319647138996514,
+      "loss_moe_lb": 1.9999999523162841,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.1686388492584228,
+      "loss_size": 0.3039375368505716,
+      "loss_pdi": 1.0021734952926635,
+      "loss_ee": 0.7893130667507648,
+      "loss_delivery": 0.6568111717700958,
+      "loss_biodist": 0.14740753322839736,
+      "loss_toxic": 0.3326639760110993,
+      "loss_moe_lb": 2.0,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 2.828946590423584,
+      "loss_size": 0.2690973848104477,
+      "loss_pdi": 1.0582403063774108,
+      "loss_ee": 0.7890313200652599,
+      "loss_delivery": 0.6396413296461105,
+      "loss_biodist": 0.140879487991333,
+      "loss_toxic": 0.03787025979545433,
+      "loss_moe_lb": 2.0,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.965114641189575,
+      "loss_size": 0.31102419793605807,
+      "loss_pdi": 1.0872425794601441,
+      "loss_ee": 0.7800759710371494,
+      "loss_delivery": 0.648045665025711,
+      "loss_biodist": 0.15429325103759767,
+      "loss_toxic": 0.08624297480710083,
+      "loss_moe_lb": 2.0,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.034215879440308,
+      "loss_size": 0.31441564708948133,
+      "loss_pdi": 1.1123341083526612,
+      "loss_ee": 0.7961474736221135,
+      "loss_delivery": 0.6239923983812332,
+      "loss_biodist": 0.1533489406108856,
+      "loss_toxic": 0.1322808950993931,
+      "loss_moe_lb": 2.0,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.972694444656372,
+      "loss_size": 0.30653163194656374,
+      "loss_pdi": 1.1311804890632629,
+      "loss_ee": 0.7937454930506647,
+      "loss_delivery": 0.6065988793969155,
+      "loss_biodist": 0.1408291682600975,
+      "loss_toxic": 0.09733039969723904,
+      "loss_moe_lb": 2.0,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/abl_full/s1_both/seed42/outer_fold_0/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15, 17, 18, 19, 21, 22, 23, 24, 25, 26, 27, 28, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 52, 53, 54, 55, 56, 57, 58, 61, 62, 63, 64, 65, 66, 67, 68, 71, 72, 73, 75, 76, 77, 80, 81, 85, 86, 88, 89, 90, 92, 94, 95, 96, 98, 99, 100, 101, 103, 104, 105, 106, 107, 108, 109, 111, 113, 114, 115, 116, 117, 118, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 133, 134, 135, 136, 138, 139, 141, 144, 145, 146, 147, 148, 149, 150, 151, 153, 154, 155, 156, 157, 158, 159, 160, 163, 164, 167, 168, 169, 170, 171, 172, 175, 176, 178, 179, 180, 182, 183, 184, 185, 188, 189, 190, 191, 192, 193, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 206, 208, 210, 211, 212, 213, 216, 217, 218, 219, 220, 222, 224, 225, 227, 228, 229, 230, 231, 232, 233, 234, 236, 237, 239, 241, 242, 244, 245, 247, 248, 249, 250, 252, 253, 255, 256, 257, 258, 259, 260, 261, 262, 265, 267, 268, 269, 270, 271, 272, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 285, 286, 287, 288, 291, 292, 294, 296, 297, 298, 299, 300, 301, 303, 304, 306, 307, 308, 309, 310, 311, 312, 313, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 332, 333, 335, 336, 337, 338, 339, 340, 342, 343, 344, 345, 346, 347, 348, 349, 351, 352, 354, 355, 356, 357, 358, 359, 360, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 374, 375, 377, 379, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389, 390, 391, 393, 395, 396, 397, 398, 399, 400, 401, 403, 404, 405, 406, 407, 408, 409, 411, 412, 413, 414, 415, 417, 418, 419], "outer_test_idx": [7, 16, 20, 29, 51, 59, 60, 69, 70, 74, 78, 79, 82, 83, 84, 87, 91, 93, 97, 102, 110, 112, 119, 120, 121, 132, 137, 140, 142, 143, 152, 161, 162, 165, 166, 173, 174, 177, 181, 186, 187, 194, 205, 207, 209, 214, 215, 221, 223, 226, 235, 238, 240, 243, 246, 251, 254, 263, 264, 266, 273, 284, 289, 290, 293, 295, 302, 305, 314, 315, 331, 334, 341, 350, 353, 361, 373, 376, 378, 392, 394, 402, 410, 416]}

models/abl_full/s1_both/seed42/outer_fold_0/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 83,
+    "mse": 1.65548933352636,
+    "rmse": 1.2866582038468337,
+    "mae": 0.5746945666081934,
+    "r2": -0.013200221826040748
+  },
+  "delivery": {
+    "n_samples": 58,
+    "mse": 0.5966502866071148,
+    "rmse": 0.7724314122348436,
+    "mae": 0.600902550468414,
+    "r2": 0.24162148150681173
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7619047619047619,
+    "precision": 0.7006414825374199,
+    "recall": 0.7006414825374199,
+    "f1": 0.7006414825374199
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.7261904761904762,
+    "precision": 0.717012617012617,
+    "recall": 0.7409523809523809,
+    "f1": 0.7039699318769087
+  },
+  "toxic": {
+    "n_samples": 58,
+    "accuracy": 0.9827586206896551,
+    "precision": 0.9912280701754386,
+    "recall": 0.75,
+    "f1": 0.8289085545722714
+  },
+  "biodist": {
+    "n_samples": 58,
+    "kl_divergence": 0.16024890157851573,
+    "js_divergence": 0.032467622854324446
+  }
+}

models/abl_full/s1_both/seed42/outer_fold_1/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.24482591245144014,
+  "lr": 0.00042458936788614755,
+  "weight_decay": 0.00014902327669828248,
+  "backbone_lr_ratio": 0.24808069123277332,
+  "moe_n_experts": 2,
+  "moe_top_k": 1,
+  "moe_expert_hidden_mult": 1,
+  "llm_lora_r": 8,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_both/seed42/outer_fold_1/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 7}

models/abl_full/s1_both/seed42/outer_fold_1/history.json

@@ -0,0 +1,167 @@
+{
+  "train": [
+    {
+      "loss": 5.600405757491653,
+      "loss_size": 1.2000684311261047,
+      "loss_pdi": 0.6828988838840175,
+      "loss_ee": 1.0604768153783437,
+      "loss_delivery": 0.8591992943673521,
+      "loss_biodist": 1.1980015538834237,
+      "loss_toxic": 0.6019431312341947,
+      "loss_moe_lb": 1.0064576300414834
+    },
+    {
+      "loss": 4.742360991400641,
+      "loss_size": 1.0796409442618087,
+      "loss_pdi": 0.6460769281194017,
+      "loss_ee": 0.9675167895652153,
+      "loss_delivery": 0.8959972065546222,
+      "loss_biodist": 0.7482811027281994,
+      "loss_toxic": 0.41472573658904516,
+      "loss_moe_lb": 1.0050648241429716
+    },
+    {
+      "loss": 4.254786182094264,
+      "loss_size": 1.1664461914349247,
+      "loss_pdi": 0.6028806623574849,
+      "loss_ee": 0.9201007727030162,
+      "loss_delivery": 0.8821316807455307,
+      "loss_biodist": 0.4595999757985811,
+      "loss_toxic": 0.23364325109365824,
+      "loss_moe_lb": 1.0039539433814384
+    },
+    {
+      "loss": 3.7494558546994186,
+      "loss_size": 1.1098671801186897,
+      "loss_pdi": 0.5598839118673995,
+      "loss_ee": 0.8385547004841469,
+      "loss_delivery": 0.7759733399426615,
+      "loss_biodist": 0.3402977386439169,
+      "loss_toxic": 0.13544689199408969,
+      "loss_moe_lb": 1.0069729911314476
+    },
+    {
+      "loss": 3.5292408047495663,
+      "loss_size": 1.080890902997674,
+      "loss_pdi": 0.5336942487471813,
+      "loss_ee": 0.8007064501981478,
+      "loss_delivery": 0.7067823821326366,
+      "loss_biodist": 0.28604381209289703,
+      "loss_toxic": 0.13082139748714966,
+      "loss_moe_lb": 1.0036791095862518
+    },
+    {
+      "loss": 3.280233860015869,
+      "loss_size": 1.0479851067066193,
+      "loss_pdi": 0.4890121052394042,
+      "loss_ee": 0.7543886859674711,
+      "loss_delivery": 0.6291292678262737,
+      "loss_biodist": 0.26939965200585286,
+      "loss_toxic": 0.09769973801952359,
+      "loss_moe_lb": 1.0032514800896515
+    },
+    {
+      "loss": 3.3657465367703825,
+      "loss_size": 1.0007865732988797,
+      "loss_pdi": 0.5091070810685286,
+      "loss_ee": 0.733278718349096,
+      "loss_delivery": 0.769810107108709,
+      "loss_biodist": 0.263388701991455,
+      "loss_toxic": 0.09867539737260321,
+      "loss_moe_lb": 1.004341889072109
+    }
+  ],
+  "val": [
+    {
+      "loss": 3.6868560791015623,
+      "loss_size": 0.5512764737010002,
+      "loss_pdi": 0.6596759557723999,
+      "loss_ee": 0.9739189267158508,
+      "loss_delivery": 0.20791610032320024,
+      "loss_biodist": 0.8941343069076538,
+      "loss_toxic": 0.40256049036979674,
+      "loss_moe_lb": 1.0218846678733826,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.7243490695953367,
+      "loss_size": 0.5449142195284367,
+      "loss_pdi": 0.6435503244400025,
+      "loss_ee": 0.863824188709259,
+      "loss_delivery": 0.2041028030216694,
+      "loss_biodist": 0.3623921275138855,
+      "loss_toxic": 0.09837098270654679,
+      "loss_moe_lb": 1.0318222165107727,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.696969696969697,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.5498147010803223,
+      "loss_size": 0.5333899289369584,
+      "loss_pdi": 0.6478025674819946,
+      "loss_ee": 0.7340908885002136,
+      "loss_delivery": 0.27008581161499023,
+      "loss_biodist": 0.34742590188980105,
+      "loss_toxic": 0.013959044264629484,
+      "loss_moe_lb": 1.022853970527649,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.701027178764343,
+      "loss_size": 0.5317124128341675,
+      "loss_pdi": 0.6931547701358796,
+      "loss_ee": 0.6460031867027283,
+      "loss_delivery": 0.5111799478530884,
+      "loss_biodist": 0.318424254655838,
+      "loss_toxic": 0.003251659439411014,
+      "loss_moe_lb": 1.003060233592987,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.322146010398865,
+      "loss_size": 0.4438261866569519,
+      "loss_pdi": 0.6595357298851013,
+      "loss_ee": 0.6595715552568435,
+      "loss_delivery": 0.2990713447332382,
+      "loss_biodist": 0.25748171210289,
+      "loss_toxic": 0.0028742717579007147,
+      "loss_moe_lb": 1.0017499089241029,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.2509952545166017,
+      "loss_size": 0.43093053102493284,
+      "loss_pdi": 0.6347038090229035,
+      "loss_ee": 0.6181163638830185,
+      "loss_delivery": 0.31103180050849916,
+      "loss_biodist": 0.2523608267307281,
+      "loss_toxic": 0.0028693197993561624,
+      "loss_moe_lb": 1.0072155237197875,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.3216882705688477,
+      "loss_size": 0.43943057060241697,
+      "loss_pdi": 0.6587613582611084,
+      "loss_ee": 0.622439831495285,
+      "loss_delivery": 0.33744184076786043,
+      "loss_biodist": 0.26239258944988253,
+      "loss_toxic": 0.002155938057694584,
+      "loss_moe_lb": 1.0101934671401978,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/abl_full/s1_both/seed42/outer_fold_1/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 3, 5, 6, 7, 10, 11, 12, 13, 15, 16, 18, 19, 20, 21, 24, 26, 27, 29, 30, 31, 32, 33, 35, 36, 38, 39, 40, 42, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 65, 66, 67, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 97, 98, 99, 100, 101, 102, 103, 105, 106, 107, 108, 109, 110, 112, 114, 115, 118, 119, 120, 121, 123, 124, 125, 126, 128, 129, 130, 131, 132, 136, 137, 140, 141, 142, 143, 144, 145, 146, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 161, 162, 164, 165, 166, 167, 168, 169, 171, 172, 173, 174, 175, 176, 177, 178, 179, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 194, 195, 196, 197, 198, 199, 201, 202, 204, 205, 206, 207, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 225, 226, 228, 229, 231, 232, 233, 234, 235, 236, 237, 238, 240, 242, 243, 244, 246, 247, 248, 250, 251, 252, 254, 255, 256, 258, 259, 260, 261, 262, 263, 264, 266, 268, 269, 270, 271, 273, 274, 275, 276, 277, 278, 279, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 295, 296, 297, 300, 301, 302, 303, 304, 305, 306, 307, 308, 310, 311, 313, 314, 315, 316, 317, 318, 319, 321, 322, 323, 324, 325, 327, 328, 329, 330, 331, 333, 334, 335, 337, 338, 339, 340, 341, 342, 343, 344, 345, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 361, 362, 364, 365, 366, 367, 368, 369, 373, 374, 375, 376, 378, 379, 380, 381, 382, 384, 387, 388, 389, 391, 392, 393, 394, 396, 397, 398, 399, 402, 403, 404, 405, 406, 408, 409, 410, 411, 412, 413, 414, 416, 417, 419], "outer_test_idx": [1, 2, 4, 8, 9, 14, 17, 22, 23, 25, 28, 34, 37, 41, 43, 45, 63, 64, 68, 85, 96, 104, 111, 113, 116, 117, 122, 127, 133, 134, 135, 138, 139, 147, 160, 163, 170, 180, 191, 192, 193, 200, 203, 208, 224, 227, 230, 239, 241, 245, 249, 253, 257, 265, 267, 272, 280, 294, 298, 299, 309, 312, 320, 326, 332, 336, 346, 347, 360, 363, 370, 371, 372, 377, 383, 385, 386, 390, 395, 400, 401, 407, 415, 418]}

models/abl_full/s1_both/seed42/outer_fold_1/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.3445300972529715,
+    "rmse": 0.5869668621421242,
+    "mae": 0.4555579432108927,
+    "r2": -0.05207242874852769
+  },
+  "delivery": {
+    "n_samples": 61,
+    "mse": 1.0775295165871506,
+    "rmse": 1.0380411921437176,
+    "mae": 0.6340284446773471,
+    "r2": 0.18438783389531754
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.6904761904761905,
+    "precision": 0.5972915181753385,
+    "recall": 0.6031746031746031,
+    "f1": 0.5997067448680352
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.5476190476190477,
+    "precision": 0.47222222222222215,
+    "recall": 0.4779744065458351,
+    "f1": 0.46562819203268635
+  },
+  "toxic": {
+    "n_samples": 61,
+    "accuracy": 0.9508196721311475,
+    "precision": 0.75,
+    "recall": 0.9741379310344828,
+    "f1": 0.8200589970501475
+  },
+  "biodist": {
+    "n_samples": 61,
+    "kl_divergence": 0.24783434314436234,
+    "js_divergence": 0.0666961611545551
+  }
+}

models/abl_full/s1_both/seed42/outer_fold_2/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.330147494187667,
+  "lr": 0.0007391626870428583,
+  "weight_decay": 0.08248214185502709,
+  "backbone_lr_ratio": 0.3321485864522685,
+  "moe_n_experts": 4,
+  "moe_top_k": 1,
+  "moe_expert_hidden_mult": 1,
+  "llm_lora_r": 8,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_both/seed42/outer_fold_2/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 10}

models/abl_full/s1_both/seed42/outer_fold_2/history.json

@@ -0,0 +1,236 @@
+{
+  "train": [
+    {
+      "loss": 5.152713775634766,
+      "loss_size": 0.8612007083925041,
+      "loss_pdi": 0.6876318486961158,
+      "loss_ee": 1.0617679599169139,
+      "loss_delivery": 0.9683949065369528,
+      "loss_biodist": 1.0637048104324855,
+      "loss_toxic": 0.5132733631778408,
+      "loss_moe_lb": 1.0177504887451996
+    },
+    {
+      "loss": 4.2314953997328475,
+      "loss_size": 0.8170983128853746,
+      "loss_pdi": 0.6610435166874448,
+      "loss_ee": 0.9547245228612745,
+      "loss_delivery": 0.8227231446150187,
+      "loss_biodist": 0.6459103758270676,
+      "loss_toxic": 0.3464304328568884,
+      "loss_moe_lb": 1.0108283081570186
+    },
+    {
+      "loss": 3.750632695249609,
+      "loss_size": 0.7929803255844761,
+      "loss_pdi": 0.593320668549151,
+      "loss_ee": 0.9024315122011546,
+      "loss_delivery": 0.8082113729135411,
+      "loss_biodist": 0.4675642582210335,
+      "loss_toxic": 0.20891104418331305,
+      "loss_moe_lb": 1.0160694025658272
+    },
+    {
+      "loss": 3.809421558637877,
+      "loss_size": 0.8187293212961506,
+      "loss_pdi": 0.6209371025497848,
+      "loss_ee": 0.8522474620793317,
+      "loss_delivery": 0.971913101705345,
+      "loss_biodist": 0.37334723005423676,
+      "loss_toxic": 0.2086577709588046,
+      "loss_moe_lb": 1.0108439245739498
+    },
+    {
+      "loss": 3.5215162812052547,
+      "loss_size": 0.7850020717124682,
+      "loss_pdi": 0.5633785048046628,
+      "loss_ee": 0.8283515363126188,
+      "loss_delivery": 0.8907835570138853,
+      "loss_biodist": 0.2965107239984177,
+      "loss_toxic": 0.19142280418954388,
+      "loss_moe_lb": 1.0072002684747852
+    },
+    {
+      "loss": 3.198103382780745,
+      "loss_size": 0.738623615253616,
+      "loss_pdi": 0.5287548682174167,
+      "loss_ee": 0.8442864297209559,
+      "loss_delivery": 0.7745753369621329,
+      "loss_biodist": 0.23759525510910395,
+      "loss_toxic": 0.1121889049426426,
+      "loss_moe_lb": 1.0110071913616077
+    },
+    {
+      "loss": 3.0494969567737065,
+      "loss_size": 0.7034210447926779,
+      "loss_pdi": 0.5016100636205157,
+      "loss_ee": 0.7678506028007817,
+      "loss_delivery": 0.653912561366687,
+      "loss_biodist": 0.23800816628578547,
+      "loss_toxic": 0.2103397584255986,
+      "loss_moe_lb": 1.0078407703219234
+    },
+    {
+      "loss": 2.7809914350509644,
+      "loss_size": 0.6419068256745467,
+      "loss_pdi": 0.5048889320444416,
+      "loss_ee": 0.7783925090287183,
+      "loss_delivery": 0.5467462676602441,
+      "loss_biodist": 0.23822186907400955,
+      "loss_toxic": 0.10388507783834545,
+      "loss_moe_lb": 1.0072059244722933
+    },
+    {
+      "loss": 2.8093391431344523,
+      "loss_size": 0.6376639669005936,
+      "loss_pdi": 0.4999076144115345,
+      "loss_ee": 0.7432701736688614,
+      "loss_delivery": 0.6609656624495983,
+      "loss_biodist": 0.20728998631238937,
+      "loss_toxic": 0.09592864762148443,
+      "loss_moe_lb": 1.0073097573744285
+    },
+    {
+      "loss": 2.763327427812525,
+      "loss_size": 0.6375405836749721,
+      "loss_pdi": 0.4994223685683431,
+      "loss_ee": 0.746296464591413,
+      "loss_delivery": 0.6282920932044854,
+      "loss_biodist": 0.1821774451716526,
+      "loss_toxic": 0.10435734552526343,
+      "loss_moe_lb": 1.0079195080576717
+    }
+  ],
+  "val": [
+    {
+      "loss": 7.724205446243286,
+      "loss_size": 3.4674210086464883,
+      "loss_pdi": 0.6508873462677002,
+      "loss_ee": 1.0635061860084534,
+      "loss_delivery": 1.9029634714126586,
+      "loss_biodist": 0.5229972720146179,
+      "loss_toxic": 0.12300186306238174,
+      "loss_moe_lb": 1.0304514646530152,
+      "acc_pdi": 0.48484848484848486,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 7.370360422134399,
+      "loss_size": 3.467023845016956,
+      "loss_pdi": 0.6449822664260865,
+      "loss_ee": 0.9673564076423645,
+      "loss_delivery": 1.9023066341876984,
+      "loss_biodist": 0.3448571443557739,
+      "loss_toxic": 0.0791855052113533,
+      "loss_moe_lb": 1.038755679130554,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 7.128422689437866,
+      "loss_size": 3.4471570923924446,
+      "loss_pdi": 0.6080391645431519,
+      "loss_ee": 1.0527662873268127,
+      "loss_delivery": 1.7938429474830628,
+      "loss_biodist": 0.23303972482681273,
+      "loss_toxic": 0.061180700571276246,
+      "loss_moe_lb": 1.0203991174697875,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.809806776046753,
+      "loss_size": 3.255460247397423,
+      "loss_pdi": 0.6362621784210205,
+      "loss_ee": 0.9918401598930359,
+      "loss_delivery": 1.7496135592460633,
+      "loss_biodist": 0.19902592599391938,
+      "loss_toxic": 0.05351642244495451,
+      "loss_moe_lb": 1.0118896961212158,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.5712010860443115,
+      "loss_size": 3.3117038667201997,
+      "loss_pdi": 0.6905999302864074,
+      "loss_ee": 0.9391629934310913,
+      "loss_delivery": 1.5461100190877914,
+      "loss_biodist": 0.12436708956956863,
+      "loss_toxic": 0.03854238981730305,
+      "loss_moe_lb": 1.0100578308105468,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.109503126144409,
+      "loss_size": 3.024673792719841,
+      "loss_pdi": 0.6201795876026154,
+      "loss_ee": 0.92969890832901,
+      "loss_delivery": 1.4628742456436157,
+      "loss_biodist": 0.12257246896624566,
+      "loss_toxic": 0.028115913399960845,
+      "loss_moe_lb": 1.0112135887145997,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.1416528701782225,
+      "loss_size": 2.9167025655508043,
+      "loss_pdi": 0.6443653285503388,
+      "loss_ee": 1.0066563487052917,
+      "loss_delivery": 1.4425293743610381,
+      "loss_biodist": 0.1540301412343979,
+      "loss_toxic": 0.06369881336577236,
+      "loss_moe_lb": 1.010964584350586,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.189666986465454,
+      "loss_size": 2.9956638380885123,
+      "loss_pdi": 0.6736391186714172,
+      "loss_ee": 1.046317207813263,
+      "loss_delivery": 1.395979829132557,
+      "loss_biodist": 0.12292753309011459,
+      "loss_toxic": 0.04645174202742055,
+      "loss_moe_lb": 1.015622067451477,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.373043394088745,
+      "loss_size": 3.116128545999527,
+      "loss_pdi": 0.6364896595478058,
+      "loss_ee": 1.0502041697502136,
+      "loss_delivery": 1.45866639316082,
+      "loss_biodist": 0.12851292118430138,
+      "loss_toxic": 0.07333654393441975,
+      "loss_moe_lb": 1.0123448848724366,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.115008068084717,
+      "loss_size": 2.921285180747509,
+      "loss_pdi": 0.6241236627101898,
+      "loss_ee": 1.0314197301864625,
+      "loss_delivery": 1.4467920392751694,
+      "loss_biodist": 0.12439441233873368,
+      "loss_toxic": 0.05612772112363018,
+      "loss_moe_lb": 1.014252781867981,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.95
+    }
+  ]
+}

models/abl_full/s1_both/seed42/outer_fold_2/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 11, 13, 14, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 32, 33, 34, 36, 37, 38, 39, 40, 41, 43, 44, 45, 47, 48, 49, 50, 51, 55, 56, 57, 58, 59, 60, 61, 63, 64, 65, 68, 69, 70, 72, 74, 76, 78, 79, 80, 82, 83, 84, 85, 86, 87, 88, 91, 93, 94, 95, 96, 97, 100, 101, 102, 103, 104, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 142, 143, 145, 147, 148, 149, 150, 151, 152, 153, 156, 157, 158, 159, 160, 161, 162, 163, 165, 166, 167, 168, 170, 171, 172, 173, 174, 175, 177, 179, 180, 181, 183, 184, 185, 186, 187, 188, 191, 192, 193, 194, 196, 197, 199, 200, 202, 203, 204, 205, 207, 208, 209, 210, 211, 212, 213, 214, 215, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 232, 233, 235, 236, 237, 238, 239, 240, 241, 242, 243, 245, 246, 249, 250, 251, 252, 253, 254, 255, 256, 257, 259, 260, 261, 263, 264, 265, 266, 267, 268, 269, 271, 272, 273, 274, 275, 276, 280, 281, 282, 283, 284, 285, 287, 289, 290, 293, 294, 295, 296, 297, 298, 299, 300, 302, 303, 305, 307, 308, 309, 310, 311, 312, 313, 314, 315, 317, 318, 320, 321, 322, 323, 324, 326, 327, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 358, 360, 361, 363, 364, 365, 366, 368, 369, 370, 371, 372, 373, 376, 377, 378, 379, 380, 381, 382, 383, 385, 386, 387, 388, 389, 390, 392, 394, 395, 398, 399, 400, 401, 402, 404, 405, 407, 408, 409, 410, 412, 413, 414, 415, 416, 417, 418], "outer_test_idx": [10, 12, 15, 31, 35, 42, 46, 52, 53, 54, 62, 66, 67, 71, 73, 75, 77, 81, 89, 90, 92, 98, 99, 105, 123, 124, 141, 144, 146, 154, 155, 164, 169, 176, 178, 182, 189, 190, 195, 198, 201, 206, 216, 217, 218, 219, 231, 234, 244, 247, 248, 258, 262, 270, 277, 278, 279, 286, 288, 291, 292, 301, 304, 306, 316, 319, 325, 328, 343, 357, 359, 362, 367, 374, 375, 384, 391, 393, 396, 397, 403, 406, 411, 419]}

models/abl_full/s1_both/seed42/outer_fold_2/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.393067832783495,
+    "rmse": 0.6269512204179006,
+    "mae": 0.48754836206457447,
+    "r2": 0.22369805318261293
+  },
+  "delivery": {
+    "n_samples": 60,
+    "mse": 0.6293676153394357,
+    "rmse": 0.7933269284093637,
+    "mae": 0.5411118950306748,
+    "r2": 0.16220261953292037
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.6547619047619048,
+    "precision": 0.6282346175963197,
+    "recall": 0.6634897360703812,
+    "f1": 0.6212097651998134
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6547619047619048,
+    "precision": 0.6078622482131254,
+    "recall": 0.6720354808590102,
+    "f1": 0.6212949538882847
+  },
+  "toxic": {
+    "n_samples": 61,
+    "accuracy": 0.9344262295081968,
+    "precision": 0.7142857142857143,
+    "recall": 0.9655172413793103,
+    "f1": 0.7821428571428571
+  },
+  "biodist": {
+    "n_samples": 60,
+    "kl_divergence": 0.18339582654236564,
+    "js_divergence": 0.04379028402728257
+  }
+}

models/abl_full/s1_both/seed42/outer_fold_3/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.2277771320023231,
+  "lr": 0.0009696339767409517,
+  "weight_decay": 0.0003933823864828228,
+  "backbone_lr_ratio": 0.4073135560471746,
+  "moe_n_experts": 4,
+  "moe_top_k": 2,
+  "moe_expert_hidden_mult": 2,
+  "llm_lora_r": 8,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_both/seed42/outer_fold_3/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 12}

models/abl_full/s1_both/seed42/outer_fold_3/history.json

@@ -0,0 +1,282 @@
+{
+  "train": [
+    {
+      "loss": 4.569545784512082,
+      "loss_size": 0.51101915538311,
+      "loss_pdi": 0.6594580573004645,
+      "loss_ee": 1.0284118587906297,
+      "loss_delivery": 0.942825937895356,
+      "loss_biodist": 0.9342355389852781,
+      "loss_toxic": 0.488486493761475,
+      "loss_moe_lb": 2.134706339320621
+    },
+    {
+      "loss": 3.37780880283665,
+      "loss_size": 0.47524130847808477,
+      "loss_pdi": 0.6199308311617052,
+      "loss_ee": 0.9292109125369304,
+      "loss_delivery": 0.8265461688106125,
+      "loss_biodist": 0.3766151179332991,
+      "loss_toxic": 0.15161741836424414,
+      "loss_moe_lb": 2.238398210422413
+    },
+    {
+      "loss": 3.0872007962819694,
+      "loss_size": 0.4445550103445311,
+      "loss_pdi": 0.5805639914564185,
+      "loss_ee": 0.8806364053004497,
+      "loss_delivery": 0.802079865960656,
+      "loss_biodist": 0.28683041116675817,
+      "loss_toxic": 0.1158963359871327,
+      "loss_moe_lb": 2.271572357899434
+    },
+    {
+      "loss": 3.1701637506484985,
+      "loss_size": 0.5009561811347265,
+      "loss_pdi": 0.5923459135197304,
+      "loss_ee": 0.8842877270402135,
+      "loss_delivery": 0.8306604349935377,
+      "loss_biodist": 0.23747707822838346,
+      "loss_toxic": 0.16886883977221678,
+      "loss_moe_lb": 2.1245123921213924
+    },
+    {
+      "loss": 2.801817297935486,
+      "loss_size": 0.4014007754221156,
+      "loss_pdi": 0.5715829125932745,
+      "loss_ee": 0.8310141136517396,
+      "loss_delivery": 0.7431407951422639,
+      "loss_biodist": 0.2067834861777924,
+      "loss_toxic": 0.10440129139051596,
+      "loss_moe_lb": 2.195342672837747
+    },
+    {
+      "loss": 2.632911917325613,
+      "loss_size": 0.40186479264819946,
+      "loss_pdi": 0.5458749492426176,
+      "loss_ee": 0.7789210926842045,
+      "loss_delivery": 0.727295874995557,
+      "loss_biodist": 0.20429579923684532,
+      "loss_toxic": 0.044171577916622976,
+      "loss_moe_lb": 2.305464100193333
+    },
+    {
+      "loss": 2.601001562298955,
+      "loss_size": 0.40228870470781586,
+      "loss_pdi": 0.4918990400997368,
+      "loss_ee": 0.7510785339651881,
+      "loss_delivery": 0.6971234280515362,
+      "loss_biodist": 0.22326555626617894,
+      "loss_toxic": 0.1021743980983612,
+      "loss_moe_lb": 2.2410518672015214
+    },
+    {
+      "loss": 2.460212440104098,
+      "loss_size": 0.3730814474055896,
+      "loss_pdi": 0.5026364958769566,
+      "loss_ee": 0.7198781137530869,
+      "loss_delivery": 0.676776959686666,
+      "loss_biodist": 0.20135338423219887,
+      "loss_toxic": 0.06579114676148763,
+      "loss_moe_lb": 2.2458660634788306
+    },
+    {
+      "loss": 2.144416024556031,
+      "loss_size": 0.29755874320461945,
+      "loss_pdi": 0.483699955448911,
+      "loss_ee": 0.6816173324713836,
+      "loss_delivery": 0.5690391652185369,
+      "loss_biodist": 0.14466739153942546,
+      "loss_toxic": 0.05192437559595124,
+      "loss_moe_lb": 2.250629737570479
+    },
+    {
+      "loss": 2.1561906949893848,
+      "loss_size": 0.3470781880456048,
+      "loss_pdi": 0.48051307974634944,
+      "loss_ee": 0.6305385231971741,
+      "loss_delivery": 0.5764511368564658,
+      "loss_biodist": 0.16402359800161542,
+      "loss_toxic": 0.04081225873921318,
+      "loss_moe_lb": 2.2769794818517326
+    },
+    {
+      "loss": 2.147863330067815,
+      "loss_size": 0.3142247807133842,
+      "loss_pdi": 0.45880394932386037,
+      "loss_ee": 0.6344889354061436,
+      "loss_delivery": 0.5934505949954729,
+      "loss_biodist": 0.1279825611493072,
+      "loss_toxic": 0.09801903067135434,
+      "loss_moe_lb": 2.289356534545486
+    },
+    {
+      "loss": 2.158412657879494,
+      "loss_size": 0.28873550851602814,
+      "loss_pdi": 0.4422310614908064,
+      "loss_ee": 0.6521380088619284,
+      "loss_delivery": 0.646223750871581,
+      "loss_biodist": 0.18561462544508883,
+      "loss_toxic": 0.0354608225309651,
+      "loss_moe_lb": 2.2619870224514522
+    }
+  ],
+  "val": [
+    {
+      "loss": 5.1847528457641605,
+      "loss_size": 2.5386213794350625,
+      "loss_pdi": 0.5081074416637421,
+      "loss_ee": 0.8049288988113403,
+      "loss_delivery": 0.8281367182731628,
+      "loss_biodist": 0.3098487854003906,
+      "loss_toxic": 0.1897735118865967,
+      "loss_moe_lb": 2.2047547817230226,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.109101510047912,
+      "loss_size": 2.5453438371419908,
+      "loss_pdi": 0.5378552436828613,
+      "loss_ee": 0.7293476760387421,
+      "loss_delivery": 0.7961073100566864,
+      "loss_biodist": 0.1953234553337097,
+      "loss_toxic": 0.2989704024977982,
+      "loss_moe_lb": 2.330738067626953,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.08425362110138,
+      "loss_size": 2.5647606000304224,
+      "loss_pdi": 0.5694735109806061,
+      "loss_ee": 0.7568024694919586,
+      "loss_delivery": 0.8387516289949417,
+      "loss_biodist": 0.22799606323242189,
+      "loss_toxic": 0.13167614876292646,
+      "loss_moe_lb": 2.4430124759674072,
+      "acc_pdi": 0.8484848484848485,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.366823887825012,
+      "loss_size": 2.4924990072846414,
+      "loss_pdi": 0.5948927760124206,
+      "loss_ee": 0.8248458862304687,
+      "loss_delivery": 0.9075531393289566,
+      "loss_biodist": 0.10045709162950515,
+      "loss_toxic": 0.4450825572013855,
+      "loss_moe_lb": 2.324281406402588,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.696969696969697,
+      "acc_toxic": 0.8260869565217391
+    },
+    {
+      "loss": 4.645944857597351,
+      "loss_size": 2.4076855540275575,
+      "loss_pdi": 0.5596936345100403,
+      "loss_ee": 0.7976774662733078,
+      "loss_delivery": 0.7818437159061432,
+      "loss_biodist": 0.12929453700780869,
+      "loss_toxic": 0.006101569114252925,
+      "loss_moe_lb": 2.4937999725341795,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.697916913032532,
+      "loss_size": 2.4932262122631075,
+      "loss_pdi": 0.45069290697574615,
+      "loss_ee": 0.8065684795379638,
+      "loss_delivery": 0.8043429493904114,
+      "loss_biodist": 0.10116395205259324,
+      "loss_toxic": 0.0682488625869155,
+      "loss_moe_lb": 2.503563070297241,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.9565217391304348
+    },
+    {
+      "loss": 3.907198989391327,
+      "loss_size": 1.7621939539909364,
+      "loss_pdi": 0.5401753157377243,
+      "loss_ee": 0.7660549342632293,
+      "loss_delivery": 0.8160712867975235,
+      "loss_biodist": 0.06947531178593636,
+      "loss_toxic": 0.008720425050705671,
+      "loss_moe_lb": 2.3909900188446045,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.236720025539398,
+      "loss_size": 1.9035443931818008,
+      "loss_pdi": 0.5434452652931213,
+      "loss_ee": 0.8244865417480469,
+      "loss_delivery": 0.7739835530519485,
+      "loss_biodist": 0.08717839270830155,
+      "loss_toxic": 0.14043191839009522,
+      "loss_moe_lb": 2.5653082847595217,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.250156939029694,
+      "loss_size": 2.0814326971769335,
+      "loss_pdi": 0.5545522332191467,
+      "loss_ee": 0.7970958948135376,
+      "loss_delivery": 0.7647629588842392,
+      "loss_biodist": 0.08868260681629181,
+      "loss_toxic": 0.014030593819916249,
+      "loss_moe_lb": 2.4910370349884032,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.149299383163452,
+      "loss_size": 2.0315980494022368,
+      "loss_pdi": 0.5397499829530716,
+      "loss_ee": 0.7687326192855835,
+      "loss_delivery": 0.7619759574532509,
+      "loss_biodist": 0.08105204850435258,
+      "loss_toxic": 0.014889129810035229,
+      "loss_moe_lb": 2.4313738346099854,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.155077862739563,
+      "loss_size": 2.040539914369583,
+      "loss_pdi": 0.5315758883953094,
+      "loss_ee": 0.7787100493907928,
+      "loss_delivery": 0.7530814379453659,
+      "loss_biodist": 0.08297589719295502,
+      "loss_toxic": 0.016451931186020376,
+      "loss_moe_lb": 2.41774845123291,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.18879998922348,
+      "loss_size": 2.0906561225652696,
+      "loss_pdi": 0.5260936617851257,
+      "loss_ee": 0.7939959585666656,
+      "loss_delivery": 0.7198820084333419,
+      "loss_biodist": 0.08342803865671158,
+      "loss_toxic": 0.019746014662086964,
+      "loss_moe_lb": 2.5341875076293947,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/abl_full/s1_both/seed42/outer_fold_3/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 1, 2, 4, 5, 7, 8, 9, 10, 11, 12, 14, 15, 16, 17, 20, 21, 22, 23, 24, 25, 27, 28, 29, 30, 31, 32, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 45, 46, 47, 51, 52, 53, 54, 55, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 73, 74, 75, 76, 77, 78, 79, 81, 82, 83, 84, 85, 86, 87, 89, 90, 91, 92, 93, 95, 96, 97, 98, 99, 100, 102, 104, 105, 106, 109, 110, 111, 112, 113, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 127, 129, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 146, 147, 148, 150, 151, 152, 153, 154, 155, 156, 158, 159, 160, 161, 162, 163, 164, 165, 166, 168, 169, 170, 172, 173, 174, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 189, 190, 191, 192, 193, 194, 195, 198, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 212, 213, 214, 215, 216, 217, 218, 219, 221, 223, 224, 226, 227, 228, 230, 231, 233, 234, 235, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 251, 252, 253, 254, 255, 256, 257, 258, 259, 262, 263, 264, 265, 266, 267, 268, 269, 270, 272, 273, 274, 275, 276, 277, 278, 279, 280, 282, 284, 286, 288, 289, 290, 291, 292, 293, 294, 295, 297, 298, 299, 301, 302, 304, 305, 306, 307, 308, 309, 310, 312, 314, 315, 316, 317, 319, 320, 321, 325, 326, 327, 328, 331, 332, 334, 336, 340, 341, 343, 344, 346, 347, 348, 350, 352, 353, 357, 359, 360, 361, 362, 363, 364, 366, 367, 368, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 382, 383, 384, 385, 386, 387, 389, 390, 391, 392, 393, 394, 395, 396, 397, 398, 400, 401, 402, 403, 404, 406, 407, 410, 411, 412, 414, 415, 416, 418, 419], "outer_test_idx": [3, 6, 13, 18, 19, 26, 33, 44, 48, 49, 50, 56, 57, 72, 80, 88, 94, 101, 103, 107, 108, 114, 126, 128, 130, 145, 149, 157, 167, 171, 175, 188, 196, 197, 199, 211, 220, 222, 225, 229, 232, 236, 250, 260, 261, 271, 281, 283, 285, 287, 296, 300, 303, 311, 313, 318, 322, 323, 324, 329, 330, 333, 335, 337, 338, 339, 342, 345, 349, 351, 354, 355, 356, 358, 365, 369, 381, 388, 399, 405, 408, 409, 413, 417]}

models/abl_full/s1_both/seed42/outer_fold_3/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 83,
+    "mse": 1.732527790200732,
+    "rmse": 1.316255214690803,
+    "mae": 0.5408693312566324,
+    "r2": 0.020069076700442334
+  },
+  "delivery": {
+    "n_samples": 59,
+    "mse": 0.5845268150974521,
+    "rmse": 0.7645435338144271,
+    "mae": 0.48413719012701917,
+    "r2": 0.43177048756483927
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7142857142857143,
+    "precision": 0.6445623342175066,
+    "recall": 0.6598240469208211,
+    "f1": 0.6499999999999999
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6428571428571429,
+    "precision": 0.5734711041691698,
+    "recall": 0.5891402714932127,
+    "f1": 0.5722830194879885
+  },
+  "toxic": {
+    "n_samples": 60,
+    "accuracy": 1.0,
+    "precision": 1.0,
+    "recall": 1.0,
+    "f1": 1.0
+  },
+  "biodist": {
+    "n_samples": 60,
+    "kl_divergence": 0.17380496676972265,
+    "js_divergence": 0.04216868552091467
+  }
+}

models/abl_full/s1_both/seed42/outer_fold_4/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.3336459802315832,
+  "lr": 0.0005171929496634874,
+  "weight_decay": 0.00014235198757319405,
+  "backbone_lr_ratio": 0.23125566110577128,
+  "moe_n_experts": 4,
+  "moe_top_k": 1,
+  "moe_expert_hidden_mult": 2,
+  "llm_lora_r": 8,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_both/seed42/outer_fold_4/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 15}

models/abl_full/s1_both/seed42/outer_fold_4/history.json

@@ -0,0 +1,351 @@
+{
+  "train": [
+    {
+      "loss": 5.680578869742316,
+      "loss_size": 1.2022036297901257,
+      "loss_pdi": 0.6911368241181245,
+      "loss_ee": 1.0734195048744615,
+      "loss_delivery": 0.9642648278055964,
+      "loss_biodist": 1.1189431129275142,
+      "loss_toxic": 0.6355154474039335,
+      "loss_moe_lb": 1.014740307588835
+    },
+    {
+      "loss": 4.765572709006232,
+      "loss_size": 1.1344646789334916,
+      "loss_pdi": 0.6499590873718262,
+      "loss_ee": 0.9887724277135488,
+      "loss_delivery": 0.9160404706726203,
+      "loss_biodist": 0.7015261118476456,
+      "loss_toxic": 0.3907836750149727,
+      "loss_moe_lb": 1.0150076376425254
+    },
+    {
+      "loss": 4.156623608357197,
+      "loss_size": 1.0709048078672305,
+      "loss_pdi": 0.6151477767003549,
+      "loss_ee": 0.9343209894927772,
+      "loss_delivery": 0.8577545055666486,
+      "loss_biodist": 0.45242498412325577,
+      "loss_toxic": 0.24410111336289225,
+      "loss_moe_lb": 1.0136639266400724
+    },
+    {
+      "loss": 3.9063501454688407,
+      "loss_size": 1.089860464169367,
+      "loss_pdi": 0.5839312841763368,
+      "loss_ee": 0.8959324150472074,
+      "loss_delivery": 0.8326679516482998,
+      "loss_biodist": 0.38330864543850357,
+      "loss_toxic": 0.14612759270931822,
+      "loss_moe_lb": 1.014538325167991
+    },
+    {
+      "loss": 3.6442219631091968,
+      "loss_size": 1.046444315966722,
+      "loss_pdi": 0.5602038599349357,
+      "loss_ee": 0.8457985858659487,
+      "loss_delivery": 0.7992855243988939,
+      "loss_biodist": 0.28028731778062677,
+      "loss_toxic": 0.13772701998587344,
+      "loss_moe_lb": 1.0089692811708193
+    },
+    {
+      "loss": 3.381216609800184,
+      "loss_size": 0.9339583967183087,
+      "loss_pdi": 0.5592660098462492,
+      "loss_ee": 0.7949836165518374,
+      "loss_delivery": 0.7507806774732229,
+      "loss_biodist": 0.2502414948634199,
+      "loss_toxic": 0.11885321380356578,
+      "loss_moe_lb": 1.0067246524063316
+    },
+    {
+      "loss": 3.2430329612783484,
+      "loss_size": 0.9114083891382089,
+      "loss_pdi": 0.5375848372240324,
+      "loss_ee": 0.7590855496960718,
+      "loss_delivery": 0.7241743039843198,
+      "loss_biodist": 0.26210401146798523,
+      "loss_toxic": 0.07838027146519036,
+      "loss_moe_lb": 1.0080170518643148
+    },
+    {
+      "loss": 3.1356499968348324,
+      "loss_size": 0.9067966339958681,
+      "loss_pdi": 0.5147047103256792,
+      "loss_ee": 0.7055981916350287,
+      "loss_delivery": 0.6798053444237322,
+      "loss_biodist": 0.21082492353948387,
+      "loss_toxic": 0.14361976250074804,
+      "loss_moe_lb": 1.009807647885503
+    },
+    {
+      "loss": 3.023964817459519,
+      "loss_size": 0.8821804201280748,
+      "loss_pdi": 0.5139715494336309,
+      "loss_ee": 0.70211508789578,
+      "loss_delivery": 0.6402809840400476,
+      "loss_biodist": 0.23446231858955846,
+      "loss_toxic": 0.08247374256992575,
+      "loss_moe_lb": 1.008427167260969
+    },
+    {
+      "loss": 2.856579084654112,
+      "loss_size": 0.7948662580267803,
+      "loss_pdi": 0.4746360078051284,
+      "loss_ee": 0.6806415633575337,
+      "loss_delivery": 0.6573987685748048,
+      "loss_biodist": 0.18960444025091222,
+      "loss_toxic": 0.09040604041247452,
+      "loss_moe_lb": 1.0075585326632939
+    },
+    {
+      "loss": 2.8019687742800325,
+      "loss_size": 0.8278479398907842,
+      "loss_pdi": 0.460802086301752,
+      "loss_ee": 0.6500890544942908,
+      "loss_delivery": 0.6351628176666595,
+      "loss_biodist": 0.1869376974331366,
+      "loss_toxic": 0.07301888717149661,
+      "loss_moe_lb": 1.0091518373102755
+    },
+    {
+      "loss": 2.7244543255986393,
+      "loss_size": 0.7427939452432297,
+      "loss_pdi": 0.4803756282136247,
+      "loss_ee": 0.6142273867452467,
+      "loss_delivery": 0.7031263080803124,
+      "loss_biodist": 0.17056602118788539,
+      "loss_toxic": 0.04867922334233299,
+      "loss_moe_lb": 1.0074441594046515
+    },
+    {
+      "loss": 2.6336539951530664,
+      "loss_size": 0.7012463568432911,
+      "loss_pdi": 0.48251094246232834,
+      "loss_ee": 0.6156406793239955,
+      "loss_delivery": 0.5994453307341885,
+      "loss_biodist": 0.18510015973368207,
+      "loss_toxic": 0.0808833715461849,
+      "loss_moe_lb": 1.0071548481245298
+    },
+    {
+      "loss": 2.6584195768510974,
+      "loss_size": 0.7399039272520993,
+      "loss_pdi": 0.44409773116176193,
+      "loss_ee": 0.6209098680599315,
+      "loss_delivery": 0.6400474600896642,
+      "loss_biodist": 0.17861945927143097,
+      "loss_toxic": 0.06789067981141694,
+      "loss_moe_lb": 1.0057777069710396
+    },
+    {
+      "loss": 2.552385117556598,
+      "loss_size": 0.7010315557992136,
+      "loss_pdi": 0.45412879940625783,
+      "loss_ee": 0.6076305202535681,
+      "loss_delivery": 0.5897340273132196,
+      "loss_biodist": 0.16656335986949303,
+      "loss_toxic": 0.06468826295128344,
+      "loss_moe_lb": 1.009352684020996
+    }
+  ],
+  "val": [
+    {
+      "loss": 4.351338243484497,
+      "loss_size": 0.5283156603574752,
+      "loss_pdi": 0.6431764841079712,
+      "loss_ee": 0.9840112090110779,
+      "loss_delivery": 0.8141101490706205,
+      "loss_biodist": 0.8709415674209595,
+      "loss_toxic": 0.5209557950496674,
+      "loss_moe_lb": 1.0300985097885131,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.414322566986084,
+      "loss_size": 0.40915088057518006,
+      "loss_pdi": 0.57725510597229,
+      "loss_ee": 0.8539698958396912,
+      "loss_delivery": 0.7813994272772107,
+      "loss_biodist": 0.4494430422782898,
+      "loss_toxic": 0.34718665331602094,
+      "loss_moe_lb": 1.01757972240448,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 3.14884934425354,
+      "loss_size": 0.36170082092285155,
+      "loss_pdi": 0.5500259160995483,
+      "loss_ee": 0.840156090259552,
+      "loss_delivery": 0.7431909930077382,
+      "loss_biodist": 0.31414745151996615,
+      "loss_toxic": 0.3457974471151829,
+      "loss_moe_lb": 1.0173659086227418,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.9835078001022337,
+      "loss_size": 0.3360332906246185,
+      "loss_pdi": 0.5084079593420029,
+      "loss_ee": 0.8543301045894622,
+      "loss_delivery": 0.81021169424057,
+      "loss_biodist": 0.20707301348447799,
+      "loss_toxic": 0.27709647375158963,
+      "loss_moe_lb": 1.0128045320510863,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.7777624845504763,
+      "loss_size": 0.27276079952716825,
+      "loss_pdi": 0.48096767291426656,
+      "loss_ee": 0.7886343747377396,
+      "loss_delivery": 0.8131667863577604,
+      "loss_biodist": 0.1389363333582878,
+      "loss_toxic": 0.2908394466154277,
+      "loss_moe_lb": 1.0190117120742799,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 3.1891578674316405,
+      "loss_size": 0.3145947799086571,
+      "loss_pdi": 0.4607068568468094,
+      "loss_ee": 0.8832492291927337,
+      "loss_delivery": 0.8955045580863953,
+      "loss_biodist": 0.17286321073770522,
+      "loss_toxic": 0.47499156957492233,
+      "loss_moe_lb": 1.0209539175033568,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8333333333333334
+    },
+    {
+      "loss": 2.962861680984497,
+      "loss_size": 0.3459408402442932,
+      "loss_pdi": 0.4526503711938858,
+      "loss_ee": 0.8540104065090418,
+      "loss_delivery": 0.934278239775449,
+      "loss_biodist": 0.1188371554017067,
+      "loss_toxic": 0.27977683106437323,
+      "loss_moe_lb": 1.0299945592880249,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.1770113468170167,
+      "loss_size": 0.34748868346214296,
+      "loss_pdi": 0.49443322122097016,
+      "loss_ee": 0.8128330364823342,
+      "loss_delivery": 0.9533780105412006,
+      "loss_biodist": 0.20197727531194687,
+      "loss_toxic": 0.3883689950918779,
+      "loss_moe_lb": 1.0118905305862427,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.1561219692230225,
+      "loss_size": 0.36988612413406374,
+      "loss_pdi": 0.45731739699840546,
+      "loss_ee": 0.8377117313444614,
+      "loss_delivery": 0.9307792767882347,
+      "loss_biodist": 0.1836371697485447,
+      "loss_toxic": 0.3993400900159031,
+      "loss_moe_lb": 1.010878896713257,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.257073348760605,
+      "loss_size": 0.3538277596235275,
+      "loss_pdi": 0.46371321603655813,
+      "loss_ee": 0.8333129465579987,
+      "loss_delivery": 0.906852000951767,
+      "loss_biodist": 0.14975143745541572,
+      "loss_toxic": 0.5621432390296832,
+      "loss_moe_lb": 1.0157330274581908,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.3074571251869203,
+      "loss_size": 0.418862184882164,
+      "loss_pdi": 0.44678124189376833,
+      "loss_ee": 0.8873284153640271,
+      "loss_delivery": 0.9675619143992662,
+      "loss_biodist": 0.14277694821357728,
+      "loss_toxic": 0.46906649945303797,
+      "loss_moe_lb": 1.0167634725570678,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.1844396233558654,
+      "loss_size": 0.3912849068641663,
+      "loss_pdi": 0.4490723855793476,
+      "loss_ee": 0.8633422922343016,
+      "loss_delivery": 0.9755209524009842,
+      "loss_biodist": 0.1369444653391838,
+      "loss_toxic": 0.39789355217944833,
+      "loss_moe_lb": 1.012496280670166,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.3202404141426087,
+      "loss_size": 0.43286535888910294,
+      "loss_pdi": 0.46595956236124036,
+      "loss_ee": 0.8726272135972977,
+      "loss_delivery": 0.892349891923368,
+      "loss_biodist": 0.17729640007019043,
+      "loss_toxic": 0.5007592872483656,
+      "loss_moe_lb": 1.0153326749801637,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.3145878672599793,
+      "loss_size": 0.4091766238212585,
+      "loss_pdi": 0.4476019456982613,
+      "loss_ee": 0.8899889543652535,
+      "loss_delivery": 1.010519946552813,
+      "loss_biodist": 0.13079079613089561,
+      "loss_toxic": 0.45538732953136785,
+      "loss_moe_lb": 1.0142157435417176,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.330686628818512,
+      "loss_size": 0.4302508354187012,
+      "loss_pdi": 0.46410428658127784,
+      "loss_ee": 0.884337130561471,
+      "loss_delivery": 1.0025697177276016,
+      "loss_biodist": 0.16711878553032875,
+      "loss_toxic": 0.4145805165870115,
+      "loss_moe_lb": 1.0122879505157472,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    }
+  ]
+}

models/abl_full/s1_both/seed42/outer_fold_4/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 23, 25, 26, 28, 29, 31, 33, 34, 35, 37, 41, 42, 43, 44, 45, 46, 48, 49, 50, 51, 52, 53, 54, 56, 57, 59, 60, 62, 63, 64, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 77, 78, 79, 80, 81, 82, 83, 84, 85, 87, 88, 89, 90, 91, 92, 93, 94, 96, 97, 98, 99, 101, 102, 103, 104, 105, 107, 108, 110, 111, 112, 113, 114, 116, 117, 119, 120, 121, 122, 123, 124, 126, 127, 128, 130, 132, 133, 134, 135, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 149, 152, 154, 155, 157, 160, 161, 162, 163, 164, 165, 166, 167, 169, 170, 171, 173, 174, 175, 176, 177, 178, 180, 181, 182, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 203, 205, 206, 207, 208, 209, 211, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 229, 230, 231, 232, 234, 235, 236, 238, 239, 240, 241, 243, 244, 245, 246, 247, 248, 249, 250, 251, 253, 254, 257, 258, 260, 261, 262, 263, 264, 265, 266, 267, 270, 271, 272, 273, 277, 278, 279, 280, 281, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 298, 299, 300, 301, 302, 303, 304, 305, 306, 309, 311, 312, 313, 314, 315, 316, 318, 319, 320, 322, 323, 324, 325, 326, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 341, 342, 343, 345, 346, 347, 349, 350, 351, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 365, 367, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 381, 383, 384, 385, 386, 388, 390, 391, 392, 393, 394, 395, 396, 397, 399, 400, 401, 402, 403, 405, 406, 407, 408, 409, 410, 411, 413, 415, 416, 417, 418, 419], "outer_test_idx": [0, 5, 11, 21, 24, 27, 30, 32, 36, 38, 39, 40, 47, 55, 58, 61, 65, 76, 86, 95, 100, 106, 109, 115, 118, 125, 129, 131, 136, 148, 150, 151, 153, 156, 158, 159, 168, 172, 179, 183, 184, 185, 202, 204, 210, 212, 213, 228, 233, 237, 242, 252, 255, 256, 259, 268, 269, 274, 275, 276, 282, 297, 307, 308, 310, 317, 321, 327, 340, 344, 348, 352, 364, 366, 368, 379, 380, 382, 387, 389, 398, 404, 412, 414]}

models/abl_full/s1_both/seed42/outer_fold_4/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.8359511180893293,
+    "rmse": 0.9143036246725315,
+    "mae": 0.5420850133467909,
+    "r2": -0.11860643443263852
+  },
+  "delivery": {
+    "n_samples": 58,
+    "mse": 1.0796427545219585,
+    "rmse": 1.03905859051449,
+    "mae": 0.649985611872298,
+    "r2": -0.04399672077007333
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7261904761904762,
+    "precision": 0.6918367346938775,
+    "recall": 0.7412023460410557,
+    "f1": 0.6946420104314841
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6785714285714286,
+    "precision": 0.6383091787439613,
+    "recall": 0.7337606837606838,
+    "f1": 0.6521164021164022
+  },
+  "toxic": {
+    "n_samples": 59,
+    "accuracy": 1.0,
+    "precision": 1.0,
+    "recall": 1.0,
+    "f1": 1.0
+  },
+  "biodist": {
+    "n_samples": 58,
+    "kl_divergence": 0.23837358595644248,
+    "js_divergence": 0.06386943316370278
+  }
+}

models/abl_full/s1_both/seed42/pip_freeze.txt

@@ -0,0 +1,153 @@
+accelerate==1.0.1
+alabaster==0.7.13
+alembic==1.14.1
+altair==5.4.1
+annotated-doc==0.0.4
+annotated-types==0.7.0
+anyio==4.5.2
+attrs==25.3.0
+babel==2.18.0
+bitsandbytes==0.45.5
+blinker==1.8.2
+Brotli @ file:///croot/brotli-split_1714483155106/work
+cachetools==5.5.2
+captum==0.7.0
+certifi @ file:///croot/certifi_1725551672989/work/certifi
+charset-normalizer @ file:///croot/charset-normalizer_1721748349566/work
+chemprop==1.7.0
+click==8.1.8
+cloudpickle==3.1.2
+colorlog==6.10.1
+contourpy==1.1.1
+cramjam==2.11.0
+cycler==0.12.1
+descriptastorus==2.8.0
+docstring_parser==0.18.0
+docutils==0.20.1
+et_xmlfile==2.0.0
+exceptiongroup==1.3.1
+fastapi==0.124.4
+fastparquet==2024.2.0
+filelock @ file:///croot/filelock_1700591183607/work
+Flask==2.1.3
+fonttools==4.57.0
+fsspec==2025.3.0
+future==1.0.0
+gitdb==4.0.12
+GitPython==3.1.50
+gmpy2 @ file:///tmp/build/80754af9/gmpy2_1645455532332/work
+greenlet==3.1.1
+h11==0.16.0
+hf-xet==1.5.1
+hf_transfer==0.1.9
+httpcore==1.0.9
+httpx==0.28.1
+huggingface_hub==0.36.2
+hyperopt==0.2.7
+idna @ file:///croot/idna_1714398848350/work
+imagesize==1.5.0
+importlib_metadata==8.5.0
+importlib_resources==6.4.5
+iniconfig==2.1.0
+inquirerpy==0.3.4
+itsdangerous==2.2.0
+Jinja2 @ file:///croot/jinja2_1716993405101/work
+joblib==1.4.2
+jsonschema==4.23.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.7
+-e git+ssh://git@github.com/RYDE-WORK/lnp_ml.git@104dfef94c6d6f03eb63369b4c4edc2f6ed4437d#egg=lnp_ml
+loguru==0.7.3
+Mako==1.3.12
+markdown-it-py==3.0.0
+MarkupSafe @ file:///croot/markupsafe_1704205993651/work
+matplotlib==3.7.5
+mdurl==0.1.2
+mkl-fft==1.3.1
+mkl-random @ file:///tmp/build/80754af9/mkl_random_1626186064646/work
+mkl-service==2.4.0
+modelscope==1.31.0
+mpmath @ file:///croot/mpmath_1690848262763/work
+mypy_extensions==1.1.0
+narwhals==1.42.1
+networkx @ file:///croot/networkx_1690561992265/work
+numpy @ file:///croot/numpy_and_numpy_base_1682520569166/work
+openpyxl==3.1.5
+optuna==4.5.0
+packaging==24.2
+pandas==2.0.3
+pandas_flavor==0.7.0
+peft==0.13.2
+pfzy==0.3.4
+pillow @ file:///croot/pillow_1721059439630/work
+pkgutil_resolve_name==1.3.10
+pluggy==1.5.0
+prompt_toolkit==3.0.52
+protobuf==5.29.6
+psutil==7.2.2
+py4j==0.10.9.9
+pyarrow==17.0.0
+pydantic==2.10.6
+pydantic_core==2.27.2
+pydeck==0.9.2
+Pygments==2.19.2
+pyparsing==3.1.4
+PySocks @ file:///tmp/build/80754af9/pysocks_1605305779399/work
+pytest==8.3.5
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+pytz==2026.2
+PyYAML @ file:///croot/pyyaml_1728657952215/work
+rdkit==2024.3.5
+referencing==0.35.1
+regex==2024.11.6
+requests @ file:///croot/requests_1721410876868/work
+rich==13.9.4
+rpds-py==0.20.1
+safetensors==0.5.3
+scikit-learn==1.3.2
+scipy==1.10.1
+selfies==2.2.0
+sentencepiece==0.2.0
+shellingham==1.5.4
+six @ file:///tmp/build/80754af9/six_1644875935023/work
+smmap==5.0.3
+sniffio==1.3.1
+snowballstemmer==3.1.1
+Sphinx==7.1.2
+sphinx_rtd_theme==3.1.0
+sphinxcontrib-applehelp==1.0.4
+sphinxcontrib-devhelp==1.0.2
+sphinxcontrib-htmlhelp==2.0.1
+sphinxcontrib-jquery==4.1
+sphinxcontrib-jsmath==1.0.1
+sphinxcontrib-qthelp==1.0.3
+sphinxcontrib-serializinghtml==1.1.5
+SQLAlchemy==2.0.51
+starlette==0.44.0
+streamlit==1.40.1
+sympy @ file:///croot/sympy_1734622612703/work
+tenacity==9.0.0
+tensorboardX==2.6.2.2
+threadpoolctl==3.5.0
+tokenizers==0.20.3
+toml==0.10.2
+tomli==2.4.1
+torch==2.4.1
+torchvision==0.20.0
+tornado==6.4.2
+tqdm==4.68.2
+transformers==4.46.3
+triton==3.0.0
+typed-argument-parser==1.10.1
+typer==0.20.1
+typing-inspect==0.9.0
+typing_extensions==4.13.2
+tzdata==2026.2
+urllib3 @ file:///croot/urllib3_1727769808118/work
+uvicorn==0.33.0
+watchdog==4.0.2
+wcwidth==0.8.2
+Werkzeug==2.3.8
+xarray==2023.1.0
+zipp==3.20.2

models/abl_full/s1_both/seed42/repro.txt

@@ -0,0 +1,5 @@
+date: 2026-07-18T12:53:18+00:00
+git_commit: 104dfef94c6d6f03eb63369b4c4edc2f6ed4437d
+git_dirty_files: 23
+seed: 42  pythonhashseed: 42
+cmd: GPU=1 --batch-size 8 --use-mpnn --use-moe --use-llm --use-rag --rag-top-k 4 --use-soft-prompt --no-llm-freeze --llm-use-qlora --llm-model-path models/qwen2.5-7b-instruct

models/abl_full/s1_both/seed42/strata_info.json

@@ -0,0 +1,60 @@
+{
+  "original_strata_counts": {
+    "T0|P0|E0": "5",
+    "T0|P0|E1": "58",
+    "T0|P0|E2": "169",
+    "T0|P1|E0": "1",
+    "T0|P1|E1": "17",
+    "T0|P1|E2": "32",
+    "T0|P2|E2": "3",
+    "T1|P0|E2": "9",
+    "T1|P1|E2": "5",
+    "TNA|P0|E0": "28",
+    "TNA|P0|E1": "21",
+    "TNA|P0|E2": "20",
+    "TNA|P1|E0": "29",
+    "TNA|P1|E1": "7",
+    "TNA|P1|E2": "14",
+    "TNA|P2|E2": "1",
+    "TNA|P3|E0": "1"
+  },
+  "rare_strata": [
+    "T0|P1|E0",
+    "T0|P2|E2",
+    "TNA|P2|E2",
+    "TNA|P3|E0"
+  ],
+  "final_strata": [
+    "RARE",
+    "T0|P0|E0",
+    "T0|P0|E1",
+    "T0|P0|E2",
+    "T0|P1|E1",
+    "T0|P1|E2",
+    "T1|P0|E2",
+    "T1|P1|E2",
+    "TNA|P0|E0",
+    "TNA|P0|E1",
+    "TNA|P0|E2",
+    "TNA|P1|E0",
+    "TNA|P1|E1",
+    "TNA|P1|E2"
+  ],
+  "final_strata_counts": {
+    "RARE": "6",
+    "T0|P0|E0": "5",
+    "T0|P0|E1": "58",
+    "T0|P0|E2": "169",
+    "T0|P1|E1": "17",
+    "T0|P1|E2": "32",
+    "T1|P0|E2": "9",
+    "T1|P1|E2": "5",
+    "TNA|P0|E0": "28",
+    "TNA|P0|E1": "21",
+    "TNA|P0|E2": "20",
+    "TNA|P1|E0": "29",
+    "TNA|P1|E1": "7",
+    "TNA|P1|E2": "14"
+  },
+  "n_rare_merged": "6"
+}

models/abl_full/s1_both/seed42/summary.json

@@ -0,0 +1,372 @@
+{
+  "fold_results": [
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.4108948503325973,
+        "lr": 0.0007340472227986773,
+        "weight_decay": 0.0007067135047679004,
+        "backbone_lr_ratio": 0.7321173243252591,
+        "moe_n_experts": 2,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 1,
+        "llm_lora_r": 32,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 14,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.65548933352636,
+          "rmse": 1.2866582038468337,
+          "mae": 0.5746945666081934,
+          "r2": -0.013200221826040748
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.5966502866071148,
+          "rmse": 0.7724314122348436,
+          "mae": 0.600902550468414,
+          "r2": 0.24162148150681173
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7619047619047619,
+          "precision": 0.7006414825374199,
+          "recall": 0.7006414825374199,
+          "f1": 0.7006414825374199
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.717012617012617,
+          "recall": 0.7409523809523809,
+          "f1": 0.7039699318769087
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9827586206896551,
+          "precision": 0.9912280701754386,
+          "recall": 0.75,
+          "f1": 0.8289085545722714
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.16024890157851573,
+          "js_divergence": 0.032467622854324446
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.24482591245144014,
+        "lr": 0.00042458936788614755,
+        "weight_decay": 0.00014902327669828248,
+        "backbone_lr_ratio": 0.24808069123277332,
+        "moe_n_experts": 2,
+        "moe_top_k": 1,
+        "moe_expert_hidden_mult": 1,
+        "llm_lora_r": 8,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 7,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.3445300972529715,
+          "rmse": 0.5869668621421242,
+          "mae": 0.4555579432108927,
+          "r2": -0.05207242874852769
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.0775295165871506,
+          "rmse": 1.0380411921437176,
+          "mae": 0.6340284446773471,
+          "r2": 0.18438783389531754
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.5972915181753385,
+          "recall": 0.6031746031746031,
+          "f1": 0.5997067448680352
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5476190476190477,
+          "precision": 0.47222222222222215,
+          "recall": 0.4779744065458351,
+          "f1": 0.46562819203268635
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.24783434314436234,
+          "js_divergence": 0.0666961611545551
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.330147494187667,
+        "lr": 0.0007391626870428583,
+        "weight_decay": 0.08248214185502709,
+        "backbone_lr_ratio": 0.3321485864522685,
+        "moe_n_experts": 4,
+        "moe_top_k": 1,
+        "moe_expert_hidden_mult": 1,
+        "llm_lora_r": 8,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 10,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.393067832783495,
+          "rmse": 0.6269512204179006,
+          "mae": 0.48754836206457447,
+          "r2": 0.22369805318261293
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6293676153394357,
+          "rmse": 0.7933269284093637,
+          "mae": 0.5411118950306748,
+          "r2": 0.16220261953292037
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.6282346175963197,
+          "recall": 0.6634897360703812,
+          "f1": 0.6212097651998134
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.6078622482131254,
+          "recall": 0.6720354808590102,
+          "f1": 0.6212949538882847
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.18339582654236564,
+          "js_divergence": 0.04379028402728257
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.2277771320023231,
+        "lr": 0.0009696339767409517,
+        "weight_decay": 0.0003933823864828228,
+        "backbone_lr_ratio": 0.4073135560471746,
+        "moe_n_experts": 4,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 2,
+        "llm_lora_r": 8,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 12,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.732527790200732,
+          "rmse": 1.316255214690803,
+          "mae": 0.5408693312566324,
+          "r2": 0.020069076700442334
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.5845268150974521,
+          "rmse": 0.7645435338144271,
+          "mae": 0.48413719012701917,
+          "r2": 0.43177048756483927
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6445623342175066,
+          "recall": 0.6598240469208211,
+          "f1": 0.6499999999999999
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5734711041691698,
+          "recall": 0.5891402714932127,
+          "f1": 0.5722830194879885
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 1.0,
+          "precision": 1.0,
+          "recall": 1.0,
+          "f1": 1.0
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.17380496676972265,
+          "js_divergence": 0.04216868552091467
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.3336459802315832,
+        "lr": 0.0005171929496634874,
+        "weight_decay": 0.00014235198757319405,
+        "backbone_lr_ratio": 0.23125566110577128,
+        "moe_n_experts": 4,
+        "moe_top_k": 1,
+        "moe_expert_hidden_mult": 2,
+        "llm_lora_r": 8,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 15,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.8359511180893293,
+          "rmse": 0.9143036246725315,
+          "mae": 0.5420850133467909,
+          "r2": -0.11860643443263852
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.0796427545219585,
+          "rmse": 1.03905859051449,
+          "mae": 0.649985611872298,
+          "r2": -0.04399672077007333
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6918367346938775,
+          "recall": 0.7412023460410557,
+          "f1": 0.6946420104314841
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6383091787439613,
+          "recall": 0.7337606837606838,
+          "f1": 0.6521164021164022
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 1.0,
+          "precision": 1.0,
+          "recall": 1.0,
+          "f1": 1.0
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.23837358595644248,
+          "js_divergence": 0.06386943316370278
+        }
+      }
+    }
+  ],
+  "summary_stats": {
+    "size": {
+      "mse_mean": 0.9923132343705776,
+      "mse_std": 0.5984792866289451,
+      "rmse_mean": 0.9462270251540387,
+      "rmse_std": 0.31139629291100435,
+      "mae_mean": 0.5201510432974168,
+      "mae_std": 0.04270959104369673,
+      "r2_mean": 0.011977608975169663,
+      "r2_std": 0.11548120300234146
+    },
+    "delivery": {
+      "mse_mean": 0.7935433976306223,
+      "mse_std": 0.23319925995033153,
+      "rmse_mean": 0.8814803314233683,
+      "rmse_std": 0.12859169080609734,
+      "mae_mean": 0.5820331384351507,
+      "mae_std": 0.061539044654339216,
+      "r2_mean": 0.19519714034596314,
+      "r2_std": 0.15267046261263886
+    },
+    "pdi": {
+      "accuracy_mean": 0.7095238095238096,
+      "accuracy_std": 0.03579356280564977,
+      "precision_mean": 0.6525133374440925,
+      "precision_std": 0.038896714345804356,
+      "recall_mean": 0.6736664429488562,
+      "recall_std": 0.045941327217941516,
+      "f1_mean": 0.6532400006073505,
+      "f1_std": 0.03965670173096385
+    },
+    "ee": {
+      "accuracy_mean": 0.65,
+      "accuracy_std": 0.058612064881386504,
+      "precision_mean": 0.6017754740722191,
+      "precision_std": 0.08029513301572835,
+      "recall_mean": 0.6427726447222246,
+      "recall_std": 0.09880211353738708,
+      "f1_mean": 0.6030584998804541,
+      "f1_std": 0.08094020366006734
+    },
+    "toxic": {
+      "accuracy_mean": 0.9736009044657997,
+      "accuracy_std": 0.026575657889142013,
+      "precision_mean": 0.8911027568922305,
+      "precision_std": 0.1303200307180046,
+      "recall_mean": 0.9379310344827585,
+      "recall_std": 0.09496854726841662,
+      "f1_mean": 0.8862220817530553,
+      "f1_std": 0.09421858471870986
+    },
+    "biodist": {
+      "kl_divergence_mean": 0.20073152479828177,
+      "kl_divergence_std": 0.03549650407693031,
+      "js_divergence_mean": 0.04979843734415591,
+      "js_divergence_std": 0.013252880014325175
+    }
+  }
+}

models/abl_full/s1_llm/seed42/local_uncommitted.patch

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

models/abl_full/s1_llm/seed42/outer_fold_0/best_params.json

@@ -0,0 +1,13 @@
+{
+  "dropout": 0.36120126405768543,
+  "lr": 0.0009943285546676244,
+  "weight_decay": 0.011860461070657894,
+  "backbone_lr_ratio": 0.09819347379144983,
+  "llm_lora_r": 8,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_llm/seed42/outer_fold_0/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 15}

models/abl_full/s1_llm/seed42/outer_fold_0/history.json

@@ -0,0 +1,321 @@
+{
+  "train": [
+    {
+      "loss": 5.147642432032405,
+      "loss_size": 0.8445845882634859,
+      "loss_pdi": 0.6722060796376821,
+      "loss_ee": 1.0654495648435645,
+      "loss_delivery": 0.9773198162784448,
+      "loss_biodist": 1.0411984002267993,
+      "loss_toxic": 0.5701748221307188
+    },
+    {
+      "loss": 4.256988151653393,
+      "loss_size": 0.8105902712087374,
+      "loss_pdi": 0.6404960864299053,
+      "loss_ee": 0.9829456693417317,
+      "loss_delivery": 0.9474382493141535,
+      "loss_biodist": 0.5263094587905986,
+      "loss_toxic": 0.38402967517440384
+    },
+    {
+      "loss": 3.7657996448310644,
+      "loss_size": 0.7923663322990006,
+      "loss_pdi": 0.6140706152529329,
+      "loss_ee": 0.9468127779058508,
+      "loss_delivery": 0.8142388411923438,
+      "loss_biodist": 0.4201673510912302,
+      "loss_toxic": 0.22732971095152804
+    },
+    {
+      "loss": 3.55830628485293,
+      "loss_size": 0.7685929627434628,
+      "loss_pdi": 0.5733279328088503,
+      "loss_ee": 0.9001653307193035,
+      "loss_delivery": 0.8628780795512973,
+      "loss_biodist": 0.3162229747788326,
+      "loss_toxic": 0.197189921890763
+    },
+    {
+      "loss": 3.348226389369449,
+      "loss_size": 0.7652994575532707,
+      "loss_pdi": 0.5595339947455639,
+      "loss_ee": 0.8787843343373891,
+      "loss_delivery": 0.8344920879682979,
+      "loss_biodist": 0.2531829160210249,
+      "loss_toxic": 0.12379072737452146
+    },
+    {
+      "loss": 3.418956476288873,
+      "loss_size": 0.7354503776374701,
+      "loss_pdi": 0.5472372681707949,
+      "loss_ee": 0.8670600813788336,
+      "loss_delivery": 0.9814510761281928,
+      "loss_biodist": 0.2594552078359836,
+      "loss_toxic": 0.10849214721521414
+    },
+    {
+      "loss": 3.1935991467656315,
+      "loss_size": 0.6718364231087066,
+      "loss_pdi": 0.542599725562173,
+      "loss_ee": 0.8166635447257274,
+      "loss_delivery": 0.8651735969894642,
+      "loss_biodist": 0.19670551441408493,
+      "loss_toxic": 0.1735602368166469
+    },
+    {
+      "loss": 3.0354659428467623,
+      "loss_size": 0.6521219149030544,
+      "loss_pdi": 0.5337435707852647,
+      "loss_ee": 0.7919033918831799,
+      "loss_delivery": 0.8176206077977612,
+      "loss_biodist": 0.19351129233837128,
+      "loss_toxic": 0.12709067441598587
+    },
+    {
+      "loss": 2.8783513468665047,
+      "loss_size": 0.6472237689068189,
+      "loss_pdi": 0.4928460322521828,
+      "loss_ee": 0.7864117219641402,
+      "loss_delivery": 0.7536062706966657,
+      "loss_biodist": 0.19095820067702113,
+      "loss_toxic": 0.09220971775278952
+    },
+    {
+      "loss": 2.8236329040011845,
+      "loss_size": 0.5908906914897867,
+      "loss_pdi": 0.4986982571112143,
+      "loss_ee": 0.7679622253856143,
+      "loss_delivery": 0.7320205288360248,
+      "loss_biodist": 0.19950280549961166,
+      "loss_toxic": 0.11956951313815709
+    },
+    {
+      "loss": 2.8038443810230977,
+      "loss_size": 0.6493027063237654,
+      "loss_pdi": 0.5011362546199077,
+      "loss_ee": 0.7827469851519611,
+      "loss_delivery": 0.7255703931724703,
+      "loss_biodist": 0.16047916933894157,
+      "loss_toxic": 0.0820891916031345
+    },
+    {
+      "loss": 2.8201219842240617,
+      "loss_size": 0.6174872288027325,
+      "loss_pdi": 0.4953894518517159,
+      "loss_ee": 0.7677139180737573,
+      "loss_delivery": 0.6763768926061489,
+      "loss_biodist": 0.20327536741624008,
+      "loss_toxic": 0.1466148302188722
+    },
+    {
+      "loss": 2.739313692659945,
+      "loss_size": 0.5882780187436052,
+      "loss_pdi": 0.48220539737392115,
+      "loss_ee": 0.7422905459597304,
+      "loss_delivery": 0.7619928093375387,
+      "loss_biodist": 0.14747168405635938,
+      "loss_toxic": 0.11162694219524998
+    },
+    {
+      "loss": 2.755237666336266,
+      "loss_size": 0.5910107342375291,
+      "loss_pdi": 0.4860636547610566,
+      "loss_ee": 0.7459143122305741,
+      "loss_delivery": 0.8056776328264056,
+      "loss_biodist": 0.1888150733870429,
+      "loss_toxic": 0.0446966980218988
+    },
+    {
+      "loss": 2.5650166079804704,
+      "loss_size": 0.5478415583839288,
+      "loss_pdi": 0.4664379881845938,
+      "loss_ee": 0.7518987720077103,
+      "loss_delivery": 0.6986862499568913,
+      "loss_biodist": 0.163654702438696,
+      "loss_toxic": 0.039902404693232196
+    }
+  ],
+  "val": [
+    {
+      "loss": 4.159415626525879,
+      "loss_size": 0.5259412823244929,
+      "loss_pdi": 0.7670119047164917,
+      "loss_ee": 1.074426245689392,
+      "loss_delivery": 0.7721726939082145,
+      "loss_biodist": 0.6637675046920777,
+      "loss_toxic": 0.389207723736763,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.3939393939393939,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.7557023048400877,
+      "loss_size": 0.5065902419388294,
+      "loss_pdi": 0.8063032388687134,
+      "loss_ee": 1.0656041860580445,
+      "loss_delivery": 0.6498562395572662,
+      "loss_biodist": 0.4285761207342148,
+      "loss_toxic": 0.3483713656663895,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.3939393939393939,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.476029634475708,
+      "loss_size": 0.4818432584404945,
+      "loss_pdi": 0.8822407007217408,
+      "loss_ee": 0.9630263566970825,
+      "loss_delivery": 0.6383201897144317,
+      "loss_biodist": 0.31057562232017516,
+      "loss_toxic": 0.26946860235184433,
+      "acc_pdi": 0.3939393939393939,
+      "acc_ee": 0.42424242424242425,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 4.863578939437867,
+      "loss_size": 0.541789197921753,
+      "loss_pdi": 1.2852459728717804,
+      "loss_ee": 0.8531787037849426,
+      "loss_delivery": 1.150505816936493,
+      "loss_biodist": 0.24585407972335815,
+      "loss_toxic": 0.8766378812491894,
+      "acc_pdi": 0.5454545454545454,
+      "acc_ee": 0.42424242424242425,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 4.034706544876099,
+      "loss_size": 0.4557234764099121,
+      "loss_pdi": 1.1182718336582185,
+      "loss_ee": 0.9067777767777443,
+      "loss_delivery": 0.7660846829414367,
+      "loss_biodist": 0.18823758214712144,
+      "loss_toxic": 0.6773600649088621,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.71524977684021,
+      "loss_size": 0.46545302867889404,
+      "loss_pdi": 1.232257217168808,
+      "loss_ee": 0.8113873779773713,
+      "loss_delivery": 0.7154026985168457,
+      "loss_biodist": 0.19339187294244767,
+      "loss_toxic": 0.3950491169467568,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.212270641326904,
+      "loss_size": 0.35887460708618163,
+      "loss_pdi": 1.260631912946701,
+      "loss_ee": 0.7949246689677238,
+      "loss_delivery": 0.6219677656888962,
+      "loss_biodist": 0.15977896004915237,
+      "loss_toxic": 0.13749083373695611,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.252304458618164,
+      "loss_size": 0.43110665678977966,
+      "loss_pdi": 1.1778656125068665,
+      "loss_ee": 0.8264414265751838,
+      "loss_delivery": 0.7276042275130749,
+      "loss_biodist": 0.16155126690864563,
+      "loss_toxic": 0.06569738076068461,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.449024963378906,
+      "loss_size": 0.4646922707557678,
+      "loss_pdi": 1.280164796113968,
+      "loss_ee": 0.8561057016253472,
+      "loss_delivery": 0.5933363705873489,
+      "loss_biodist": 0.1324687570333481,
+      "loss_toxic": 0.25103175402618944,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.059537887573242,
+      "loss_size": 0.35314797461032865,
+      "loss_pdi": 1.2232667624950408,
+      "loss_ee": 0.7745084583759307,
+      "loss_delivery": 0.5932506948709488,
+      "loss_biodist": 0.14098457098007203,
+      "loss_toxic": 0.10889038797467947,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.950475645065308,
+      "loss_size": 0.35260245501995086,
+      "loss_pdi": 1.092390513420105,
+      "loss_ee": 0.8254107091575861,
+      "loss_delivery": 0.6297225534915925,
+      "loss_biodist": 0.13818495273590087,
+      "loss_toxic": 0.05596545431762934,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.9905550956726072,
+      "loss_size": 0.36834086775779723,
+      "loss_pdi": 1.1431765377521514,
+      "loss_ee": 0.8174747616052628,
+      "loss_delivery": 0.6127920657396316,
+      "loss_biodist": 0.1217936098575592,
+      "loss_toxic": 0.07282344694249332,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.014424228668213,
+      "loss_size": 0.37132986187934874,
+      "loss_pdi": 1.1549767553806305,
+      "loss_ee": 0.8054959446191787,
+      "loss_delivery": 0.5878907024860383,
+      "loss_biodist": 0.14284245520830155,
+      "loss_toxic": 0.09392537409439683,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.075956106185913,
+      "loss_size": 0.3624262988567352,
+      "loss_pdi": 1.2138631284236907,
+      "loss_ee": 0.8189812514930963,
+      "loss_delivery": 0.580648323893547,
+      "loss_biodist": 0.13009978979825973,
+      "loss_toxic": 0.11513978778384626,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.0229308128356935,
+      "loss_size": 0.3652563512325287,
+      "loss_pdi": 1.1822814345359802,
+      "loss_ee": 0.8041247572749853,
+      "loss_delivery": 0.5846836477518081,
+      "loss_biodist": 0.12438895255327224,
+      "loss_toxic": 0.10542423813603818,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/abl_full/s1_llm/seed42/outer_fold_0/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15, 17, 18, 19, 21, 22, 23, 24, 25, 26, 27, 28, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 52, 53, 54, 55, 56, 57, 58, 61, 62, 63, 64, 65, 66, 67, 68, 71, 72, 73, 75, 76, 77, 80, 81, 85, 86, 88, 89, 90, 92, 94, 95, 96, 98, 99, 100, 101, 103, 104, 105, 106, 107, 108, 109, 111, 113, 114, 115, 116, 117, 118, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 133, 134, 135, 136, 138, 139, 141, 144, 145, 146, 147, 148, 149, 150, 151, 153, 154, 155, 156, 157, 158, 159, 160, 163, 164, 167, 168, 169, 170, 171, 172, 175, 176, 178, 179, 180, 182, 183, 184, 185, 188, 189, 190, 191, 192, 193, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 206, 208, 210, 211, 212, 213, 216, 217, 218, 219, 220, 222, 224, 225, 227, 228, 229, 230, 231, 232, 233, 234, 236, 237, 239, 241, 242, 244, 245, 247, 248, 249, 250, 252, 253, 255, 256, 257, 258, 259, 260, 261, 262, 265, 267, 268, 269, 270, 271, 272, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 285, 286, 287, 288, 291, 292, 294, 296, 297, 298, 299, 300, 301, 303, 304, 306, 307, 308, 309, 310, 311, 312, 313, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 332, 333, 335, 336, 337, 338, 339, 340, 342, 343, 344, 345, 346, 347, 348, 349, 351, 352, 354, 355, 356, 357, 358, 359, 360, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 374, 375, 377, 379, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389, 390, 391, 393, 395, 396, 397, 398, 399, 400, 401, 403, 404, 405, 406, 407, 408, 409, 411, 412, 413, 414, 415, 417, 418, 419], "outer_test_idx": [7, 16, 20, 29, 51, 59, 60, 69, 70, 74, 78, 79, 82, 83, 84, 87, 91, 93, 97, 102, 110, 112, 119, 120, 121, 132, 137, 140, 142, 143, 152, 161, 162, 165, 166, 173, 174, 177, 181, 186, 187, 194, 205, 207, 209, 214, 215, 221, 223, 226, 235, 238, 240, 243, 246, 251, 254, 263, 264, 266, 273, 284, 289, 290, 293, 295, 302, 305, 314, 315, 331, 334, 341, 350, 353, 361, 373, 376, 378, 392, 394, 402, 410, 416]}

models/abl_full/s1_llm/seed42/outer_fold_0/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 83,
+    "mse": 1.6152706135152473,
+    "rmse": 1.2709329697176195,
+    "mae": 0.5560500422066235,
+    "r2": 0.011414624800616835
+  },
+  "delivery": {
+    "n_samples": 58,
+    "mse": 0.569237390911561,
+    "rmse": 0.7544782242792438,
+    "mae": 0.5994908674138373,
+    "r2": 0.27646492613737106
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.75,
+    "precision": 0.6964285714285714,
+    "recall": 0.7195295794725588,
+    "f1": 0.7043740573152337
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.7261904761904762,
+    "precision": 0.684408546477512,
+    "recall": 0.7380952380952381,
+    "f1": 0.6942196418797405
+  },
+  "toxic": {
+    "n_samples": 58,
+    "accuracy": 0.9827586206896551,
+    "precision": 0.9912280701754386,
+    "recall": 0.75,
+    "f1": 0.8289085545722714
+  },
+  "biodist": {
+    "n_samples": 58,
+    "kl_divergence": 0.2234068484038443,
+    "js_divergence": 0.04400725882151094
+  }
+}

models/abl_full/s1_llm/seed42/outer_fold_1/best_params.json

@@ -0,0 +1,13 @@
+{
+  "dropout": 0.2307845821368613,
+  "lr": 0.000958470938363808,
+  "weight_decay": 7.743828905178866e-05,
+  "backbone_lr_ratio": 0.011104773529240554,
+  "llm_lora_r": 8,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_llm/seed42/outer_fold_1/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 11}

models/abl_full/s1_llm/seed42/outer_fold_1/history.json

@@ -0,0 +1,237 @@
+{
+  "train": [
+    {
+      "loss": 5.31889381924191,
+      "loss_size": 1.2377011441701167,
+      "loss_pdi": 0.6591631512384157,
+      "loss_ee": 1.053300973531362,
+      "loss_delivery": 0.8986851844835926,
+      "loss_biodist": 0.9206196327467222,
+      "loss_toxic": 0.5676593901337804
+    },
+    {
+      "loss": 4.242555192998938,
+      "loss_size": 1.1428092319417644,
+      "loss_pdi": 0.5818789673818124,
+      "loss_ee": 0.9055582446021002,
+      "loss_delivery": 0.860341256534731,
+      "loss_biodist": 0.45966394124804316,
+      "loss_toxic": 0.3206170860778641
+    },
+    {
+      "loss": 3.945884408177556,
+      "loss_size": 1.1150637850970835,
+      "loss_pdi": 0.5731848647465577,
+      "loss_ee": 0.8864123740711728,
+      "loss_delivery": 0.7410666408973772,
+      "loss_biodist": 0.3720009846461786,
+      "loss_toxic": 0.2981052903106084
+    },
+    {
+      "loss": 3.771327904752783,
+      "loss_size": 1.0933466932258091,
+      "loss_pdi": 0.5638367139004372,
+      "loss_ee": 0.8369277715682983,
+      "loss_delivery": 0.808194908137257,
+      "loss_biodist": 0.3217562503508619,
+      "loss_toxic": 0.20463744457811117
+    },
+    {
+      "loss": 3.5512277210080945,
+      "loss_size": 1.0661521545938544,
+      "loss_pdi": 0.5355479205782349,
+      "loss_ee": 0.8346078138093691,
+      "loss_delivery": 0.7495328228215914,
+      "loss_biodist": 0.24907558391223084,
+      "loss_toxic": 0.17771735889333728
+    },
+    {
+      "loss": 3.080665900900557,
+      "loss_size": 0.735050067909666,
+      "loss_pdi": 0.5247787269386085,
+      "loss_ee": 0.8116669840103871,
+      "loss_delivery": 0.7121920490989814,
+      "loss_biodist": 0.24619822848487544,
+      "loss_toxic": 0.11244597721442177
+    },
+    {
+      "loss": 3.3846415925670312,
+      "loss_size": 1.0514244253973704,
+      "loss_pdi": 0.5351601130253559,
+      "loss_ee": 0.8008756734229423,
+      "loss_delivery": 0.7044348140826097,
+      "loss_biodist": 0.25983165345481923,
+      "loss_toxic": 0.10588467437648089
+    },
+    {
+      "loss": 3.40125953506779,
+      "loss_size": 0.9813242455994761,
+      "loss_pdi": 0.5397726861206261,
+      "loss_ee": 0.7782152754229468,
+      "loss_delivery": 0.8498910393666577,
+      "loss_biodist": 0.22946381991779483,
+      "loss_toxic": 0.10234788009846532
+    },
+    {
+      "loss": 3.3028472243128597,
+      "loss_size": 1.0344087632524002,
+      "loss_pdi": 0.5138594145710403,
+      "loss_ee": 0.7748850149077338,
+      "loss_delivery": 0.7315809285721263,
+      "loss_biodist": 0.2475472209421364,
+      "loss_toxic": 0.08030925857255587
+    },
+    {
+      "loss": 3.321901753142073,
+      "loss_size": 1.0006535474110294,
+      "loss_pdi": 0.5102769502916852,
+      "loss_ee": 0.7776487961008742,
+      "loss_delivery": 0.7848934556382734,
+      "loss_biodist": 0.24158037957307454,
+      "loss_toxic": 0.08913449681288488
+    },
+    {
+      "loss": 3.2876750778507544,
+      "loss_size": 0.9780543958013123,
+      "loss_pdi": 0.5238370359749407,
+      "loss_ee": 0.7781052114190282,
+      "loss_delivery": 0.7209929026461936,
+      "loss_biodist": 0.2487740989875149,
+      "loss_toxic": 0.11638363839655712
+    }
+  ],
+  "val": [
+    {
+      "loss": 2.9308527231216432,
+      "loss_size": 0.519802288711071,
+      "loss_pdi": 0.6647319912910461,
+      "loss_ee": 0.8960770487785339,
+      "loss_delivery": 0.21736948788166047,
+      "loss_biodist": 0.45979435443878175,
+      "loss_toxic": 0.1797662064433098,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.696969696969697,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.673656702041626,
+      "loss_size": 0.49040685296058656,
+      "loss_pdi": 0.6421751320362091,
+      "loss_ee": 0.6792589068412781,
+      "loss_delivery": 0.4068264842033386,
+      "loss_biodist": 0.4101373553276062,
+      "loss_toxic": 0.05604152362793684,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.6190099716186523,
+      "loss_size": 0.5263730645179748,
+      "loss_pdi": 0.6318800538778305,
+      "loss_ee": 0.6358642637729645,
+      "loss_delivery": 0.44621982574462893,
+      "loss_biodist": 0.37990050911903384,
+      "loss_toxic": 0.02586401244625449,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.497333312034607,
+      "loss_size": 0.4611061841249466,
+      "loss_pdi": 0.6644555628299713,
+      "loss_ee": 0.6429485023021698,
+      "loss_delivery": 0.451871919631958,
+      "loss_biodist": 0.28512915670871736,
+      "loss_toxic": 0.027778143156319858,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.459736967086792,
+      "loss_size": 0.46326924562454225,
+      "loss_pdi": 0.6687620878219604,
+      "loss_ee": 0.5824844896793365,
+      "loss_delivery": 0.4660391449928284,
+      "loss_biodist": 0.29676556289196016,
+      "loss_toxic": 0.0203330360352993,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.3264069318771363,
+      "loss_size": 0.46273518800735475,
+      "loss_pdi": 0.651288640499115,
+      "loss_ee": 0.5867357015609741,
+      "loss_delivery": 0.38362357914447787,
+      "loss_biodist": 0.27055719792842864,
+      "loss_toxic": 0.012285364512354135,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.431403112411499,
+      "loss_size": 0.4746141493320465,
+      "loss_pdi": 0.6360470652580261,
+      "loss_ee": 0.6429702758789062,
+      "loss_delivery": 0.4433935463428497,
+      "loss_biodist": 0.2768822431564331,
+      "loss_toxic": 0.010033991513773799,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.3240861892700195,
+      "loss_size": 0.4379171907901764,
+      "loss_pdi": 0.6249558329582214,
+      "loss_ee": 0.608390498161316,
+      "loss_delivery": 0.4470574378967285,
+      "loss_biodist": 0.24718185961246492,
+      "loss_toxic": 0.010378359304741024,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.307906484603882,
+      "loss_size": 0.4245007187128067,
+      "loss_pdi": 0.619099599123001,
+      "loss_ee": 0.6267234861850739,
+      "loss_delivery": 0.44055399894714353,
+      "loss_biodist": 0.24240740835666658,
+      "loss_toxic": 0.009858347801491618,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.3521193981170656,
+      "loss_size": 0.4394673764705658,
+      "loss_pdi": 0.6199764013290405,
+      "loss_ee": 0.6198862552642822,
+      "loss_delivery": 0.4753984332084656,
+      "loss_biodist": 0.2460687130689621,
+      "loss_toxic": 0.008661644207313657,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.2861815452575684,
+      "loss_size": 0.41219576001167296,
+      "loss_pdi": 0.6171818912029267,
+      "loss_ee": 0.6221594095230103,
+      "loss_delivery": 0.4256651043891907,
+      "loss_biodist": 0.25570506155490874,
+      "loss_toxic": 0.008512460021302104,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/abl_full/s1_llm/seed42/outer_fold_1/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 3, 5, 6, 7, 10, 11, 12, 13, 15, 16, 18, 19, 20, 21, 24, 26, 27, 29, 30, 31, 32, 33, 35, 36, 38, 39, 40, 42, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 65, 66, 67, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 97, 98, 99, 100, 101, 102, 103, 105, 106, 107, 108, 109, 110, 112, 114, 115, 118, 119, 120, 121, 123, 124, 125, 126, 128, 129, 130, 131, 132, 136, 137, 140, 141, 142, 143, 144, 145, 146, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 161, 162, 164, 165, 166, 167, 168, 169, 171, 172, 173, 174, 175, 176, 177, 178, 179, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 194, 195, 196, 197, 198, 199, 201, 202, 204, 205, 206, 207, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 225, 226, 228, 229, 231, 232, 233, 234, 235, 236, 237, 238, 240, 242, 243, 244, 246, 247, 248, 250, 251, 252, 254, 255, 256, 258, 259, 260, 261, 262, 263, 264, 266, 268, 269, 270, 271, 273, 274, 275, 276, 277, 278, 279, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 295, 296, 297, 300, 301, 302, 303, 304, 305, 306, 307, 308, 310, 311, 313, 314, 315, 316, 317, 318, 319, 321, 322, 323, 324, 325, 327, 328, 329, 330, 331, 333, 334, 335, 337, 338, 339, 340, 341, 342, 343, 344, 345, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 361, 362, 364, 365, 366, 367, 368, 369, 373, 374, 375, 376, 378, 379, 380, 381, 382, 384, 387, 388, 389, 391, 392, 393, 394, 396, 397, 398, 399, 402, 403, 404, 405, 406, 408, 409, 410, 411, 412, 413, 414, 416, 417, 419], "outer_test_idx": [1, 2, 4, 8, 9, 14, 17, 22, 23, 25, 28, 34, 37, 41, 43, 45, 63, 64, 68, 85, 96, 104, 111, 113, 116, 117, 122, 127, 133, 134, 135, 138, 139, 147, 160, 163, 170, 180, 191, 192, 193, 200, 203, 208, 224, 227, 230, 239, 241, 245, 249, 253, 257, 265, 267, 272, 280, 294, 298, 299, 309, 312, 320, 326, 332, 336, 346, 347, 360, 363, 370, 371, 372, 377, 383, 385, 386, 390, 395, 400, 401, 407, 415, 418]}

models/abl_full/s1_llm/seed42/outer_fold_1/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.3497929868728069,
+    "rmse": 0.5914329944066419,
+    "mae": 0.4466874681501871,
+    "r2": -0.06814342257090522
+  },
+  "delivery": {
+    "n_samples": 61,
+    "mse": 1.1250146367822282,
+    "rmse": 1.0606670716026911,
+    "mae": 0.6891685391730461,
+    "r2": 0.14844502105923318
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.6190476190476191,
+    "precision": 0.5575167376749848,
+    "recall": 0.5714285714285714,
+    "f1": 0.5543766578249336
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.5476190476190477,
+    "precision": 0.4739422084623323,
+    "recall": 0.4779744065458351,
+    "f1": 0.4646469622331691
+  },
+  "toxic": {
+    "n_samples": 61,
+    "accuracy": 0.9836065573770492,
+    "precision": 0.9915254237288136,
+    "recall": 0.8333333333333333,
+    "f1": 0.8957264957264958
+  },
+  "biodist": {
+    "n_samples": 61,
+    "kl_divergence": 0.22230263127128363,
+    "js_divergence": 0.05931753635423146
+  }
+}

models/abl_full/s1_llm/seed42/outer_fold_2/best_params.json

@@ -0,0 +1,13 @@
+{
+  "dropout": 0.3632267382224968,
+  "lr": 0.0005436310282210407,
+  "weight_decay": 0.001502346334594783,
+  "backbone_lr_ratio": 0.08436478652740312,
+  "llm_lora_r": 32,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_llm/seed42/outer_fold_2/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 13}

models/abl_full/s1_llm/seed42/outer_fold_2/history.json

@@ -0,0 +1,279 @@
+{
+  "train": [
+    {
+      "loss": 5.36462926864624,
+      "loss_size": 0.8706878608948475,
+      "loss_pdi": 0.6885033813682763,
+      "loss_ee": 1.0611503881377142,
+      "loss_delivery": 1.0215906931741818,
+      "loss_biodist": 1.1023758311529417,
+      "loss_toxic": 0.6337891567397762
+    },
+    {
+      "loss": 4.529071704761402,
+      "loss_size": 0.8401351039474075,
+      "loss_pdi": 0.6459278867051408,
+      "loss_ee": 0.9848497203878455,
+      "loss_delivery": 0.914910903653583,
+      "loss_biodist": 0.7195982079248171,
+      "loss_toxic": 0.4474188113534773
+    },
+    {
+      "loss": 3.951905495411641,
+      "loss_size": 0.813644140757419,
+      "loss_pdi": 0.6092791750624373,
+      "loss_ee": 0.9192135785077069,
+      "loss_delivery": 0.8252965711258553,
+      "loss_biodist": 0.4887416153340726,
+      "loss_toxic": 0.3196989362304275
+    },
+    {
+      "loss": 3.6365428873010583,
+      "loss_size": 0.7937702816483136,
+      "loss_pdi": 0.5979801311686233,
+      "loss_ee": 0.8533155418731071,
+      "loss_delivery": 0.8274937568081392,
+      "loss_biodist": 0.36762599707455246,
+      "loss_toxic": 0.22278225973147797
+    },
+    {
+      "loss": 3.4147697074993237,
+      "loss_size": 0.7715920066914043,
+      "loss_pdi": 0.5473445915692562,
+      "loss_ee": 0.8263872604112368,
+      "loss_delivery": 0.795418934745563,
+      "loss_biodist": 0.32617241587187795,
+      "loss_toxic": 0.17444259638117776
+    },
+    {
+      "loss": 3.2687396294361837,
+      "loss_size": 0.7419949951204093,
+      "loss_pdi": 0.5502567855087487,
+      "loss_ee": 0.7747293452958803,
+      "loss_delivery": 0.8301814322334689,
+      "loss_biodist": 0.25826014558205734,
+      "loss_toxic": 0.1417507005704416
+    },
+    {
+      "loss": 3.1720887680311463,
+      "loss_size": 0.7395644131544474,
+      "loss_pdi": 0.5217430712403478,
+      "loss_ee": 0.7548408459972691,
+      "loss_delivery": 0.8303363386642288,
+      "loss_biodist": 0.2547873252549687,
+      "loss_toxic": 0.10291569284792687
+    },
+    {
+      "loss": 3.075252732715091,
+      "loss_size": 0.7378537817984014,
+      "loss_pdi": 0.5076718390793413,
+      "loss_ee": 0.7184092974340593,
+      "loss_delivery": 0.7718015240656363,
+      "loss_biodist": 0.24027213152195956,
+      "loss_toxic": 0.13018597279851502
+    },
+    {
+      "loss": 2.884459128250947,
+      "loss_size": 0.7027379934852188,
+      "loss_pdi": 0.4766775667667389,
+      "loss_ee": 0.7148908223654773,
+      "loss_delivery": 0.7258990050570385,
+      "loss_biodist": 0.21120643394218908,
+      "loss_toxic": 0.08622548082289663
+    },
+    {
+      "loss": 2.8339545920088485,
+      "loss_size": 0.7234651190606324,
+      "loss_pdi": 0.4769008582508242,
+      "loss_ee": 0.6739020033462627,
+      "loss_delivery": 0.6441003712447914,
+      "loss_biodist": 0.20530315810764158,
+      "loss_toxic": 0.13904721330146533
+    },
+    {
+      "loss": 2.8291303434887447,
+      "loss_size": 0.6923665050920602,
+      "loss_pdi": 0.4679024296032416,
+      "loss_ee": 0.6651387975828068,
+      "loss_delivery": 0.7329641821215281,
+      "loss_biodist": 0.1865290553384536,
+      "loss_toxic": 0.11662771172720839
+    },
+    {
+      "loss": 2.725346584577818,
+      "loss_size": 0.6627247101954512,
+      "loss_pdi": 0.4644476014214593,
+      "loss_ee": 0.6552067705102869,
+      "loss_delivery": 0.688079476658557,
+      "loss_biodist": 0.20521912643232862,
+      "loss_toxic": 0.08315514632802759
+    },
+    {
+      "loss": 2.7874236944559456,
+      "loss_size": 0.6820027900708688,
+      "loss_pdi": 0.4862427932990564,
+      "loss_ee": 0.6529161462912688,
+      "loss_delivery": 0.6929996894622171,
+      "loss_biodist": 0.20486629175374638,
+      "loss_toxic": 0.10216211944164054
+    }
+  ],
+  "val": [
+    {
+      "loss": 7.804502153396607,
+      "loss_size": 3.3283013671636583,
+      "loss_pdi": 0.6776975989341736,
+      "loss_ee": 1.0560791611671447,
+      "loss_delivery": 1.734302707016468,
+      "loss_biodist": 0.7140680432319642,
+      "loss_toxic": 0.3228553652763367,
+      "acc_pdi": 0.42424242424242425,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 7.4086761474609375,
+      "loss_size": 3.451000216603279,
+      "loss_pdi": 0.6704203605651855,
+      "loss_ee": 0.9652581334114074,
+      "loss_delivery": 1.7616875678300858,
+      "loss_biodist": 0.42040573954582217,
+      "loss_toxic": 0.18553604185581207,
+      "acc_pdi": 0.48484848484848486,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 7.2073790550231935,
+      "loss_size": 3.3938671171665193,
+      "loss_pdi": 0.6683620095252991,
+      "loss_ee": 0.9469543337821961,
+      "loss_delivery": 1.8101840883493423,
+      "loss_biodist": 0.31719583868980405,
+      "loss_toxic": 0.1326232858002186,
+      "acc_pdi": 0.45454545454545453,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 7.006436920166015,
+      "loss_size": 3.343517942726612,
+      "loss_pdi": 0.674301701784134,
+      "loss_ee": 0.9568410634994506,
+      "loss_delivery": 1.7925197899341583,
+      "loss_biodist": 0.22353795021772385,
+      "loss_toxic": 0.08939621523022652,
+      "acc_pdi": 0.48484848484848486,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.760958242416382,
+      "loss_size": 3.2666656225919724,
+      "loss_pdi": 0.639546251296997,
+      "loss_ee": 0.9238664388656617,
+      "loss_delivery": 1.7377906262874603,
+      "loss_biodist": 0.18202168345451356,
+      "loss_toxic": 0.08490984290838241,
+      "acc_pdi": 0.48484848484848486,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.776791954040528,
+      "loss_size": 3.2788748323917387,
+      "loss_pdi": 0.6404872536659241,
+      "loss_ee": 0.94258354306221,
+      "loss_delivery": 1.7482545793056488,
+      "loss_biodist": 0.18415743559598924,
+      "loss_toxic": 0.06798670999705791,
+      "acc_pdi": 0.5151515151515151,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.630231809616089,
+      "loss_size": 3.1941952049732207,
+      "loss_pdi": 0.6442001104354859,
+      "loss_ee": 0.9295283854007721,
+      "loss_delivery": 1.701593428850174,
+      "loss_biodist": 0.1511024206876755,
+      "loss_toxic": 0.09654238410294055,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.538211107254028,
+      "loss_size": 3.0761350959539415,
+      "loss_pdi": 0.6581819236278534,
+      "loss_ee": 0.9308917582035064,
+      "loss_delivery": 1.738530421257019,
+      "loss_biodist": 0.14440711140632628,
+      "loss_toxic": 0.08246456542983652,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.4778365135192875,
+      "loss_size": 3.127567061781883,
+      "loss_pdi": 0.6294267117977143,
+      "loss_ee": 0.9284797489643097,
+      "loss_delivery": 1.7013780117034911,
+      "loss_biodist": 0.12876486480236055,
+      "loss_toxic": 0.055887991283088925,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.5424212455749515,
+      "loss_size": 3.1456640899181365,
+      "loss_pdi": 0.6438539206981659,
+      "loss_ee": 0.9560280680656433,
+      "loss_delivery": 1.7197022199630738,
+      "loss_biodist": 0.12746439576148988,
+      "loss_toxic": 0.048482421692460774,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.444079160690308,
+      "loss_size": 3.0503403902053834,
+      "loss_pdi": 0.6505418002605439,
+      "loss_ee": 0.9497699677944184,
+      "loss_delivery": 1.7166073799133301,
+      "loss_biodist": 0.13068773299455644,
+      "loss_toxic": 0.045501152984797955,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.4499735832214355,
+      "loss_size": 3.053487080335617,
+      "loss_pdi": 0.6503841936588287,
+      "loss_ee": 0.9494261622428894,
+      "loss_delivery": 1.7213685989379883,
+      "loss_biodist": 0.12877217978239058,
+      "loss_toxic": 0.04628004720434546,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.596410894393921,
+      "loss_size": 3.2027933269739153,
+      "loss_pdi": 0.6455744743347168,
+      "loss_ee": 0.94369837641716,
+      "loss_delivery": 1.714592695236206,
+      "loss_biodist": 0.13734084889292716,
+      "loss_toxic": 0.05228173770010471,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.95
+    }
+  ]
+}

models/abl_full/s1_llm/seed42/outer_fold_2/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 11, 13, 14, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 32, 33, 34, 36, 37, 38, 39, 40, 41, 43, 44, 45, 47, 48, 49, 50, 51, 55, 56, 57, 58, 59, 60, 61, 63, 64, 65, 68, 69, 70, 72, 74, 76, 78, 79, 80, 82, 83, 84, 85, 86, 87, 88, 91, 93, 94, 95, 96, 97, 100, 101, 102, 103, 104, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 142, 143, 145, 147, 148, 149, 150, 151, 152, 153, 156, 157, 158, 159, 160, 161, 162, 163, 165, 166, 167, 168, 170, 171, 172, 173, 174, 175, 177, 179, 180, 181, 183, 184, 185, 186, 187, 188, 191, 192, 193, 194, 196, 197, 199, 200, 202, 203, 204, 205, 207, 208, 209, 210, 211, 212, 213, 214, 215, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 232, 233, 235, 236, 237, 238, 239, 240, 241, 242, 243, 245, 246, 249, 250, 251, 252, 253, 254, 255, 256, 257, 259, 260, 261, 263, 264, 265, 266, 267, 268, 269, 271, 272, 273, 274, 275, 276, 280, 281, 282, 283, 284, 285, 287, 289, 290, 293, 294, 295, 296, 297, 298, 299, 300, 302, 303, 305, 307, 308, 309, 310, 311, 312, 313, 314, 315, 317, 318, 320, 321, 322, 323, 324, 326, 327, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 358, 360, 361, 363, 364, 365, 366, 368, 369, 370, 371, 372, 373, 376, 377, 378, 379, 380, 381, 382, 383, 385, 386, 387, 388, 389, 390, 392, 394, 395, 398, 399, 400, 401, 402, 404, 405, 407, 408, 409, 410, 412, 413, 414, 415, 416, 417, 418], "outer_test_idx": [10, 12, 15, 31, 35, 42, 46, 52, 53, 54, 62, 66, 67, 71, 73, 75, 77, 81, 89, 90, 92, 98, 99, 105, 123, 124, 141, 144, 146, 154, 155, 164, 169, 176, 178, 182, 189, 190, 195, 198, 201, 206, 216, 217, 218, 219, 231, 234, 244, 247, 248, 258, 262, 270, 277, 278, 279, 286, 288, 291, 292, 301, 304, 306, 316, 319, 325, 328, 343, 357, 359, 362, 367, 374, 375, 384, 391, 393, 396, 397, 403, 406, 411, 419]}

models/abl_full/s1_llm/seed42/outer_fold_2/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.3808630997006998,
+    "rmse": 0.6171410695300547,
+    "mae": 0.4970361513440453,
+    "r2": 0.24780218295956957
+  },
+  "delivery": {
+    "n_samples": 60,
+    "mse": 0.6823172211803693,
+    "rmse": 0.826024951911484,
+    "mae": 0.5678542086932187,
+    "r2": 0.09171751672639383
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.6904761904761905,
+    "precision": 0.6693363844393593,
+    "recall": 0.717008797653959,
+    "f1": 0.662962962962963
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6547619047619048,
+    "precision": 0.5985449735449735,
+    "recall": 0.6507936507936507,
+    "f1": 0.6140979689366786
+  },
+  "toxic": {
+    "n_samples": 61,
+    "accuracy": 0.9344262295081968,
+    "precision": 0.7142857142857143,
+    "recall": 0.9655172413793103,
+    "f1": 0.7821428571428571
+  },
+  "biodist": {
+    "n_samples": 60,
+    "kl_divergence": 0.2095544756354361,
+    "js_divergence": 0.05129503689646577
+  }
+}

models/abl_full/s1_llm/seed42/outer_fold_3/best_params.json

@@ -0,0 +1,13 @@
+{
+  "dropout": 0.2834072840481449,
+  "lr": 0.0009817126260852267,
+  "weight_decay": 0.02817395190380758,
+  "backbone_lr_ratio": 0.07975174608349046,
+  "llm_lora_r": 8,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_llm/seed42/outer_fold_3/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 10}

models/abl_full/s1_llm/seed42/outer_fold_3/history.json

@@ -0,0 +1,216 @@
+{
+  "train": [
+    {
+      "loss": 4.703645029583493,
+      "loss_size": 0.5132228625384537,
+      "loss_pdi": 0.6786868137282294,
+      "loss_ee": 1.0606231979421668,
+      "loss_delivery": 0.88484697909774,
+      "loss_biodist": 1.031302827435571,
+      "loss_toxic": 0.5537116310080966
+    },
+    {
+      "loss": 3.7224996605434932,
+      "loss_size": 0.48492095800670415,
+      "loss_pdi": 0.6347826719284058,
+      "loss_ee": 0.9606194512264149,
+      "loss_delivery": 0.8909749190348226,
+      "loss_biodist": 0.5026347186114337,
+      "loss_toxic": 0.2761031350976712
+    },
+    {
+      "loss": 3.4812703841441386,
+      "loss_size": 0.45948930688806483,
+      "loss_pdi": 0.6280171420123126,
+      "loss_ee": 0.9439059093191817,
+      "loss_delivery": 0.8247231158956483,
+      "loss_biodist": 0.39155318205421036,
+      "loss_toxic": 0.27610667973656106
+    },
+    {
+      "loss": 3.151053058134543,
+      "loss_size": 0.42764424834702464,
+      "loss_pdi": 0.6216636799477242,
+      "loss_ee": 0.8778461346755156,
+      "loss_delivery": 0.751379218757958,
+      "loss_biodist": 0.32537463065740224,
+      "loss_toxic": 0.20028542850569292
+    },
+    {
+      "loss": 2.8697496620384424,
+      "loss_size": 0.4269665209828196,
+      "loss_pdi": 0.5621184295899159,
+      "loss_ee": 0.8231808357947582,
+      "loss_delivery": 0.761883601445604,
+      "loss_biodist": 0.2445085850116369,
+      "loss_toxic": 0.1063569859047798
+    },
+    {
+      "loss": 2.7854691131694898,
+      "loss_size": 0.39138617799491493,
+      "loss_pdi": 0.523831074302261,
+      "loss_ee": 0.8215621967573423,
+      "loss_delivery": 0.8069386293205458,
+      "loss_biodist": 0.23227983008365374,
+      "loss_toxic": 0.0727254229132086
+    },
+    {
+      "loss": 2.644588808755617,
+      "loss_size": 0.3799504083958832,
+      "loss_pdi": 0.5460223128666749,
+      "loss_ee": 0.8235893032035312,
+      "loss_delivery": 0.6982345264908429,
+      "loss_biodist": 0.19347638072999748,
+      "loss_toxic": 0.07254853126643276
+    },
+    {
+      "loss": 2.447996841894614,
+      "loss_size": 0.3444129961567956,
+      "loss_pdi": 0.5091852706831854,
+      "loss_ee": 0.7676352334989084,
+      "loss_delivery": 0.6081535461786631,
+      "loss_biodist": 0.2014586832877752,
+      "loss_toxic": 0.08061650964884541
+    },
+    {
+      "loss": 2.5595616069999902,
+      "loss_size": 0.401192020423509,
+      "loss_pdi": 0.5187365811418843,
+      "loss_ee": 0.7414932194593791,
+      "loss_delivery": 0.6709390034949457,
+      "loss_biodist": 0.20140804048325564,
+      "loss_toxic": 0.08921324366049187
+    },
+    {
+      "loss": 2.514197868269843,
+      "loss_size": 0.37954199555757884,
+      "loss_pdi": 0.5400866350612125,
+      "loss_ee": 0.7406274740760391,
+      "loss_delivery": 0.6182017973750025,
+      "loss_biodist": 0.19105441717279925,
+      "loss_toxic": 0.10805481451412512
+    }
+  ],
+  "val": [
+    {
+      "loss": 6.006975531578064,
+      "loss_size": 2.7494787633419038,
+      "loss_pdi": 0.6572931408882141,
+      "loss_ee": 0.9051256060600281,
+      "loss_delivery": 0.9734304904937744,
+      "loss_biodist": 0.4705023288726807,
+      "loss_toxic": 0.26852029263973237,
+      "acc_pdi": 0.5151515151515151,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.215833234786987,
+      "loss_size": 2.7748794183135033,
+      "loss_pdi": 0.622614860534668,
+      "loss_ee": 0.7935141026973724,
+      "loss_delivery": 0.7108514219522476,
+      "loss_biodist": 0.21298229992389678,
+      "loss_toxic": 0.1014245718717575,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.9565217391304348
+    },
+    {
+      "loss": 5.176661276817322,
+      "loss_size": 2.6878706961870193,
+      "loss_pdi": 0.577756530046463,
+      "loss_ee": 0.7635396361351013,
+      "loss_delivery": 0.8958718746900558,
+      "loss_biodist": 0.17920992374420167,
+      "loss_toxic": 0.08617440611124039,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.696969696969697,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.082661437988281,
+      "loss_size": 2.627828113734722,
+      "loss_pdi": 0.6501552999019623,
+      "loss_ee": 0.7462156653404236,
+      "loss_delivery": 0.7903401762247085,
+      "loss_biodist": 0.14461142122745513,
+      "loss_toxic": 0.1362426983192563,
+      "acc_pdi": 0.5454545454545454,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.142218232154846,
+      "loss_size": 2.715782304108143,
+      "loss_pdi": 0.6310679972171783,
+      "loss_ee": 0.7377537488937378,
+      "loss_delivery": 0.7469718545675278,
+      "loss_biodist": 0.14446477591991425,
+      "loss_toxic": 0.17223090548068284,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.191743612289429,
+      "loss_size": 2.7231680706143377,
+      "loss_pdi": 0.62291060090065,
+      "loss_ee": 0.7647787451744079,
+      "loss_delivery": 0.7376459240913391,
+      "loss_biodist": 0.1462427243590355,
+      "loss_toxic": 0.2010015532374382,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.971458315849304,
+      "loss_size": 2.6256699919700623,
+      "loss_pdi": 0.6229885518550873,
+      "loss_ee": 0.7519989430904388,
+      "loss_delivery": 0.72496417760849,
+      "loss_biodist": 0.13869253545999527,
+      "loss_toxic": 0.11672328729182482,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.8839984893798825,
+      "loss_size": 2.5973687767982483,
+      "loss_pdi": 0.6202696621417999,
+      "loss_ee": 0.7469191431999207,
+      "loss_delivery": 0.7374312400817871,
+      "loss_biodist": 0.13846834599971772,
+      "loss_toxic": 0.057255149632692334,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 5.073031163215637,
+      "loss_size": 2.7210794240236282,
+      "loss_pdi": 0.62847341299057,
+      "loss_ee": 0.7558909296989441,
+      "loss_delivery": 0.7458723217248917,
+      "loss_biodist": 0.14287284761667252,
+      "loss_toxic": 0.0880776752717793,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.985839319229126,
+      "loss_size": 2.670389157533646,
+      "loss_pdi": 0.6253753185272217,
+      "loss_ee": 0.7514976680278778,
+      "loss_delivery": 0.7372064888477325,
+      "loss_biodist": 0.13751591145992278,
+      "loss_toxic": 0.07459766920655966,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.9565217391304348
+    }
+  ]
+}

models/abl_full/s1_llm/seed42/outer_fold_3/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [0, 1, 2, 4, 5, 7, 8, 9, 10, 11, 12, 14, 15, 16, 17, 20, 21, 22, 23, 24, 25, 27, 28, 29, 30, 31, 32, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 45, 46, 47, 51, 52, 53, 54, 55, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 73, 74, 75, 76, 77, 78, 79, 81, 82, 83, 84, 85, 86, 87, 89, 90, 91, 92, 93, 95, 96, 97, 98, 99, 100, 102, 104, 105, 106, 109, 110, 111, 112, 113, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 127, 129, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 146, 147, 148, 150, 151, 152, 153, 154, 155, 156, 158, 159, 160, 161, 162, 163, 164, 165, 166, 168, 169, 170, 172, 173, 174, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 189, 190, 191, 192, 193, 194, 195, 198, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 212, 213, 214, 215, 216, 217, 218, 219, 221, 223, 224, 226, 227, 228, 230, 231, 233, 234, 235, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 251, 252, 253, 254, 255, 256, 257, 258, 259, 262, 263, 264, 265, 266, 267, 268, 269, 270, 272, 273, 274, 275, 276, 277, 278, 279, 280, 282, 284, 286, 288, 289, 290, 291, 292, 293, 294, 295, 297, 298, 299, 301, 302, 304, 305, 306, 307, 308, 309, 310, 312, 314, 315, 316, 317, 319, 320, 321, 325, 326, 327, 328, 331, 332, 334, 336, 340, 341, 343, 344, 346, 347, 348, 350, 352, 353, 357, 359, 360, 361, 362, 363, 364, 366, 367, 368, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 382, 383, 384, 385, 386, 387, 389, 390, 391, 392, 393, 394, 395, 396, 397, 398, 400, 401, 402, 403, 404, 406, 407, 410, 411, 412, 414, 415, 416, 418, 419], "outer_test_idx": [3, 6, 13, 18, 19, 26, 33, 44, 48, 49, 50, 56, 57, 72, 80, 88, 94, 101, 103, 107, 108, 114, 126, 128, 130, 145, 149, 157, 167, 171, 175, 188, 196, 197, 199, 211, 220, 222, 225, 229, 232, 236, 250, 260, 261, 271, 281, 283, 285, 287, 296, 300, 303, 311, 313, 318, 322, 323, 324, 329, 330, 333, 335, 337, 338, 339, 342, 345, 349, 351, 354, 355, 356, 358, 365, 369, 381, 388, 399, 405, 408, 409, 413, 417]}

models/abl_full/s1_llm/seed42/outer_fold_3/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 83,
+    "mse": 1.771712610404338,
+    "rmse": 1.3310569523519036,
+    "mae": 0.5784341517768531,
+    "r2": -0.002094156269689673
+  },
+  "delivery": {
+    "n_samples": 59,
+    "mse": 0.7036047486492218,
+    "rmse": 0.8388115096070283,
+    "mae": 0.5862117674107016,
+    "r2": 0.3160125883953585
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.75,
+    "precision": 0.7014607425441266,
+    "recall": 0.7426686217008798,
+    "f1": 0.7105824446267432
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6547619047619048,
+    "precision": 0.5854030501089325,
+    "recall": 0.6147812971342383,
+    "f1": 0.5895045852492662
+  },
+  "toxic": {
+    "n_samples": 60,
+    "accuracy": 0.9666666666666667,
+    "precision": 0.8,
+    "recall": 0.9824561403508771,
+    "f1": 0.8660714285714286
+  },
+  "biodist": {
+    "n_samples": 60,
+    "kl_divergence": 0.20841596050219596,
+    "js_divergence": 0.05264348216657474
+  }
+}

models/abl_full/s1_llm/seed42/outer_fold_4/best_params.json

@@ -0,0 +1,13 @@
+{
+  "dropout": 0.3555361701988309,
+  "lr": 0.0008600917157970105,
+  "weight_decay": 0.04751323754945044,
+  "backbone_lr_ratio": 0.0916633676039248,
+  "llm_lora_r": 32,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
+}

models/abl_full/s1_llm/seed42/outer_fold_4/epoch_mean.json

@@ -0,0 +1 @@
+{"epoch_mean": 12}

models/abl_full/s1_llm/seed42/outer_fold_4/history.json

@@ -0,0 +1,258 @@
+{
+  "train": [
+    {
+      "loss": 5.394337982744784,
+      "loss_size": 1.1257124124346554,
+      "loss_pdi": 0.6762782705796732,
+      "loss_ee": 1.0582815134847485,
+      "loss_delivery": 0.9602544831263052,
+      "loss_biodist": 1.0551355674460128,
+      "loss_toxic": 0.5365343851012152
+    },
+    {
+      "loss": 4.523071179518828,
+      "loss_size": 1.0805974346963134,
+      "loss_pdi": 0.6548431483474938,
+      "loss_ee": 0.9541445181176469,
+      "loss_delivery": 0.9257775768637657,
+      "loss_biodist": 0.5845872131553856,
+      "loss_toxic": 0.3552327539671112
+    },
+    {
+      "loss": 4.188256134858003,
+      "loss_size": 1.0770723203549515,
+      "loss_pdi": 0.6466329935434703,
+      "loss_ee": 0.9437865614891052,
+      "loss_delivery": 0.8513611442333943,
+      "loss_biodist": 0.46490952453097784,
+      "loss_toxic": 0.25296762661152594
+    },
+    {
+      "loss": 3.84365147191125,
+      "loss_size": 1.0425346574267826,
+      "loss_pdi": 0.618281994317029,
+      "loss_ee": 0.8794757723808289,
+      "loss_delivery": 0.8202072083950043,
+      "loss_biodist": 0.3390743921737413,
+      "loss_toxic": 0.2001333768051621
+    },
+    {
+      "loss": 3.5898072912886336,
+      "loss_size": 0.9976331996353897,
+      "loss_pdi": 0.5824355587765977,
+      "loss_ee": 0.8188786997988418,
+      "loss_delivery": 0.7497271782643086,
+      "loss_biodist": 0.32773358717158035,
+      "loss_toxic": 0.16699596401304007
+    },
+    {
+      "loss": 3.5503548641462586,
+      "loss_size": 1.0286318790670987,
+      "loss_pdi": 0.5576911544477617,
+      "loss_ee": 0.856319905938329,
+      "loss_delivery": 0.7199911005996369,
+      "loss_biodist": 0.28663341952739535,
+      "loss_toxic": 0.16064182748212605
+    },
+    {
+      "loss": 3.3900407649375297,
+      "loss_size": 0.9947644510381931,
+      "loss_pdi": 0.5822967455999272,
+      "loss_ee": 0.8434135639989698,
+      "loss_delivery": 0.6596610352799699,
+      "loss_biodist": 0.2585667297244072,
+      "loss_toxic": 0.11662859046781385
+    },
+    {
+      "loss": 3.2567197471051603,
+      "loss_size": 0.9685135679067792,
+      "loss_pdi": 0.5588751040600441,
+      "loss_ee": 0.79942778316704,
+      "loss_delivery": 0.6265396055941647,
+      "loss_biodist": 0.24395358562469482,
+      "loss_toxic": 0.12165178422740586
+    },
+    {
+      "loss": 3.234087083790753,
+      "loss_size": 0.9572132026424279,
+      "loss_pdi": 0.5511158140929969,
+      "loss_ee": 0.8025100319772154,
+      "loss_delivery": 0.6379532136425778,
+      "loss_biodist": 0.21870556433458585,
+      "loss_toxic": 0.13136480687614027
+    },
+    {
+      "loss": 3.0502883646939254,
+      "loss_size": 0.8992311735008214,
+      "loss_pdi": 0.5489526101060815,
+      "loss_ee": 0.7788462799948614,
+      "loss_delivery": 0.6070568097906338,
+      "loss_biodist": 0.20949353901921092,
+      "loss_toxic": 0.0742840920417645
+    },
+    {
+      "loss": 3.173892449688267,
+      "loss_size": 0.9275565492744381,
+      "loss_pdi": 0.5518025916975897,
+      "loss_ee": 0.8022452979474455,
+      "loss_delivery": 0.6698433639028588,
+      "loss_biodist": 0.20099210326333303,
+      "loss_toxic": 0.09389892752546616
+    },
+    {
+      "loss": 3.090608519476813,
+      "loss_size": 0.9454825053343902,
+      "loss_pdi": 0.5328148651767421,
+      "loss_ee": 0.809598571545369,
+      "loss_delivery": 0.5869329646229744,
+      "loss_biodist": 0.230247830015582,
+      "loss_toxic": 0.05867864513404768
+    }
+  ],
+  "val": [
+    {
+      "loss": 4.145606851577758,
+      "loss_size": 0.5787453949451447,
+      "loss_pdi": 0.6668040752410889,
+      "loss_ee": 0.9050501465797425,
+      "loss_delivery": 0.8079444321803748,
+      "loss_biodist": 0.7703335046768188,
+      "loss_toxic": 0.4415561527013779,
+      "acc_pdi": 0.45454545454545453,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.4442593097686767,
+      "loss_size": 0.41342132091522216,
+      "loss_pdi": 0.643909078836441,
+      "loss_ee": 0.8091854572296142,
+      "loss_delivery": 0.8259047922125319,
+      "loss_biodist": 0.4396663784980774,
+      "loss_toxic": 0.33677796721458436,
+      "acc_pdi": 0.48484848484848486,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.148733377456665,
+      "loss_size": 0.39056936502456663,
+      "loss_pdi": 0.5834551930427552,
+      "loss_ee": 0.7725244641304017,
+      "loss_delivery": 0.8219722658395767,
+      "loss_biodist": 0.30380556881427767,
+      "loss_toxic": 0.3053280901163816,
+      "acc_pdi": 0.48484848484848486,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.9237773418426514,
+      "loss_size": 0.26635559648275375,
+      "loss_pdi": 0.5148423090577126,
+      "loss_ee": 0.7525633133947849,
+      "loss_delivery": 0.8397242486476898,
+      "loss_biodist": 0.17795073464512826,
+      "loss_toxic": 0.3923140401020646,
+      "acc_pdi": 0.5454545454545454,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.8045753479003905,
+      "loss_size": 0.25571008771657944,
+      "loss_pdi": 0.466125363111496,
+      "loss_ee": 0.7717210814356804,
+      "loss_delivery": 0.8160159582272172,
+      "loss_biodist": 0.12289684787392616,
+      "loss_toxic": 0.3907925594598055,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.79259135723114,
+      "loss_size": 0.2913134977221489,
+      "loss_pdi": 0.45867272838950157,
+      "loss_ee": 0.8326532807201147,
+      "loss_delivery": 0.8062762513756752,
+      "loss_biodist": 0.11587586551904679,
+      "loss_toxic": 0.3189347217790782,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.815433645248413,
+      "loss_size": 0.3001242786645889,
+      "loss_pdi": 0.45082499235868456,
+      "loss_ee": 0.7799030734226108,
+      "loss_delivery": 0.7875660380348564,
+      "loss_biodist": 0.11556445211172103,
+      "loss_toxic": 0.4047958321869373,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.69639847278595,
+      "loss_size": 0.2592822894454002,
+      "loss_pdi": 0.45043728351593015,
+      "loss_ee": 0.7782680813223124,
+      "loss_delivery": 0.7862008877098561,
+      "loss_biodist": 0.11447847336530685,
+      "loss_toxic": 0.33743340661749244,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.754783272743225,
+      "loss_size": 0.2988031134009361,
+      "loss_pdi": 0.4530772313475609,
+      "loss_ee": 0.8197335122153163,
+      "loss_delivery": 0.816107339411974,
+      "loss_biodist": 0.10424908772110938,
+      "loss_toxic": 0.3031175131909549,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.6354313254356385,
+      "loss_size": 0.26673090904951097,
+      "loss_pdi": 0.4474984996020794,
+      "loss_ee": 0.7827383052557707,
+      "loss_delivery": 0.7797933895606548,
+      "loss_biodist": 0.10235336273908616,
+      "loss_toxic": 0.29335027942433956,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.5953806400299073,
+      "loss_size": 0.249676850438118,
+      "loss_pdi": 0.44969420284032824,
+      "loss_ee": 0.7588210146874189,
+      "loss_delivery": 0.7748311610892415,
+      "loss_biodist": 0.10375507175922394,
+      "loss_toxic": 0.29373050848953425,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.6621654272079467,
+      "loss_size": 0.2608113825321198,
+      "loss_pdi": 0.4512368381023407,
+      "loss_ee": 0.8041598474606871,
+      "loss_delivery": 0.7632074613822624,
+      "loss_biodist": 0.09309678822755814,
+      "loss_toxic": 0.326570615102537,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9444444444444444
+    }
+  ]
+}

models/abl_full/s1_llm/seed42/outer_fold_4/splits.json

@@ -0,0 +1 @@
+{"outer_train_idx": [1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 23, 25, 26, 28, 29, 31, 33, 34, 35, 37, 41, 42, 43, 44, 45, 46, 48, 49, 50, 51, 52, 53, 54, 56, 57, 59, 60, 62, 63, 64, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 77, 78, 79, 80, 81, 82, 83, 84, 85, 87, 88, 89, 90, 91, 92, 93, 94, 96, 97, 98, 99, 101, 102, 103, 104, 105, 107, 108, 110, 111, 112, 113, 114, 116, 117, 119, 120, 121, 122, 123, 124, 126, 127, 128, 130, 132, 133, 134, 135, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 149, 152, 154, 155, 157, 160, 161, 162, 163, 164, 165, 166, 167, 169, 170, 171, 173, 174, 175, 176, 177, 178, 180, 181, 182, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 203, 205, 206, 207, 208, 209, 211, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 229, 230, 231, 232, 234, 235, 236, 238, 239, 240, 241, 243, 244, 245, 246, 247, 248, 249, 250, 251, 253, 254, 257, 258, 260, 261, 262, 263, 264, 265, 266, 267, 270, 271, 272, 273, 277, 278, 279, 280, 281, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 298, 299, 300, 301, 302, 303, 304, 305, 306, 309, 311, 312, 313, 314, 315, 316, 318, 319, 320, 322, 323, 324, 325, 326, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 341, 342, 343, 345, 346, 347, 349, 350, 351, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 365, 367, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 381, 383, 384, 385, 386, 388, 390, 391, 392, 393, 394, 395, 396, 397, 399, 400, 401, 402, 403, 405, 406, 407, 408, 409, 410, 411, 413, 415, 416, 417, 418, 419], "outer_test_idx": [0, 5, 11, 21, 24, 27, 30, 32, 36, 38, 39, 40, 47, 55, 58, 61, 65, 76, 86, 95, 100, 106, 109, 115, 118, 125, 129, 131, 136, 148, 150, 151, 153, 156, 158, 159, 168, 172, 179, 183, 184, 185, 202, 204, 210, 212, 213, 228, 233, 237, 242, 252, 255, 256, 259, 268, 269, 274, 275, 276, 282, 297, 307, 308, 310, 317, 321, 327, 340, 344, 348, 352, 364, 366, 368, 379, 380, 382, 387, 389, 398, 404, 412, 414]}

models/abl_full/s1_llm/seed42/outer_fold_4/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.7148602454419013,
+    "rmse": 0.8454940836232394,
+    "mae": 0.5174528707229045,
+    "r2": 0.0434281945825975
+  },
+  "delivery": {
+    "n_samples": 58,
+    "mse": 0.8905622140496255,
+    "rmse": 0.9436960390134238,
+    "mae": 0.6175072552186662,
+    "r2": 0.138841040505834
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7976190476190477,
+    "precision": 0.7387740555951532,
+    "recall": 0.7456011730205279,
+    "f1": 0.7420054200542006
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.7142857142857143,
+    "precision": 0.6509661835748792,
+    "recall": 0.6675213675213675,
+    "f1": 0.6549590850833087
+  },
+  "toxic": {
+    "n_samples": 59,
+    "accuracy": 1.0,
+    "precision": 1.0,
+    "recall": 1.0,
+    "f1": 1.0
+  },
+  "biodist": {
+    "n_samples": 58,
+    "kl_divergence": 0.1907601176792889,
+    "js_divergence": 0.047605114707691666
+  }
+}

models/abl_full/s1_llm/seed42/pip_freeze.txt

@@ -0,0 +1,153 @@
+accelerate==1.0.1
+alabaster==0.7.13
+alembic==1.14.1
+altair==5.4.1
+annotated-doc==0.0.4
+annotated-types==0.7.0
+anyio==4.5.2
+attrs==25.3.0
+babel==2.18.0
+bitsandbytes==0.45.5
+blinker==1.8.2
+Brotli @ file:///croot/brotli-split_1714483155106/work
+cachetools==5.5.2
+captum==0.7.0
+certifi @ file:///croot/certifi_1725551672989/work/certifi
+charset-normalizer @ file:///croot/charset-normalizer_1721748349566/work
+chemprop==1.7.0
+click==8.1.8
+cloudpickle==3.1.2
+colorlog==6.10.1
+contourpy==1.1.1
+cramjam==2.11.0
+cycler==0.12.1
+descriptastorus==2.8.0
+docstring_parser==0.18.0
+docutils==0.20.1
+et_xmlfile==2.0.0
+exceptiongroup==1.3.1
+fastapi==0.124.4
+fastparquet==2024.2.0
+filelock @ file:///croot/filelock_1700591183607/work
+Flask==2.1.3
+fonttools==4.57.0
+fsspec==2025.3.0
+future==1.0.0
+gitdb==4.0.12
+GitPython==3.1.50
+gmpy2 @ file:///tmp/build/80754af9/gmpy2_1645455532332/work
+greenlet==3.1.1
+h11==0.16.0
+hf-xet==1.5.1
+hf_transfer==0.1.9
+httpcore==1.0.9
+httpx==0.28.1
+huggingface_hub==0.36.2
+hyperopt==0.2.7
+idna @ file:///croot/idna_1714398848350/work
+imagesize==1.5.0
+importlib_metadata==8.5.0
+importlib_resources==6.4.5
+iniconfig==2.1.0
+inquirerpy==0.3.4
+itsdangerous==2.2.0
+Jinja2 @ file:///croot/jinja2_1716993405101/work
+joblib==1.4.2
+jsonschema==4.23.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.7
+-e git+ssh://git@github.com/RYDE-WORK/lnp_ml.git@104dfef94c6d6f03eb63369b4c4edc2f6ed4437d#egg=lnp_ml
+loguru==0.7.3
+Mako==1.3.12
+markdown-it-py==3.0.0
+MarkupSafe @ file:///croot/markupsafe_1704205993651/work
+matplotlib==3.7.5
+mdurl==0.1.2
+mkl-fft==1.3.1
+mkl-random @ file:///tmp/build/80754af9/mkl_random_1626186064646/work
+mkl-service==2.4.0
+modelscope==1.31.0
+mpmath @ file:///croot/mpmath_1690848262763/work
+mypy_extensions==1.1.0
+narwhals==1.42.1
+networkx @ file:///croot/networkx_1690561992265/work
+numpy @ file:///croot/numpy_and_numpy_base_1682520569166/work
+openpyxl==3.1.5
+optuna==4.5.0
+packaging==24.2
+pandas==2.0.3
+pandas_flavor==0.7.0
+peft==0.13.2
+pfzy==0.3.4
+pillow @ file:///croot/pillow_1721059439630/work
+pkgutil_resolve_name==1.3.10
+pluggy==1.5.0
+prompt_toolkit==3.0.52
+protobuf==5.29.6
+psutil==7.2.2
+py4j==0.10.9.9
+pyarrow==17.0.0
+pydantic==2.10.6
+pydantic_core==2.27.2
+pydeck==0.9.2
+Pygments==2.19.2
+pyparsing==3.1.4
+PySocks @ file:///tmp/build/80754af9/pysocks_1605305779399/work
+pytest==8.3.5
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+pytz==2026.2
+PyYAML @ file:///croot/pyyaml_1728657952215/work
+rdkit==2024.3.5
+referencing==0.35.1
+regex==2024.11.6
+requests @ file:///croot/requests_1721410876868/work
+rich==13.9.4
+rpds-py==0.20.1
+safetensors==0.5.3
+scikit-learn==1.3.2
+scipy==1.10.1
+selfies==2.2.0
+sentencepiece==0.2.0
+shellingham==1.5.4
+six @ file:///tmp/build/80754af9/six_1644875935023/work
+smmap==5.0.3
+sniffio==1.3.1
+snowballstemmer==3.1.1
+Sphinx==7.1.2
+sphinx_rtd_theme==3.1.0
+sphinxcontrib-applehelp==1.0.4
+sphinxcontrib-devhelp==1.0.2
+sphinxcontrib-htmlhelp==2.0.1
+sphinxcontrib-jquery==4.1
+sphinxcontrib-jsmath==1.0.1
+sphinxcontrib-qthelp==1.0.3
+sphinxcontrib-serializinghtml==1.1.5
+SQLAlchemy==2.0.51
+starlette==0.44.0
+streamlit==1.40.1
+sympy @ file:///croot/sympy_1734622612703/work
+tenacity==9.0.0
+tensorboardX==2.6.2.2
+threadpoolctl==3.5.0
+tokenizers==0.20.3
+toml==0.10.2
+tomli==2.4.1
+torch==2.4.1
+torchvision==0.20.0
+tornado==6.4.2
+tqdm==4.68.2
+transformers==4.46.3
+triton==3.0.0
+typed-argument-parser==1.10.1
+typer==0.20.1
+typing-inspect==0.9.0
+typing_extensions==4.13.2
+tzdata==2026.2
+urllib3 @ file:///croot/urllib3_1727769808118/work
+uvicorn==0.33.0
+watchdog==4.0.2
+wcwidth==0.8.2
+Werkzeug==2.3.8
+xarray==2023.1.0
+zipp==3.20.2