fix: 修正bug

app/SCORE.md

@@ -12,8 +12,6 @@ encapsulation_efficiency_2: score = weight, where weight=0.08
 
 pdi_0: score = weight, where weight=0.08
 pdi_1: score = weight, where weight=0.02
-pdi_2: score = weight, where weight=0
-pdi_3: score = weight, where weight=0
 
 toxicity_0: score=weight, where weight=0.2
 toxicity_1: score=weight, where weight=0

app/api.py

@@ -25,23 +25,28 @@ from app.optimize import (
     TARGET_BIODIST,
     CompRanges,
     ScoringWeights,
+    HELPER_LIPID_OPTIONS,
+    ROUTE_OPTIONS,
+    create_dataframe_from_formulations,
+    predict_all,
 )
 
 
 # ============ Pydantic Models ============
 
 class CompRangesRequest(BaseModel):
-    """组分范围配置（mol 比例为百分数 0-100）"""
-    weight_ratio_min: float = Field(default=5.0, ge=1.0, le=50.0, description="阳离子脂质/mRNA 重量比最小值")
-    weight_ratio_max: float = Field(default=30.0, ge=1.0, le=50.0, description="阳离子脂质/mRNA 重量比最大值")
-    cationic_mol_min: float = Field(default=5.0, ge=0.0, le=100.0, description="阳离子脂质 mol 比例最小值 (%)")
-    cationic_mol_max: float = Field(default=80.0, ge=0.0, le=100.0, description="阳离子脂质 mol 比例最大值 (%)")
-    phospholipid_mol_min: float = Field(default=0.0, ge=0.0, le=100.0, description="磷脂 mol 比例最小值 (%)")
-    phospholipid_mol_max: float = Field(default=80.0, ge=0.0, le=100.0, description="磷脂 mol 比例最大值 (%)")
-    cholesterol_mol_min: float = Field(default=0.0, ge=0.0, le=100.0, description="胆固醇 mol 比例最小值 (%)")
-    cholesterol_mol_max: float = Field(default=80.0, ge=0.0, le=100.0, description="胆固醇 mol 比例最大值 (%)")
-    peg_mol_min: float = Field(default=0.0, ge=0.0, le=20.0, description="PEG 脂质 mol 比例最小值 (%)")
-    peg_mol_max: float = Field(default=5.0, ge=0.0, le=20.0, description="PEG 脂质 mol 比例最大值 (%)")
+    """组分范围配置（mol 比例为百分数 0-100）
+    """
+    weight_ratio_min: float = Field(default=7.0, ge=1.0, le=50.0, description="阳离子脂质/mRNA 重量比最小值")
+    weight_ratio_max: float = Field(default=20.0, ge=1.0, le=50.0, description="阳离子脂质/mRNA 重量比最大值")
+    cationic_mol_min: float = Field(default=22.0, ge=0.0, le=100.0, description="阳离子脂质 mol 比例最小值 (%)")
+    cationic_mol_max: float = Field(default=55.0, ge=0.0, le=100.0, description="阳离子脂质 mol 比例最大值 (%)")
+    phospholipid_mol_min: float = Field(default=7.0, ge=0.0, le=100.0, description="磷脂 mol 比例最小值 (%)")
+    phospholipid_mol_max: float = Field(default=42.0, ge=0.0, le=100.0, description="磷脂 mol 比例最大值 (%)")
+    cholesterol_mol_min: float = Field(default=15.0, ge=0.0, le=100.0, description="胆固醇 mol 比例最小值 (%)")
+    cholesterol_mol_max: float = Field(default=46.0, ge=0.0, le=100.0, description="胆固醇 mol 比例最大值 (%)")
+    peg_mol_min: float = Field(default=1.0, ge=0.0, le=20.0, description="PEG 脂质 mol 比例最小值 (%)")
+    peg_mol_max: float = Field(default=6.0, ge=0.0, le=20.0, description="PEG 脂质 mol 比例最大值 (%)")
     
     def to_comp_ranges(self) -> CompRanges:
         """转换为 CompRanges 对象"""
@@ -65,7 +70,7 @@ class ScoringWeightsRequest(BaseModel):
     delivery_weight: float = Field(default=0.0, ge=0.0, description="量化递送权重")
     size_weight: float = Field(default=0.0, ge=0.0, description="粒径权重 (80-150nm)")
     ee_class_weights: List[float] = Field(default=[0.0, 0.0, 0.0], description="EE 分类权重 [class0, class1, class2]")
-    pdi_class_weights: List[float] = Field(default=[0.0, 0.0, 0.0, 0.0], description="PDI 分类权重 [class0, class1, class2, class3]")
+    pdi_class_weights: List[float] = Field(default=[0.0, 0.0], description="PDI 分类权重 [class0(<0.2), class1(>=0.2)]")
     toxic_class_weights: List[float] = Field(default=[0.0, 0.0], description="毒性分类权重 [无毒, 有毒]")
     
     def to_scoring_weights(self) -> ScoringWeights:
@@ -87,6 +92,7 @@ class OptimizeRequest(BaseModel):
     top_k: int = Field(default=20, ge=1, le=100, description="Number of top formulations to return")
     num_seeds: Optional[int] = Field(default=None, ge=1, le=500, description="Number of seed points from first iteration (default: top_k * 5)")
     top_per_seed: int = Field(default=1, ge=1, le=10, description="Number of local best to keep per seed in refinement")
+    rerank_top_n: int = Field(default=200, ge=0, le=1000, description="两阶段推理：粗筛后用完整模型重排的候选数。0 表示关闭，此时 LLM 会在粗筛阶段跑满全部候选")
     step_sizes: Optional[List[float]] = Field(default=None, description="Mol ratio step sizes for each iteration (default: [10, 2, 1])")
     wr_step_sizes: Optional[List[float]] = Field(default=None, description="Weight ratio step sizes for each iteration (default: [5, 2, 1])")
     comp_ranges: Optional[CompRangesRequest] = Field(default=None, description="组分范围配置（默认使用标准范围）")
@@ -126,7 +132,7 @@ class FormulationResult(BaseModel):
     quantified_delivery: Optional[float] = None
     unnormalized_delivery: Optional[float] = None  # 反推的原始递送值（z-score 逆变换）
     size: Optional[float] = None
-    pdi_class: Optional[int] = None   # PDI 分类 (0: <0.2, 1: 0.2-0.3, 2: 0.3-0.4, 3: >0.4)
+    pdi_class: Optional[int] = None   # PDI 分类 (0: <0.2, 1: ≥0.2)
     ee_class: Optional[int] = None    # EE 分类 (0: <80%, 1: 80-90%, 2: >90%)
     toxic_class: Optional[int] = None # 毒性分类 (0: 无毒, 1: 有毒)
 
@@ -145,6 +151,73 @@ class HealthResponse(BaseModel):
     model_loaded: bool
     device: str
     available_organs: List[str]
+    use_moe: bool = False
+    use_llm: bool = False
+    use_rag: bool = False
+
+
+class PredictRequest(BaseModel):
+    """单配方预测请求"""
+    smiles: str = Field(..., description="Cationic lipid SMILES")
+    cationic_lipid_to_mrna_ratio: float = Field(..., gt=0, le=50, description="阳离子脂质/mRNA 重量比")
+    cationic_lipid_mol_ratio: float = Field(..., ge=0, le=100, description="阳离子脂质 mol 比例 (%)")
+    phospholipid_mol_ratio: float = Field(..., ge=0, le=100, description="磷脂 mol 比例 (%)")
+    cholesterol_mol_ratio: float = Field(..., ge=0, le=100, description="胆固醇 mol 比例 (%)")
+    peg_lipid_mol_ratio: float = Field(..., ge=0, le=20, description="PEG 脂质 mol 比例 (%)")
+    helper_lipid: str = Field(default="DOPE", description=f"辅助脂质，可选 {HELPER_LIPID_OPTIONS}")
+    route: str = Field(default="intravenous", description=f"给药途径，可选 {ROUTE_OPTIONS}")
+
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "smiles": "CC(C)NCCNC(C)C",
+                "cationic_lipid_to_mrna_ratio": 10.0,
+                "cationic_lipid_mol_ratio": 50.0,
+                "phospholipid_mol_ratio": 10.0,
+                "cholesterol_mol_ratio": 38.5,
+                "peg_lipid_mol_ratio": 1.5,
+                "helper_lipid": "DOPE",
+                "route": "intravenous",
+            }
+        }
+
+
+class PredictResponse(BaseModel):
+    """单配方预测响应"""
+    smiles: str
+    helper_lipid: str
+    route: str
+    biodist: Dict[str, float]
+    size: Optional[float] = None
+    quantified_delivery: Optional[float] = None
+    unnormalized_delivery: Optional[float] = None
+    pdi_class: Optional[int] = None
+    ee_class: Optional[int] = None
+    toxic_class: Optional[int] = None
+
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "smiles": "CC(C)NCCNC(C)C",
+                "helper_lipid": "DOPE",
+                "route": "intravenous",
+                "biodist": {
+                    "lymph_nodes": 0.0048,
+                    "heart": 0.0044,
+                    "liver": 0.6591,
+                    "spleen": 0.2817,
+                    "lung": 0.0245,
+                    "kidney": 0.0052,
+                    "muscle": 0.0203,
+                },
+                "size": 100.7,
+                "quantified_delivery": 0.1169,
+                "unnormalized_delivery": 0.3432,
+                "pdi_class": 0,
+                "ee_class": 2,
+                "toxic_class": 0,
+            }
+        }
 
 
 # ============ Global State ============
@@ -188,6 +261,20 @@ async def lifespan(app: FastAPI):
     try:
         state.model = load_model(model_path, state.device)
         logger.success("Model loaded successfully!")
+        # RAG 检索池：服务期用全部内部数据（无泄漏顾虑，查询分子会被 _retrieve_topk 自动排除）
+        _llm = getattr(state.model, "llm_prompt", None)
+        if _llm is not None and getattr(_llm, "use_rag", False):
+            import numpy as np
+            import pandas as pd
+            from lnp_ml.dataset import LNPDataset, process_dataframe
+            from lnp_ml.modeling.nested_cv_optuna import _build_rag_pool
+
+            rag_csv = Path(os.environ.get("RAG_POOL_CSV", "data/interim/internal.csv"))
+            logger.info(f"Building RAG retrieval pool from {rag_csv}...")
+            _ds = LNPDataset(process_dataframe(pd.read_csv(rag_csv)))
+            _s, _d, _ex = _build_rag_pool(_ds, np.arange(len(_ds)))
+            _llm.set_retrieval_pool(_s, _d, pool_id="serve", extra_labels=_ex)
+            logger.success(f"RAG pool ready: {len(_s)} molecules")
     except Exception as e:
         logger.error(f"Failed to load model: {e}")
         raise
@@ -224,11 +311,16 @@ app.add_middleware(
 @app.get("/", response_model=HealthResponse)
 async def health_check():
     """健康检查"""
+    _m = state.model
+    _llm = getattr(_m, "llm_prompt", None) if _m is not None else None
     return HealthResponse(
-        status="healthy" if state.model is not None else "model_not_loaded",
-        model_loaded=state.model is not None,
+        status="healthy" if _m is not None else "model_not_loaded",
+        model_loaded=_m is not None,
         device=str(state.device),
         available_organs=AVAILABLE_ORGANS,
+        use_moe=getattr(_m, "moe", None) is not None,
+        use_llm=_llm is not None,
+        use_rag=bool(getattr(_llm, "use_rag", False)),
     )
 
 
@@ -238,6 +330,85 @@ async def get_available_organs():
     return AVAILABLE_ORGANS
 
 
+@app.post("/predict", response_model=PredictResponse)
+async def predict_formulation(request: PredictRequest):
+    """对单个指定配方做属性预测（不搜索，约 0.2 秒）。"""
+    import pandas as pd
+    from rdkit import Chem, RDLogger
+
+    if state.model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    if request.helper_lipid not in HELPER_LIPID_OPTIONS:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid helper_lipid: {request.helper_lipid}. Available: {HELPER_LIPID_OPTIONS}",
+        )
+    if request.route not in ROUTE_OPTIONS:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid route: {request.route}. Available: {ROUTE_OPTIONS}",
+        )
+
+    mol_sum = (
+        request.cationic_lipid_mol_ratio
+        + request.phospholipid_mol_ratio
+        + request.cholesterol_mol_ratio
+        + request.peg_lipid_mol_ratio
+    )
+    if abs(mol_sum - 100.0) > 0.5:
+        raise HTTPException(
+            status_code=400, detail=f"四项 mol 比例之和须为 100（当前 {mol_sum:.2f}）"
+        )
+
+    RDLogger.DisableLog("rdApp.*")
+    if Chem.MolFromSmiles(request.smiles) is None:
+        raise HTTPException(status_code=400, detail=f"无法解析的 SMILES: {request.smiles[:80]}")
+
+    logger.info(f"Predict request: helper={request.helper_lipid}, route={request.route}, "
+                f"smiles={request.smiles[:50]}...")
+
+    df = create_dataframe_from_formulations(
+        request.smiles,
+        [(
+            request.cationic_lipid_to_mrna_ratio,
+            request.cationic_lipid_mol_ratio,
+            request.phospholipid_mol_ratio,
+            request.cholesterol_mol_ratio,
+            request.peg_lipid_mol_ratio,
+        )],
+        [request.helper_lipid],
+        [request.route],
+    )
+
+    try:
+        if getattr(state.model, "llm_prompt", None) is not None:
+            state.model.set_llm_enabled(True)
+        df = predict_all(state.model, df, state.device, batch_size=1)
+    except Exception as e:
+        logger.error(f"Prediction failed: {e}")
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        raise HTTPException(status_code=500, detail=str(e))
+
+    row = df.iloc[0]
+    _size, _unnorm = row.get("pred_size"), row.get("pred_unnorm_delivery")
+    return PredictResponse(
+        smiles=request.smiles,
+        helper_lipid=request.helper_lipid,
+        route=request.route,
+        biodist={
+            c.replace("Biodistribution_", ""): float(row[f"pred_{c}"]) for c in TARGET_BIODIST
+        },
+        size=float(_size) if pd.notna(_size) else None,
+        quantified_delivery=float(row["pred_delivery"]),
+        unnormalized_delivery=float(_unnorm) if pd.notna(_unnorm) else None,
+        pdi_class=int(row["pred_pdi_class"]),
+        ee_class=int(row["pred_ee_class"]),
+        toxic_class=int(row["pred_toxic_class"]),
+    )
+
+
 @app.post("/optimize", response_model=OptimizeResponse)
 async def optimize_formulation(request: OptimizeRequest):
     """
@@ -305,6 +476,7 @@ async def optimize_formulation(request: OptimizeRequest):
             comp_ranges=comp_ranges,
             routes=request.routes,
             scoring_weights=scoring_weights,
+            rerank_top_n=request.rerank_top_n,
             batch_size=256,
         )
         
@@ -350,6 +522,10 @@ async def optimize_formulation(request: OptimizeRequest):
         
     except Exception as e:
         logger.error(f"Optimization failed: {e}")
+        if getattr(state.model, "llm_prompt", None) is not None:
+            state.model.set_llm_enabled(True)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         raise HTTPException(status_code=500, detail=str(e))
 
 

app/app.py

@@ -143,6 +143,7 @@ def call_optimize_api(
     top_k: int = 20,
     num_seeds: int = None,
     top_per_seed: int = 1,
+    rerank_top_n: int = 200,
     step_sizes: list = None,
     wr_step_sizes: list = None,
     comp_ranges: dict = None,
@@ -156,6 +157,7 @@ def call_optimize_api(
         "top_k": top_k,
         "num_seeds": num_seeds,
         "top_per_seed": top_per_seed,
+        "rerank_top_n": rerank_top_n,
         "step_sizes": step_sizes,
         "wr_step_sizes": wr_step_sizes,
         "comp_ranges": comp_ranges,
@@ -175,9 +177,7 @@ def call_optimize_api(
 # PDI 分类标签
 PDI_CLASS_LABELS = {
     0: "<0.2 (优)",
-    1: "0.2-0.3 (良)",
-    2: "0.3-0.4 (中)",
-    3: ">0.4 (差)",
+    1: "≥0.2 (欠佳)",
 }
 
 # EE 分类标签
@@ -272,9 +272,18 @@ def main():
         # API 状态
         if api_online:
             st.success("🟢 API 服务在线")
+            try:
+                with httpx.Client(timeout=5) as _c:
+                    _info = _c.get(f"{API_URL}/").json()
+                _caps = [n for n, k in (("MoE", "use_moe"), ("LLM", "use_llm"), ("RAG", "use_rag"))
+                        if _info.get(k)]
+                st.caption(f"模型: {' + '.join(_caps) if _caps else '仅 backbone'}｜{_info.get('device', '?')}")
+            except Exception:
+                pass
         else:
             st.error("🔴 API 服务离线")
-            st.info("请先启动 API 服务:\n```\nuvicorn app.api:app --port 8000\n```")
+            st.info(f"请先启动 API 服务:\n```\nuvicorn app.api:app --port 8010\n```\n"
+                    f"当前 API_URL: {API_URL}")
         
         # st.divider()
         
@@ -354,6 +363,18 @@ def main():
                 help="后续迭代中，每个种子点邻域保留的局部最优数量",
             )
 
+            st.markdown("**LLM 推理精度**")
+            rerank_top_n = st.slider(
+                "LLM 重排候选数 (rerank_top_n)",
+                min_value=20,
+                max_value=500,
+                value=200,
+                step=20,
+                help="粗筛阶段只用 backbone 打分，再用完整模型对前 N 个"
+                    "候选重新预测并排序。"
+                    "调大更接近完整模型的判断，但耗时线性增加。",
+            )
+            
             st.markdown("**迭代步长与轮数**")
             use_custom_steps = st.checkbox(
                 "自定义迭代步长",
@@ -439,37 +460,37 @@ def main():
                 st.caption("阳离子脂质/mRNA 重量比")
                 col1, col2 = st.columns(2)
                 with col1:
-                    weight_ratio_min = st.number_input("最小", min_value=1.0, max_value=50.0, value=5.0, step=1.0, format="%.1f", key="wr_min")
+                    weight_ratio_min = st.number_input("最小", min_value=1.0, max_value=50.0, value=7.0, step=1.0, format="%.1f", key="wr_min")
                 with col2:
-                    weight_ratio_max = st.number_input("最大", min_value=1.0, max_value=50.0, value=30.0, step=1.0, format="%.1f", key="wr_max")
+                    weight_ratio_max = st.number_input("最大", min_value=1.0, max_value=50.0, value=20.0, step=1.0, format="%.1f", key="wr_max")
 
                 st.caption("阳离子脂质 mol 比例 (%)")
                 col1, col2 = st.columns(2)
                 with col1:
-                    cationic_mol_min = st.number_input("最小", min_value=0.0, max_value=100.0, value=5.0, step=5.0, format="%.1f", key="cat_min")
+                    cationic_mol_min = st.number_input("最小", min_value=0.0, max_value=100.0, value=22.0, step=5.0, format="%.1f", key="cat_min")
                 with col2:
-                    cationic_mol_max = st.number_input("最大", min_value=0.0, max_value=100.0, value=80.0, step=5.0, format="%.1f", key="cat_max")
+                    cationic_mol_max = st.number_input("最大", min_value=0.0, max_value=100.0, value=55.0, step=5.0, format="%.1f", key="cat_max")
 
                 st.caption("磷脂 mol 比例 (%)")
                 col1, col2 = st.columns(2)
                 with col1:
-                    phospholipid_mol_min = st.number_input("最小", min_value=0.0, max_value=100.0, value=0.0, step=5.0, format="%.1f", key="phos_min")
+                    phospholipid_mol_min = st.number_input("最小", min_value=0.0, max_value=100.0, value=7.0, step=5.0, format="%.1f", key="phos_min")
                 with col2:
-                    phospholipid_mol_max = st.number_input("最大", min_value=0.0, max_value=100.0, value=80.0, step=5.0, format="%.1f", key="phos_max")
+                    phospholipid_mol_max = st.number_input("最大", min_value=0.0, max_value=100.0, value=42.0, step=5.0, format="%.1f", key="phos_max")
 
                 st.caption("胆固醇 mol 比例 (%)")
                 col1, col2 = st.columns(2)
                 with col1:
-                    cholesterol_mol_min = st.number_input("最小", min_value=0.0, max_value=100.0, value=0.0, step=5.0, format="%.1f", key="chol_min")
+                    cholesterol_mol_min = st.number_input("最小", min_value=0.0, max_value=100.0, value=15.0, step=5.0, format="%.1f", key="chol_min")
                 with col2:
-                    cholesterol_mol_max = st.number_input("最大", min_value=0.0, max_value=100.0, value=80.0, step=5.0, format="%.1f", key="chol_max")
+                    cholesterol_mol_max = st.number_input("最大", min_value=0.0, max_value=100.0, value=46.0, step=5.0, format="%.1f", key="chol_max")
 
                 st.caption("PEG 脂质 mol 比例 (%)")
                 col1, col2 = st.columns(2)
                 with col1:
-                    peg_mol_min = st.number_input("最小", min_value=0.0, max_value=20.0, value=0.0, step=1.0, format="%.1f", key="peg_min")
+                    peg_mol_min = st.number_input("最小", min_value=0.0, max_value=20.0, value=1.0, step=1.0, format="%.1f", key="peg_min")
                 with col2:
-                    peg_mol_max = st.number_input("最大", min_value=0.0, max_value=20.0, value=5.0, step=1.0, format="%.1f", key="peg_max")
+                    peg_mol_max = st.number_input("最大", min_value=0.0, max_value=20.0, value=6.0, step=1.0, format="%.1f", key="peg_max")
                 
                 comp_ranges = {
                     "weight_ratio_min": weight_ratio_min,
@@ -514,7 +535,7 @@ def main():
                     help="score = normalize(delivery, route) × weight",
                 )
                 sw_size = st.number_input(
-                    "粒径 (Size, 80-150nm)",
+                    "粒径 (Size, 60-150nm)",
                     min_value=0.00, max_value=10.00, value=0.05,
                     step=0.05, format="%.2f", key="sw_size",
                     help="score = (1 if 60≤size≤150 else 0) × weight",
@@ -530,15 +551,11 @@ def main():
                     sw_ee2 = st.number_input(">80% (高)", min_value=0.00, max_value=1.00, value=0.08, step=0.01, format="%.2f", key="sw_ee2")
                 
                 st.caption("**PDI 分类权重**")
-                col1, col2, col3, col4 = st.columns(4)
+                col1, col2 = st.columns(2)
                 with col1:
                     sw_pdi0 = st.number_input("<0.2 (优)", min_value=0.00, max_value=1.00, value=0.08, step=0.01, format="%.2f", key="sw_pdi0")
                 with col2:
-                    sw_pdi1 = st.number_input("0.2-0.3 (良)", min_value=0.00, max_value=1.00, value=0.02, step=0.01, format="%.2f", key="sw_pdi1")
-                with col3:
-                    sw_pdi2 = st.number_input("0.3-0.4 (中)", min_value=0.00, max_value=1.00, value=0.00, step=0.01, format="%.2f", key="sw_pdi2")
-                with col4:
-                    sw_pdi3 = st.number_input(">0.4 (差)", min_value=0.00, max_value=1.00, value=0.00, step=0.01, format="%.2f", key="sw_pdi3")
+                    sw_pdi1 = st.number_input("≥0.2 (欠佳)", min_value=0.00, max_value=1.00, value=0.00, step=0.01, format="%.2f", key="sw_pdi1")
                 
                 st.caption("**毒性分类权重**")
                 col1, col2 = st.columns(2)
@@ -552,7 +569,7 @@ def main():
                     "delivery_weight": sw_delivery,
                     "size_weight": sw_size,
                     "ee_class_weights": [sw_ee0, sw_ee1, sw_ee2],
-                    "pdi_class_weights": [sw_pdi0, sw_pdi1, sw_pdi2, sw_pdi3],
+                    "pdi_class_weights": [sw_pdi0, sw_pdi1],
                     "toxic_class_weights": [sw_toxic0, sw_toxic1],
                 }
             else:
@@ -604,6 +621,7 @@ def main():
                         top_k=top_k,
                         num_seeds=num_seeds,
                         top_per_seed=top_per_seed,
+                        rerank_top_n=rerank_top_n,
                         step_sizes=step_sizes,
                         wr_step_sizes=wr_step_sizes_val,
                         comp_ranges=comp_ranges,
@@ -612,9 +630,7 @@ def main():
                     )
                     all_results.append({"smiles": smiles, "results": results})
                     
-                    # 为多 SMILES 模式添加 SMILES 标签
-                    smiles_label = smiles[:30] + "..." if len(smiles) > 30 else smiles
-                    df = format_results_dataframe(results, smiles_label if is_multi_smiles else None)
+                    df = format_results_dataframe(results, smiles if is_multi_smiles else None)
                     all_dfs.append(df)
                     
                 except httpx.HTTPStatusError as e:
@@ -709,7 +725,7 @@ def main():
         with col_export:
             smiles_used = st.session_state.get("smiles_used", "")
             if isinstance(smiles_used, list):
-                smiles_used = ",".join(smiles_used)
+                smiles_used = " | ".join(smiles_used)
             
             csv_content = create_export_csv(
                 df,
@@ -736,6 +752,13 @@ def main():
             use_container_width=True,
             hide_index=True,
             height=600,
+            column_config={
+                "SMILES": st.column_config.TextColumn(
+                    "SMILES",
+                    width="medium",
+                    help="界面上按列宽省略显示，可拖动列头加宽；导出的 CSV 中为完整字符串",
+                ),
+            },
         )
         
         # 详细信息

app/optimize.py

@@ -45,22 +45,23 @@ AVAILABLE_ORGANS = ["lymph_nodes", "heart", "liver", "spleen", "lung", "kidney",
 
 @dataclass
 class CompRanges:
-    """组分参数范围配置（mol 比例为百分数 0-100）"""
-    # 阳离子脂质/mRNA 重量比
-    weight_ratio_min: float = 5.0
-    weight_ratio_max: float = 30.0
-    # 阳离子脂质 mol 比例 (%)
-    cationic_mol_min: float = 5.0
-    cationic_mol_max: float = 80.0
-    # 磷脂 mol 比例 (%)
-    phospholipid_mol_min: float = 0.0
-    phospholipid_mol_max: float = 80.0
-    # 胆固醇 mol 比例 (%)
-    cholesterol_mol_min: float = 0.0
-    cholesterol_mol_max: float = 80.0
-    # PEG 脂质 mol 比例 (%)
-    peg_mol_min: float = 0.0
-    peg_mol_max: float = 5.0
+    """组分参数范围配置（mol 比例为百分数 0-100）
+    """
+    # 阳离子脂质/mRNA 重量比（数据 1%–99%: 7.52–19.35）
+    weight_ratio_min: float = 7.0
+    weight_ratio_max: float = 20.0
+    # 阳离子脂质 mol 比例 (%)（数据 1%–99%: 22.38–52.83）
+    cationic_mol_min: float = 22.0
+    cationic_mol_max: float = 55.0
+    # 磷脂 mol 比例 (%)（数据 1%–99%: 7.50–40.50）
+    phospholipid_mol_min: float = 7.0
+    phospholipid_mol_max: float = 42.0
+    # 胆固醇 mol 比例 (%)（数据 1%–99%: 16.00–45.00）
+    cholesterol_mol_min: float = 15.0
+    cholesterol_mol_max: float = 46.0
+    # PEG 脂质 mol 比例 (%)（数据 1%–99%: 1.0–6.0）
+    peg_mol_min: float = 1.0
+    peg_mol_max: float = 6.0
     
     def to_dict(self) -> Dict:
         """转换为字典"""
@@ -146,6 +147,16 @@ if not DELIVERY_NORM:
     logger.warning("DELIVERY_NORM is empty — scoring normalization for delivery will be disabled")
 
 
+_SIZE_STATS_PATH = Path(__file__).resolve().parent / "size_zscore_stats.json"
+if _SIZE_STATS_PATH.exists():
+    with open(_SIZE_STATS_PATH) as _f:
+        SIZE_ZSCORE_STATS: Dict[str, float] = json.load(_f)
+    logger.info(f"Loaded size stats from {_SIZE_STATS_PATH}")
+else:
+    SIZE_ZSCORE_STATS = {}
+    logger.warning(f"size_zscore_stats.json not found at {_SIZE_STATS_PATH}, "
+                   "pred_size 将置为 NaN，run 'make preprocess' to generate it")
+
 @dataclass
 class ScoringWeights:
     """
@@ -161,7 +172,7 @@ class ScoringWeights:
     size_weight: float = 0.0          # score = (1 if 80<=size<=150 else 0) * weight
     # 分类任务：per-class 权重（预测为该类时，得分 = 对应权重）
     ee_class_weights: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0])        # EE class 0, 1, 2
-    pdi_class_weights: List[float] = field(default_factory=lambda: [0.0, 0.0, 0.0, 0.0])  # PDI class 0, 1, 2, 3
+    pdi_class_weights: List[float] = field(default_factory=lambda: [0.0, 0.0])  # PDI class 0(<0.2), 1(>=0.2)
     toxic_class_weights: List[float] = field(default_factory=lambda: [0.0, 0.0])           # Toxic class 0, 1
 
 
@@ -299,7 +310,7 @@ class Formulation:
     quantified_delivery: Optional[float] = None
     unnormalized_delivery: Optional[float] = None  # 反推的原始递送值（z-score 逆变换）
     size: Optional[float] = None
-    pdi_class: Optional[int] = None   # PDI 分类 (0-3)
+    pdi_class: Optional[int] = None   # PDI 分类 (0: <0.2, 1: ≥0.2)
     ee_class: Optional[int] = None    # EE 分类 (0-2)
     toxic_class: Optional[int] = None # 毒性分类 (0: 无毒, 1: 有毒)
     
@@ -597,8 +608,12 @@ def predict_all(
     for i, col in enumerate(TARGET_BIODIST):
         df[f"pred_{col}"] = biodist_preds[:, i]
 
-    # size 模型输出为 log(size)，转换回真实粒径 (nm)
-    df["pred_size"] = np.exp(size_preds)
+    if SIZE_ZSCORE_STATS:
+        df["pred_size"] = np.exp(
+            size_preds * SIZE_ZSCORE_STATS["std"] + SIZE_ZSCORE_STATS["mean"]
+        )
+    else:
+        df["pred_size"] = np.nan
     df["pred_delivery"] = delivery_preds
     df["pred_pdi_class"] = pdi_preds
     df["pred_ee_class"] = ee_preds
@@ -688,7 +703,7 @@ def select_top_k(
                 # 额外预测值
                 quantified_delivery=row.get("pred_delivery"),
                 unnormalized_delivery=unnorm_delivery,
-                size=row.get("pred_size"),
+                size=float(row["pred_size"]) if pd.notna(row.get("pred_size")) else None,
                 pdi_class=int(row.get("pred_pdi_class")) if row.get("pred_pdi_class") is not None else None,
                 ee_class=int(row.get("pred_ee_class")) if row.get("pred_ee_class") is not None else None,
                 toxic_class=int(row.get("pred_toxic_class")) if row.get("pred_toxic_class") is not None else None,
@@ -701,6 +716,43 @@ def select_top_k(
     return formulations
 
 
+def create_dataframe_from_candidates(
+    smiles: str,
+    candidates: List[Formulation],
+) -> pd.DataFrame:
+    """为一批已确定 helper_lipid / route 的候选配方构建 DataFrame（一行一个候选）。"""
+    frames = []
+    for f in candidates:
+        comp = (
+            f.cationic_lipid_to_mrna_ratio,
+            f.cationic_lipid_mol_ratio,
+            f.phospholipid_mol_ratio,
+            f.cholesterol_mol_ratio,
+            f.peg_lipid_mol_ratio,
+        )
+        frames.append(
+            create_dataframe_from_formulations(smiles, [comp], [f.helper_lipid], [f.route])
+        )
+    return pd.concat(frames, ignore_index=True)
+
+
+def rerank_with_llm(
+    smiles: str,
+    organ: str,
+    candidates: List[Formulation],
+    model: torch.nn.Module,
+    device: torch.device,
+    scoring_weights: ScoringWeights,
+    batch_size: int = 16,
+) -> List[Formulation]:
+    """第二阶段：用完整模型（含 LLM 旁路）对候选集重新预测并排序。"""
+    if not candidates:
+        return list(candidates)
+    df = create_dataframe_from_candidates(smiles, candidates)
+    df = predict_all(model, df, device, batch_size)
+    return select_top_k(df, organ, k=len(df), scoring_weights=scoring_weights)
+
+
 def generate_single_seed_grid(
     seed: Formulation,
     mol_step: float,
@@ -780,6 +832,8 @@ def optimize(
     routes: Optional[List[str]] = None,
     scoring_weights: Optional[ScoringWeights] = None,
     batch_size: int = 256,
+    rerank_top_n: int = 200,
+    rerank_batch_size: int = 16,
 ) -> List[Formulation]:
     """
     执行配方优化（层级搜索策略）。
@@ -843,6 +897,15 @@ def optimize(
     logger.info(f"Comp ranges: {comp_ranges.to_dict()}")
     
     seeds = None
+    global_pool: List[Formulation] = []
+
+    _has_llm = getattr(model, "llm_prompt", None) is not None
+    _do_rerank = bool(_has_llm and rerank_top_n and rerank_top_n > 0)
+    if _has_llm:
+        model.set_llm_enabled(not _do_rerank)
+    if _do_rerank:
+        logger.info(f"Two-stage inference: coarse search w/o LLM, rerank top-{rerank_top_n} w/ LLM")
+
     
     for iteration, (mol_step, wr_step) in enumerate(zip(step_sizes, wr_step_sizes)):
         logger.info(f"\n{'='*60}")
@@ -872,7 +935,12 @@ def optimize(
             # 选择 top num_seeds 个种子点
             seeds = select_top_k(df, organ, num_seeds, scoring_weights)
 
-            logger.info(f"Selected {len(seeds)} seeds for next iteration")
+            global_pool = select_top_k(
+                df, organ, max(num_seeds, top_k * 10), scoring_weights
+            )
+
+            logger.info(f"Selected {len(seeds)} seeds for next iteration "
+                        f"(global pool: {len(global_pool)})")
             
         else:
             # ==================== 后续迭代：层级局部搜索 ====================
@@ -924,10 +992,9 @@ def optimize(
         logger.info(f"Current best score: {_score(best):.4f} (biodist_{organ}={best.get_biodist(organ):.4f})")
         logger.info(f"Best formulation: {best.to_dict()}")
     
-    # 最终去重、按综合评分排序并返回 top_k
+    # 粗筛结果：去重、按综合评分排序
     seeds_sorted = sorted(seeds, key=_score, reverse=True)
 
-    # 去重：保留每个唯一配方中得分最高的（已排序，所以第一个出现的就是最高的）
     seen_keys = set()
     unique_results = []
     for f in seeds_sorted:
@@ -936,10 +1003,44 @@ def optimize(
             seen_keys.add(key)
             unique_results.append(f)
 
-    logger.info(f"Final results: {len(unique_results)} unique formulations (from {len(seeds)} candidates)")
+    logger.info(f"Stage-1: {len(unique_results)} unique formulations (from {len(seeds)} candidates)")
 
+    if len(unique_results) < top_k and global_pool:
+        n_before = len(unique_results)
+        for f in global_pool:
+            if len(unique_results) >= top_k:
+                break
+            key = f.unique_key()
+            if key not in seen_keys:
+                seen_keys.add(key)
+                unique_results.append(f)
+        if len(unique_results) > n_before:
+            logger.warning(
+                f"局部搜索收敛至 {n_before} 个不同配方（请求 {top_k} 个），"
+                f"已从粗筛池回填 {len(unique_results) - n_before} 个未细化候选"
+            )
         return unique_results[:top_k]
 
+    # ==================== 第二阶段：完整模型重排 ====================
+    pool = unique_results[: max(rerank_top_n, top_k)]
+    logger.info(f"Stage-2: reranking {len(pool)} candidates with LLM enabled...")
+    model.set_llm_enabled(True)
+    reranked = rerank_with_llm(
+        smiles, organ, pool, model, device, scoring_weights, rerank_batch_size,
+    )
+
+    # 监控：粗筛与重排的秩相关。若长期接近 1，说明 LLM 旁路对排序无贡献，可直接关闭
+    try:
+        from scipy.stats import spearmanr
+        before = {f.unique_key(): i for i, f in enumerate(pool)}
+        after = [before[f.unique_key()] for f in reranked if f.unique_key() in before]
+        rho = spearmanr(after, range(len(after))).correlation
+        logger.info(f"Stage-2 rank correlation with stage-1: rho={rho:.4f}")
+    except Exception:
+        pass
+
+    return reranked[:top_k]
+
 
 def format_results(formulations: List[Formulation], organ: str) -> pd.DataFrame:
     """格式化结果为 DataFrame"""

app/size_zscore_stats.json

@@ -0,0 +1,4 @@
+{
+  "mean": 4.593287493148074,
+  "std": 0.42933069758929093
+}

lnp_ml/modeling/final_train_optuna_cv.py

@@ -53,6 +53,7 @@ from lnp_ml.modeling.trainer_balanced import (
     train_fixed_epochs,
 )
 from lnp_ml.modeling.visualization import plot_multitask_loss_curves
+from lnp_ml.modeling.nested_cv_optuna import _build_rag_pool
 
 # MPNN ensemble 默认路径
 DEFAULT_MPNN_ENSEMBLE_DIR = MODELS_DIR / "mpnn" / "all_amine_split_for_LiON"
@@ -178,20 +179,36 @@ def create_model(
     mpnn_device: str = "cpu",
     chemeleon_cache: Optional[str] = None,
     unimol_cache: Optional[str] = None,
-    # ============ MoE 相关（新增） ============
+    moleculestm_cache: Optional[str] = None,
+    mole_cache: Optional[str] = None,
+    set_transformer_block: str = "sab",
+    # MoE
     use_moe: bool = False,
     moe_n_experts: int = 4,
     moe_top_k: int = 2,
     moe_expert_hidden_mult: int = 2,
     moe_jitter_noise: float = 0.0,
+    # LLM（含 reg_bypass / use_rag / soft prompt / QLoRA）
+    llm_kwargs: Optional[Dict] = None,
+    # 检索增强
+    use_retrieval: bool = False,
+    retr_feature_dim: int = 0,
 ) -> Union[LNPModel, LNPModelWithoutMPNN]:
-    """创建模型"""
-    moe_kwargs = dict(
+    """创建模型。llm_kwargs 含 use_llm / llm_model_path / llm_freeze / llm_use_lora / llm_lora_* / reg_bypass。"""
+    extra_kwargs = dict(
+        set_transformer_block=set_transformer_block,
         use_moe=use_moe,
         moe_n_experts=moe_n_experts,
         moe_top_k=moe_top_k,
         moe_expert_hidden_mult=moe_expert_hidden_mult,
         moe_jitter_noise=moe_jitter_noise,
+        use_retrieval=use_retrieval,
+        retr_feature_dim=retr_feature_dim,
+        chemeleon_cache_path=chemeleon_cache,
+        unimol_cache_path=unimol_cache,
+        moleculestm_cache_path=moleculestm_cache,
+        mole_cache_path=mole_cache,
+        **(llm_kwargs or {}),
     )
 
     if use_mpnn:
@@ -205,9 +222,7 @@ 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,
+            **extra_kwargs,
         )
     else:
         return LNPModelWithoutMPNN(
@@ -217,9 +232,7 @@ 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,
+            **extra_kwargs,
         )
 
 
@@ -277,6 +290,10 @@ def run_optuna_cv(
     moe_top_k: int = 2,
     moe_expert_hidden_mult: int = 2,
     moe_jitter_noise: float = 0.0,
+    set_transformer_block: str = "sab",
+    llm_kwargs: Optional[Dict] = None,
+    use_retrieval: bool = False,
+    freeze_backbone_epochs: int = 3,
 ) -> Tuple[Dict, int, optuna.Study]:
     """
     使用全量数据做 3-fold CV Optuna 超参搜索。
@@ -332,6 +349,23 @@ def run_optuna_cv(
         weight_decay = trial.suggest_float("weight_decay", 1e-5, 1e-1, log=True)
         backbone_lr_ratio = trial.suggest_float("backbone_lr_ratio", 0.01, 1.0, log=True)
 
+        # MoE / LoRA 结构超参：搜索空间与 nested_cv_optuna 保持一致，便于横向比较
+        moe_t = dict(
+            moe_n_experts=moe_n_experts,
+            moe_top_k=moe_top_k,
+            moe_expert_hidden_mult=moe_expert_hidden_mult,
+        )
+        if use_moe:
+            moe_t["moe_n_experts"] = trial.suggest_categorical("moe_n_experts", [2, 4, 8])
+            moe_t["moe_top_k"] = trial.suggest_int("moe_top_k", 1, 2)
+            moe_t["moe_expert_hidden_mult"] = trial.suggest_categorical(
+                "moe_expert_hidden_mult", [1, 2])
+
+        llm_kwargs_t = dict(llm_kwargs or {})
+        if llm_kwargs_t.get("use_llm", False):
+            llm_kwargs_t["llm_use_lora"] = True
+            llm_kwargs_t["llm_lora_r"] = trial.suggest_categorical("llm_lora_r", [8, 16, 32])
+
         # 3-fold CV
         cv = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=seed)
 
@@ -339,7 +373,6 @@ def run_optuna_cv(
         fold_best_epochs = []
 
         for fold, (train_idx, val_idx) in enumerate(cv.split(indices, strata)):
-            # 创建 DataLoader
             train_subset = Subset(full_dataset, train_idx.tolist())
             val_subset = Subset(full_dataset, val_idx.tolist())
 
@@ -350,10 +383,8 @@ def run_optuna_cv(
                 val_subset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn
             )
 
-            # 计算类权重
             class_weights = compute_class_weights_from_loader(train_loader)
 
-            # 创建模型
             model = create_model(
                 d_model=d_model,
                 num_heads=num_heads,
@@ -365,22 +396,29 @@ def run_optuna_cv(
                 mpnn_device=device.type,
                 chemeleon_cache=chemeleon_cache,
                 unimol_cache=unimol_cache,
+                set_transformer_block=set_transformer_block,
                 use_moe=use_moe,
-                moe_n_experts=moe_n_experts,
-                moe_top_k=moe_top_k,
-                moe_expert_hidden_mult=moe_expert_hidden_mult,
                 moe_jitter_noise=moe_jitter_noise,
+                llm_kwargs=llm_kwargs_t,
+                use_retrieval=use_retrieval,
+                retr_feature_dim=(3 if use_retrieval else 0),
+                **moe_t,
             )
             if rdkit_cache is not None:
                 model.rdkit_encoder._cache = rdkit_cache
 
-            # 加载预训练权重
+            # 检索池只用当前 train fold：用全量会让验证损失被自己的标签污染，
+            # 选出的超参和 epoch_mean 都会有偏
+            if llm_kwargs_t.get("use_rag", False) and getattr(model, "llm_prompt", None) is not None:
+                _s, _d, _ex = _build_rag_pool(full_dataset, train_idx)
+                model.llm_prompt.set_retrieval_pool(
+                    _s, _d, pool_id=f"opt{fold}", extra_labels=_ex)
+
             if pretrain_state_dict is not None and pretrain_config is not None:
                 load_pretrain_weights_to_model(
                     model, pretrain_state_dict, d_model, pretrain_config, load_delivery_head
                 )
 
-            # 训练（带早停）
             result = train_with_early_stopping(
                 model=model,
                 train_loader=train_loader,
@@ -392,6 +430,7 @@ def run_optuna_cv(
                 patience=patience,
                 class_weights=class_weights,
                 backbone_lr_ratio=backbone_lr_ratio,
+                freeze_backbone_epochs=freeze_backbone_epochs,
             )
 
             fold_val_losses.append(result["best_val_loss"])
@@ -427,6 +466,7 @@ def run_optuna_cv(
         "n_attn_layers": fixed_n_attn_layers,
         "fusion_strategy": fixed_fusion_strategy,
         "head_hidden_dim": fixed_head_hidden_dim,
+        "set_transformer_block": set_transformer_block,
     })
     epoch_mean = study.best_trial.user_attrs.get("epoch_mean", epochs_per_trial)
     
@@ -451,6 +491,8 @@ def main(
     n_trials: int = 20,
     epochs_per_trial: int = 30,
     patience: int = 10,
+    fixed_params_json: Optional[Path] = None,
+    fixed_epoch_mean: Optional[int] = None,
     # 训练参数
     batch_size: int = 32,
     # 最终训练参数
@@ -472,6 +514,27 @@ def main(
     moe_top_k: int = 2,
     moe_expert_hidden_mult: int = 2,
     moe_jitter_noise: float = 0.0,
+    # Set Transformer
+    set_transformer_block: str = "sab",
+    # 回归旁路
+    reg_bypass: str = "on",
+    # LLM
+    use_llm: bool = False,
+    use_rag: bool = False,
+    rag_top_k: int = 4,
+    use_soft_prompt: bool = False,
+    llm_model_path: str = "models/qwen2.5-7b-instruct",
+    llm_freeze: bool = True,
+    llm_use_lora: bool = False,
+    llm_use_qlora: bool = False,
+    llm_lora_r: int = 8,
+    llm_lora_alpha: int = 16,
+    llm_lora_dropout: float = 0.05,
+    llm_max_length: int = 1536,
+    # 检索增强
+    use_retrieval: bool = False,
+    # backbone 冻结轮数（必须与 Optuna 阶段一致）
+    freeze_backbone_epochs: int = 3,
     # 设备
     device: str = "cuda" if torch.cuda.is_available() else "cpu",
 ):
@@ -492,6 +555,26 @@ def main(
     logger.info(f"Using device: {device}")
     device = torch.device(device)
 
+    if use_swa and llm_use_qlora:
+        logger.error("use_swa 与 QLoRA 不兼容：SWA 要对全部参数做滑动平均，4bit 量化基座无法参与")
+        raise typer.Exit(1)
+
+    llm_kwargs = dict(
+        reg_bypass=reg_bypass,
+        use_rag=use_rag,
+        rag_top_k=rag_top_k,
+        llm_max_length=llm_max_length,
+        use_llm=use_llm,
+        llm_model_path=llm_model_path,
+        llm_freeze=llm_freeze,
+        llm_use_lora=llm_use_lora,
+        llm_use_qlora=llm_use_qlora,
+        use_soft_prompt=use_soft_prompt,
+        llm_lora_r=llm_lora_r,
+        llm_lora_alpha=llm_lora_alpha,
+        llm_lora_dropout=llm_lora_dropout,
+    )
+
     # 加载预训练权重（如果指定）
     pretrain_state_dict = None
     pretrain_config = None
@@ -532,7 +615,17 @@ def main(
     # 预热 RDKit 缓存（在整个训练流程中共享）
     rdkit_cache = warmup_rdkit_cache(full_dataset.smiles)
     
-    # 运行 Optuna 调参
+    if fixed_params_json is not None:
+        if fixed_epoch_mean is None:
+            raise typer.BadParameter("--fixed-params-json 必须配合 --fixed-epoch-mean 使用")
+        logger.info(f"Skipping Optuna, loading fixed params from {fixed_params_json}")
+        with open(fixed_params_json) as f:
+            best_params = json.load(f)
+        epoch_mean = fixed_epoch_mean
+        study = None
+        logger.info(f"Fixed params: {best_params}")
+        logger.info(f"Fixed epoch_mean: {epoch_mean}")
+    else:
         logger.info(f"\nRunning {n_folds}-fold Optuna with {n_trials} trials...")
         study_path = output_dir / "optuna_study.sqlite3"
 
@@ -559,6 +652,10 @@ def main(
             moe_top_k=moe_top_k,
             moe_expert_hidden_mult=moe_expert_hidden_mult,
             moe_jitter_noise=moe_jitter_noise,
+            set_transformer_block=set_transformer_block,
+            llm_kwargs=llm_kwargs,
+            use_retrieval=use_retrieval,
+            freeze_backbone_epochs=freeze_backbone_epochs,
         )
 
     # 保存最佳参数
@@ -569,6 +666,7 @@ def main(
         json.dump({"epoch_mean": epoch_mean}, f)
 
     # 保存 Optuna 试验历史
+    if study is not None:
         trials_history = []
         for trial in study.trials:
             trials_history.append({
@@ -606,7 +704,19 @@ def main(
     with open(output_dir / "class_weights.json", "w") as f:
         json.dump(class_weights_info, f, indent=2)
 
-    # 创建模型
+    arch = {
+        "moe_n_experts": best_params.get("moe_n_experts", moe_n_experts),
+        "moe_top_k": best_params.get("moe_top_k", moe_top_k),
+        "moe_expert_hidden_mult": best_params.get("moe_expert_hidden_mult", moe_expert_hidden_mult),
+        "moe_jitter_noise": moe_jitter_noise,
+        "set_transformer_block": best_params.get("set_transformer_block", set_transformer_block),
+    }
+    llm_kwargs_resolved = {
+        **llm_kwargs,
+        **({"llm_use_lora": True, "llm_lora_r": best_params["llm_lora_r"]}
+           if (use_llm and "llm_lora_r" in best_params) else {}),
+    }
+
     model = create_model(
         d_model=best_params["d_model"],
         num_heads=best_params["num_heads"],
@@ -619,13 +729,21 @@ def main(
         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,
-        moe_expert_hidden_mult=moe_expert_hidden_mult,
-        moe_jitter_noise=moe_jitter_noise,
+        llm_kwargs=llm_kwargs_resolved,
+        use_retrieval=use_retrieval,
+        retr_feature_dim=(3 if use_retrieval else 0),
+        **arch,
     )
     model.rdkit_encoder._cache = rdkit_cache
 
+    # 全量训练阶段检索池用全量：_retrieve_topk 按 SMILES 精确匹配排除查询分子自己
+    # （含它的所有重复行），所以不存在抄自己标签的退化解。这也让训练期的池子规模
+    # 与服务期完全一致。
+    if use_rag and getattr(model, "llm_prompt", None) is not None:
+        _s, _d, _ex = _build_rag_pool(full_dataset, np.arange(len(full_dataset)))
+        model.llm_prompt.set_retrieval_pool(_s, _d, pool_id="full", extra_labels=_ex)
+        logger.info(f"[RAG] 检索池={len(_s)} 全量分子")
+
     # 加载预训练权重
     if pretrain_state_dict is not None and pretrain_config is not None:
         loaded = load_pretrain_weights_to_model(
@@ -635,18 +753,16 @@ def main(
         if loaded:
             logger.info("Loaded pretrain weights for final training")
 
-    # 打印模型信息
     n_params_total = sum(p.numel() for p in model.parameters())
     n_params_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
     logger.info(f"Model parameters: {n_params_total:,} total, {n_params_trainable:,} trainable")
 
-    # 训练（固定 epoch，不 early-stop）
     swa_start = int(epoch_mean * swa_start_ratio) if use_swa else None
 
     train_result = train_fixed_epochs(
         model=model,
         train_loader=full_loader,
-        val_loader=None,  # 全量训练，无验证集
+        val_loader=None,
         device=device,
         lr=best_params["lr"],
         weight_decay=best_params["weight_decay"],
@@ -656,12 +772,12 @@ def main(
         use_swa=use_swa,
         swa_start_epoch=swa_start,
         backbone_lr_ratio=best_params.get("backbone_lr_ratio", 1.0),
+        freeze_backbone_epochs=freeze_backbone_epochs,
     )
 
-    # 加载最终权重
-    model.load_state_dict(train_result["final_state"])
+    # QLoRA 4-bit 基座的量化元数据键不在 final_state 里，必须 strict=False
+    model.load_state_dict(train_result["final_state"], strict=False)
 
-    # 保存模型
     config = {
         "d_model": best_params["d_model"],
         "num_heads": best_params["num_heads"],
@@ -675,14 +791,23 @@ def main(
         "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,
-        "moe_expert_hidden_mult": moe_expert_hidden_mult,
-        "moe_jitter_noise": moe_jitter_noise,
+        "use_retrieval": use_retrieval,
+        "retr_feature_dim": (3 if use_retrieval else 0),
+        **arch,
+        **llm_kwargs_resolved,
     }
 
+    # 只丢 LLM 基座权重（几 GB，加载时从 llm_model_path 现读），LoRA 适配器必须留下
+    _full_state = train_result["final_state"]
+    _slim_state = {
+        k: v for k, v in _full_state.items()
+        if (not k.startswith("llm_prompt.encoder")) or ("lora_" in k)
+    }
+    _n_lora = sum(1 for k in _slim_state if "lora_" in k)
+    logger.info(f"保存 checkpoint: {len(_slim_state)} 个张量，其中 LoRA 适配器 {_n_lora} 个")
+
     torch.save({
-        "model_state_dict": train_result["final_state"],
+        "model_state_dict": _slim_state,
         "config": config,
         "best_params": best_params,
         "epoch_mean": epoch_mean,

lnp_ml/modeling/heads.py

@@ -80,7 +80,7 @@ class MultiTaskHead(nn.Module):
         size_dropout = min(0.5, dropout + 0.2)
         self.size_head = RegressionHead(in_dim, hidden_dim, size_dropout)
 
-        # PDI: 2 分类
+        # PDI: 2 分类（dataset.py 已将 4 段 one-hot 折叠为 pdi_4 >= 1）
         self.pdi_head = ClassificationHead(in_dim, num_classes=2, hidden_dim=hidden_dim, dropout=dropout)
 
         # Encapsulation Efficiency: 3 分类

lnp_ml/modeling/layers/fusion.py

@@ -114,8 +114,10 @@ class FusionLayer(nn.Module):
 class ResidualConcatFusion(nn.Module):
     """对真实 token 做 attention pooling，再用零初始化门把 MoE/LLM 旁路以残差方式加入。
 
-    g_moe / g_llm 初始为 0 → +moe/+llm 起点严格等于 baseline；
-    旁路只有确实有用时才会被训练打开，从机制上保证“加了不会更差”。
+    门为逐维向量（初始全零）：
+      - 起点仍严格等于 baseline，"加了不会更差"的保证不变；
+      - 标量门只能整体调音量，最优值是"有用维度的收益"与"噪声维度的损害"之间的妥协；
+      - 标量门的梯度 ∂L/∂g = Σ_i u_i f_i 会自我抵消，逐维门 ∂L/∂g_i = u_i f_i 不会。
     """
 
     def __init__(self, d_model: int, strategy: PoolingStrategy = "attention") -> None:
@@ -125,10 +127,18 @@ class ResidualConcatFusion(nn.Module):
         self.d_model = d_model
         self.pool = FusionLayer(d_model=d_model, n_tokens=1, strategy=strategy)
         self.fusion_dim = self.pool.fusion_dim
-        # 零初始化门控（可学习标量），旁路初始不参与
-        self.g_moe = nn.Parameter(torch.zeros(()))
-        self.g_llm = nn.Parameter(torch.zeros(()))
-        self.g_retr = nn.Parameter(torch.zeros(()))  # 检索旁路零初始化门控
+        # 零初始化逐维门控，旁路初始不参与
+        self.g_moe = nn.Parameter(torch.zeros(d_model))
+        self.g_llm = nn.Parameter(torch.zeros(d_model))
+        self.g_retr = nn.Parameter(torch.zeros(d_model))
+
+    def gate_norms(self) -> Dict[str, float]:
+        """诊断用。标量门时代的 |g| 对应这里的 norm / sqrt(d_model)。"""
+        return {
+            "g_moe": float(self.g_moe.norm()),
+            "g_llm": float(self.g_llm.norm()),
+            "g_retr": float(self.g_retr.norm()),
+        }
 
     def forward(
         self,
@@ -145,13 +155,14 @@ class ResidualConcatFusion(nn.Module):
         if return_attn_weights:
             pooled, attn = pooled
 
+        # [d_model] 与 [B, d_model] 自动广播
         out = pooled
         if f_moe is not None:
-            out = out + self.g_moe * f_moe  # 残差 + 零初始化门
+            out = out + self.g_moe * f_moe
         if f_llm is not None:
             out = out + self.g_llm * f_llm
         if f_retr is not None:
-            out = out + self.g_retr * f_retr  # 检索旁路，零初始化门保证起点=不开
+            out = out + self.g_retr * f_retr
 
         # 回归旁路：额外返回 pooled（纯 chem+tab，不含 f_llm/f_moe/f_retr）
         if return_attn_weights:

lnp_ml/modeling/layers/llm_prompt.py

@@ -12,6 +12,15 @@ DEFAULT_MOLT5_PATH = os.environ.get("MOLT5_PATH", "models/molt5-base")
 # 检索池支持的额外多任务（除 delivery 外）
 _EXTRA_TASKS = ["size", "pdi", "ee", "toxic", "biodist"]
 
+# 连续量的分位桶数。边界只从训练检索池统计，不引入测试集分布。
+_N_QBINS = 5
+
+# 离散标签的语义，避免 LLM 只看到裸整数
+_PDI_LABELS = {0: "<0.2", 1: ">=0.2"}
+_EE_LABELS = {0: "<50%", 1: "50-80%", 2: ">80%"}
+_TOX_LABELS = {0: "non-toxic", 1: "toxic"}
+_ORGANS = ["lymph_nodes", "heart", "liver", "spleen", "lung", "kidney", "muscle"]
+
 
 class LLMPromptEncoder(nn.Module):
     """用 LLM 编码分子，输出 F_llm [B, d_model]。
@@ -35,7 +44,7 @@ class LLMPromptEncoder(nn.Module):
         lora_r: int = 8,
         lora_alpha: int = 16,
         lora_dropout: float = 0.05,
-        max_length: int = 256,
+        max_length: int = 1536,
         use_rag: bool = False,
         rag_top_k: int = 4,
         use_soft_prompt: bool = False,
@@ -64,6 +73,9 @@ class LLMPromptEncoder(nn.Module):
         )
         if _is_qwen and self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
+        if use_soft_prompt:
+            # soft token 后置要求文本右对齐，否则 soft 会落在 PAD 之后
+            self.tokenizer.padding_side = "left"
 
         if _is_t5:
             self.encoder = T5EncoderModel.from_pretrained(model_name_or_path)
@@ -128,6 +140,7 @@ class LLMPromptEncoder(nn.Module):
         self._rag_pool_extra = None           # dict: task -> (values, valid)
         self._rag_pool_fps = None
         self._rag_pool_id: str = "none"
+        self._rag_pool_qedges: Dict[str, List[float]] = {}
 
     def _apply_lora(self, r, alpha, dropout, prepare_kbit=False):
         from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
@@ -167,6 +180,19 @@ class LLMPromptEncoder(nn.Module):
         self._rag_pool_labels = np.asarray(labels, dtype=np.float32).reshape(-1)
         self._rag_pool_extra = extra_labels
         self._rag_pool_fps = [_smiles_to_fp(s) for s in self._rag_pool_smiles]
+
+        self._rag_pool_qedges = {}
+        _probs = np.linspace(0.0, 1.0, _N_QBINS + 1)[1:-1]
+        _cont = {"delivery": self._rag_pool_labels}
+        if extra_labels is not None and "size" in extra_labels:
+            _v, _ok = extra_labels["size"]
+            _cont["size"] = np.asarray(_v, dtype=np.float32)[np.asarray(_ok, dtype=bool)]
+        for _name, _vals in _cont.items():
+            _vals = np.asarray(_vals, dtype=np.float32)
+            _vals = _vals[np.isfinite(_vals)]
+            if _vals.size >= _N_QBINS:
+                self._rag_pool_qedges[_name] = np.quantile(_vals, _probs).tolist()
+
         if pool_id != self._rag_pool_id:
             self._cache = {k: v for k, v in self._cache.items() if not k.startswith("RAG::")}
             self._prompt_cache.clear()
@@ -213,6 +239,21 @@ class LLMPromptEncoder(nn.Module):
             return "[" + ", ".join(f"{x:.3f}" for x in v) + "]"
         return f"{float(v):.3f}"
 
+    def _qbin(self, name: str, v) -> str:
+        """v 落在训练池经验分布的第几桶，如 '(Q3/5)'；无边界或缺失时返回空串。"""
+        import bisect
+        edges = self._rag_pool_qedges.get(name)
+        if v is None or not edges:
+            return ""
+        return f"(Q{bisect.bisect_right(edges, float(v)) + 1}/{len(edges) + 1})"
+
+    @staticmethod
+    def _fmt_class(v, labels: Dict[int, str]) -> str:
+        """离散标签附带语义，例如 '1(>=0.2)'。"""
+        if v is None:
+            return "unknown"
+        return f"{int(v)}({labels.get(int(v), '?')})"
+
     def _fmt_mol(self, smiles: str) -> str:
         """按 backbone 期望格式化分子。
         BioT5：SMILES -> SELFIES，用 <bom>...<eom> 紧贴包裹（官方格式，token 间无空格）；
@@ -227,36 +268,54 @@ class LLMPromptEncoder(nn.Module):
         return f"<bom>{sfs}<eom>"
         
     def _build_rag_prompt(self, target_smiles: str, neighbors) -> str:
-        """构造 RAG prompt：原始 SMILES + 邻居多任务结果（numeric）。"""
+        """构造 RAG prompt：原始 SMILES + 邻居多任务结果（数值 + 训练池分位桶）。"""
         blocks = []
         for rank, nb in enumerate(neighbors, 1):
             ex = nb["extra"]
+            bio = ex.get("biodist")
+            if bio is not None:
+                _top = max(range(len(bio)), key=lambda i: bio[i])
+                bio_s = (f"{_ORGANS[_top]}-dominant ["
+                         + ",".join(f"{x:.2f}" for x in bio) + "]")
+            else:
+                bio_s = "unknown"
             blocks.append(
-                f"Retrieved sample {rank}:\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"
-                f"pdi_class: {self._fmt(ex.get('pdi'), 'int')}\n"
-                f"ee_class: {self._fmt(ex.get('ee'), 'int')}\n"
-                f"toxic: {self._fmt(ex.get('toxic'), 'int')}\n"
-                f"biodist: {self._fmt(ex.get('biodist'), 'vec')}"
+                f"#{rank} sim={nb['sim']:.2f} {self._fmt_mol(nb['smiles'])} "
+                f"deliv={self._fmt(nb['delivery'])}{self._qbin('delivery', nb['delivery'])} "
+                f"size={self._fmt(ex.get('size'))}{self._qbin('size', ex.get('size'))} "
+                f"pdi={self._fmt_class(ex.get('pdi'), _PDI_LABELS)} "
+                f"ee={self._fmt_class(ex.get('ee'), _EE_LABELS)} "
+                f"tox={self._fmt_class(ex.get('toxic'), _TOX_LABELS)} bio={bio_s}"
             )
-        retrieved_block = "\n\n".join(blocks) if blocks else "(no retrieved samples)"
+        retrieved_block = "\n".join(blocks) if blocks else "(none)"
         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]\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; "
-            "'unknown' means the measurement is missing):\n\n"
-            f"{retrieved_block}\n\n"
-            "[Encoding Instructions]\n"
-            "Capture the target structure and the retrieval evidence (structural similarity "
-            "and consistency of retrieved outcomes) into your internal representation."
+            "Encode the target LNP molecule into a retrieval-aware representation "
+            "for multi-task property prediction. Do not output predictions.\n"
+            f"[Target] {self._fmt_mol(target_smiles)}\n"
+            "[Retrieved] Nearest training molecules by fingerprint similarity, with "
+            "known outcomes. deliv=delivery_log and size=size_z are raw values, each "
+            f"followed by (Qi/{_N_QBINS}) = which {_N_QBINS}-quantile bin it falls into "
+            "among training molecules, Q1=lowest. pdi/ee/tox give the class index with "
+            "its meaning. bio=fraction in [lymph_nodes,heart,liver,spleen,lung,kidney,"
+            "muscle] with the dominant organ named. 'unknown' = missing:\n"
+            f"{retrieved_block}\n"
+            "[Instruction] Capture the target structure and the retrieval evidence "
+            "(similarity, and both the level and the consistency of retrieved outcomes) "
+            "into your representation."
         )
 
+    def _warn_if_truncated(self, enc) -> None:
+        if getattr(self, "_trunc_warned", False):
+            return
+        if int(enc["attention_mask"].sum(1).max()) >= self.max_length:
+            from loguru import logger
+            logger.warning(
+                f"RAG prompt 触达 max_length={self.max_length} 被截断："
+                f"尾部邻居与 [Instruction] 丢失，读出位置偏移。"
+                f"请调大 max_length 或减小 rag_top_k（当前 {self.rag_top_k}）。"
+            )
+            self._trunc_warned = True
+
     def _get_prompt(self, s: str) -> str:
         if not self.use_rag:
             return self._fmt_mol(s)
@@ -273,6 +332,7 @@ class LLMPromptEncoder(nn.Module):
                 bt, padding=True, truncation=True,
                 max_length=self.max_length, return_tensors="pt",
             ).to(device)
+            self._warn_if_truncated(enc)
             out = self.encoder(**enc).last_hidden_state          # [B,L,H]
             lengths = enc["attention_mask"].sum(1) - 1
             b = torch.arange(out.size(0), device=device)
@@ -318,18 +378,33 @@ class LLMPromptEncoder(nn.Module):
 
         if soft_list:
             soft = torch.cat(soft_list, dim=1).to(text_embeds.dtype)
-            inputs_embeds = torch.cat([soft, text_embeds], dim=1)
+            # soft 后置：因果注意力下文本表示不受右侧影响，文本段因此只依赖 SMILES，
+            # 对同一分子的所有配方保持不变（可缓存）；读出点落在最后一个 soft token，
+            # 它同时看到全部文本与全部 soft，整合能力不损失。
+            inputs_embeds = torch.cat([text_embeds, soft], dim=1)
             soft_mask = torch.ones(soft.size(0), soft.size(1),
                                    device=device, dtype=text_mask.dtype)
-            attn_mask = torch.cat([soft_mask, text_mask], dim=1)
+            attn_mask = torch.cat([text_mask, soft_mask], dim=1)
         else:
             inputs_embeds = text_embeds
             attn_mask = text_mask
 
-        out = self.encoder(inputs_embeds=inputs_embeds, attention_mask=attn_mask).last_hidden_state
+        # 左填充下位置编码必须由 mask 推出，否则 PAD 会把真实 token 的位置顶偏
+        position_ids = attn_mask.long().cumsum(-1) - 1
+        position_ids = position_ids.masked_fill(attn_mask == 0, 1)
+
+        out = self.encoder(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attn_mask,
+            position_ids=position_ids,
+        ).last_hidden_state
+
+        if soft_list:
+            feat = out[:, -1, :]                        # 后置时末位必为最后一个 soft token
+        else:
             lengths = attn_mask.sum(1) - 1
             b = torch.arange(out.size(0), device=device)
-        feat = out[b, lengths.long(), :]                # [B, H] 最后有效 token
+            feat = out[b, lengths.long(), :]
         return self.proj_down(feat.float())
 
     # ---------- 旧路径（保留，向后兼容）----------

lnp_ml/modeling/models.py

@@ -20,7 +20,7 @@ from lnp_ml.modeling.layers import (
     LLMPromptEncoder,
 )
 from lnp_ml.modeling.layers.llm_prompt import DEFAULT_MOLT5_PATH
-from lnp_ml.modeling.heads import MultiTaskHead
+from lnp_ml.modeling.heads import MultiTaskHead, RegressionHead
 
 
 PoolingStrategy = Literal["attention", "avg", "max"]
@@ -128,6 +128,7 @@ class LNPModel(nn.Module):
         llm_lora_dropout: float = 0.05,
         use_rag: bool = False,
         rag_top_k: int = 4,
+        llm_max_length: int = 1536,
         use_retrieval: bool = False,
         retr_feature_dim: int = 0,
     ) -> None:
@@ -246,6 +247,7 @@ class LNPModel(nn.Module):
                 lora_dropout=llm_lora_dropout,
                 use_rag=use_rag,
                 rag_top_k=rag_top_k,
+                max_length=llm_max_length,
                 use_soft_prompt=use_soft_prompt,
             )
         else:
@@ -269,6 +271,14 @@ class LNPModel(nn.Module):
             dropout=dropout,
         )
 
+        # ============ 分支辅助头（仅训练期参与 loss，不进 forward）============
+        # f_moe / f_llm 要经过零初始化门 g 才进入主路径，g=0 时这两条分支梯度为零。
+        # 辅助头提供一条绕开 g 的通路：先把分支练成对 delivery 有判别力的方向，
+        # ⟨u, f⟩ 从随机噪声变成明确的正数后，g 才拿得到强梯度。
+        # 它不参与 forward()，所以 g=0 时主路径输出仍逐比特等于 baseline。
+        self.aux_head_moe = RegressionHead(d_model, head_hidden_dim, dropout) if use_moe else None
+        self.aux_head_llm = RegressionHead(d_model, head_hidden_dim, dropout) if use_llm else None
+
     def _encode_and_project(
         self,
         smiles: List[str],
@@ -324,7 +334,8 @@ class LNPModel(nn.Module):
             self._last_moe_extras = None
 
         f_llm = None
-        if self.llm_prompt is not None and smiles is not None:
+        if (self.llm_prompt is not None and smiles is not None
+                and getattr(self, "_llm_enabled", True)):
             if getattr(self.llm_prompt, "use_soft_prompt", False):
                 f_llm = self.llm_prompt(smiles, chem=chem, tab=tab)
             else:
@@ -336,10 +347,27 @@ class LNPModel(nn.Module):
             _feats_t = torch.as_tensor(_feats, dtype=chem.dtype, device=chem.device)
             f_retr = self.retr_proj(_feats_t)
 
+        # 供辅助监督使用（训练期），推理期不消费
+        self._last_f_moe = f_moe
+        self._last_f_llm = f_llm
+
         fused, pooled = self.fusion(chem, tab, f_moe=f_moe, f_llm=f_llm, f_retr=f_retr)
         self._last_pooled = pooled  # 纯数值向量，供回归 head 使用
         return fused
 
+    def get_aux_outputs(self) -> Dict[str, torch.Tensor]:
+        """返回各旁路分支对 delivery 的辅助预测。仅训练期由 loss 消费。"""
+        out: Dict[str, torch.Tensor] = {}
+        if not self.training:
+            return out
+        f_moe = getattr(self, "_last_f_moe", None)
+        f_llm = getattr(self, "_last_f_llm", None)
+        if getattr(self, "aux_head_moe", None) is not None and f_moe is not None:
+            out["aux_moe"] = self.aux_head_moe(f_moe)
+        if getattr(self, "aux_head_llm", None) is not None and f_llm is not None:
+            out["aux_llm"] = self.aux_head_llm(f_llm)
+        return out
+
     def forward_from_projected(
         self,
         stacked: torch.Tensor,
@@ -431,12 +459,20 @@ class LNPModel(nn.Module):
         完整的多任务 forward。
 
         Returns:
-            Dict[str, Tensor]: size [B,1], pdi [B,4], ee [B,3],
+            Dict[str, Tensor]: size [B,1], pdi [B,2], ee [B,3],
             delivery [B,1], biodist [B,7], toxic [B,2]
         """
         fused = self.forward_backbone(smiles, tabular)
         return self.head(fused, getattr(self, "_last_pooled", None) if getattr(self, "reg_bypass", True) else None)
 
+    def set_llm_enabled(self, enabled: bool) -> None:
+        """开关 LLM 旁路，供两阶段推理在粗筛阶段跳过 7B 前向。
+
+        forward 用 getattr(self, "_llm_enabled", True) 读取，因此没调用过
+        本方法的实例默认开启，与改动前行为一致。
+        """
+        self._llm_enabled = bool(enabled)
+
     def clear_cache(self) -> None:
         """清空所有 encoder 的缓存"""
         self.rdkit_encoder.clear_cache()
@@ -492,13 +528,22 @@ class LNPModel(nn.Module):
         }
 
         unexpected = []
+        loaded = []
         model_state = self.state_dict()
         for k, v in filtered_state_dict.items():
             if k in model_state and model_state[k].shape == v.shape:
                 model_state[k] = v
+                loaded.append(k)
             else:
                 unexpected.append(k)
 
+        from loguru import logger
+        logger.info(
+            f"预训练权重：迁移 {len(loaded)} 个张量"
+            f"（{sum(model_state[k].numel() for k in loaded) / 1e6:.2f}M 参数），"
+            f"跳过 {len(unexpected)} 个"
+            + (f"，例如 {unexpected[:3]}" if unexpected else "")
+        )
 
         self.load_state_dict(model_state, strict=False)
 
@@ -543,6 +588,7 @@ class LNPModelWithoutMPNN(LNPModel):
         llm_lora_dropout: float = 0.05,
         use_rag: bool = False,
         rag_top_k: int = 4,
+        llm_max_length: int = 1536,
         use_retrieval: bool = False,
         retr_feature_dim: int = 0,
     ) -> None:
@@ -573,6 +619,7 @@ class LNPModelWithoutMPNN(LNPModel):
             use_llm=use_llm,
             use_rag=use_rag,
             rag_top_k=rag_top_k,
+            llm_max_length=llm_max_length,
             use_retrieval=use_retrieval,
             retr_feature_dim=retr_feature_dim,
             llm_model_path=llm_model_path,

lnp_ml/modeling/nested_cv_optuna.py

@@ -804,6 +804,19 @@ def _run_single_outer_fold(
 
     class_weights = compute_class_weights_from_loader(train_loader)
 
+    arch = {
+        "moe_n_experts": best_params.get("moe_n_experts", moe_n_experts),
+        "moe_top_k": best_params.get("moe_top_k", moe_top_k),
+        "moe_expert_hidden_mult": best_params.get("moe_expert_hidden_mult", moe_expert_hidden_mult),
+        "moe_jitter_noise": moe_jitter_noise,
+        "set_transformer_block": best_params.get("set_transformer_block", "sab"),
+    }
+    llm_kwargs_resolved = {
+        **(llm_kwargs or {}),
+        **({"llm_use_lora": True, "llm_lora_r": best_params["llm_lora_r"]}
+           if ((llm_kwargs or {}).get("use_llm", False) and "llm_lora_r" in best_params) else {}),
+    }
+
     model = create_model(
         d_model=best_params["d_model"],
         num_heads=best_params["num_heads"],
@@ -818,15 +831,10 @@ def _run_single_outer_fold(
         moleculestm_cache=moleculestm_cache,
         mole_cache=mole_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),
-        moe_expert_hidden_mult=best_params.get("moe_expert_hidden_mult", moe_expert_hidden_mult),
-        moe_jitter_noise=moe_jitter_noise,
-        set_transformer_block=best_params.get("set_transformer_block", "sab"),
-        llm_kwargs={**llm_kwargs, **({"llm_use_lora": True, "llm_lora_r": best_params["llm_lora_r"]}
-                                     if (llm_kwargs.get("use_llm", False) and "llm_lora_r" in best_params) else {})},
+        llm_kwargs=llm_kwargs_resolved,
         use_retrieval=use_retrieval,
         retr_feature_dim=(3 if use_retrieval else 0),
+        **arch,
     )
     model.rdkit_encoder._cache = rdkit_cache
     if use_retrieval:
@@ -878,7 +886,6 @@ def _run_single_outer_fold(
         "n_attn_layers": best_params["n_attn_layers"],
         "fusion_strategy": best_params["fusion_strategy"],
         "head_hidden_dim": best_params["head_hidden_dim"],
-        "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,
@@ -890,16 +897,19 @@ def _run_single_outer_fold(
         "use_mole": mole_cache is not None,
         "mole_cache": mole_cache,
         "use_moe": use_moe,
-        "moe_n_experts": moe_n_experts,
-        "moe_top_k": moe_top_k,
-        "moe_expert_hidden_mult": moe_expert_hidden_mult,
-        "moe_jitter_noise": moe_jitter_noise,
-        **(llm_kwargs or {}),
+        "use_retrieval": use_retrieval,
+        "retr_feature_dim": (3 if use_retrieval else 0),
+        **arch,
+        **llm_kwargs_resolved,
     }
 
     _full_state = train_result["final_state"]
-    _slim_state = {k: v for k, v in _full_state.items()
-                   if not k.startswith("llm_prompt.encoder")}
+    _slim_state = {
+        k: v for k, v in _full_state.items()
+        if (not k.startswith("llm_prompt.encoder")) or ("lora_" in k)
+    }
+    _n_lora = sum(1 for k in _slim_state if "lora_" in k)
+    logger.info(f"保存 checkpoint: {len(_slim_state)} 个张量，其中 LoRA 适配器 {_n_lora} 个")
     torch.save({
         "model_state_dict": _slim_state,
         "config": config,
@@ -981,6 +991,7 @@ def main(
     use_llm: bool = False,
     use_rag: bool = False,
     rag_top_k: int = 4,
+    llm_max_length: int = 1536,
     use_retrieval: bool = False,
     retrieval_source: str = "internal",
     llm_model_path: str = DEFAULT_MOLT5_PATH,
@@ -991,6 +1002,8 @@ def main(
     llm_lora_r: int = 8,
     llm_lora_alpha: int = 16,
     llm_lora_dropout: float = 0.05,
+    fix_hparams_json: Optional[Path] = None,
+    fix_epoch_mean: int = 12,
     # 并行
     parallel: bool = False,
     # 设备
@@ -1017,6 +1030,7 @@ def main(
         reg_bypass=reg_bypass,
         use_rag=use_rag,
         rag_top_k=rag_top_k,
+        llm_max_length=llm_max_length,
         use_llm=use_llm,
         llm_model_path=llm_model_path,
         llm_freeze=llm_freeze,
@@ -1028,7 +1042,17 @@ def main(
         llm_lora_dropout=llm_lora_dropout,
     )
 
-    # 加载预训练权重（如果指定）
+    _fixed_bp = None
+    _fixed_em = None
+    if fix_hparams_json is not None:
+        _fixed_bp = json.loads(Path(fix_hparams_json).read_text())
+        _fixed_em = int(fix_epoch_mean)
+        if use_moe:
+            _fixed_bp["moe_n_experts"] = moe_n_experts
+            _fixed_bp["moe_top_k"] = moe_top_k
+        logger.info(f"[FIX] 固定超参={_fixed_bp}, epoch_mean={_fixed_em}, 跳过内层 Optuna")
+
+    # 加载预训练权重
     pretrain_state_dict = None
     pretrain_config = None
     if init_from_pretrain is not None:
@@ -1113,6 +1137,8 @@ def main(
             moe_jitter_noise=moe_jitter_noise,
             set_transformer_block=set_transformer_block,
             llm_kwargs=llm_kwargs,
+            precomputed_best_params=_fixed_bp,
+            precomputed_epoch_mean=_fixed_em,
         ))
     
     if parallel:

lnp_ml/modeling/predict.py

@@ -36,57 +36,78 @@ def load_model(
     """
     加载训练好的模型。
 
-    自动根据 checkpoint 的 config.use_mpnn 选择模型类型。
+    根据 checkpoint 的 config 决定模型类型与全部结构开关
+    (MoE / LLM / RAG / soft-prompt / reg_bypass)。
     """
     checkpoint = torch.load(model_path, map_location=device, weights_only=False)
     config = checkpoint["config"]
     use_mpnn = config.get("use_mpnn", False)
 
+    common_kwargs = dict(
+        d_model=config["d_model"],
+        num_heads=config["num_heads"],
+        n_attn_layers=config["n_attn_layers"],
+        set_transformer_block=config.get("set_transformer_block", "sab"),
+        fusion_strategy=config["fusion_strategy"],
+        head_hidden_dim=config["head_hidden_dim"],
+        dropout=config["dropout"],
+        reg_bypass=config.get("reg_bypass", "on"),
+        use_moe=config.get("use_moe", False),
+        moe_n_experts=config.get("moe_n_experts", 4),
+        moe_top_k=config.get("moe_top_k", 2),
+        moe_expert_hidden_mult=config.get("moe_expert_hidden_mult", 2),
+        moe_jitter_noise=config.get("moe_jitter_noise", 0.0),
+        use_llm=config.get("use_llm", False),
+        llm_model_path=config.get("llm_model_path", "models/molt5-base"),
+        llm_freeze=config.get("llm_freeze", True),
+        llm_use_lora=config.get("llm_use_lora", False),
+        llm_use_qlora=config.get("llm_use_qlora", False),
+        use_soft_prompt=config.get("use_soft_prompt", False),
+        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),
+        use_rag=config.get("use_rag", False),
+        rag_top_k=config.get("rag_top_k", 4),
+        llm_max_length=config.get("llm_max_length", 256),
+        use_retrieval=config.get("use_retrieval", False),
+        retr_feature_dim=config.get("retr_feature_dim", 0),
+        chemeleon_cache_path=config.get("chemeleon_cache"),
+        unimol_cache_path=config.get("unimol_cache"),
+    )
+
     if use_mpnn:
         # 总是自动查找 MPNN ensemble，避免使用 checkpoint 中的旧绝对路径（可能来自其他机器）
         logger.info("Model was trained with MPNN, auto-detecting ensemble...")
         ensemble_paths = find_mpnn_ensemble_paths()
         logger.info(f"Found {len(ensemble_paths)} MPNN models")
-        
         model = LNPModel(
-            d_model=config["d_model"],
-            num_heads=config["num_heads"],
-            n_attn_layers=config["n_attn_layers"],
-            fusion_strategy=config["fusion_strategy"],
-            head_hidden_dim=config["head_hidden_dim"],
-            dropout=config["dropout"],
-            use_llm=config.get("use_llm", False),
-            llm_model_path=config.get("llm_model_path", "models/molt5-base"),
-            llm_freeze=config.get("llm_freeze", True),
-            llm_use_lora=config.get("llm_use_lora", False),
-            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),
             mpnn_ensemble_paths=ensemble_paths,
             mpnn_device=mpnn_device,
-            chemeleon_cache_path=config.get("chemeleon_cache"),
-            unimol_cache_path=config.get("unimol_cache"),
+            **common_kwargs,
         )
     else:
-        model = LNPModelWithoutMPNN(
-            d_model=config["d_model"],
-            num_heads=config["num_heads"],
-            n_attn_layers=config["n_attn_layers"],
-            fusion_strategy=config["fusion_strategy"],
-            head_hidden_dim=config["head_hidden_dim"],
-            dropout=config["dropout"],
-            use_llm=config.get("use_llm", False),
-            llm_model_path=config.get("llm_model_path", "models/molt5-base"),
-            llm_freeze=config.get("llm_freeze", True),
-            llm_use_lora=config.get("llm_use_lora", False),
-            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 = LNPModelWithoutMPNN(**common_kwargs)
+
+    # 兼容旧 checkpoint：fusion 门从标量升级为逐维向量后，需把 0 维张量展开
+    _sd = checkpoint["model_state_dict"]
+    for _k in ("fusion.g_moe", "fusion.g_llm", "fusion.g_retr"):
+        if _k in _sd and _sd[_k].dim() == 0:
+            _sd[_k] = _sd[_k].reshape(1).expand(model.fusion.d_model).clone()
+
+    missing, unexpected = model.load_state_dict(
+        checkpoint["model_state_dict"], strict=False
+    )
+    # strict=False 不会因结构不匹配报错，这里手动兜底
+    if unexpected:
+        raise RuntimeError(
+            f"checkpoint 中有 {len(unexpected)} 个权重找不到对应模块，"
+            f"结构开关可能未对齐: {unexpected[:5]}"
+        )
+    if missing:
+        logger.warning(
+            f"{len(missing)} 个参数未从 checkpoint 恢复，将使用随机初始化: {missing[:5]}"
         )
 
-    model.load_state_dict(checkpoint["model_state_dict"], strict=False)
     model.to(device)
     model.eval()
 
@@ -152,7 +173,7 @@ def predictions_to_dataframe(predictions: Dict) -> pd.DataFrame:
     })
     
     # PDI 类别映射
-    pdi_labels = ["0_0to0_2", "0_2to0_3", "0_3to0_4", "0_4to0_5"]
+    pdi_labels = ["PDI<0.2", "PDI>=0.2"]
     df["pred_pdi_label"] = df["pred_pdi_class"].map(lambda x: pdi_labels[x])
     
     # EE 类别映射
@@ -311,10 +332,10 @@ def test(
                 "r2": float(r2_score(y_true, y_pred)),
             }
     
-    # 分类指标：PDI
+    # 分类指标：PDI（真值需与 dataset.py 同样折叠为二分类）
     pdi_cols = ["PDI_0_0to0_2", "PDI_0_2to0_3", "PDI_0_3to0_4", "PDI_0_4to0_5"]
     if all(c in test_df.columns for c in pdi_cols):
-        pdi_true = test_df[pdi_cols].values.argmax(axis=1)
+        pdi_true = (test_df[pdi_cols].values.argmax(axis=1) >= 1).astype(np.int64)
         mask = test_df[pdi_cols].sum(axis=1) > 0
         if mask.any():
             y_true = pdi_true[mask]

lnp_ml/modeling/trainer_balanced.py

@@ -1,7 +1,7 @@
 """带类权重的训练器：处理分类任务的数据不均衡问题"""
 
 from typing import Dict, List, Optional, Tuple
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 
 import numpy as np
 import torch
@@ -14,7 +14,7 @@ from tqdm import tqdm
 @dataclass
 class ClassWeights:
     """分类任务的类权重"""
-    pdi: Optional[torch.Tensor] = None  # [4] for 4 PDI classes
+    pdi: Optional[torch.Tensor] = None  # [2] 二分类 (0: PDI<0.2, 1: PDI>=0.2)
     ee: Optional[torch.Tensor] = None   # [3] for 3 EE classes
     toxic: Optional[torch.Tensor] = None  # [2] for binary toxic
 
@@ -29,6 +29,10 @@ class LossWeightsBalanced:
     biodist: float = 1.0
     toxic: float = 1.0
     moe_lb: float = 0.01  # MoE load-balancing 系数（仅在 use_moe=True 时生效）
+    # 旁路分支辅助监督系数。GoogLeNet 辅助分类器用的也是 0.3。
+    # 训练后期应退火到 0：辅助 loss 优化的是"f 单独能预测 delivery"，
+    # 而主 loss 要的是"f 对 pooled 有增量价值"，两者并不等价。
+    aux_branch: float = 0.3
 
 
 def compute_class_weights_from_loader(
@@ -188,6 +192,16 @@ def compute_multitask_loss_balanced(
             losses["moe_lb"] = extras["lb_loss"]
             total_loss = total_loss + task_weights.moe_lb * losses["moe_lb"]
 
+    # 旁路分支辅助监督：绕开零初始化门，直接用 f_moe / f_llm 预测 delivery。
+    # 不纳入不确定性加权——它是优化脚手架，不是一个真实任务。
+    if (model is not None and hasattr(model, "get_aux_outputs")
+            and task_weights.aux_branch > 0
+            and "delivery" in targets and mask["delivery"].any()):
+        m = mask["delivery"]
+        for name, pred in model.get_aux_outputs().items():
+            losses[name] = F.mse_loss(pred[m].squeeze(-1), targets["delivery"][m])
+            total_loss = total_loss + task_weights.aux_branch * losses[name]
+
     return total_loss, losses
 
 
@@ -202,7 +216,8 @@ def train_epoch_balanced(
     """带类权重的训练一个 epoch"""
     model.train()
     total_loss = 0.0
-    task_losses = {k: 0.0 for k in ["size", "pdi", "ee", "delivery", "biodist", "toxic", "moe_lb"]}
+    task_losses = {k: 0.0 for k in ["size", "pdi", "ee", "delivery", "biodist", "toxic",
+                                    "moe_lb", "aux_moe", "aux_llm"]}
     n_batches = 0
     
     for batch in tqdm(loader, desc="Training", leave=False):
@@ -394,10 +409,14 @@ def train_with_early_stopping(
     task_weights: Optional[LossWeightsBalanced] = None,
     class_weights: Optional[ClassWeights] = None,
     backbone_lr_ratio: float = 1.0,
+    freeze_backbone_epochs: int = 0,
 ) -> Dict:
     """
     带早停的完整训练流程。
 
+    freeze_backbone_epochs 必须与 train_fixed_epochs 取同一个值：best_epoch 是在
+    这个调度下测出来的，最终训练换了调度这个数就不可迁移。
+
     Returns:
         Dict with keys: history, best_val_loss, best_epoch, best_state
     """
@@ -408,11 +427,27 @@ def train_with_early_stopping(
     )
     early_stopping = EarlyStoppingBalanced(patience=patience, min_delta=1e-3)
 
+    backbone_named = [
+        (n, p) for n, p in model.named_parameters()
+        if n.startswith(BACKBONE_PREFIXES) and not n.startswith(FROM_SCRATCH_PREFIXES)
+        and p.requires_grad
+    ]
+    if freeze_backbone_epochs > 0:
+        for _, p in backbone_named:
+            p.requires_grad_(False)
+
+    task_weights = replace(task_weights) if task_weights is not None else LossWeightsBalanced()
+    _aux_w0 = task_weights.aux_branch
+
     history = {"train": [], "val": []}
     best_val_loss = float("inf")
     best_state = None
 
     for epoch in range(epochs):
+        task_weights.aux_branch = _aux_w0 * max(0.0, 1.0 - epoch / (0.5 * epochs))
+        if freeze_backbone_epochs > 0 and epoch == freeze_backbone_epochs:
+            for _, p in backbone_named:
+                p.requires_grad_(True)
         # Train
         train_metrics = train_epoch_balanced(
             model, train_loader, optimizer, device, task_weights, class_weights
@@ -516,9 +551,13 @@ def train_fixed_epochs(
         swa_start = swa_start_epoch or int(epochs * 0.75)
         swa_scheduler = SWALR(optimizer, swa_lr=lr * 0.1)
 
+    task_weights = replace(task_weights) if task_weights is not None else LossWeightsBalanced()
+    _aux_w0 = task_weights.aux_branch
+
     history = {"train": [], "val": []}
 
     for epoch in range(epochs):
+        task_weights.aux_branch = _aux_w0 * max(0.0, 1.0 - epoch / (0.5 * epochs))
         if freeze_backbone_epochs > 0 and epoch == freeze_backbone_epochs:
             for _, p in backbone_named:
                 p.requires_grad_(True)

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

@@ -1,42 +1,105 @@
 {
   "size": {
     "n_samples": 83,
-    "mse": 1.368858521000182,
-    "rmse": 1.1699822737974204,
-    "mae": 0.48265016079386586,
-    "r2": 0.16222489685168928
+    "mse": 1.3484356700772522,
+    "rmse": 1.1612216283196124,
+    "mae": 0.4798589242061937,
+    "r2": 0.1747241783006852
   },
   "delivery": {
     "n_samples": 58,
-    "mse": 0.4106486248647463,
-    "rmse": 0.6408187145088275,
-    "mae": 0.4092221012826776,
-    "r2": 0.4780408138229023
+    "mse": 0.3976989005700695,
+    "rmse": 0.6306337293311146,
+    "mae": 0.4073375633664443,
+    "r2": 0.49450069866587976
   },
   "toxic": {
     "n_samples": 58,
-    "accuracy": 1.0,
-    "precision": 1.0,
-    "recall": 1.0,
-    "f1": 1.0
+    "accuracy": 0.9655172413793104,
+    "precision": 0.75,
+    "recall": 0.9821428571428572,
+    "f1": 0.8242424242424242,
+    "true_class_counts": [
+      56,
+      2
+    ],
+    "pred_class_counts": [
+      54,
+      4
+    ],
+    "confusion_matrix": [
+      [
+        54,
+        2
+      ],
+      [
+        0,
+        2
+      ]
+    ]
   },
   "pdi": {
     "n_samples": 84,
-    "accuracy": 0.8095238095238095,
-    "precision": 0.7797101449275362,
-    "recall": 0.7063435495367071,
-    "f1": 0.7279352226720648
+    "accuracy": 0.7857142857142857,
+    "precision": 0.7551020408163265,
+    "recall": 0.8118317890235209,
+    "f1": 0.7630094043887148,
+    "true_class_counts": [
+      61,
+      23
+    ],
+    "pred_class_counts": [
+      49,
+      35
+    ],
+    "confusion_matrix": [
+      [
+        46,
+        15
+      ],
+      [
+        3,
+        20
+      ]
+    ]
   },
   "ee": {
     "n_samples": 84,
     "accuracy": 0.75,
-    "precision": 0.6912280701754385,
-    "recall": 0.6471428571428571,
-    "f1": 0.6631611379274931
+    "precision": 0.7092916445857623,
+    "recall": 0.7614285714285715,
+    "f1": 0.7175824990730441,
+    "true_class_counts": [
+      14,
+      20,
+      50
+    ],
+    "pred_class_counts": [
+      17,
+      30,
+      37
+    ],
+    "confusion_matrix": [
+      [
+        10,
+        3,
+        1
+      ],
+      [
+        3,
+        17,
+        0
+      ],
+      [
+        4,
+        10,
+        36
+      ]
+    ]
   },
   "biodist": {
     "n_samples": 58,
-    "kl_divergence": 0.16366098804686893,
-    "js_divergence": 0.0359968761663274
+    "kl_divergence": 0.15998570647485624,
+    "js_divergence": 0.035472610690609384
   }
 }

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

@@ -4,39 +4,102 @@
     "mse": 0.30493822479834903,
     "rmse": 0.5522121193874225,
     "mae": 0.34891908283267253,
-    "r2": 0.0688270762246912
+    "r2": 0.06882707391945686
   },
   "delivery": {
     "n_samples": 61,
-    "mse": 0.7776284206466513,
-    "rmse": 0.8818324220886026,
-    "mae": 0.519796750021438,
-    "r2": 0.4113913613868899
+    "mse": 0.7752937492958416,
+    "rmse": 0.8805076656655759,
+    "mae": 0.5151286462071908,
+    "r2": 0.4131585218812348
   },
   "toxic": {
     "n_samples": 61,
-    "accuracy": 1.0,
-    "precision": 1.0,
-    "recall": 1.0,
-    "f1": 1.0
+    "accuracy": 0.9508196721311475,
+    "precision": 0.75,
+    "recall": 0.9741379310344828,
+    "f1": 0.8200589970501475,
+    "true_class_counts": [
+      58,
+      3
+    ],
+    "pred_class_counts": [
+      55,
+      6
+    ],
+    "confusion_matrix": [
+      [
+        55,
+        3
+      ],
+      [
+        0,
+        3
+      ]
+    ]
   },
   "pdi": {
     "n_samples": 84,
-    "accuracy": 0.7619047619047619,
-    "precision": 0.6711111111111111,
-    "recall": 0.5873015873015873,
-    "f1": 0.5942028985507246
+    "accuracy": 0.6547619047619048,
+    "precision": 0.5803571428571429,
+    "recall": 0.5952380952380952,
+    "f1": 0.5822328931572629,
+    "true_class_counts": [
+      63,
+      21
+    ],
+    "pred_class_counts": [
+      56,
+      28
+    ],
+    "confusion_matrix": [
+      [
+        45,
+        18
+      ],
+      [
+        11,
+        10
+      ]
+    ]
   },
   "ee": {
     "n_samples": 84,
-    "accuracy": 0.6785714285714286,
-    "precision": 0.6062341913405743,
-    "recall": 0.6071785357499643,
-    "f1": 0.6041666666666666
+    "accuracy": 0.5476190476190477,
+    "precision": 0.5001165151969107,
+    "recall": 0.4989534275248561,
+    "f1": 0.48349530196256724,
+    "true_class_counts": [
+      13,
+      22,
+      49
+    ],
+    "pred_class_counts": [
+      19,
+      31,
+      34
+    ],
+    "confusion_matrix": [
+      [
+        5,
+        7,
+        1
+      ],
+      [
+        8,
+        11,
+        3
+      ],
+      [
+        6,
+        13,
+        30
+      ]
+    ]
   },
   "biodist": {
     "n_samples": 61,
     "kl_divergence": 0.1177531067102451,
-    "js_divergence": 0.027774748827162937
+    "js_divergence": 0.027774748827162944
   }
 }

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

@@ -1,42 +1,105 @@
 {
   "size": {
     "n_samples": 84,
-    "mse": 0.3089202253953697,
-    "rmse": 0.5558059242175902,
-    "mae": 0.3994691472061615,
-    "r2": 0.3898880719888205
+    "mse": 0.3091245824692921,
+    "rmse": 0.5559897323416073,
+    "mae": 0.4046506446335209,
+    "r2": 0.3894844727418765
   },
   "delivery": {
     "n_samples": 60,
-    "mse": 0.6215887874238653,
-    "rmse": 0.7884090229213928,
-    "mae": 0.4654098044336172,
-    "r2": 0.17255759426449502
+    "mse": 0.623471858388545,
+    "rmse": 0.789602341934562,
+    "mae": 0.4688382081680016,
+    "r2": 0.17005089391015027
   },
   "toxic": {
     "n_samples": 61,
-    "accuracy": 1.0,
-    "precision": 1.0,
-    "recall": 1.0,
-    "f1": 1.0
+    "accuracy": 0.9344262295081968,
+    "precision": 0.7142857142857143,
+    "recall": 0.9655172413793103,
+    "f1": 0.7821428571428571,
+    "true_class_counts": [
+      58,
+      3
+    ],
+    "pred_class_counts": [
+      54,
+      7
+    ],
+    "confusion_matrix": [
+      [
+        54,
+        4
+      ],
+      [
+        0,
+        3
+      ]
+    ]
   },
   "pdi": {
     "n_samples": 84,
-    "accuracy": 0.7261904761904762,
-    "precision": 0.6428571428571428,
-    "recall": 0.6385630498533724,
-    "f1": 0.6405581395348837
+    "accuracy": 0.6904761904761905,
+    "precision": 0.6693363844393593,
+    "recall": 0.717008797653959,
+    "f1": 0.662962962962963,
+    "true_class_counts": [
+      62,
+      22
+    ],
+    "pred_class_counts": [
+      46,
+      38
+    ],
+    "confusion_matrix": [
+      [
+        41,
+        21
+      ],
+      [
+        5,
+        17
+      ]
+    ]
   },
   "ee": {
     "n_samples": 84,
-    "accuracy": 0.6547619047619048,
-    "precision": 0.5676343636869953,
-    "recall": 0.5378151260504201,
-    "f1": 0.5491025070464323
+    "accuracy": 0.6428571428571429,
+    "precision": 0.5842305037957212,
+    "recall": 0.6349206349206349,
+    "f1": 0.5925185925185925,
+    "true_class_counts": [
+      12,
+      21,
+      51
+    ],
+    "pred_class_counts": [
+      21,
+      23,
+      40
+    ],
+    "confusion_matrix": [
+      [
+        8,
+        2,
+        2
+      ],
+      [
+        5,
+        12,
+        4
+      ],
+      [
+        8,
+        9,
+        34
+      ]
+    ]
   },
   "biodist": {
     "n_samples": 60,
-    "kl_divergence": 0.11869463263891707,
-    "js_divergence": 0.02929511833384903
+    "kl_divergence": 0.11799692057231032,
+    "js_divergence": 0.029125166530752162
   }
 }

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

@@ -1,42 +1,105 @@
 {
   "size": {
     "n_samples": 83,
-    "mse": 1.6827545971608378,
-    "rmse": 1.2972103133882484,
-    "mae": 0.5269114175759023,
-    "r2": 0.048221123148072476
+    "mse": 1.6732277890580922,
+    "rmse": 1.2935330645399414,
+    "mae": 0.5211672051452704,
+    "r2": 0.05360960932568781
   },
   "delivery": {
     "n_samples": 59,
-    "mse": 0.8527625246024667,
+    "mse": 0.8527625246024666,
     "rmse": 0.9234514197306032,
     "mae": 0.5525073854104011,
-    "r2": 0.1710134699931437
+    "r2": 0.17101350859831832
   },
   "toxic": {
     "n_samples": 60,
-    "accuracy": 1.0,
-    "precision": 1.0,
-    "recall": 1.0,
-    "f1": 1.0
+    "accuracy": 0.9333333333333333,
+    "precision": 0.7142857142857143,
+    "recall": 0.9649122807017544,
+    "f1": 0.7818181818181817,
+    "true_class_counts": [
+      57,
+      3
+    ],
+    "pred_class_counts": [
+      53,
+      7
+    ],
+    "confusion_matrix": [
+      [
+        53,
+        4
+      ],
+      [
+        0,
+        3
+      ]
+    ]
   },
   "pdi": {
     "n_samples": 84,
-    "accuracy": 0.75,
-    "precision": 0.6676909569798068,
-    "recall": 0.6400293255131965,
-    "f1": 0.6493738819320214
+    "accuracy": 0.7380952380952381,
+    "precision": 0.6612903225806452,
+    "recall": 0.6612903225806452,
+    "f1": 0.6612903225806452,
+    "true_class_counts": [
+      62,
+      22
+    ],
+    "pred_class_counts": [
+      62,
+      22
+    ],
+    "confusion_matrix": [
+      [
+        51,
+        11
+      ],
+      [
+        11,
+        11
+      ]
+    ]
   },
   "ee": {
     "n_samples": 84,
-    "accuracy": 0.6309523809523809,
-    "precision": 0.5132964178288214,
-    "recall": 0.49489693313222727,
-    "f1": 0.4975934917795383
+    "accuracy": 0.6666666666666666,
+    "precision": 0.5872039589961117,
+    "recall": 0.6381096028154852,
+    "f1": 0.5938808373590982,
+    "true_class_counts": [
+      13,
+      20,
+      51
+    ],
+    "pred_class_counts": [
+      23,
+      20,
+      41
+    ],
+    "confusion_matrix": [
+      [
+        10,
+        3,
+        0
+      ],
+      [
+        9,
+        8,
+        3
+      ],
+      [
+        4,
+        9,
+        38
+      ]
+    ]
   },
   "biodist": {
     "n_samples": 60,
-    "kl_divergence": 0.1458100103551393,
-    "js_divergence": 0.034398655034422854
+    "kl_divergence": 0.14886390502458902,
+    "js_divergence": 0.03503301624445018
   }
 }

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

@@ -1,42 +1,105 @@
 {
   "size": {
     "n_samples": 84,
-    "mse": 0.539048457673242,
-    "rmse": 0.7341991948192548,
-    "mae": 0.41399459761950447,
-    "r2": 0.27868624280784415
+    "mse": 0.5439429374081691,
+    "rmse": 0.7375248723996833,
+    "mae": 0.41855190159962563,
+    "r2": 0.2721367833807379
   },
   "delivery": {
     "n_samples": 58,
-    "mse": 0.8685551446715326,
-    "rmse": 0.9319630597140278,
-    "mae": 0.48785093074168856,
-    "r2": 0.16012146519578374
+    "mse": 0.8807361070482976,
+    "rmse": 0.9384754163260206,
+    "mae": 0.4923390269539427,
+    "r2": 0.14834272376574187
   },
   "toxic": {
     "n_samples": 59,
-    "accuracy": 1.0,
-    "precision": 1.0,
-    "recall": 1.0,
-    "f1": 1.0
+    "accuracy": 0.9661016949152542,
+    "precision": 0.8,
+    "recall": 0.9821428571428572,
+    "f1": 0.865909090909091,
+    "true_class_counts": [
+      56,
+      3
+    ],
+    "pred_class_counts": [
+      54,
+      5
+    ],
+    "confusion_matrix": [
+      [
+        54,
+        2
+      ],
+      [
+        0,
+        3
+      ]
+    ]
   },
   "pdi": {
     "n_samples": 84,
-    "accuracy": 0.7380952380952381,
-    "precision": 0.6515151515151515,
-    "recall": 0.6319648093841642,
-    "f1": 0.6390625
+    "accuracy": 0.6785714285714286,
+    "precision": 0.6428571428571428,
+    "recall": 0.6796187683284457,
+    "f1": 0.6415362731152205,
+    "true_class_counts": [
+      62,
+      22
+    ],
+    "pred_class_counts": [
+      49,
+      35
+    ],
+    "confusion_matrix": [
+      [
+        42,
+        20
+      ],
+      [
+        7,
+        15
+      ]
+    ]
   },
   "ee": {
     "n_samples": 84,
-    "accuracy": 0.6785714285714286,
-    "precision": 0.6048840048840048,
-    "recall": 0.6175213675213675,
-    "f1": 0.6107062011796588
+    "accuracy": 0.6309523809523809,
+    "precision": 0.5673525820584645,
+    "recall": 0.6021367521367521,
+    "f1": 0.5701075744179193,
+    "true_class_counts": [
+      12,
+      20,
+      52
+    ],
+    "pred_class_counts": [
+      17,
+      28,
+      39
+    ],
+    "confusion_matrix": [
+      [
+        7,
+        4,
+        1
+      ],
+      [
+        6,
+        11,
+        3
+      ],
+      [
+        4,
+        13,
+        35
+      ]
+    ]
   },
   "biodist": {
     "n_samples": 58,
-    "kl_divergence": 0.12253709805850461,
-    "js_divergence": 0.03003058068886421
+    "kl_divergence": 0.12072341813237816,
+    "js_divergence": 0.029660750070876794
   }
 }

models/abl_full/rf/seed42/summary.json

@@ -5,43 +5,106 @@
       "test_metrics": {
         "size": {
           "n_samples": 83,
-          "mse": 1.368858521000182,
-          "rmse": 1.1699822737974204,
-          "mae": 0.48265016079386586,
-          "r2": 0.16222489685168928
+          "mse": 1.3484356700772522,
+          "rmse": 1.1612216283196124,
+          "mae": 0.4798589242061937,
+          "r2": 0.1747241783006852
         },
         "delivery": {
           "n_samples": 58,
-          "mse": 0.4106486248647463,
-          "rmse": 0.6408187145088275,
-          "mae": 0.4092221012826776,
-          "r2": 0.4780408138229023
+          "mse": 0.3976989005700695,
+          "rmse": 0.6306337293311146,
+          "mae": 0.4073375633664443,
+          "r2": 0.49450069866587976
         },
         "toxic": {
           "n_samples": 58,
-          "accuracy": 1.0,
-          "precision": 1.0,
-          "recall": 1.0,
-          "f1": 1.0
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242,
+          "true_class_counts": [
+            56,
+            2
+          ],
+          "pred_class_counts": [
+            54,
+            4
+          ],
+          "confusion_matrix": [
+            [
+              54,
+              2
+            ],
+            [
+              0,
+              2
+            ]
+          ]
         },
         "pdi": {
           "n_samples": 84,
-          "accuracy": 0.8095238095238095,
-          "precision": 0.7797101449275362,
-          "recall": 0.7063435495367071,
-          "f1": 0.7279352226720648
+          "accuracy": 0.7857142857142857,
+          "precision": 0.7551020408163265,
+          "recall": 0.8118317890235209,
+          "f1": 0.7630094043887148,
+          "true_class_counts": [
+            61,
+            23
+          ],
+          "pred_class_counts": [
+            49,
+            35
+          ],
+          "confusion_matrix": [
+            [
+              46,
+              15
+            ],
+            [
+              3,
+              20
+            ]
+          ]
         },
         "ee": {
           "n_samples": 84,
           "accuracy": 0.75,
-          "precision": 0.6912280701754385,
-          "recall": 0.6471428571428571,
-          "f1": 0.6631611379274931
+          "precision": 0.7092916445857623,
+          "recall": 0.7614285714285715,
+          "f1": 0.7175824990730441,
+          "true_class_counts": [
+            14,
+            20,
+            50
+          ],
+          "pred_class_counts": [
+            17,
+            30,
+            37
+          ],
+          "confusion_matrix": [
+            [
+              10,
+              3,
+              1
+            ],
+            [
+              3,
+              17,
+              0
+            ],
+            [
+              4,
+              10,
+              36
+            ]
+          ]
         },
         "biodist": {
           "n_samples": 58,
-          "kl_divergence": 0.16366098804686893,
-          "js_divergence": 0.0359968761663274
+          "kl_divergence": 0.15998570647485624,
+          "js_divergence": 0.035472610690609384
         }
       }
     },
@@ -53,40 +116,103 @@
           "mse": 0.30493822479834903,
           "rmse": 0.5522121193874225,
           "mae": 0.34891908283267253,
-          "r2": 0.0688270762246912
+          "r2": 0.06882707391945686
         },
         "delivery": {
           "n_samples": 61,
-          "mse": 0.7776284206466513,
-          "rmse": 0.8818324220886026,
-          "mae": 0.519796750021438,
-          "r2": 0.4113913613868899
+          "mse": 0.7752937492958416,
+          "rmse": 0.8805076656655759,
+          "mae": 0.5151286462071908,
+          "r2": 0.4131585218812348
         },
         "toxic": {
           "n_samples": 61,
-          "accuracy": 1.0,
-          "precision": 1.0,
-          "recall": 1.0,
-          "f1": 1.0
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475,
+          "true_class_counts": [
+            58,
+            3
+          ],
+          "pred_class_counts": [
+            55,
+            6
+          ],
+          "confusion_matrix": [
+            [
+              55,
+              3
+            ],
+            [
+              0,
+              3
+            ]
+          ]
         },
         "pdi": {
           "n_samples": 84,
-          "accuracy": 0.7619047619047619,
-          "precision": 0.6711111111111111,
-          "recall": 0.5873015873015873,
-          "f1": 0.5942028985507246
+          "accuracy": 0.6547619047619048,
+          "precision": 0.5803571428571429,
+          "recall": 0.5952380952380952,
+          "f1": 0.5822328931572629,
+          "true_class_counts": [
+            63,
+            21
+          ],
+          "pred_class_counts": [
+            56,
+            28
+          ],
+          "confusion_matrix": [
+            [
+              45,
+              18
+            ],
+            [
+              11,
+              10
+            ]
+          ]
         },
         "ee": {
           "n_samples": 84,
-          "accuracy": 0.6785714285714286,
-          "precision": 0.6062341913405743,
-          "recall": 0.6071785357499643,
-          "f1": 0.6041666666666666
+          "accuracy": 0.5476190476190477,
+          "precision": 0.5001165151969107,
+          "recall": 0.4989534275248561,
+          "f1": 0.48349530196256724,
+          "true_class_counts": [
+            13,
+            22,
+            49
+          ],
+          "pred_class_counts": [
+            19,
+            31,
+            34
+          ],
+          "confusion_matrix": [
+            [
+              5,
+              7,
+              1
+            ],
+            [
+              8,
+              11,
+              3
+            ],
+            [
+              6,
+              13,
+              30
+            ]
+          ]
         },
         "biodist": {
           "n_samples": 61,
           "kl_divergence": 0.1177531067102451,
-          "js_divergence": 0.027774748827162937
+          "js_divergence": 0.027774748827162944
         }
       }
     },
@@ -95,43 +221,106 @@
       "test_metrics": {
         "size": {
           "n_samples": 84,
-          "mse": 0.3089202253953697,
-          "rmse": 0.5558059242175902,
-          "mae": 0.3994691472061615,
-          "r2": 0.3898880719888205
+          "mse": 0.3091245824692921,
+          "rmse": 0.5559897323416073,
+          "mae": 0.4046506446335209,
+          "r2": 0.3894844727418765
         },
         "delivery": {
           "n_samples": 60,
-          "mse": 0.6215887874238653,
-          "rmse": 0.7884090229213928,
-          "mae": 0.4654098044336172,
-          "r2": 0.17255759426449502
+          "mse": 0.623471858388545,
+          "rmse": 0.789602341934562,
+          "mae": 0.4688382081680016,
+          "r2": 0.17005089391015027
         },
         "toxic": {
           "n_samples": 61,
-          "accuracy": 1.0,
-          "precision": 1.0,
-          "recall": 1.0,
-          "f1": 1.0
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571,
+          "true_class_counts": [
+            58,
+            3
+          ],
+          "pred_class_counts": [
+            54,
+            7
+          ],
+          "confusion_matrix": [
+            [
+              54,
+              4
+            ],
+            [
+              0,
+              3
+            ]
+          ]
         },
         "pdi": {
           "n_samples": 84,
-          "accuracy": 0.7261904761904762,
-          "precision": 0.6428571428571428,
-          "recall": 0.6385630498533724,
-          "f1": 0.6405581395348837
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6693363844393593,
+          "recall": 0.717008797653959,
+          "f1": 0.662962962962963,
+          "true_class_counts": [
+            62,
+            22
+          ],
+          "pred_class_counts": [
+            46,
+            38
+          ],
+          "confusion_matrix": [
+            [
+              41,
+              21
+            ],
+            [
+              5,
+              17
+            ]
+          ]
         },
         "ee": {
           "n_samples": 84,
-          "accuracy": 0.6547619047619048,
-          "precision": 0.5676343636869953,
-          "recall": 0.5378151260504201,
-          "f1": 0.5491025070464323
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5842305037957212,
+          "recall": 0.6349206349206349,
+          "f1": 0.5925185925185925,
+          "true_class_counts": [
+            12,
+            21,
+            51
+          ],
+          "pred_class_counts": [
+            21,
+            23,
+            40
+          ],
+          "confusion_matrix": [
+            [
+              8,
+              2,
+              2
+            ],
+            [
+              5,
+              12,
+              4
+            ],
+            [
+              8,
+              9,
+              34
+            ]
+          ]
         },
         "biodist": {
           "n_samples": 60,
-          "kl_divergence": 0.11869463263891707,
-          "js_divergence": 0.02929511833384903
+          "kl_divergence": 0.11799692057231032,
+          "js_divergence": 0.029125166530752162
         }
       }
     },
@@ -140,43 +329,106 @@
       "test_metrics": {
         "size": {
           "n_samples": 83,
-          "mse": 1.6827545971608378,
-          "rmse": 1.2972103133882484,
-          "mae": 0.5269114175759023,
-          "r2": 0.048221123148072476
+          "mse": 1.6732277890580922,
+          "rmse": 1.2935330645399414,
+          "mae": 0.5211672051452704,
+          "r2": 0.05360960932568781
         },
         "delivery": {
           "n_samples": 59,
-          "mse": 0.8527625246024667,
+          "mse": 0.8527625246024666,
           "rmse": 0.9234514197306032,
           "mae": 0.5525073854104011,
-          "r2": 0.1710134699931437
+          "r2": 0.17101350859831832
         },
         "toxic": {
           "n_samples": 60,
-          "accuracy": 1.0,
-          "precision": 1.0,
-          "recall": 1.0,
-          "f1": 1.0
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817,
+          "true_class_counts": [
+            57,
+            3
+          ],
+          "pred_class_counts": [
+            53,
+            7
+          ],
+          "confusion_matrix": [
+            [
+              53,
+              4
+            ],
+            [
+              0,
+              3
+            ]
+          ]
         },
         "pdi": {
           "n_samples": 84,
-          "accuracy": 0.75,
-          "precision": 0.6676909569798068,
-          "recall": 0.6400293255131965,
-          "f1": 0.6493738819320214
+          "accuracy": 0.7380952380952381,
+          "precision": 0.6612903225806452,
+          "recall": 0.6612903225806452,
+          "f1": 0.6612903225806452,
+          "true_class_counts": [
+            62,
+            22
+          ],
+          "pred_class_counts": [
+            62,
+            22
+          ],
+          "confusion_matrix": [
+            [
+              51,
+              11
+            ],
+            [
+              11,
+              11
+            ]
+          ]
         },
         "ee": {
           "n_samples": 84,
-          "accuracy": 0.6309523809523809,
-          "precision": 0.5132964178288214,
-          "recall": 0.49489693313222727,
-          "f1": 0.4975934917795383
+          "accuracy": 0.6666666666666666,
+          "precision": 0.5872039589961117,
+          "recall": 0.6381096028154852,
+          "f1": 0.5938808373590982,
+          "true_class_counts": [
+            13,
+            20,
+            51
+          ],
+          "pred_class_counts": [
+            23,
+            20,
+            41
+          ],
+          "confusion_matrix": [
+            [
+              10,
+              3,
+              0
+            ],
+            [
+              9,
+              8,
+              3
+            ],
+            [
+              4,
+              9,
+              38
+            ]
+          ]
         },
         "biodist": {
           "n_samples": 60,
-          "kl_divergence": 0.1458100103551393,
-          "js_divergence": 0.034398655034422854
+          "kl_divergence": 0.14886390502458902,
+          "js_divergence": 0.03503301624445018
         }
       }
     },
@@ -185,103 +437,166 @@
       "test_metrics": {
         "size": {
           "n_samples": 84,
-          "mse": 0.539048457673242,
-          "rmse": 0.7341991948192548,
-          "mae": 0.41399459761950447,
-          "r2": 0.27868624280784415
+          "mse": 0.5439429374081691,
+          "rmse": 0.7375248723996833,
+          "mae": 0.41855190159962563,
+          "r2": 0.2721367833807379
         },
         "delivery": {
           "n_samples": 58,
-          "mse": 0.8685551446715326,
-          "rmse": 0.9319630597140278,
-          "mae": 0.48785093074168856,
-          "r2": 0.16012146519578374
+          "mse": 0.8807361070482976,
+          "rmse": 0.9384754163260206,
+          "mae": 0.4923390269539427,
+          "r2": 0.14834272376574187
         },
         "toxic": {
           "n_samples": 59,
-          "accuracy": 1.0,
-          "precision": 1.0,
-          "recall": 1.0,
-          "f1": 1.0
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091,
+          "true_class_counts": [
+            56,
+            3
+          ],
+          "pred_class_counts": [
+            54,
+            5
+          ],
+          "confusion_matrix": [
+            [
+              54,
+              2
+            ],
+            [
+              0,
+              3
+            ]
+          ]
         },
         "pdi": {
           "n_samples": 84,
-          "accuracy": 0.7380952380952381,
-          "precision": 0.6515151515151515,
-          "recall": 0.6319648093841642,
-          "f1": 0.6390625
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6428571428571428,
+          "recall": 0.6796187683284457,
+          "f1": 0.6415362731152205,
+          "true_class_counts": [
+            62,
+            22
+          ],
+          "pred_class_counts": [
+            49,
+            35
+          ],
+          "confusion_matrix": [
+            [
+              42,
+              20
+            ],
+            [
+              7,
+              15
+            ]
+          ]
         },
         "ee": {
           "n_samples": 84,
-          "accuracy": 0.6785714285714286,
-          "precision": 0.6048840048840048,
-          "recall": 0.6175213675213675,
-          "f1": 0.6107062011796588
+          "accuracy": 0.6309523809523809,
+          "precision": 0.5673525820584645,
+          "recall": 0.6021367521367521,
+          "f1": 0.5701075744179193,
+          "true_class_counts": [
+            12,
+            20,
+            52
+          ],
+          "pred_class_counts": [
+            17,
+            28,
+            39
+          ],
+          "confusion_matrix": [
+            [
+              7,
+              4,
+              1
+            ],
+            [
+              6,
+              11,
+              3
+            ],
+            [
+              4,
+              13,
+              35
+            ]
+          ]
         },
         "biodist": {
           "n_samples": 58,
-          "kl_divergence": 0.12253709805850461,
-          "js_divergence": 0.03003058068886421
+          "kl_divergence": 0.12072341813237816,
+          "js_divergence": 0.029660750070876794
         }
       }
     }
   ],
   "summary_stats": {
     "size": {
-      "mse_mean": 0.840904005205596,
-      "mse_std": 0.5742527703456151,
-      "rmse_mean": 0.8618819651219873,
-      "rmse_std": 0.3131508955808008,
-      "mae_mean": 0.4343888812056213,
-      "mae_std": 0.06297614136064734,
-      "r2_mean": 0.1895694822042235,
-      "r2_std": 0.12913183030644884
+      "mse_mean": 0.8359338407622309,
+      "mse_std": 0.5671793909058501,
+      "rmse_mean": 0.8600962833976535,
+      "rmse_std": 0.3101100192637678,
+      "mae_mean": 0.43462955168345657,
+      "mae_std": 0.06006325295257579,
+      "r2_mean": 0.19175642353368885,
+      "r2_std": 0.12652560874612837
     },
     "delivery": {
-      "mse_mean": 0.7062367004418524,
-      "mse_std": 0.17172737219125497,
-      "rmse_mean": 0.8332949277926908,
-      "rmse_std": 0.1088864718724354,
-      "mae_mean": 0.4869573943779645,
-      "mae_std": 0.048750584479951106,
-      "r2_mean": 0.27862494093264295,
-      "r2_std": 0.13730782732933916
+      "mse_mean": 0.705992627981044,
+      "mse_std": 0.1781763346919953,
+      "rmse_mean": 0.8325341145975752,
+      "rmse_std": 0.11348821970704938,
+      "mae_mean": 0.4872301660211962,
+      "mae_std": 0.048552308082167685,
+      "r2_mean": 0.279413269364265,
+      "r2_std": 0.14494172624677884
     },
     "toxic": {
-      "accuracy_mean": 1.0,
-      "accuracy_std": 0.0,
-      "precision_mean": 1.0,
-      "precision_std": 0.0,
-      "recall_mean": 1.0,
-      "recall_std": 0.0,
-      "f1_mean": 1.0,
-      "f1_std": 0.0
+      "accuracy_mean": 0.9500396342534485,
+      "accuracy_std": 0.014290090778642328,
+      "precision_mean": 0.7457142857142858,
+      "precision_std": 0.03149343955006944,
+      "recall_mean": 0.9737706334802525,
+      "recall_std": 0.007575117962473199,
+      "f1_mean": 0.8148343102325404,
+      "f1_std": 0.031251364844790935
     },
     "pdi": {
-      "accuracy_mean": 0.7571428571428571,
-      "accuracy_std": 0.02876915707998708,
-      "precision_mean": 0.6825769014781496,
-      "precision_std": 0.04965967017816626,
-      "recall_mean": 0.6408404643178055,
-      "recall_std": 0.03805702915009694,
-      "f1_mean": 0.6502265285379389,
-      "f1_std": 0.04335003701571735
+      "accuracy_mean": 0.7095238095238094,
+      "accuracy_std": 0.04677829215330594,
+      "precision_mean": 0.6617886067101234,
+      "precision_std": 0.05613216163532689,
+      "recall_mean": 0.6929975545649333,
+      "recall_std": 0.07132933119437145,
+      "f1_mean": 0.6622063712409613,
+      "f1_std": 0.05828494354685521
     },
     "ee": {
-      "accuracy_mean": 0.6785714285714286,
-      "accuracy_std": 0.03984095364447979,
-      "precision_mean": 0.5966554095831669,
-      "precision_std": 0.05812812328518918,
-      "recall_mean": 0.5809109639193673,
-      "recall_std": 0.055986847376187095,
-      "f1_mean": 0.5849460009199577,
-      "f1_std": 0.056684824199185184
+      "accuracy_mean": 0.6476190476190476,
+      "accuracy_std": 0.0650309537321317,
+      "precision_mean": 0.5896390409265941,
+      "precision_std": 0.06762084933247267,
+      "recall_mean": 0.62710979776526,
+      "recall_std": 0.08398957802940274,
+      "f1_mean": 0.5915169610662443,
+      "f1_std": 0.07486808628755862
     },
     "biodist": {
-      "kl_divergence_mean": 0.13369116716193502,
-      "kl_divergence_std": 0.01815702398766046,
-      "js_divergence_mean": 0.031499195810125294,
-      "js_divergence_std": 0.0031471347478191324
+      "kl_divergence_mean": 0.13306461138287579,
+      "kl_divergence_std": 0.017822152356421223,
+      "js_divergence_mean": 0.031413258472770295,
+      "js_divergence_std": 0.003197698848519836
     }
   }
 }

models/cv5_sample/baseline/seed42/outer_fold_0/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "moe_n_experts": 2,
+  "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/cv5_sample/baseline/seed42/outer_fold_0/epoch_mean.json

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

models/cv5_sample/baseline/seed42/outer_fold_0/history.json

@@ -0,0 +1,237 @@
+{
+  "train": [
+    {
+      "loss": 5.055936246304898,
+      "loss_size": 0.8606430002160974,
+      "loss_pdi": 0.6768325905542116,
+      "loss_ee": 1.0347085675677739,
+      "loss_delivery": 1.0286872987006161,
+      "loss_biodist": 0.9479167284192266,
+      "loss_toxic": 0.52346063224045
+    },
+    {
+      "loss": 3.99745157602671,
+      "loss_size": 0.7864497557685182,
+      "loss_pdi": 0.6331288669560406,
+      "loss_ee": 0.9356951681343285,
+      "loss_delivery": 0.8529001963702408,
+      "loss_biodist": 0.4785909547999099,
+      "loss_toxic": 0.33874212809511134
+    },
+    {
+      "loss": 3.6507615302060104,
+      "loss_size": 0.7465558921968615,
+      "loss_pdi": 0.5953955263704866,
+      "loss_ee": 0.8915211453631118,
+      "loss_delivery": 0.8915189435151784,
+      "loss_biodist": 0.3364428563697918,
+      "loss_toxic": 0.23094557698916746
+    },
+    {
+      "loss": 3.6875334591478914,
+      "loss_size": 0.7882098569660574,
+      "loss_pdi": 0.5826210242670935,
+      "loss_ee": 0.8787951453312023,
+      "loss_delivery": 1.0553698157136504,
+      "loss_biodist": 0.24477992005444862,
+      "loss_toxic": 0.19957500592981642
+    },
+    {
+      "loss": 3.2242523431777954,
+      "loss_size": 0.7434807216396203,
+      "loss_pdi": 0.5325888967191851,
+      "loss_ee": 0.8109645078311095,
+      "loss_delivery": 0.8456436809454415,
+      "loss_biodist": 0.2152977317974374,
+      "loss_toxic": 0.13414243228632855
+    },
+    {
+      "loss": 2.72972969428913,
+      "loss_size": 0.41291164667219726,
+      "loss_pdi": 0.5447176517666997,
+      "loss_ee": 0.7833192106839773,
+      "loss_delivery": 0.7865552172266148,
+      "loss_biodist": 0.18164325535700127,
+      "loss_toxic": 0.08577694238251629
+    },
+    {
+      "loss": 2.7861866564364046,
+      "loss_size": 0.5654866753800495,
+      "loss_pdi": 0.4972686191668382,
+      "loss_ee": 0.748628426242519,
+      "loss_delivery": 0.7679579907373802,
+      "loss_biodist": 0.2135133976871903,
+      "loss_toxic": 0.05895325155123263
+    },
+    {
+      "loss": 2.7905788904911764,
+      "loss_size": 0.5844308375506788,
+      "loss_pdi": 0.5185149493249687,
+      "loss_ee": 0.7592554019915091,
+      "loss_delivery": 0.7543600313365459,
+      "loss_biodist": 0.16718868860924566,
+      "loss_toxic": 0.0781663300664592
+    },
+    {
+      "loss": 2.7408436472351485,
+      "loss_size": 0.4887551304859084,
+      "loss_pdi": 0.5093961752749778,
+      "loss_ee": 0.7876343791549271,
+      "loss_delivery": 0.7577200457453728,
+      "loss_biodist": 0.18930391161828428,
+      "loss_toxic": 0.08456783540660164
+    },
+    {
+      "loss": 2.4752910588238692,
+      "loss_size": 0.4774871263149622,
+      "loss_pdi": 0.46845414630464605,
+      "loss_ee": 0.7078113322322434,
+      "loss_delivery": 0.6665648940849949,
+      "loss_biodist": 0.16393172116698446,
+      "loss_toxic": 0.05852524563040886
+    },
+    {
+      "loss": 2.5551916940792188,
+      "loss_size": 0.5141892947659299,
+      "loss_pdi": 0.4798985337083404,
+      "loss_ee": 0.7492339151936609,
+      "loss_delivery": 0.6760205814640063,
+      "loss_biodist": 0.14918651633165977,
+      "loss_toxic": 0.060266225545264375
+    }
+  ],
+  "val": [
+    {
+      "loss": 3.9572858810424805,
+      "loss_size": 0.5028525212779641,
+      "loss_pdi": 0.7401092886924744,
+      "loss_ee": 0.9889105439186097,
+      "loss_delivery": 0.6921922609210014,
+      "loss_biodist": 0.7097499012947083,
+      "loss_toxic": 0.35348092019557953,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 3.4830109119415282,
+      "loss_size": 0.4342193953692913,
+      "loss_pdi": 0.81399564743042,
+      "loss_ee": 0.8462728559970856,
+      "loss_delivery": 0.598291178047657,
+      "loss_biodist": 0.5577272236347198,
+      "loss_toxic": 0.2763485822826624,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.4455579280853272,
+      "loss_size": 0.39207540452480316,
+      "loss_pdi": 0.8946391701698303,
+      "loss_ee": 0.8172489017248153,
+      "loss_delivery": 0.6878725722432136,
+      "loss_biodist": 0.3735358417034149,
+      "loss_toxic": 0.3407814013771713,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 4.0154601573944095,
+      "loss_size": 0.4502963304519653,
+      "loss_pdi": 1.2591865181922912,
+      "loss_ee": 0.8082401096820832,
+      "loss_delivery": 0.6105723276734352,
+      "loss_biodist": 0.29027132838964465,
+      "loss_toxic": 0.6681207884568721,
+      "acc_pdi": 0.3939393939393939,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.85
+    },
+    {
+      "loss": 3.352661943435669,
+      "loss_size": 0.32106279730796816,
+      "loss_pdi": 1.0119264781475068,
+      "loss_ee": 0.9331188946962357,
+      "loss_delivery": 0.5763307940214872,
+      "loss_biodist": 0.21531529426574708,
+      "loss_toxic": 0.37886948613449933,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.444174289703369,
+      "loss_size": 0.35679062008857726,
+      "loss_pdi": 1.155838769674301,
+      "loss_ee": 0.8481447480618953,
+      "loss_delivery": 0.6408442467451095,
+      "loss_biodist": 0.17143512666225433,
+      "loss_toxic": 0.36500417003408075,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.0015304565429686,
+      "loss_size": 0.3286006212234497,
+      "loss_pdi": 1.0518546760082246,
+      "loss_ee": 0.8692154731601477,
+      "loss_delivery": 0.5558455713093281,
+      "loss_biodist": 0.15108025521039964,
+      "loss_toxic": 0.15213317363522946,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.3050911903381346,
+      "loss_size": 0.3655768930912018,
+      "loss_pdi": 1.2608545422554016,
+      "loss_ee": 0.8624444365501404,
+      "loss_delivery": 0.5735122933983803,
+      "loss_biodist": 0.1565265588462353,
+      "loss_toxic": 0.20591226250398903,
+      "acc_pdi": 0.5454545454545454,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.1276957035064696,
+      "loss_size": 0.30587572753429415,
+      "loss_pdi": 1.2177553713321685,
+      "loss_ee": 0.8245017357170582,
+      "loss_delivery": 0.5470387808978557,
+      "loss_biodist": 0.13348158523440362,
+      "loss_toxic": 0.2128354288637638,
+      "acc_pdi": 0.45454545454545453,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.1956005096435547,
+      "loss_size": 0.3409146010875702,
+      "loss_pdi": 1.250557976961136,
+      "loss_ee": 0.8290754735469819,
+      "loss_delivery": 0.5627829760313035,
+      "loss_biodist": 0.12690359950065613,
+      "loss_toxic": 0.20395291885361075,
+      "acc_pdi": 0.5151515151515151,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.1633092880249025,
+      "loss_size": 0.35122549831867217,
+      "loss_pdi": 1.2853083610534668,
+      "loss_ee": 0.7996505998075009,
+      "loss_delivery": 0.579404155164957,
+      "loss_biodist": 0.13538600280880927,
+      "loss_toxic": 0.13744947910308838,
+      "acc_pdi": 0.5151515151515151,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/cv5_sample/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/cv5_sample/baseline/seed42/outer_fold_0/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 83,
+    "mse": 1.5968842484584136,
+    "rmse": 1.263678854954222,
+    "mae": 0.5682635198246284,
+    "r2": 0.022667532793975265
+  },
+  "delivery": {
+    "n_samples": 58,
+    "mse": 0.5483822645024997,
+    "rmse": 0.740528368465719,
+    "mae": 0.5250802371843621,
+    "r2": 0.3029730502833111
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7142857142857143,
+    "precision": 0.6737967914438503,
+    "recall": 0.7084818246614397,
+    "f1": 0.6785714285714286
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.7380952380952381,
+    "precision": 0.7073879551820729,
+    "recall": 0.7647619047619049,
+    "f1": 0.7124896608767578
+  },
+  "toxic": {
+    "n_samples": 58,
+    "accuracy": 0.9827586206896551,
+    "precision": 0.9912280701754386,
+    "recall": 0.75,
+    "f1": 0.8289085545722714
+  },
+  "biodist": {
+    "n_samples": 58,
+    "kl_divergence": 0.1971223395298808,
+    "js_divergence": 0.043043902163233595
+  }
+}

models/cv5_sample/baseline/seed42/outer_fold_1/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "moe_n_experts": 2,
+  "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/cv5_sample/baseline/seed42/outer_fold_1/epoch_mean.json

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

models/cv5_sample/baseline/seed42/outer_fold_1/history.json

@@ -0,0 +1,237 @@
+{
+  "train": [
+    {
+      "loss": 5.3353071212768555,
+      "loss_size": 1.2256073724176433,
+      "loss_pdi": 0.6600620666065732,
+      "loss_ee": 1.020295173735232,
+      "loss_delivery": 0.9249667110475334,
+      "loss_biodist": 0.9719328582286835,
+      "loss_toxic": 0.5454153837384405
+    },
+    {
+      "loss": 4.285548815856108,
+      "loss_size": 1.111898030179578,
+      "loss_pdi": 0.6168666175893835,
+      "loss_ee": 0.9040256867537627,
+      "loss_delivery": 0.9542275908831004,
+      "loss_biodist": 0.4325796188534917,
+      "loss_toxic": 0.29354191521132317
+    },
+    {
+      "loss": 3.56095738024325,
+      "loss_size": 0.7991758820575636,
+      "loss_pdi": 0.5580040844711097,
+      "loss_ee": 0.8332083772968601,
+      "loss_delivery": 0.9466991460806614,
+      "loss_biodist": 0.27789403236395605,
+      "loss_toxic": 0.18792380020022392
+    },
+    {
+      "loss": 3.5710522806322254,
+      "loss_size": 1.046761174661082,
+      "loss_pdi": 0.5474661910856092,
+      "loss_ee": 0.8386256227622161,
+      "loss_delivery": 0.758317036201825,
+      "loss_biodist": 0.23423250090028788,
+      "loss_toxic": 0.19732182759892297
+    },
+    {
+      "loss": 3.344043564152073,
+      "loss_size": 0.9940445900366113,
+      "loss_pdi": 0.5329670773164646,
+      "loss_ee": 0.7994254424765304,
+      "loss_delivery": 0.7578543800910985,
+      "loss_biodist": 0.2094646818130403,
+      "loss_toxic": 0.1107913074335335
+    },
+    {
+      "loss": 2.9736292458869316,
+      "loss_size": 0.9058008608785836,
+      "loss_pdi": 0.5148797063408671,
+      "loss_ee": 0.7447971218341106,
+      "loss_delivery": 0.6041139666670682,
+      "loss_biodist": 0.19747554249054677,
+      "loss_toxic": 0.06619761599152274
+    },
+    {
+      "loss": 3.091302794379157,
+      "loss_size": 0.8345884137056969,
+      "loss_pdi": 0.4844536020143612,
+      "loss_ee": 0.7188561582887495,
+      "loss_delivery": 0.6844744932007145,
+      "loss_biodist": 0.16666732728481293,
+      "loss_toxic": 0.24889929544784734
+    },
+    {
+      "loss": 2.474903399879868,
+      "loss_size": 0.4893314063146308,
+      "loss_pdi": 0.4681552906294127,
+      "loss_ee": 0.699250686007577,
+      "loss_delivery": 0.6240661612434967,
+      "loss_biodist": 0.181533552202824,
+      "loss_toxic": 0.07646923640906508
+    },
+    {
+      "loss": 2.7190622445699333,
+      "loss_size": 0.7405592789118355,
+      "loss_pdi": 0.4723592585808522,
+      "loss_ee": 0.6855430079473032,
+      "loss_delivery": 0.6662846245878452,
+      "loss_biodist": 0.1865380845359854,
+      "loss_toxic": 0.03832938376738614
+    },
+    {
+      "loss": 2.7083647992159867,
+      "loss_size": 0.7273675418786101,
+      "loss_pdi": 0.46286967477282964,
+      "loss_ee": 0.6985194336723637,
+      "loss_delivery": 0.634993327781558,
+      "loss_biodist": 0.16010415433226405,
+      "loss_toxic": 0.09129857794834755
+    },
+    {
+      "loss": 2.5737238510234937,
+      "loss_size": 0.7433734733510662,
+      "loss_pdi": 0.46628197262416016,
+      "loss_ee": 0.668087279474413,
+      "loss_delivery": 0.5395700915439708,
+      "loss_biodist": 0.1628361604100949,
+      "loss_toxic": 0.05862165063720297
+    }
+  ],
+  "val": [
+    {
+      "loss": 3.031777000427246,
+      "loss_size": 0.5401369854807854,
+      "loss_pdi": 0.6682288587093353,
+      "loss_ee": 0.9218176007270813,
+      "loss_delivery": 0.2300173431634903,
+      "loss_biodist": 0.5083122372627258,
+      "loss_toxic": 0.17078153491020204,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.696969696969697,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.502068042755127,
+      "loss_size": 0.48314091861248015,
+      "loss_pdi": 0.639093154668808,
+      "loss_ee": 0.6983676433563233,
+      "loss_delivery": 0.3520810127258301,
+      "loss_biodist": 0.2895366489887238,
+      "loss_toxic": 0.052681690081954,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.417718243598938,
+      "loss_size": 0.44972713887691496,
+      "loss_pdi": 0.6729098677635192,
+      "loss_ee": 0.6924970507621765,
+      "loss_delivery": 0.36018016636371614,
+      "loss_biodist": 0.26438169181346893,
+      "loss_toxic": 0.019849006040021776,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.685157823562622,
+      "loss_size": 0.4716562867164612,
+      "loss_pdi": 0.5998583734035492,
+      "loss_ee": 0.601423978805542,
+      "loss_delivery": 0.7582527756690979,
+      "loss_biodist": 0.2994823604822159,
+      "loss_toxic": 0.007456492912024259,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.0070614576339723,
+      "loss_size": 0.44010728001594546,
+      "loss_pdi": 0.6203509151935578,
+      "loss_ee": 0.6338537752628326,
+      "loss_delivery": 0.1313400834798813,
+      "loss_biodist": 0.22033900320529937,
+      "loss_toxic": 0.004034112859517336,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.1099953651428223,
+      "loss_size": 0.4530618578195572,
+      "loss_pdi": 0.6262755393981934,
+      "loss_ee": 0.597742722928524,
+      "loss_delivery": 0.30005890130996704,
+      "loss_biodist": 0.18362512215971946,
+      "loss_toxic": 0.0021228168159723284,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.2114261865615843,
+      "loss_size": 0.5167493343353271,
+      "loss_pdi": 0.5672660648822785,
+      "loss_ee": 0.6098539769649506,
+      "loss_delivery": 0.3752159357070923,
+      "loss_biodist": 0.19700948819518088,
+      "loss_toxic": 0.0033008580910973253,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.7878787878787878,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.0603028774261474,
+      "loss_size": 0.4996716260910034,
+      "loss_pdi": 0.5430601418018342,
+      "loss_ee": 0.623375853896141,
+      "loss_delivery": 0.24462970793247224,
+      "loss_biodist": 0.19848335832357406,
+      "loss_toxic": 0.007659496087580919,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 1.9372182607650756,
+      "loss_size": 0.5049063771963119,
+      "loss_pdi": 0.5282617956399918,
+      "loss_ee": 0.6546011686325073,
+      "loss_delivery": 0.11980208605527878,
+      "loss_biodist": 0.18347663804888725,
+      "loss_toxic": 0.004195089940913022,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.696969696969697,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 1.9230464458465577,
+      "loss_size": 0.4861533761024475,
+      "loss_pdi": 0.4957872599363327,
+      "loss_ee": 0.6811205118894577,
+      "loss_delivery": 0.12559136152267455,
+      "loss_biodist": 0.19170617461204528,
+      "loss_toxic": 0.003053023712709546,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 1.9245877981185913,
+      "loss_size": 0.47571042776107786,
+      "loss_pdi": 0.5140088766813278,
+      "loss_ee": 0.6643553763628006,
+      "loss_delivery": 0.13551929593086243,
+      "loss_biodist": 0.19236155450344086,
+      "loss_toxic": 0.00303305524867028,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/cv5_sample/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/cv5_sample/baseline/seed42/outer_fold_1/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.30847479651898757,
+    "rmse": 0.5554050742647096,
+    "mae": 0.41110294015102444,
+    "r2": 0.05802764121606785
+  },
+  "delivery": {
+    "n_samples": 61,
+    "mse": 0.9177343265746907,
+    "rmse": 0.9579845127008529,
+    "mae": 0.589223776898179,
+    "r2": 0.30534127327020133
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.6785714285714286,
+    "precision": 0.5875,
+    "recall": 0.5952380952380952,
+    "f1": 0.5902439024390244
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.5833333333333334,
+    "precision": 0.5230357877416701,
+    "recall": 0.5150801579373008,
+    "f1": 0.49626374349858454
+  },
+  "toxic": {
+    "n_samples": 61,
+    "accuracy": 0.9836065573770492,
+    "precision": 0.9915254237288136,
+    "recall": 0.8333333333333333,
+    "f1": 0.8957264957264958
+  },
+  "biodist": {
+    "n_samples": 61,
+    "kl_divergence": 0.13675394519783235,
+    "js_divergence": 0.03248669854336822
+  }
+}

models/cv5_sample/baseline/seed42/outer_fold_2/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "moe_n_experts": 2,
+  "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/cv5_sample/baseline/seed42/outer_fold_2/epoch_mean.json

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

models/cv5_sample/baseline/seed42/outer_fold_2/history.json

@@ -0,0 +1,237 @@
+{
+  "train": [
+    {
+      "loss": 4.981144982415277,
+      "loss_size": 0.8744071665647868,
+      "loss_pdi": 0.6796890819394911,
+      "loss_ee": 1.0396443009376526,
+      "loss_delivery": 0.8872663096801655,
+      "loss_biodist": 0.9659920004574029,
+      "loss_toxic": 0.5519921908507476
+    },
+    {
+      "loss": 3.534698486328125,
+      "loss_size": 0.502043268769174,
+      "loss_pdi": 0.6077678847957302,
+      "loss_ee": 0.9173400402069092,
+      "loss_delivery": 0.7972146357233459,
+      "loss_biodist": 0.47308555164852656,
+      "loss_toxic": 0.26515947262177597
+    },
+    {
+      "loss": 3.369716721612054,
+      "loss_size": 0.8003431667347212,
+      "loss_pdi": 0.5898223374340985,
+      "loss_ee": 0.8692510144130604,
+      "loss_delivery": 0.6483961460558144,
+      "loss_biodist": 0.33141391241067164,
+      "loss_toxic": 0.16803750934431683
+    },
+    {
+      "loss": 3.485442135785077,
+      "loss_size": 0.8031057560363332,
+      "loss_pdi": 0.581869965469515,
+      "loss_ee": 0.891696730981002,
+      "loss_delivery": 0.7873844331180727,
+      "loss_biodist": 0.30878648705579137,
+      "loss_toxic": 0.1734091161146156
+    },
+    {
+      "loss": 3.1509077484543258,
+      "loss_size": 0.7464953818031259,
+      "loss_pdi": 0.5425146785942284,
+      "loss_ee": 0.8121652844789866,
+      "loss_delivery": 0.7709234127321759,
+      "loss_biodist": 0.23068079734976227,
+      "loss_toxic": 0.11076271264008372
+    },
+    {
+      "loss": 2.9153835676811837,
+      "loss_size": 0.7114814908520596,
+      "loss_pdi": 0.5249752043872267,
+      "loss_ee": 0.7787095694928556,
+      "loss_delivery": 0.6547171264081388,
+      "loss_biodist": 0.21912198936617053,
+      "loss_toxic": 0.08922003149180799
+    },
+    {
+      "loss": 2.7910783709706486,
+      "loss_size": 0.6730798660097895,
+      "loss_pdi": 0.48799943561489517,
+      "loss_ee": 0.7486014172837541,
+      "loss_delivery": 0.6180728810864526,
+      "loss_biodist": 0.20473402114333333,
+      "loss_toxic": 0.11729335879302911
+    },
+    {
+      "loss": 2.6673463969617277,
+      "loss_size": 0.6288768926987777,
+      "loss_pdi": 0.4827019000375593,
+      "loss_ee": 0.7593149873050483,
+      "loss_delivery": 0.5716850237065071,
+      "loss_biodist": 0.1800537146627903,
+      "loss_toxic": 0.10778728406876326
+    },
+    {
+      "loss": 2.5156604889276863,
+      "loss_size": 0.5605991873386744,
+      "loss_pdi": 0.4511222758808652,
+      "loss_ee": 0.7027070796167528,
+      "loss_delivery": 0.5257743572866594,
+      "loss_biodist": 0.18494447743570483,
+      "loss_toxic": 0.14391961756376298
+    },
+    {
+      "loss": 2.589712107503736,
+      "loss_size": 0.6075560066949677,
+      "loss_pdi": 0.4832051074182665,
+      "loss_ee": 0.742237118450371,
+      "loss_delivery": 0.5646769211098954,
+      "loss_biodist": 0.20049028380497083,
+      "loss_toxic": 0.06180886638577323
+    },
+    {
+      "loss": 2.605213010633314,
+      "loss_size": 0.6251636191799834,
+      "loss_pdi": 0.46899717240720185,
+      "loss_ee": 0.701141885003528,
+      "loss_delivery": 0.5535404454614665,
+      "loss_biodist": 0.163242267193021,
+      "loss_toxic": 0.14924246485571604
+    }
+  ],
+  "val": [
+    {
+      "loss": 7.663464832305908,
+      "loss_size": 3.313545787334442,
+      "loss_pdi": 0.6900765538215637,
+      "loss_ee": 1.0373500823974608,
+      "loss_delivery": 1.899225541949272,
+      "loss_biodist": 0.5267939090728759,
+      "loss_toxic": 0.2363927185535431,
+      "acc_pdi": 0.48484848484848486,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 7.12261176109314,
+      "loss_size": 3.4245932430028914,
+      "loss_pdi": 0.7202669441699981,
+      "loss_ee": 0.9270447134971619,
+      "loss_delivery": 1.7875666916370392,
+      "loss_biodist": 0.239187753200531,
+      "loss_toxic": 0.07735542804002762,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.7173912525177,
+      "loss_size": 3.356651684641838,
+      "loss_pdi": 0.6973654150962829,
+      "loss_ee": 0.9874879479408264,
+      "loss_delivery": 1.495903843641281,
+      "loss_biodist": 0.1953073710203171,
+      "loss_toxic": 0.04306442430242896,
+      "acc_pdi": 0.5151515151515151,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.669378137588501,
+      "loss_size": 3.177801153063774,
+      "loss_pdi": 0.6953718304634094,
+      "loss_ee": 0.9659196972846985,
+      "loss_delivery": 1.685699325799942,
+      "loss_biodist": 0.1821014702320099,
+      "loss_toxic": 0.04902317344676703,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.0091077327728275,
+      "loss_size": 2.877595055103302,
+      "loss_pdi": 0.6217953264713287,
+      "loss_ee": 1.0046721696853638,
+      "loss_delivery": 1.47297545671463,
+      "loss_biodist": 0.09909218698740005,
+      "loss_toxic": 0.02542670485563576,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 5.837002372741699,
+      "loss_size": 2.829563948512077,
+      "loss_pdi": 0.6526294827461243,
+      "loss_ee": 0.9968664407730102,
+      "loss_delivery": 1.3078757762908935,
+      "loss_biodist": 0.09006003737449646,
+      "loss_toxic": 0.050850869109854105,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 5.693418264389038,
+      "loss_size": 2.7425797268748284,
+      "loss_pdi": 0.6487438917160034,
+      "loss_ee": 1.0159194469451904,
+      "loss_delivery": 1.2614001840353013,
+      "loss_biodist": 0.08654961585998536,
+      "loss_toxic": 0.03482890441082418,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 5.848294258117676,
+      "loss_size": 2.9319378077983855,
+      "loss_pdi": 0.6582076787948609,
+      "loss_ee": 1.0355260014533996,
+      "loss_delivery": 1.2072753131389617,
+      "loss_biodist": 0.09763360768556595,
+      "loss_toxic": 0.015172032522968948,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 5.659598445892334,
+      "loss_size": 2.7996189504861833,
+      "loss_pdi": 0.650446206331253,
+      "loss_ee": 0.9975044012069703,
+      "loss_delivery": 1.197677844762802,
+      "loss_biodist": 0.0949702337384224,
+      "loss_toxic": 0.015282840328291058,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 5.686920213699341,
+      "loss_size": 2.7936870113015173,
+      "loss_pdi": 0.6796894907951355,
+      "loss_ee": 1.025844794511795,
+      "loss_delivery": 1.1696001946926118,
+      "loss_biodist": 0.09400582611560822,
+      "loss_toxic": 0.023353599617257713,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 5.827069425582886,
+      "loss_size": 2.8803999423980713,
+      "loss_pdi": 0.6822858631610871,
+      "loss_ee": 1.0088962197303772,
+      "loss_delivery": 1.2416288018226624,
+      "loss_biodist": 0.09613501951098442,
+      "loss_toxic": 0.017144279484637082,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/cv5_sample/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/cv5_sample/baseline/seed42/outer_fold_2/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.3450456742192763,
+    "rmse": 0.5874058854142307,
+    "mae": 0.46413040039173903,
+    "r2": 0.31854095833662033
+  },
+  "delivery": {
+    "n_samples": 60,
+    "mse": 0.633114832363136,
+    "rmse": 0.7956851339337289,
+    "mae": 0.5481160653211797,
+    "r2": 0.15721442419210252
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7142857142857143,
+    "precision": 0.6517857142857143,
+    "recall": 0.6744868035190617,
+    "f1": 0.6583050847457628
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6428571428571429,
+    "precision": 0.5853276353276353,
+    "recall": 0.6349206349206349,
+    "f1": 0.5970747928194737
+  },
+  "toxic": {
+    "n_samples": 61,
+    "accuracy": 1.0,
+    "precision": 1.0,
+    "recall": 1.0,
+    "f1": 1.0
+  },
+  "biodist": {
+    "n_samples": 60,
+    "kl_divergence": 0.1542080131383728,
+    "js_divergence": 0.03808539778255157
+  }
+}

models/cv5_sample/baseline/seed42/outer_fold_3/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "moe_n_experts": 2,
+  "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/cv5_sample/baseline/seed42/outer_fold_3/epoch_mean.json

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

models/cv5_sample/baseline/seed42/outer_fold_3/history.json

@@ -0,0 +1,237 @@
+{
+  "train": [
+    {
+      "loss": 4.500459426158184,
+      "loss_size": 0.5179172538824983,
+      "loss_pdi": 0.6786781702492688,
+      "loss_ee": 1.0435632644472896,
+      "loss_delivery": 0.8494648298097623,
+      "loss_biodist": 0.9263661117167086,
+      "loss_toxic": 0.5011754672269564
+    },
+    {
+      "loss": 3.5351730681754447,
+      "loss_size": 0.4758393716167759,
+      "loss_pdi": 0.6050793609103641,
+      "loss_ee": 0.9199129455798382,
+      "loss_delivery": 0.8140221713362513,
+      "loss_biodist": 0.46644948382635376,
+      "loss_toxic": 0.28111603362737475
+    },
+    {
+      "loss": 3.054368998553302,
+      "loss_size": 0.4495431559311377,
+      "loss_pdi": 0.5573401378618704,
+      "loss_ee": 0.8661556098912213,
+      "loss_delivery": 0.7387672813760268,
+      "loss_biodist": 0.3471993232095564,
+      "loss_toxic": 0.13824860089634722
+    },
+    {
+      "loss": 3.0508435287991085,
+      "loss_size": 0.44718711551379514,
+      "loss_pdi": 0.5715750842481047,
+      "loss_ee": 0.8271728351309493,
+      "loss_delivery": 0.8101292844559695,
+      "loss_biodist": 0.30609360800401586,
+      "loss_toxic": 0.14899400996698722
+    },
+    {
+      "loss": 2.663577508282017,
+      "loss_size": 0.41014224129754145,
+      "loss_pdi": 0.5433303649361069,
+      "loss_ee": 0.7740156594160441,
+      "loss_delivery": 0.6729272297910742,
+      "loss_biodist": 0.21821164641831373,
+      "loss_toxic": 0.10631358106872318
+    },
+    {
+      "loss": 2.557415685138187,
+      "loss_size": 0.4118609937864381,
+      "loss_pdi": 0.5263543410881145,
+      "loss_ee": 0.7567270873366175,
+      "loss_delivery": 0.6311398055907842,
+      "loss_biodist": 0.2105274689761368,
+      "loss_toxic": 0.0864857490224814
+    },
+    {
+      "loss": 2.3681529148204907,
+      "loss_size": 0.34351927102417557,
+      "loss_pdi": 0.4974673124583992,
+      "loss_ee": 0.7299669586323403,
+      "loss_delivery": 0.5682296187088296,
+      "loss_biodist": 0.18714645907685562,
+      "loss_toxic": 0.10800810708865724
+    },
+    {
+      "loss": 2.3271716382052445,
+      "loss_size": 0.3356305380930772,
+      "loss_pdi": 0.4922029255209742,
+      "loss_ee": 0.7045397452405981,
+      "loss_delivery": 0.6236170645098429,
+      "loss_biodist": 0.1928157155980935,
+      "loss_toxic": 0.05377595526845874
+    },
+    {
+      "loss": 2.3417651943258337,
+      "loss_size": 0.33690569489388855,
+      "loss_pdi": 0.49544163652368495,
+      "loss_ee": 0.7062375666321935,
+      "loss_delivery": 0.6490140618706072,
+      "loss_biodist": 0.17358206846826785,
+      "loss_toxic": 0.060061405567024405
+    },
+    {
+      "loss": 2.245423777683361,
+      "loss_size": 0.35165876744164004,
+      "loss_pdi": 0.48304893237513463,
+      "loss_ee": 0.635595595111718,
+      "loss_delivery": 0.5934803685223734,
+      "loss_biodist": 0.1646844497000849,
+      "loss_toxic": 0.08570577780640609
+    },
+    {
+      "loss": 2.2503104499868445,
+      "loss_size": 0.3315257770386902,
+      "loss_pdi": 0.4818291132514541,
+      "loss_ee": 0.7003061497533644,
+      "loss_delivery": 0.57664611291241,
+      "loss_biodist": 0.16049769621443105,
+      "loss_toxic": 0.07698796648686324
+    }
+  ],
+  "val": [
+    {
+      "loss": 5.7637759923934935,
+      "loss_size": 2.7772069811820983,
+      "loss_pdi": 0.6040945172309875,
+      "loss_ee": 0.8834547519683837,
+      "loss_delivery": 0.8618888959288598,
+      "loss_biodist": 0.40433400869369507,
+      "loss_toxic": 0.22842871248722077,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 0.782608695652174
+    },
+    {
+      "loss": 5.360143804550171,
+      "loss_size": 2.749151438474655,
+      "loss_pdi": 0.5780697166919708,
+      "loss_ee": 0.7454802870750428,
+      "loss_delivery": 0.8378195464611053,
+      "loss_biodist": 0.2229689359664917,
+      "loss_toxic": 0.21747921109199525,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.201284193992615,
+      "loss_size": 2.6817896857857706,
+      "loss_pdi": 0.6185303628444672,
+      "loss_ee": 0.7450856685638427,
+      "loss_delivery": 0.773091234266758,
+      "loss_biodist": 0.20598501861095428,
+      "loss_toxic": 0.16629037708044053,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.911232304573059,
+      "loss_size": 2.483246049284935,
+      "loss_pdi": 0.5181228935718536,
+      "loss_ee": 0.8096402883529663,
+      "loss_delivery": 0.7650228798389435,
+      "loss_biodist": 0.17875694036483764,
+      "loss_toxic": 0.14957769010216,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.086032152175903,
+      "loss_size": 2.546382114291191,
+      "loss_pdi": 0.5509662330150604,
+      "loss_ee": 0.8134036421775818,
+      "loss_delivery": 0.778993546962738,
+      "loss_biodist": 0.16294443011283874,
+      "loss_toxic": 0.21995164193212985,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.123322796821594,
+      "loss_size": 2.5980996996164323,
+      "loss_pdi": 0.587708693742752,
+      "loss_ee": 0.7929542601108551,
+      "loss_delivery": 0.8301725268363953,
+      "loss_biodist": 0.1425508365035057,
+      "loss_toxic": 0.1622565558180213,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.9730717420578,
+      "loss_size": 2.5694424003362655,
+      "loss_pdi": 0.5331394076347351,
+      "loss_ee": 0.8161154568195343,
+      "loss_delivery": 0.771108528971672,
+      "loss_biodist": 0.16225484609603882,
+      "loss_toxic": 0.11304349759593606,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.696969696969697,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.881006336212158,
+      "loss_size": 2.4189414471387862,
+      "loss_pdi": 0.5496642231941223,
+      "loss_ee": 0.8132100462913513,
+      "loss_delivery": 0.8486790001392365,
+      "loss_biodist": 0.15106511116027832,
+      "loss_toxic": 0.10017689382657409,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.704700064659119,
+      "loss_size": 2.282734674215317,
+      "loss_pdi": 0.5450537085533143,
+      "loss_ee": 0.8220642626285553,
+      "loss_delivery": 0.7426864296197891,
+      "loss_biodist": 0.15404343008995056,
+      "loss_toxic": 0.15513503681868315,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.758828639984131,
+      "loss_size": 2.3555777221918106,
+      "loss_pdi": 0.5239109575748444,
+      "loss_ee": 0.8412651300430298,
+      "loss_delivery": 0.7657725661993027,
+      "loss_biodist": 0.13986787497997283,
+      "loss_toxic": 0.13041157741099596,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.7822669506072994,
+      "loss_size": 2.370918756723404,
+      "loss_pdi": 0.5238461703062057,
+      "loss_ee": 0.8227214813232422,
+      "loss_delivery": 0.77011858522892,
+      "loss_biodist": 0.14329722970724107,
+      "loss_toxic": 0.1449897738173604,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.8695652173913043
+    }
+  ]
+}

models/cv5_sample/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/cv5_sample/baseline/seed42/outer_fold_3/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 83,
+    "mse": 1.763092308495774,
+    "rmse": 1.3278148622815509,
+    "mae": 0.5504888193233006,
+    "r2": 0.0027815521929155373
+  },
+  "delivery": {
+    "n_samples": 59,
+    "mse": 0.6925828062408207,
+    "rmse": 0.8322156008155703,
+    "mae": 0.5598128846357182,
+    "r2": 0.32672722594328685
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7380952380952381,
+    "precision": 0.6612903225806452,
+    "recall": 0.6612903225806452,
+    "f1": 0.6612903225806452
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.5952380952380952,
+    "precision": 0.5301311591634171,
+    "recall": 0.5438914027149321,
+    "f1": 0.5247639796659405
+  },
+  "toxic": {
+    "n_samples": 60,
+    "accuracy": 0.9333333333333333,
+    "precision": 0.7142857142857143,
+    "recall": 0.9649122807017544,
+    "f1": 0.7818181818181817
+  },
+  "biodist": {
+    "n_samples": 60,
+    "kl_divergence": 0.15787098167920866,
+    "js_divergence": 0.04007834983720861
+  }
+}

models/cv5_sample/baseline/seed42/outer_fold_4/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "moe_n_experts": 2,
+  "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/cv5_sample/baseline/seed42/outer_fold_4/epoch_mean.json

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

models/cv5_sample/baseline/seed42/outer_fold_4/history.json

@@ -0,0 +1,237 @@
+{
+  "train": [
+    {
+      "loss": 5.102151986714956,
+      "loss_size": 1.1567820576397148,
+      "loss_pdi": 0.6515501222094974,
+      "loss_ee": 1.0294628127201184,
+      "loss_delivery": 0.8426589999045875,
+      "loss_biodist": 0.9110404494646434,
+      "loss_toxic": 0.5230041408055538
+    },
+    {
+      "loss": 4.055108985385379,
+      "loss_size": 1.0688643268234022,
+      "loss_pdi": 0.5984339689886248,
+      "loss_ee": 0.9147920866270323,
+      "loss_delivery": 0.798665193689836,
+      "loss_biodist": 0.4416384439210634,
+      "loss_toxic": 0.2584042294403991
+    },
+    {
+      "loss": 3.8014610168096183,
+      "loss_size": 1.0324330774916184,
+      "loss_pdi": 0.5652096682303661,
+      "loss_ee": 0.9040436511104172,
+      "loss_delivery": 0.7405963850182455,
+      "loss_biodist": 0.3845312357754321,
+      "loss_toxic": 0.21792648017809196
+    },
+    {
+      "loss": 3.3098922484629862,
+      "loss_size": 0.6985325795170423,
+      "loss_pdi": 0.5925048146698926,
+      "loss_ee": 0.846102629158948,
+      "loss_delivery": 0.7961102732129999,
+      "loss_biodist": 0.2757947507339555,
+      "loss_toxic": 0.15982103052020474
+    },
+    {
+      "loss": 3.176285627725962,
+      "loss_size": 0.9075968438709104,
+      "loss_pdi": 0.5314604060069935,
+      "loss_ee": 0.8254372380875252,
+      "loss_delivery": 0.6657844270906739,
+      "loss_biodist": 0.21802176696223183,
+      "loss_toxic": 0.0905371925390854
+    },
+    {
+      "loss": 3.274916900170816,
+      "loss_size": 0.9856320102069829,
+      "loss_pdi": 0.5265752342907158,
+      "loss_ee": 0.7516567006304458,
+      "loss_delivery": 0.7083130106329918,
+      "loss_biodist": 0.2212505919506421,
+      "loss_toxic": 0.1400652697681122
+    },
+    {
+      "loss": 2.9465325329754806,
+      "loss_size": 0.8816265802125673,
+      "loss_pdi": 0.5074402512730779,
+      "loss_ee": 0.757787544179607,
+      "loss_delivery": 0.6106935801940996,
+      "loss_biodist": 0.19856224893718152,
+      "loss_toxic": 0.05817761752011909
+    },
+    {
+      "loss": 2.957754476650341,
+      "loss_size": 0.854996056572811,
+      "loss_pdi": 0.48137593349894964,
+      "loss_ee": 0.7679745611306783,
+      "loss_delivery": 0.6376201487272173,
+      "loss_biodist": 0.17370915422971184,
+      "loss_toxic": 0.10946399376671007
+    },
+    {
+      "loss": 2.7943139656169995,
+      "loss_size": 0.7861643652054103,
+      "loss_pdi": 0.4862389399393185,
+      "loss_ee": 0.6960344652871828,
+      "loss_delivery": 0.620618630301308,
+      "loss_biodist": 0.19384471256587957,
+      "loss_toxic": 0.07682121185762053
+    },
+    {
+      "loss": 2.7292510528822205,
+      "loss_size": 0.736359557187235,
+      "loss_pdi": 0.5096078466724705,
+      "loss_ee": 0.7071713265534993,
+      "loss_delivery": 0.5689382711677132,
+      "loss_biodist": 0.2055618997562576,
+      "loss_toxic": 0.07013601859609873
+    },
+    {
+      "loss": 2.688862778045036,
+      "loss_size": 0.7182406856602913,
+      "loss_pdi": 0.4632230537968713,
+      "loss_ee": 0.7033636360555082,
+      "loss_delivery": 0.6424240871156389,
+      "loss_biodist": 0.17127765292251432,
+      "loss_toxic": 0.05990550624886276
+    }
+  ],
+  "val": [
+    {
+      "loss": 3.6000691413879395,
+      "loss_size": 0.3701790302991867,
+      "loss_pdi": 0.5664441585540771,
+      "loss_ee": 0.8832380712032318,
+      "loss_delivery": 0.8911014348268509,
+      "loss_biodist": 0.5329773306846619,
+      "loss_toxic": 0.3749651610851288,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.0433499336242678,
+      "loss_size": 0.3372409701347351,
+      "loss_pdi": 0.4964510649442673,
+      "loss_ee": 0.7701501965522766,
+      "loss_delivery": 0.8530217319726944,
+      "loss_biodist": 0.2348724514245987,
+      "loss_toxic": 0.36663210168480875,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.779363822937012,
+      "loss_size": 0.30666323006153107,
+      "loss_pdi": 0.4720880776643753,
+      "loss_ee": 0.7368923991918563,
+      "loss_delivery": 0.828652124106884,
+      "loss_biodist": 0.16220578402280808,
+      "loss_toxic": 0.3004120915196836,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 3.281400537490845,
+      "loss_size": 0.5272440910339355,
+      "loss_pdi": 0.5226441502571106,
+      "loss_ee": 0.8246899724006653,
+      "loss_delivery": 0.8534515619277954,
+      "loss_biodist": 0.28649153709411623,
+      "loss_toxic": 0.31357612535357476,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.9561454892158507,
+      "loss_size": 0.2726425051689148,
+      "loss_pdi": 0.45198590606451033,
+      "loss_ee": 0.7590681105852127,
+      "loss_delivery": 0.8566071918583475,
+      "loss_biodist": 0.10773630291223527,
+      "loss_toxic": 0.5303979313932359,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.8890976190567015,
+      "loss_size": 0.34611437618732455,
+      "loss_pdi": 0.4896680176258087,
+      "loss_ee": 0.7619812674820423,
+      "loss_delivery": 0.866328160604462,
+      "loss_biodist": 0.09204872399568557,
+      "loss_toxic": 0.37474240330047903,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.7058202266693114,
+      "loss_size": 0.3040961235761642,
+      "loss_pdi": 0.44319443553686144,
+      "loss_ee": 0.7749164622277022,
+      "loss_delivery": 0.7846140109002591,
+      "loss_biodist": 0.09509205557405949,
+      "loss_toxic": 0.34775193249806763,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.8554219126701357,
+      "loss_size": 0.24815200865268708,
+      "loss_pdi": 0.4591174952685833,
+      "loss_ee": 0.7767300330102443,
+      "loss_delivery": 0.75640877969563,
+      "loss_biodist": 0.09491635411977768,
+      "loss_toxic": 0.5468520291149617,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.8406914949417112,
+      "loss_size": 0.2561314105987549,
+      "loss_pdi": 0.46083473414182663,
+      "loss_ee": 0.7739205928519368,
+      "loss_delivery": 0.784732368728146,
+      "loss_biodist": 0.08942811377346516,
+      "loss_toxic": 0.5080395077820867,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.8140346884727476,
+      "loss_size": 0.2685915812849998,
+      "loss_pdi": 0.4499114632606506,
+      "loss_ee": 0.8004547510296106,
+      "loss_delivery": 0.8297218790045008,
+      "loss_biodist": 0.08691173046827316,
+      "loss_toxic": 0.4234957183711231,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.763965880870819,
+      "loss_size": 0.2510989636182785,
+      "loss_pdi": 0.45986512303352356,
+      "loss_ee": 0.7952149830758571,
+      "loss_delivery": 0.8199413266032934,
+      "loss_biodist": 0.0973357006907463,
+      "loss_toxic": 0.3888999660965055,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    }
+  ]
+}

models/cv5_sample/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/cv5_sample/baseline/seed42/outer_fold_4/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.5697698999186622,
+    "rmse": 0.75483104064331,
+    "mae": 0.45241403512634515,
+    "r2": 0.23757709942203908
+  },
+  "delivery": {
+    "n_samples": 58,
+    "mse": 0.8494501546811344,
+    "rmse": 0.9216562019978677,
+    "mae": 0.6002610280852894,
+    "r2": 0.17859572323309747
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7380952380952381,
+    "precision": 0.6724137931034483,
+    "recall": 0.6906158357771262,
+    "f1": 0.6791666666666667
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6785714285714286,
+    "precision": 0.6324786324786325,
+    "recall": 0.6799145299145298,
+    "f1": 0.6406077872744539
+  },
+  "toxic": {
+    "n_samples": 59,
+    "accuracy": 1.0,
+    "precision": 1.0,
+    "recall": 1.0,
+    "f1": 1.0
+  },
+  "biodist": {
+    "n_samples": 58,
+    "kl_divergence": 0.17354171533022728,
+    "js_divergence": 0.046613494801719065
+  }
+}

models/cv5_sample/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/cv5_sample/baseline/seed42/summary.json

@@ -0,0 +1,372 @@
+{
+  "fold_results": [
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.19323529834582587,
+        "lr": 0.000980484886752915,
+        "weight_decay": 0.00014367509113739864,
+        "backbone_lr_ratio": 0.13093310334961422,
+        "moe_n_experts": 2,
+        "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": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.5968842484584136,
+          "rmse": 1.263678854954222,
+          "mae": 0.5682635198246284,
+          "r2": 0.022667532793975265
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.5483822645024997,
+          "rmse": 0.740528368465719,
+          "mae": 0.5250802371843621,
+          "r2": 0.3029730502833111
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6737967914438503,
+          "recall": 0.7084818246614397,
+          "f1": 0.6785714285714286
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7380952380952381,
+          "precision": 0.7073879551820729,
+          "recall": 0.7647619047619049,
+          "f1": 0.7124896608767578
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9827586206896551,
+          "precision": 0.9912280701754386,
+          "recall": 0.75,
+          "f1": 0.8289085545722714
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.1971223395298808,
+          "js_divergence": 0.043043902163233595
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.19323529834582587,
+        "lr": 0.000980484886752915,
+        "weight_decay": 0.00014367509113739864,
+        "backbone_lr_ratio": 0.13093310334961422,
+        "moe_n_experts": 2,
+        "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": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.30847479651898757,
+          "rmse": 0.5554050742647096,
+          "mae": 0.41110294015102444,
+          "r2": 0.05802764121606785
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 0.9177343265746907,
+          "rmse": 0.9579845127008529,
+          "mae": 0.589223776898179,
+          "r2": 0.30534127327020133
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.5875,
+          "recall": 0.5952380952380952,
+          "f1": 0.5902439024390244
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5833333333333334,
+          "precision": 0.5230357877416701,
+          "recall": 0.5150801579373008,
+          "f1": 0.49626374349858454
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9836065573770492,
+          "precision": 0.9915254237288136,
+          "recall": 0.8333333333333333,
+          "f1": 0.8957264957264958
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.13675394519783235,
+          "js_divergence": 0.03248669854336822
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.19323529834582587,
+        "lr": 0.000980484886752915,
+        "weight_decay": 0.00014367509113739864,
+        "backbone_lr_ratio": 0.13093310334961422,
+        "moe_n_experts": 2,
+        "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": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.3450456742192763,
+          "rmse": 0.5874058854142307,
+          "mae": 0.46413040039173903,
+          "r2": 0.31854095833662033
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.633114832363136,
+          "rmse": 0.7956851339337289,
+          "mae": 0.5481160653211797,
+          "r2": 0.15721442419210252
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6517857142857143,
+          "recall": 0.6744868035190617,
+          "f1": 0.6583050847457628
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5853276353276353,
+          "recall": 0.6349206349206349,
+          "f1": 0.5970747928194737
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 1.0,
+          "precision": 1.0,
+          "recall": 1.0,
+          "f1": 1.0
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.1542080131383728,
+          "js_divergence": 0.03808539778255157
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.19323529834582587,
+        "lr": 0.000980484886752915,
+        "weight_decay": 0.00014367509113739864,
+        "backbone_lr_ratio": 0.13093310334961422,
+        "moe_n_experts": 2,
+        "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": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.763092308495774,
+          "rmse": 1.3278148622815509,
+          "mae": 0.5504888193233006,
+          "r2": 0.0027815521929155373
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.6925828062408207,
+          "rmse": 0.8322156008155703,
+          "mae": 0.5598128846357182,
+          "r2": 0.32672722594328685
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7380952380952381,
+          "precision": 0.6612903225806452,
+          "recall": 0.6612903225806452,
+          "f1": 0.6612903225806452
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5952380952380952,
+          "precision": 0.5301311591634171,
+          "recall": 0.5438914027149321,
+          "f1": 0.5247639796659405
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.15787098167920866,
+          "js_divergence": 0.04007834983720861
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.19323529834582587,
+        "lr": 0.000980484886752915,
+        "weight_decay": 0.00014367509113739864,
+        "backbone_lr_ratio": 0.13093310334961422,
+        "moe_n_experts": 2,
+        "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": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.5697698999186622,
+          "rmse": 0.75483104064331,
+          "mae": 0.45241403512634515,
+          "r2": 0.23757709942203908
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.8494501546811344,
+          "rmse": 0.9216562019978677,
+          "mae": 0.6002610280852894,
+          "r2": 0.17859572323309747
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7380952380952381,
+          "precision": 0.6724137931034483,
+          "recall": 0.6906158357771262,
+          "f1": 0.6791666666666667
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6324786324786325,
+          "recall": 0.6799145299145298,
+          "f1": 0.6406077872744539
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 1.0,
+          "precision": 1.0,
+          "recall": 1.0,
+          "f1": 1.0
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.17354171533022728,
+          "js_divergence": 0.046613494801719065
+        }
+      }
+    }
+  ],
+  "summary_stats": {
+    "size": {
+      "mse_mean": 0.9166533855222226,
+      "mse_std": 0.6318413727420791,
+      "rmse_mean": 0.8978271435116048,
+      "rmse_std": 0.33250534716905716,
+      "mae_mean": 0.4892799429634075,
+      "mae_std": 0.06014677111861511,
+      "r2_mean": 0.12791895679232362,
+      "r2_std": 0.1264782068138924
+    },
+    "delivery": {
+      "mse_mean": 0.7282528768724562,
+      "mse_std": 0.13657894316164607,
+      "rmse_mean": 0.8496139635827478,
+      "rmse_std": 0.08005616626887421,
+      "mae_mean": 0.5644987984249458,
+      "mae_std": 0.027329238310088152,
+      "r2_mean": 0.2541703393843998,
+      "r2_std": 0.07124134045642266
+    },
+    "pdi": {
+      "accuracy_mean": 0.7166666666666667,
+      "accuracy_std": 0.021821789023599242,
+      "precision_mean": 0.6493573242827317,
+      "precision_std": 0.03194842845400596,
+      "recall_mean": 0.6660225763552736,
+      "recall_std": 0.038761053417086286,
+      "f1_mean": 0.6535154810007054,
+      "f1_std": 0.03277945989198496
+    },
+    "ee": {
+      "accuracy_mean": 0.6476190476190476,
+      "accuracy_std": 0.05664465358747772,
+      "precision_mean": 0.5956722339786855,
+      "precision_std": 0.06857668680486774,
+      "recall_mean": 0.6277137260498604,
+      "recall_std": 0.09085325454039572,
+      "f1_mean": 0.5942399928270421,
+      "f1_std": 0.07818855962865172
+    },
+    "toxic": {
+      "accuracy_mean": 0.9799397022800076,
+      "accuracy_std": 0.02448827002026046,
+      "precision_mean": 0.9394078416379934,
+      "precision_std": 0.11262714591271072,
+      "recall_mean": 0.9096491228070175,
+      "recall_std": 0.10069013931425785,
+      "f1_mean": 0.9012906464233896,
+      "f1_std": 0.0883525548692155
+    },
+    "biodist": {
+      "kl_divergence_mean": 0.16389939897510436,
+      "kl_divergence_std": 0.020316882039832528,
+      "js_divergence_mean": 0.0400615686256162,
+      "js_divergence_std": 0.0047561799696747425
+    }
+  }
+}

models/cv5_sample/moe_llm/seed42/outer_fold_0/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "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/cv5_sample/moe_llm/seed42/outer_fold_0/epoch_mean.json

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

models/cv5_sample/moe_llm/seed42/outer_fold_0/history.json

@@ -0,0 +1,271 @@
+{
+  "train": [
+    {
+      "loss": 5.9101777463345915,
+      "loss_size": 0.8219048421125155,
+      "loss_pdi": 0.6955033202429075,
+      "loss_ee": 1.060509283800383,
+      "loss_delivery": 1.1355218297323666,
+      "loss_biodist": 0.907981518958066,
+      "loss_toxic": 0.5266113526917793,
+      "loss_moe_lb": 2.0565945038924345,
+      "loss_aux_moe": 1.1926589547782331,
+      "loss_aux_llm": 1.3497599073358484
+    },
+    {
+      "loss": 4.279001815898998,
+      "loss_size": 0.5204920379092565,
+      "loss_pdi": 0.643413982681326,
+      "loss_ee": 0.9418714610305993,
+      "loss_delivery": 0.8884327468034383,
+      "loss_biodist": 0.43965512856438355,
+      "loss_toxic": 0.3611150317498156,
+      "loss_moe_lb": 2.0773819910513387,
+      "loss_aux_moe": 0.9263223323064882,
+      "loss_aux_llm": 1.095245057666624
+    },
+    {
+      "loss": 4.1005129363085775,
+      "loss_size": 0.8035221945595097,
+      "loss_pdi": 0.6013445999171283,
+      "loss_ee": 0.8792690650836842,
+      "loss_delivery": 0.9753083366951024,
+      "loss_biodist": 0.260869637333058,
+      "loss_toxic": 0.19401404719699072,
+      "loss_moe_lb": 2.0695742400916846,
+      "loss_aux_moe": 0.9956301312792946,
+      "loss_aux_llm": 1.1593363244791288
+    },
+    {
+      "loss": 3.8045457988171965,
+      "loss_size": 0.7318794419233864,
+      "loss_pdi": 0.5873883729045456,
+      "loss_ee": 0.8806033150569813,
+      "loss_delivery": 0.8927270299880892,
+      "loss_biodist": 0.24839287150550532,
+      "loss_toxic": 0.2283072332486611,
+      "loss_moe_lb": 2.1437587512506022,
+      "loss_aux_moe": 0.9384613774112753,
+      "loss_aux_llm": 1.0804562792181969
+    },
+    {
+      "loss": 3.240251125516118,
+      "loss_size": 0.6515310636243304,
+      "loss_pdi": 0.538264666860168,
+      "loss_ee": 0.8188068826456327,
+      "loss_delivery": 0.7787331332628792,
+      "loss_biodist": 0.17340928879944054,
+      "loss_toxic": 0.17186429866312727,
+      "loss_moe_lb": 2.1789169665929435,
+      "loss_aux_moe": 0.7798387190780124,
+      "loss_aux_llm": 1.0161442956006206
+    },
+    {
+      "loss": 2.756747068585576,
+      "loss_size": 0.6201924379612949,
+      "loss_pdi": 0.5108480352807689,
+      "loss_ee": 0.7891972757674552,
+      "loss_delivery": 0.5784409728807371,
+      "loss_biodist": 0.1625980019972131,
+      "loss_toxic": 0.09638489400842101,
+      "loss_moe_lb": 2.1672985134898006,
+      "loss_aux_moe": 0.6525389511842985,
+      "loss_aux_llm": 0.866097974414761
+    },
+    {
+      "loss": 2.8097413456117786,
+      "loss_size": 0.6104585203165943,
+      "loss_pdi": 0.49535534148280685,
+      "loss_ee": 0.7758352208781887,
+      "loss_delivery": 0.6814873780853845,
+      "loss_biodist": 0.15574933598573143,
+      "loss_toxic": 0.13788988649203265,
+      "loss_moe_lb": 2.149112511325527
+    },
+    {
+      "loss": 2.5547807119988106,
+      "loss_size": 0.49592660246668635,
+      "loss_pdi": 0.4421862122174856,
+      "loss_ee": 0.7361203488465902,
+      "loss_delivery": 0.6629904395825154,
+      "loss_biodist": 0.14514514481699145,
+      "loss_toxic": 0.12097700466590657,
+      "loss_moe_lb": 2.127687138480109
+    },
+    {
+      "loss": 2.496382748758471,
+      "loss_size": 0.49584831820951925,
+      "loss_pdi": 0.43293014209012726,
+      "loss_ee": 0.7089789686976252,
+      "loss_delivery": 0.7198760602119807,
+      "loss_biodist": 0.1322829470642515,
+      "loss_toxic": 0.06398353036304281,
+      "loss_moe_lb": 2.137723809963948
+    },
+    {
+      "loss": 2.3917723153088546,
+      "loss_size": 0.5447976748282845,
+      "loss_pdi": 0.47060893274642324,
+      "loss_ee": 0.6911834728073429,
+      "loss_delivery": 0.5664171859420635,
+      "loss_biodist": 0.12129703466151212,
+      "loss_toxic": 0.05351947958272806,
+      "loss_moe_lb": 2.1298333728635632
+    },
+    {
+      "loss": 2.1267837833713843,
+      "loss_size": 0.2846768913236824,
+      "loss_pdi": 0.44730447313270055,
+      "loss_ee": 0.6802649554368612,
+      "loss_delivery": 0.6072367915027851,
+      "loss_biodist": 0.12430884417247128,
+      "loss_toxic": 0.04373953997698688,
+      "loss_moe_lb": 2.14098026623597
+    }
+  ],
+  "val": [
+    {
+      "loss": 4.0116489887237545,
+      "loss_size": 0.4896596247330308,
+      "loss_pdi": 0.7048518657684326,
+      "loss_ee": 0.9747519016265869,
+      "loss_delivery": 0.7728810399770737,
+      "loss_biodist": 0.706810986995697,
+      "loss_toxic": 0.37354564368724824,
+      "loss_moe_lb": 2.0459502696990968,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.3040961265563964,
+      "loss_size": 0.445052208006382,
+      "loss_pdi": 0.7240784287452697,
+      "loss_ee": 0.8502145409584045,
+      "loss_delivery": 0.6241754226386547,
+      "loss_biodist": 0.35449219942092897,
+      "loss_toxic": 0.3185396656394005,
+      "loss_moe_lb": 2.1541085243225098,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.4476874828338624,
+      "loss_size": 0.4868860960006714,
+      "loss_pdi": 0.7957421660423278,
+      "loss_ee": 0.8902200803160667,
+      "loss_delivery": 0.712676465883851,
+      "loss_biodist": 0.18424639031291007,
+      "loss_toxic": 0.4164220209233463,
+      "loss_moe_lb": 2.1585488319396973,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 3.6464715957641602,
+      "loss_size": 0.40821063220500947,
+      "loss_pdi": 0.9250725567340851,
+      "loss_ee": 0.8542257234454155,
+      "loss_delivery": 0.7186837211251259,
+      "loss_biodist": 0.1707709789276123,
+      "loss_toxic": 0.6113252578303218,
+      "loss_moe_lb": 2.218433880805969,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 2.98254132270813,
+      "loss_size": 0.3915577203035355,
+      "loss_pdi": 0.9889228403568268,
+      "loss_ee": 0.8030516915023327,
+      "loss_delivery": 0.6232511460781097,
+      "loss_biodist": 0.18688103817403318,
+      "loss_toxic": 0.06858860440552235,
+      "loss_moe_lb": 2.2559383392333983,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.17673282623291,
+      "loss_size": 0.45360138416290285,
+      "loss_pdi": 1.150715070962906,
+      "loss_ee": 0.8275567952543497,
+      "loss_delivery": 0.613179263472557,
+      "loss_biodist": 0.1491141177713871,
+      "loss_toxic": 0.08308818615041673,
+      "loss_moe_lb": 2.248614454269409,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.5197216987609865,
+      "loss_size": 0.44484867453575133,
+      "loss_pdi": 1.3462523818016052,
+      "loss_ee": 0.8313753753900528,
+      "loss_delivery": 0.6039341509342193,
+      "loss_biodist": 0.14939168095588684,
+      "loss_toxic": 0.24860251196660102,
+      "loss_moe_lb": 2.2827333211898804,
+      "acc_pdi": 0.5454545454545454,
+      "acc_ee": 0.45454545454545453,
+      "acc_toxic": 0.9
+    },
+    {
+      "loss": 2.7240954875946044,
+      "loss_size": 0.4011200428009033,
+      "loss_pdi": 0.9446926414966583,
+      "loss_ee": 0.8310217900201678,
+      "loss_delivery": 0.4591434508562088,
+      "loss_biodist": 0.14846671670675277,
+      "loss_toxic": 0.03622314129024744,
+      "loss_moe_lb": 2.3130435943603516,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.0855964183807374,
+      "loss_size": 0.4515984505414963,
+      "loss_pdi": 1.174130117893219,
+      "loss_ee": 0.8735770538449288,
+      "loss_delivery": 0.500551962852478,
+      "loss_biodist": 0.12562411203980445,
+      "loss_toxic": 0.07660643630661071,
+      "loss_moe_lb": 2.2815173625946046,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.9979569911956787,
+      "loss_size": 0.4335483551025391,
+      "loss_pdi": 1.2359270095825194,
+      "loss_ee": 0.7991622228175401,
+      "loss_delivery": 0.481366953253746,
+      "loss_biodist": 0.12457970529794693,
+      "loss_toxic": 0.04048144668340683,
+      "loss_moe_lb": 2.2985286951065063,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 3.0956326007843016,
+      "loss_size": 0.4739636182785034,
+      "loss_pdi": 1.278123664855957,
+      "loss_ee": 0.8378281531855464,
+      "loss_delivery": 0.46590901277959346,
+      "loss_biodist": 0.12750579938292503,
+      "loss_toxic": 0.03617289457470178,
+      "loss_moe_lb": 2.2634385585784913,
+      "acc_pdi": 0.5757575757575758,
+      "acc_ee": 0.48484848484848486,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/cv5_sample/moe_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/cv5_sample/moe_llm/seed42/outer_fold_0/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 83,
+    "mse": 1.4858819540561683,
+    "rmse": 1.2189675771144073,
+    "mae": 0.5309677058494234,
+    "r2": 0.0906036692787614
+  },
+  "delivery": {
+    "n_samples": 58,
+    "mse": 0.4578579988449137,
+    "rmse": 0.6766520515338099,
+    "mae": 0.44371995239936074,
+    "r2": 0.41803485452290234
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7857142857142857,
+    "precision": 0.7322033898305085,
+    "recall": 0.7441197434069851,
+    "f1": 0.7375
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.7380952380952381,
+    "precision": 0.7146208708708709,
+    "recall": 0.7547619047619047,
+    "f1": 0.7160330091364574
+  },
+  "toxic": {
+    "n_samples": 58,
+    "accuracy": 0.9827586206896551,
+    "precision": 0.9912280701754386,
+    "recall": 0.75,
+    "f1": 0.8289085545722714
+  },
+  "biodist": {
+    "n_samples": 58,
+    "kl_divergence": 0.2124931399005794,
+    "js_divergence": 0.0436771104854127
+  }
+}

models/cv5_sample/moe_llm/seed42/outer_fold_1/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "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/cv5_sample/moe_llm/seed42/outer_fold_1/epoch_mean.json

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

models/cv5_sample/moe_llm/seed42/outer_fold_1/history.json

@@ -0,0 +1,271 @@
+{
+  "train": [
+    {
+      "loss": 5.844873963175593,
+      "loss_size": 1.1697139462103714,
+      "loss_pdi": 0.6559584736824036,
+      "loss_ee": 0.9990178346633911,
+      "loss_delivery": 0.9326752996122515,
+      "loss_biodist": 0.9255792949650739,
+      "loss_toxic": 0.5104590783248076,
+      "loss_moe_lb": 2.0308946306641036,
+      "loss_aux_moe": 1.0037819767320477,
+      "loss_aux_llm": 1.1631469360074482
+    },
+    {
+      "loss": 4.553732060097359,
+      "loss_size": 1.1190044841653592,
+      "loss_pdi": 0.5833910236487517,
+      "loss_ee": 0.8829387386102934,
+      "loss_delivery": 0.7804326554005211,
+      "loss_biodist": 0.4103407849733894,
+      "loss_toxic": 0.30222512459432754,
+      "loss_moe_lb": 2.0332680811753145,
+      "loss_aux_moe": 0.8820628979922952,
+      "loss_aux_llm": 1.0768901111709106
+    },
+    {
+      "loss": 4.165987249967214,
+      "loss_size": 1.0680863821023219,
+      "loss_pdi": 0.5564941554456144,
+      "loss_ee": 0.8305287328926293,
+      "loss_delivery": 0.8199278615616463,
+      "loss_biodist": 0.30551903876098424,
+      "loss_toxic": 0.23546614855326511,
+      "loss_moe_lb": 2.1472520892684526,
+      "loss_aux_moe": 0.9134891435906694,
+      "loss_aux_llm": 1.0313057750463486
+    },
+    {
+      "loss": 3.887093511787621,
+      "loss_size": 1.103906466147384,
+      "loss_pdi": 0.5772485410844957,
+      "loss_ee": 0.8390914693072036,
+      "loss_delivery": 0.7768730807143289,
+      "loss_biodist": 0.21761337989890897,
+      "loss_toxic": 0.1712454239258895,
+      "loss_moe_lb": 2.1270299988823966,
+      "loss_aux_moe": 0.7972461388672929,
+      "loss_aux_llm": 1.0023344352438643
+    },
+    {
+      "loss": 3.5329479301297986,
+      "loss_size": 1.0075848493725061,
+      "loss_pdi": 0.5278141957682532,
+      "loss_ee": 0.7903922522390211,
+      "loss_delivery": 0.7493079117423779,
+      "loss_biodist": 0.216409088087243,
+      "loss_toxic": 0.14804497885366752,
+      "loss_moe_lb": 2.1215679355569788,
+      "loss_aux_moe": 0.7791867956921861,
+      "loss_aux_llm": 0.9144315647112357
+    },
+    {
+      "loss": 3.140749799238669,
+      "loss_size": 0.8968957497461422,
+      "loss_pdi": 0.4775858110672719,
+      "loss_ee": 0.7530861561362808,
+      "loss_delivery": 0.6636832944444708,
+      "loss_biodist": 0.1660661126512128,
+      "loss_toxic": 0.1779262529673508,
+      "loss_moe_lb": 2.1186544766297213,
+      "loss_aux_moe": 0.6614102075832922,
+      "loss_aux_llm": 0.9802552050835377
+    },
+    {
+      "loss": 2.982136368751526,
+      "loss_size": 0.8015535189896017,
+      "loss_pdi": 0.4799974121757456,
+      "loss_ee": 0.7338198085088987,
+      "loss_delivery": 0.6604716744575951,
+      "loss_biodist": 0.17151938232819777,
+      "loss_toxic": 0.1764316958488544,
+      "loss_moe_lb": 2.116116572070766
+    },
+    {
+      "loss": 2.6167806999103442,
+      "loss_size": 0.6552065660220545,
+      "loss_pdi": 0.4558954645653029,
+      "loss_ee": 0.6795905675437,
+      "loss_delivery": 0.5747650408664265,
+      "loss_biodist": 0.1658533130948608,
+      "loss_toxic": 0.12498937264663747,
+      "loss_moe_lb": 2.0987441733076766
+    },
+    {
+      "loss": 2.564945269275356,
+      "loss_size": 0.701622544812995,
+      "loss_pdi": 0.4624234595008799,
+      "loss_ee": 0.6582107753367037,
+      "loss_delivery": 0.58780188199032,
+      "loss_biodist": 0.16030761096123103,
+      "loss_toxic": 0.04463169629987631,
+      "loss_moe_lb": 2.111907923543775
+    },
+    {
+      "loss": 2.6338622086756938,
+      "loss_size": 0.7301617483230861,
+      "loss_pdi": 0.43557040232258876,
+      "loss_ee": 0.6678765420172665,
+      "loss_delivery": 0.6153083404375089,
+      "loss_biodist": 0.15970830488446597,
+      "loss_toxic": 0.07487160047025394,
+      "loss_moe_lb": 2.1209699076575204
+    },
+    {
+      "loss": 2.4515427640966467,
+      "loss_size": 0.6441191072399551,
+      "loss_pdi": 0.4289833254910804,
+      "loss_ee": 0.680172837263829,
+      "loss_delivery": 0.5573766311076848,
+      "loss_biodist": 0.14778587272441066,
+      "loss_toxic": 0.04419851102362815,
+      "loss_moe_lb": 2.0920171866545805
+    }
+  ],
+  "val": [
+    {
+      "loss": 2.8340189456939697,
+      "loss_size": 0.5454001704230904,
+      "loss_pdi": 0.6431448519229889,
+      "loss_ee": 0.9064770460128784,
+      "loss_delivery": 0.19316499456763267,
+      "loss_biodist": 0.4068158447742462,
+      "loss_toxic": 0.1233778491616249,
+      "loss_moe_lb": 2.032278299331665,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.3471620321273803,
+      "loss_size": 0.4725328356027603,
+      "loss_pdi": 0.6254292964935303,
+      "loss_ee": 0.7353704273700714,
+      "loss_delivery": 0.21038319170475006,
+      "loss_biodist": 0.2688927501440048,
+      "loss_toxic": 0.025847082026302814,
+      "loss_moe_lb": 2.2519073486328125,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.248750996589661,
+      "loss_size": 0.41793634593486784,
+      "loss_pdi": 0.6494660258293152,
+      "loss_ee": 0.6911145269870758,
+      "loss_delivery": 0.2551379919052124,
+      "loss_biodist": 0.23418427407741546,
+      "loss_toxic": 0.012842147471383214,
+      "loss_moe_lb": 2.597307586669922,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.2044697284698485,
+      "loss_size": 0.3363520873710513,
+      "loss_pdi": 0.6014128744602203,
+      "loss_ee": 0.7720682263374329,
+      "loss_delivery": 0.285314205288887,
+      "loss_biodist": 0.21102917790412903,
+      "loss_toxic": 0.024753990955650805,
+      "loss_moe_lb": 2.42631254196167,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.172011137008667,
+      "loss_size": 0.4511841475963593,
+      "loss_pdi": 0.6105897426605225,
+      "loss_ee": 0.7073207706212997,
+      "loss_delivery": 0.24147587269544601,
+      "loss_biodist": 0.18464208096265794,
+      "loss_toxic": 0.0065372010692954065,
+      "loss_moe_lb": 2.4683003425598145,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.072033905982971,
+      "loss_size": 0.4331493258476257,
+      "loss_pdi": 0.5731127798557282,
+      "loss_ee": 0.6349433034658432,
+      "loss_delivery": 0.2421866089105606,
+      "loss_biodist": 0.21265464127063752,
+      "loss_toxic": 0.0034104947932064532,
+      "loss_moe_lb": 2.45847373008728,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.0299159049987794,
+      "loss_size": 0.44897343069314954,
+      "loss_pdi": 0.5645411908626556,
+      "loss_ee": 0.6812305867671966,
+      "loss_delivery": 0.1939051305875182,
+      "loss_biodist": 0.17333385050296785,
+      "loss_toxic": 0.0041761530097574,
+      "loss_moe_lb": 2.4599356651306152,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.7272727272727273,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.059905481338501,
+      "loss_size": 0.48319281786680224,
+      "loss_pdi": 0.5864004135131836,
+      "loss_ee": 0.6532717317342758,
+      "loss_delivery": 0.1758626704569906,
+      "loss_biodist": 0.19589750468730927,
+      "loss_toxic": 0.003806475317105651,
+      "loss_moe_lb": 2.442372369766235,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.0781530857086183,
+      "loss_size": 0.49511486291885376,
+      "loss_pdi": 0.5853226751089096,
+      "loss_ee": 0.6847641676664352,
+      "loss_delivery": 0.14788594823330642,
+      "loss_biodist": 0.20549408048391343,
+      "loss_toxic": 0.002917639911174774,
+      "loss_moe_lb": 2.4569051265716553,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.016889500617981,
+      "loss_size": 0.4801466763019562,
+      "loss_pdi": 0.5440871745347977,
+      "loss_ee": 0.6530755490064621,
+      "loss_delivery": 0.17224864036543294,
+      "loss_biodist": 0.20416025221347808,
+      "loss_toxic": 0.00327652373816818,
+      "loss_moe_lb": 2.4228747844696046,
+      "acc_pdi": 0.8484848484848485,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 1.995576024055481,
+      "loss_size": 0.4439514249563217,
+      "loss_pdi": 0.5852887749671936,
+      "loss_ee": 0.6712759345769882,
+      "loss_delivery": 0.1613160498905927,
+      "loss_biodist": 0.17435047775506973,
+      "loss_toxic": 0.0036445653764531015,
+      "loss_moe_lb": 2.464684581756592,
+      "acc_pdi": 0.8181818181818182,
+      "acc_ee": 0.7575757575757576,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/cv5_sample/moe_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/cv5_sample/moe_llm/seed42/outer_fold_1/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.3521366695033162,
+    "rmse": 0.593411045990312,
+    "mae": 0.417994069290303,
+    "r2": -0.07530019609216088
+  },
+  "delivery": {
+    "n_samples": 61,
+    "mse": 0.8615718355352413,
+    "rmse": 0.9282089395902419,
+    "mae": 0.54733644703739,
+    "r2": 0.34785223029308054
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.6309523809523809,
+    "precision": 0.5417771883289124,
+    "recall": 0.5476190476190477,
+    "f1": 0.5421135924037278
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.5952380952380952,
+    "precision": 0.531971296677179,
+    "recall": 0.5323723895152467,
+    "f1": 0.5202020202020202
+  },
+  "toxic": {
+    "n_samples": 61,
+    "accuracy": 0.9836065573770492,
+    "precision": 0.9915254237288136,
+    "recall": 0.8333333333333333,
+    "f1": 0.8957264957264958
+  },
+  "biodist": {
+    "n_samples": 61,
+    "kl_divergence": 0.12922332899139097,
+    "js_divergence": 0.030354793549857463
+  }
+}

models/cv5_sample/moe_llm/seed42/outer_fold_2/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "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/cv5_sample/moe_llm/seed42/outer_fold_2/epoch_mean.json

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

models/cv5_sample/moe_llm/seed42/outer_fold_2/history.json

@@ -0,0 +1,271 @@
+{
+  "train": [
+    {
+      "loss": 5.448386102109342,
+      "loss_size": 0.8643624792630608,
+      "loss_pdi": 0.6676707122776959,
+      "loss_ee": 1.0198738494434871,
+      "loss_delivery": 0.8485204742063541,
+      "loss_biodist": 0.9448029020348111,
+      "loss_toxic": 0.4876952078696844,
+      "loss_moe_lb": 2.093161576503032,
+      "loss_aux_moe": 0.8641154044181913,
+      "loss_aux_llm": 1.1632047895644162
+    },
+    {
+      "loss": 4.324384463799967,
+      "loss_size": 0.825573201517801,
+      "loss_pdi": 0.5958146203208614,
+      "loss_ee": 0.9189161110568691,
+      "loss_delivery": 0.7985159166761346,
+      "loss_biodist": 0.43780089069057154,
+      "loss_toxic": 0.2720214316772448,
+      "loss_moe_lb": 2.1437764683285274,
+      "loss_aux_moe": 0.8986711423542049,
+      "loss_aux_llm": 1.0598642354881442
+    },
+    {
+      "loss": 3.8106770064379716,
+      "loss_size": 0.7928319295515885,
+      "loss_pdi": 0.5804953180454873,
+      "loss_ee": 0.8457762880905254,
+      "loss_delivery": 0.7977923632070825,
+      "loss_biodist": 0.2861963832700575,
+      "loss_toxic": 0.14996869308320251,
+      "loss_moe_lb": 2.0927669937546187,
+      "loss_aux_moe": 0.9209618949004121,
+      "loss_aux_llm": 1.0866852402687073
+    },
+    {
+      "loss": 3.590217042613674,
+      "loss_size": 0.7697639918407878,
+      "loss_pdi": 0.6194508639541833,
+      "loss_ee": 0.8710581531395784,
+      "loss_delivery": 0.7570703561241562,
+      "loss_biodist": 0.2637654389884021,
+      "loss_toxic": 0.128640679239824,
+      "loss_moe_lb": 2.2199561015979663,
+      "loss_aux_moe": 0.7979112141639799,
+      "loss_aux_llm": 0.8974349112124056
+    },
+    {
+      "loss": 3.2631049865000956,
+      "loss_size": 0.7027561263860883,
+      "loss_pdi": 0.540183317822379,
+      "loss_ee": 0.8333719823811505,
+      "loss_delivery": 0.7394266468850342,
+      "loss_biodist": 0.20225958405314265,
+      "loss_toxic": 0.15996630871794312,
+      "loss_moe_lb": 2.2460104194847315,
+      "loss_aux_moe": 0.7544193904142122,
+      "loss_aux_llm": 0.8569866848757138
+    },
+    {
+      "loss": 2.9335989533243954,
+      "loss_size": 0.6793409892030664,
+      "loss_pdi": 0.5341143374507492,
+      "loss_ee": 0.8107952403055655,
+      "loss_delivery": 0.7013196065216452,
+      "loss_biodist": 0.17505165954699387,
+      "loss_toxic": 0.0537688604703903,
+      "loss_moe_lb": 2.211098677403218,
+      "loss_aux_moe": 0.6585364743463091,
+      "loss_aux_llm": 0.8589774281591982
+    },
+    {
+      "loss": 2.572078147450009,
+      "loss_size": 0.6022349424861573,
+      "loss_pdi": 0.5041297573495556,
+      "loss_ee": 0.7460332577292984,
+      "loss_delivery": 0.5461117913191383,
+      "loss_biodist": 0.1844338054793912,
+      "loss_toxic": 0.046531276912090205,
+      "loss_moe_lb": 2.26575459016336
+    },
+    {
+      "loss": 2.613092970203709,
+      "loss_size": 0.5582844796212943,
+      "loss_pdi": 0.5014023039792035,
+      "loss_ee": 0.7368455045932049,
+      "loss_delivery": 0.5232992003092894,
+      "loss_biodist": 0.1375063386921947,
+      "loss_toxic": 0.20026905486687413,
+      "loss_moe_lb": 2.2357002980000265
+    },
+    {
+      "loss": 2.414530187039762,
+      "loss_size": 0.5297113753653862,
+      "loss_pdi": 0.47427408356924317,
+      "loss_ee": 0.7108970058930887,
+      "loss_delivery": 0.5342327124766402,
+      "loss_biodist": 0.14476421435137052,
+      "loss_toxic": 0.07928779668401222,
+      "loss_moe_lb": 2.235232720503936
+    },
+    {
+      "loss": 2.374394999968039,
+      "loss_size": 0.5287634996143548,
+      "loss_pdi": 0.46810457754779505,
+      "loss_ee": 0.6844177930741697,
+      "loss_delivery": 0.527905018496755,
+      "loss_biodist": 0.18576635964013435,
+      "loss_toxic": 0.0401537462823546,
+      "loss_moe_lb": 2.2262662423623576
+    },
+    {
+      "loss": 2.186512160945583,
+      "loss_size": 0.4162472438167881,
+      "loss_pdi": 0.4552439733131512,
+      "loss_ee": 0.6912266974513596,
+      "loss_delivery": 0.4482238549437072,
+      "loss_biodist": 0.14196190117178736,
+      "loss_toxic": 0.08895144933941648,
+      "loss_moe_lb": 2.2699929765752844
+    }
+  ],
+  "val": [
+    {
+      "loss": 7.46802167892456,
+      "loss_size": 3.2265774309635162,
+      "loss_pdi": 0.7353087663650513,
+      "loss_ee": 1.0791295766830444,
+      "loss_delivery": 1.8187558829784394,
+      "loss_biodist": 0.45695706605911257,
+      "loss_toxic": 0.15167853236198425,
+      "loss_moe_lb": 2.1358907222747803,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 7.090808343887329,
+      "loss_size": 3.410844311118126,
+      "loss_pdi": 0.7231593370437622,
+      "loss_ee": 1.0354070901870727,
+      "loss_delivery": 1.6264947444200515,
+      "loss_biodist": 0.2861800193786621,
+      "loss_toxic": 0.063536443002522,
+      "loss_moe_lb": 2.133509969711304,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.772978925704956,
+      "loss_size": 3.1503812193870546,
+      "loss_pdi": 0.7274282336235046,
+      "loss_ee": 1.085149884223938,
+      "loss_delivery": 1.600457674264908,
+      "loss_biodist": 0.2441837251186371,
+      "loss_toxic": 0.05247659794986248,
+      "loss_moe_lb": 2.238656997680664,
+      "acc_pdi": 0.6060606060606061,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.508674764633179,
+      "loss_size": 3.038449627161026,
+      "loss_pdi": 0.7216925919055939,
+      "loss_ee": 1.0778231978416444,
+      "loss_delivery": 1.5543653905391692,
+      "loss_biodist": 0.14704219698905946,
+      "loss_toxic": 0.06883886940777302,
+      "loss_moe_lb": 2.328789567947388,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 6.362945938110352,
+      "loss_size": 3.0777094930410387,
+      "loss_pdi": 0.6802633225917816,
+      "loss_ee": 1.05179603099823,
+      "loss_delivery": 1.5254353404045105,
+      "loss_biodist": 0.10500667095184327,
+      "loss_toxic": 0.028991057851817457,
+      "loss_moe_lb": 2.4010870933532713,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.95
+    },
+    {
+      "loss": 5.968770718574524,
+      "loss_size": 2.767032378911972,
+      "loss_pdi": 0.6399851977825165,
+      "loss_ee": 1.0736777305603027,
+      "loss_delivery": 1.4773260951042175,
+      "loss_biodist": 0.11008006632328034,
+      "loss_toxic": 0.014215384027920663,
+      "loss_moe_lb": 2.367847776412964,
+      "acc_pdi": 0.6363636363636364,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 6.202451419830322,
+      "loss_size": 3.0308651685714723,
+      "loss_pdi": 0.6326265096664428,
+      "loss_ee": 1.0713923573493958,
+      "loss_delivery": 1.4624701291322708,
+      "loss_biodist": 0.11829227358102798,
+      "loss_toxic": 0.0033330414851661773,
+      "loss_moe_lb": 2.3702832221984864,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5151515151515151,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 6.452511215209961,
+      "loss_size": 3.152136731147766,
+      "loss_pdi": 0.7097072839736939,
+      "loss_ee": 1.0254977822303772,
+      "loss_delivery": 1.543718020617962,
+      "loss_biodist": 0.13677483052015305,
+      "loss_toxic": 0.0033543511119205503,
+      "loss_moe_lb": 2.3397052764892576,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 6.112395906448365,
+      "loss_size": 2.9286181777715683,
+      "loss_pdi": 0.6871815204620362,
+      "loss_ee": 1.0443474650382996,
+      "loss_delivery": 1.444450181722641,
+      "loss_biodist": 0.12184233590960503,
+      "loss_toxic": 0.004058638200513087,
+      "loss_moe_lb": 2.3976557731628416,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 6.221766138076783,
+      "loss_size": 3.082172819972038,
+      "loss_pdi": 0.670418244600296,
+      "loss_ee": 1.0383009672164918,
+      "loss_delivery": 1.4141629114747047,
+      "loss_biodist": 0.1265248589217663,
+      "loss_toxic": 0.005223917227704078,
+      "loss_moe_lb": 2.367207431793213,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 6.086094331741333,
+      "loss_size": 3.0313263446092606,
+      "loss_pdi": 0.6579509675502777,
+      "loss_ee": 1.0429080128669739,
+      "loss_delivery": 1.34516292065382,
+      "loss_biodist": 0.1171148531138897,
+      "loss_toxic": 0.005285785451997071,
+      "loss_moe_lb": 2.3800329685211183,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 1.0
+    }
+  ]
+}

models/cv5_sample/moe_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/cv5_sample/moe_llm/seed42/outer_fold_2/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.4170957202870868,
+    "rmse": 0.6458294823613171,
+    "mae": 0.4985973424876907,
+    "r2": 0.1762434046684933
+  },
+  "delivery": {
+    "n_samples": 60,
+    "mse": 0.5854544655957109,
+    "rmse": 0.7651499628149445,
+    "mae": 0.49963968732239056,
+    "r2": 0.22065863304023947
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7380952380952381,
+    "precision": 0.6851851851851851,
+    "recall": 0.7199413489736071,
+    "f1": 0.693633952254642
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6666666666666666,
+    "precision": 0.613035584604212,
+    "recall": 0.6666666666666666,
+    "f1": 0.627436739505705
+  },
+  "toxic": {
+    "n_samples": 61,
+    "accuracy": 1.0,
+    "precision": 1.0,
+    "recall": 1.0,
+    "f1": 1.0
+  },
+  "biodist": {
+    "n_samples": 60,
+    "kl_divergence": 0.17218262970209255,
+    "js_divergence": 0.039882423974733884
+  }
+}

models/cv5_sample/moe_llm/seed42/outer_fold_3/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "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/cv5_sample/moe_llm/seed42/outer_fold_3/epoch_mean.json

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

models/cv5_sample/moe_llm/seed42/outer_fold_3/history.json

@@ -0,0 +1,271 @@
+{
+  "train": [
+    {
+      "loss": 5.186875691285005,
+      "loss_size": 0.5132114932545133,
+      "loss_pdi": 0.6726247926016111,
+      "loss_ee": 1.0378613858609587,
+      "loss_delivery": 0.9076983817928547,
+      "loss_biodist": 0.9157768181852393,
+      "loss_toxic": 0.5092847890950538,
+      "loss_moe_lb": 2.0349878427144645,
+      "loss_aux_moe": 1.010407390223967,
+      "loss_aux_llm": 1.083440064981177
+    },
+    {
+      "loss": 3.907044636236655,
+      "loss_size": 0.45616640432460887,
+      "loss_pdi": 0.6273065104677871,
+      "loss_ee": 0.9086244025745908,
+      "loss_delivery": 0.8223057431546418,
+      "loss_biodist": 0.37386319846720306,
+      "loss_toxic": 0.2553665132538692,
+      "loss_moe_lb": 2.043943031414135,
+      "loss_aux_moe": 0.9434679752668819,
+      "loss_aux_llm": 0.9746431202904599
+    },
+    {
+      "loss": 3.3520812343906714,
+      "loss_size": 0.43402249788915787,
+      "loss_pdi": 0.5923634042610993,
+      "loss_ee": 0.8609370621475013,
+      "loss_delivery": 0.7527698189825625,
+      "loss_biodist": 0.23749435713162292,
+      "loss_toxic": 0.14371532126254327,
+      "loss_moe_lb": 2.1175068971273063,
+      "loss_aux_moe": 0.9599171609089181,
+      "loss_aux_llm": 0.9244846647252908
+    },
+    {
+      "loss": 3.13085011855976,
+      "loss_size": 0.4183622161681588,
+      "loss_pdi": 0.5564500198170945,
+      "loss_ee": 0.8595020851573428,
+      "loss_delivery": 0.7666488471063407,
+      "loss_biodist": 0.2030809574030541,
+      "loss_toxic": 0.13739074190852005,
+      "loss_moe_lb": 2.0961369018296936,
+      "loss_aux_moe": 0.8233248805375518,
+      "loss_aux_llm": 0.9434398358335366
+    },
+    {
+      "loss": 2.7858061146091773,
+      "loss_size": 0.39161805264853145,
+      "loss_pdi": 0.566842187900801,
+      "loss_ee": 0.7711971643808726,
+      "loss_delivery": 0.6954803979527708,
+      "loss_biodist": 0.18086066509823542,
+      "loss_toxic": 0.09432240759928685,
+      "loss_moe_lb": 2.094309987248601,
+      "loss_aux_moe": 0.7374768013486991,
+      "loss_aux_llm": 0.9944169878758289
+    },
+    {
+      "loss": 2.470173813201286,
+      "loss_size": 0.3628059673953701,
+      "loss_pdi": 0.5112751488750046,
+      "loss_ee": 0.7295194646796664,
+      "loss_delivery": 0.6388976760208607,
+      "loss_biodist": 0.16272256372345462,
+      "loss_toxic": 0.0796801333057669,
+      "loss_moe_lb": 2.0946719453141496,
+      "loss_aux_moe": 0.6715661742397256,
+      "loss_aux_llm": 0.8760432788246387
+    },
+    {
+      "loss": 2.361145433541891,
+      "loss_size": 0.3760294072531365,
+      "loss_pdi": 0.5063537709616326,
+      "loss_ee": 0.7174645541487513,
+      "loss_delivery": 0.5734429028187249,
+      "loss_biodist": 0.15295277354685036,
+      "loss_toxic": 0.09341773955698905,
+      "loss_moe_lb": 2.107072810868959
+    },
+    {
+      "loss": 2.285447710269206,
+      "loss_size": 0.33821960578899124,
+      "loss_pdi": 0.5100937501804249,
+      "loss_ee": 0.7393307581141189,
+      "loss_delivery": 0.5426343622038493,
+      "loss_biodist": 0.12398734537733568,
+      "loss_toxic": 0.09918186348875531,
+      "loss_moe_lb": 2.0946988995010787
+    },
+    {
+      "loss": 2.1923618091119303,
+      "loss_size": 0.3291457059818345,
+      "loss_pdi": 0.48687964556990443,
+      "loss_ee": 0.7320111675842388,
+      "loss_delivery": 0.5369202897355363,
+      "loss_biodist": 0.10976744191469373,
+      "loss_toxic": 0.0715399857372914,
+      "loss_moe_lb": 2.107653498649597
+    },
+    {
+      "loss": 2.1388532954293327,
+      "loss_size": 0.289924700970988,
+      "loss_pdi": 0.4686911090000256,
+      "loss_ee": 0.6864962835569639,
+      "loss_delivery": 0.5863513408883197,
+      "loss_biodist": 0.12377736652018251,
+      "loss_toxic": 0.05824326928956686,
+      "loss_moe_lb": 2.0926922881925427
+    },
+    {
+      "loss": 2.0976792155085384,
+      "loss_size": 0.31547431859212954,
+      "loss_pdi": 0.47137918746149216,
+      "loss_ee": 0.6413027813305726,
+      "loss_delivery": 0.57016261661979,
+      "loss_biodist": 0.12815634000140266,
+      "loss_toxic": 0.04252448489077742,
+      "loss_moe_lb": 2.0866256404567407
+    }
+  ],
+  "val": [
+    {
+      "loss": 5.893757343292236,
+      "loss_size": 2.7006002753973006,
+      "loss_pdi": 0.6564314365386963,
+      "loss_ee": 0.8844664573669434,
+      "loss_delivery": 0.9374306380748749,
+      "loss_biodist": 0.38643685579299925,
+      "loss_toxic": 0.312559700012207,
+      "loss_moe_lb": 2.0757519721984865,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.696969696969697,
+      "acc_toxic": 0.782608695652174
+    },
+    {
+      "loss": 5.338843441009521,
+      "loss_size": 2.6443060830235483,
+      "loss_pdi": 0.6313182950019837,
+      "loss_ee": 0.7473037540912628,
+      "loss_delivery": 0.9139358669519424,
+      "loss_biodist": 0.16414627730846404,
+      "loss_toxic": 0.21320655420422555,
+      "loss_moe_lb": 2.070327711105347,
+      "acc_pdi": 0.7272727272727273,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.012914252281189,
+      "loss_size": 2.5547497063875197,
+      "loss_pdi": 0.5505357623100281,
+      "loss_ee": 0.6987002611160278,
+      "loss_delivery": 0.823940533399582,
+      "loss_biodist": 0.15457155257463456,
+      "loss_toxic": 0.19016381772235036,
+      "loss_moe_lb": 2.116373634338379,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.13802444934845,
+      "loss_size": 2.5034602016210554,
+      "loss_pdi": 0.554446917772293,
+      "loss_ee": 0.737292206287384,
+      "loss_delivery": 0.9933651030063629,
+      "loss_biodist": 0.12027869075536728,
+      "loss_toxic": 0.2056742062792182,
+      "loss_moe_lb": 2.114448404312134,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 5.1256492137908936,
+      "loss_size": 2.59128378033638,
+      "loss_pdi": 0.5555723905563354,
+      "loss_ee": 0.8101335763931274,
+      "loss_delivery": 0.8393854141235352,
+      "loss_biodist": 0.1212102085351944,
+      "loss_toxic": 0.18476905608549715,
+      "loss_moe_lb": 2.1382421493530273,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6666666666666666,
+      "acc_toxic": 0.8695652173913043
+    },
+    {
+      "loss": 4.646122789382934,
+      "loss_size": 2.2567166566848753,
+      "loss_pdi": 0.5647770702838898,
+      "loss_ee": 0.8054781496524811,
+      "loss_delivery": 0.8111683785915375,
+      "loss_biodist": 0.11808478087186813,
+      "loss_toxic": 0.08715959303081036,
+      "loss_moe_lb": 2.135502338409424,
+      "acc_pdi": 0.7878787878787878,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.9130434782608695
+    },
+    {
+      "loss": 4.743265771865845,
+      "loss_size": 2.4424239844083786,
+      "loss_pdi": 0.5591305613517761,
+      "loss_ee": 0.869558048248291,
+      "loss_delivery": 0.6986709222197532,
+      "loss_biodist": 0.11900895535945892,
+      "loss_toxic": 0.04972788635641336,
+      "loss_moe_lb": 2.146485376358032,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.451839828491211,
+      "loss_size": 2.216589081287384,
+      "loss_pdi": 0.4908704340457916,
+      "loss_ee": 0.8932434380054474,
+      "loss_delivery": 0.6621725678443908,
+      "loss_biodist": 0.11659160852432252,
+      "loss_toxic": 0.07040527649223804,
+      "loss_moe_lb": 2.1616697788238524,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9565217391304348
+    },
+    {
+      "loss": 4.532372415065765,
+      "loss_size": 2.252890150249004,
+      "loss_pdi": 0.5175323605537414,
+      "loss_ee": 0.9039882540702819,
+      "loss_delivery": 0.6627729907631874,
+      "loss_biodist": 0.12313011735677719,
+      "loss_toxic": 0.0710445934906602,
+      "loss_moe_lb": 2.1507219314575194,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 4.523290491104126,
+      "loss_size": 2.22698954641819,
+      "loss_pdi": 0.5167619049549103,
+      "loss_ee": 0.8895878970623017,
+      "loss_delivery": 0.6798219382762909,
+      "loss_biodist": 0.12791793793439865,
+      "loss_toxic": 0.07941748769953846,
+      "loss_moe_lb": 2.1426663398742676,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9130434782608695
+    },
+    {
+      "loss": 4.4532647132873535,
+      "loss_size": 2.1416719794273376,
+      "loss_pdi": 0.522536426782608,
+      "loss_ee": 0.9024902045726776,
+      "loss_delivery": 0.681278795003891,
+      "loss_biodist": 0.11594324558973312,
+      "loss_toxic": 0.091428601462394,
+      "loss_moe_lb": 2.1225868701934814,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9130434782608695
+    }
+  ]
+}

models/cv5_sample/moe_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/cv5_sample/moe_llm/seed42/outer_fold_3/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 83,
+    "mse": 1.773779594197195,
+    "rmse": 1.331833170557482,
+    "mae": 0.5606547438887408,
+    "r2": -0.003263258057283336
+  },
+  "delivery": {
+    "n_samples": 59,
+    "mse": 0.6480387215025862,
+    "rmse": 0.8050085226272988,
+    "mae": 0.5227914613403254,
+    "r2": 0.37002936863191194
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.75,
+    "precision": 0.6789023521026372,
+    "recall": 0.684017595307918,
+    "f1": 0.6813008130081302
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.6428571428571429,
+    "precision": 0.5700750584471516,
+    "recall": 0.5891402714932127,
+    "f1": 0.5736800630417652
+  },
+  "toxic": {
+    "n_samples": 60,
+    "accuracy": 0.9833333333333333,
+    "precision": 0.875,
+    "recall": 0.9912280701754386,
+    "f1": 0.9241466498103665
+  },
+  "biodist": {
+    "n_samples": 60,
+    "kl_divergence": 0.15399347135454697,
+    "js_divergence": 0.036819733414228055
+  }
+}

models/cv5_sample/moe_llm/seed42/outer_fold_4/best_params.json

@@ -0,0 +1,16 @@
+{
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "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/cv5_sample/moe_llm/seed42/outer_fold_4/epoch_mean.json

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

models/cv5_sample/moe_llm/seed42/outer_fold_4/history.json

@@ -0,0 +1,271 @@
+{
+  "train": [
+    {
+      "loss": 5.859205187977971,
+      "loss_size": 1.117435747103111,
+      "loss_pdi": 0.67151229768186,
+      "loss_ee": 1.0485850894773328,
+      "loss_delivery": 0.9128926674256453,
+      "loss_biodist": 0.9507973387434676,
+      "loss_toxic": 0.5287212576415088,
+      "loss_moe_lb": 2.1035370407877743,
+      "loss_aux_moe": 0.964274688951067,
+      "loss_aux_llm": 1.1332960209330998
+    },
+    {
+      "loss": 4.493758427130209,
+      "loss_size": 1.0583569308793224,
+      "loss_pdi": 0.6126068854654158,
+      "loss_ee": 0.9225417810517389,
+      "loss_delivery": 0.8066660902790121,
+      "loss_biodist": 0.4459190773400101,
+      "loss_toxic": 0.22362688823124846,
+      "loss_moe_lb": 2.0593858892853194,
+      "loss_aux_moe": 0.8353829867131001,
+      "loss_aux_llm": 0.9469703129819922
+    },
+    {
+      "loss": 3.938486898267591,
+      "loss_size": 0.9910998611232719,
+      "loss_pdi": 0.5878341842342067,
+      "loss_ee": 0.8390434464892825,
+      "loss_delivery": 0.7174631838460226,
+      "loss_biodist": 0.2919453759048436,
+      "loss_toxic": 0.1582938922267105,
+      "loss_moe_lb": 2.0062548792040027,
+      "loss_aux_moe": 0.9655985453644315,
+      "loss_aux_llm": 1.0280679303246576
+    },
+    {
+      "loss": 3.858765460349418,
+      "loss_size": 0.9830634577854259,
+      "loss_pdi": 0.5594958969064661,
+      "loss_ee": 0.8622410829002792,
+      "loss_delivery": 0.814423064525063,
+      "loss_biodist": 0.2729547468794359,
+      "loss_toxic": 0.16985372757237102,
+      "loss_moe_lb": 2.0284129091211267,
+      "loss_aux_moe": 0.7865037793243254,
+      "loss_aux_llm": 1.015865511588148
+    },
+    {
+      "loss": 3.2949182278401143,
+      "loss_size": 0.9350303498071593,
+      "loss_pdi": 0.5480591703105617,
+      "loss_ee": 0.7708510216828939,
+      "loss_delivery": 0.6737078446391467,
+      "loss_biodist": 0.18565125703006177,
+      "loss_toxic": 0.0994621277031665,
+      "loss_moe_lb": 2.0545133126748576,
+      "loss_aux_moe": 0.682390890713479,
+      "loss_aux_llm": 0.9088856340260119
+    },
+    {
+      "loss": 3.1421183669889294,
+      "loss_size": 0.9002627499804303,
+      "loss_pdi": 0.540199404632723,
+      "loss_ee": 0.7999422775732504,
+      "loss_delivery": 0.6205245248369269,
+      "loss_biodist": 0.20558074076433439,
+      "loss_toxic": 0.089201840867817,
+      "loss_moe_lb": 2.057482020275013,
+      "loss_aux_moe": 0.68815186116341,
+      "loss_aux_llm": 0.8820263811865369
+    },
+    {
+      "loss": 2.9004717743074573,
+      "loss_size": 0.7811062932014465,
+      "loss_pdi": 0.5034840235839019,
+      "loss_ee": 0.7158847797561336,
+      "loss_delivery": 0.6796143882580705,
+      "loss_biodist": 0.17312858465152817,
+      "loss_toxic": 0.09950823990024023,
+      "loss_moe_lb": 2.035009638683216
+    },
+    {
+      "loss": 2.9226431202244116,
+      "loss_size": 0.8814810771193053,
+      "loss_pdi": 0.5112376519151636,
+      "loss_ee": 0.7327054779271822,
+      "loss_delivery": 0.591531743874421,
+      "loss_biodist": 0.16855915958011472,
+      "loss_toxic": 0.0948821196735308,
+      "loss_moe_lb": 2.0614523597665735
+    },
+    {
+      "loss": 2.655469871856071,
+      "loss_size": 0.760228231347896,
+      "loss_pdi": 0.48637631092522593,
+      "loss_ee": 0.6979158569026638,
+      "loss_delivery": 0.5602184928752281,
+      "loss_biodist": 0.14575726085820714,
+      "loss_toxic": 0.06323230568306974,
+      "loss_moe_lb": 2.050211484367783
+    },
+    {
+      "loss": 2.509097969209826,
+      "loss_size": 0.7080223632825388,
+      "loss_pdi": 0.4473253053587836,
+      "loss_ee": 0.6626571416854858,
+      "loss_delivery": 0.5532580835593713,
+      "loss_biodist": 0.1437241429009953,
+      "loss_toxic": 0.04900663166709647,
+      "loss_moe_lb": 2.079811640687891
+    },
+    {
+      "loss": 2.484670532716287,
+      "loss_size": 0.7306602957683641,
+      "loss_pdi": 0.47313058899866567,
+      "loss_ee": 0.660492988454329,
+      "loss_delivery": 0.5034271306484133,
+      "loss_biodist": 0.13187124368709488,
+      "loss_toxic": 0.04232822387869394,
+      "loss_moe_lb": 2.062951352145221
+    }
+  ],
+  "val": [
+    {
+      "loss": 3.755842113494873,
+      "loss_size": 0.4265119343996048,
+      "loss_pdi": 0.5915934443473816,
+      "loss_ee": 0.8747140049934388,
+      "loss_delivery": 0.8283049013465643,
+      "loss_biodist": 0.6451014578342438,
+      "loss_toxic": 0.39165826290845873,
+      "loss_moe_lb": 2.1249768257141115,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.932078742980957,
+      "loss_size": 0.3444555729627609,
+      "loss_pdi": 0.5126577079296112,
+      "loss_ee": 0.7018583565950394,
+      "loss_delivery": 0.8085120553150773,
+      "loss_biodist": 0.24770227670669556,
+      "loss_toxic": 0.309023617208004,
+      "loss_moe_lb": 2.076921796798706,
+      "acc_pdi": 0.5151515151515151,
+      "acc_ee": 0.6060606060606061,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.628650116920471,
+      "loss_size": 0.31893138736486437,
+      "loss_pdi": 0.5060839474201202,
+      "loss_ee": 0.6603254437446594,
+      "loss_delivery": 0.6911659525707364,
+      "loss_biodist": 0.16578628718852997,
+      "loss_toxic": 0.28838122070301325,
+      "loss_moe_lb": 2.0901700019836427,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.6363636363636364,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.635863518714905,
+      "loss_size": 0.21646426171064376,
+      "loss_pdi": 0.4807057544589043,
+      "loss_ee": 0.7196594156324864,
+      "loss_delivery": 0.7288934513926506,
+      "loss_biodist": 0.12086161077022553,
+      "loss_toxic": 0.37542195692658425,
+      "loss_moe_lb": 2.13119113445282,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.6391834259033202,
+      "loss_size": 0.25473385751247407,
+      "loss_pdi": 0.515596741437912,
+      "loss_ee": 0.7078545540571213,
+      "loss_delivery": 0.7006262980401516,
+      "loss_biodist": 0.13355249539017677,
+      "loss_toxic": 0.3424140252172947,
+      "loss_moe_lb": 2.1053977251052856,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.8888888888888888
+    },
+    {
+      "loss": 2.775922429561615,
+      "loss_size": 0.26124624609947206,
+      "loss_pdi": 0.46465142369270324,
+      "loss_ee": 0.7462939266115427,
+      "loss_delivery": 0.699479284894187,
+      "loss_biodist": 0.10829939097166061,
+      "loss_toxic": 0.5012864658609033,
+      "loss_moe_lb": 2.1504496335983276,
+      "acc_pdi": 0.7575757575757576,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.565546917915344,
+      "loss_size": 0.23962730169296265,
+      "loss_pdi": 0.46498828530311587,
+      "loss_ee": 0.8164169616997242,
+      "loss_delivery": 0.734196837246418,
+      "loss_biodist": 0.09018271192908286,
+      "loss_toxic": 0.2623913433402777,
+      "loss_moe_lb": 2.1556874752044677,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.461030912399292,
+      "loss_size": 0.20080993920564652,
+      "loss_pdi": 0.4726409152150154,
+      "loss_ee": 0.8052330955862999,
+      "loss_delivery": 0.7270681807538495,
+      "loss_biodist": 0.08629608154296875,
+      "loss_toxic": 0.21949858171865344,
+      "loss_moe_lb": 2.152100443840027,
+      "acc_pdi": 0.696969696969697,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.446087622642517,
+      "loss_size": 0.2072613552212715,
+      "loss_pdi": 0.4464923307299614,
+      "loss_ee": 0.8108729934319854,
+      "loss_delivery": 0.7515319633297622,
+      "loss_biodist": 0.08003018759191036,
+      "loss_toxic": 0.20465249773114919,
+      "loss_moe_lb": 2.1997459888458253,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5454545454545454,
+      "acc_toxic": 1.0
+    },
+    {
+      "loss": 2.4781503915786742,
+      "loss_size": 0.19303087145090103,
+      "loss_pdi": 0.45752228796482086,
+      "loss_ee": 0.7772767193615436,
+      "loss_delivery": 0.777904954974656,
+      "loss_biodist": 0.09190915115177631,
+      "loss_toxic": 0.23136587778571993,
+      "loss_moe_lb": 2.1780391693115235,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    },
+    {
+      "loss": 2.4587207794189454,
+      "loss_size": 0.1924375042319298,
+      "loss_pdi": 0.4666991338133812,
+      "loss_ee": 0.770481814071536,
+      "loss_delivery": 0.7665328018600122,
+      "loss_biodist": 0.09240943193435669,
+      "loss_toxic": 0.22187743864487858,
+      "loss_moe_lb": 2.1950268745422363,
+      "acc_pdi": 0.6666666666666666,
+      "acc_ee": 0.5757575757575758,
+      "acc_toxic": 0.9444444444444444
+    }
+  ]
+}

models/cv5_sample/moe_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/cv5_sample/moe_llm/seed42/outer_fold_4/test_metrics.json

@@ -0,0 +1,42 @@
+{
+  "size": {
+    "n_samples": 84,
+    "mse": 0.6805742520206829,
+    "rmse": 0.8249692430755724,
+    "mae": 0.4822031688111435,
+    "r2": 0.08930711264607172
+  },
+  "delivery": {
+    "n_samples": 58,
+    "mse": 0.7669285932160084,
+    "rmse": 0.8757445935979328,
+    "mae": 0.567918715805843,
+    "r2": 0.2583927108955246
+  },
+  "pdi": {
+    "n_samples": 84,
+    "accuracy": 0.7619047619047619,
+    "precision": 0.7002652519893899,
+    "recall": 0.7214076246334311,
+    "f1": 0.7083333333333334
+  },
+  "ee": {
+    "n_samples": 84,
+    "accuracy": 0.7023809523809523,
+    "precision": 0.6327380952380953,
+    "recall": 0.6816239316239318,
+    "f1": 0.6431159420289855
+  },
+  "toxic": {
+    "n_samples": 59,
+    "accuracy": 1.0,
+    "precision": 1.0,
+    "recall": 1.0,
+    "f1": 1.0
+  },
+  "biodist": {
+    "n_samples": 58,
+    "kl_divergence": 0.16534984510664474,
+    "js_divergence": 0.041883856181021084
+  }
+}

models/cv5_sample/moe_llm/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/cv5_sample/moe_llm/seed42/summary.json

@@ -0,0 +1,372 @@
+{
+  "fold_results": [
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.19323529834582587,
+        "lr": 0.000980484886752915,
+        "weight_decay": 0.00014367509113739864,
+        "backbone_lr_ratio": 0.13093310334961422,
+        "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": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.4858819540561683,
+          "rmse": 1.2189675771144073,
+          "mae": 0.5309677058494234,
+          "r2": 0.0906036692787614
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.4578579988449137,
+          "rmse": 0.6766520515338099,
+          "mae": 0.44371995239936074,
+          "r2": 0.41803485452290234
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7857142857142857,
+          "precision": 0.7322033898305085,
+          "recall": 0.7441197434069851,
+          "f1": 0.7375
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7380952380952381,
+          "precision": 0.7146208708708709,
+          "recall": 0.7547619047619047,
+          "f1": 0.7160330091364574
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9827586206896551,
+          "precision": 0.9912280701754386,
+          "recall": 0.75,
+          "f1": 0.8289085545722714
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2124931399005794,
+          "js_divergence": 0.0436771104854127
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.19323529834582587,
+        "lr": 0.000980484886752915,
+        "weight_decay": 0.00014367509113739864,
+        "backbone_lr_ratio": 0.13093310334961422,
+        "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": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.3521366695033162,
+          "rmse": 0.593411045990312,
+          "mae": 0.417994069290303,
+          "r2": -0.07530019609216088
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 0.8615718355352413,
+          "rmse": 0.9282089395902419,
+          "mae": 0.54733644703739,
+          "r2": 0.34785223029308054
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6309523809523809,
+          "precision": 0.5417771883289124,
+          "recall": 0.5476190476190477,
+          "f1": 0.5421135924037278
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5952380952380952,
+          "precision": 0.531971296677179,
+          "recall": 0.5323723895152467,
+          "f1": 0.5202020202020202
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9836065573770492,
+          "precision": 0.9915254237288136,
+          "recall": 0.8333333333333333,
+          "f1": 0.8957264957264958
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.12922332899139097,
+          "js_divergence": 0.030354793549857463
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.19323529834582587,
+        "lr": 0.000980484886752915,
+        "weight_decay": 0.00014367509113739864,
+        "backbone_lr_ratio": 0.13093310334961422,
+        "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": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.4170957202870868,
+          "rmse": 0.6458294823613171,
+          "mae": 0.4985973424876907,
+          "r2": 0.1762434046684933
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.5854544655957109,
+          "rmse": 0.7651499628149445,
+          "mae": 0.49963968732239056,
+          "r2": 0.22065863304023947
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7380952380952381,
+          "precision": 0.6851851851851851,
+          "recall": 0.7199413489736071,
+          "f1": 0.693633952254642
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.613035584604212,
+          "recall": 0.6666666666666666,
+          "f1": 0.627436739505705
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 1.0,
+          "precision": 1.0,
+          "recall": 1.0,
+          "f1": 1.0
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.17218262970209255,
+          "js_divergence": 0.039882423974733884
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.19323529834582587,
+        "lr": 0.000980484886752915,
+        "weight_decay": 0.00014367509113739864,
+        "backbone_lr_ratio": 0.13093310334961422,
+        "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": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.773779594197195,
+          "rmse": 1.331833170557482,
+          "mae": 0.5606547438887408,
+          "r2": -0.003263258057283336
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.6480387215025862,
+          "rmse": 0.8050085226272988,
+          "mae": 0.5227914613403254,
+          "r2": 0.37002936863191194
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.75,
+          "precision": 0.6789023521026372,
+          "recall": 0.684017595307918,
+          "f1": 0.6813008130081302
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5700750584471516,
+          "recall": 0.5891402714932127,
+          "f1": 0.5736800630417652
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9833333333333333,
+          "precision": 0.875,
+          "recall": 0.9912280701754386,
+          "f1": 0.9241466498103665
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.15399347135454697,
+          "js_divergence": 0.036819733414228055
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.19323529834582587,
+        "lr": 0.000980484886752915,
+        "weight_decay": 0.00014367509113739864,
+        "backbone_lr_ratio": 0.13093310334961422,
+        "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": 11,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.6805742520206829,
+          "rmse": 0.8249692430755724,
+          "mae": 0.4822031688111435,
+          "r2": 0.08930711264607172
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.7669285932160084,
+          "rmse": 0.8757445935979328,
+          "mae": 0.567918715805843,
+          "r2": 0.2583927108955246
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7619047619047619,
+          "precision": 0.7002652519893899,
+          "recall": 0.7214076246334311,
+          "f1": 0.7083333333333334
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6327380952380953,
+          "recall": 0.6816239316239318,
+          "f1": 0.6431159420289855
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 1.0,
+          "precision": 1.0,
+          "recall": 1.0,
+          "f1": 1.0
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.16534984510664474,
+          "js_divergence": 0.041883856181021084
+        }
+      }
+    }
+  ],
+  "summary_stats": {
+    "size": {
+      "mse_mean": 0.9418936380128897,
+      "mse_std": 0.5795638860278585,
+      "rmse_mean": 0.9230021038198182,
+      "rmse_std": 0.29993458346292695,
+      "mae_mean": 0.4980834060654603,
+      "mae_std": 0.048297613647592774,
+      "r2_mean": 0.05551814648877644,
+      "r2_std": 0.0866209771152643
+    },
+    "delivery": {
+      "mse_mean": 0.6639703229388921,
+      "mse_std": 0.14039060447259435,
+      "rmse_mean": 0.8101528140328454,
+      "rmse_std": 0.08730830918963939,
+      "mae_mean": 0.5162812527810619,
+      "mae_std": 0.04293092010991082,
+      "r2_mean": 0.3229935594767318,
+      "r2_std": 0.07281328920995907
+    },
+    "pdi": {
+      "accuracy_mean": 0.7333333333333333,
+      "accuracy_std": 0.05355819942477139,
+      "precision_mean": 0.6676666734873267,
+      "precision_std": 0.06558804201410597,
+      "recall_mean": 0.6834210719881977,
+      "recall_std": 0.07057419243292298,
+      "f1_mean": 0.6725763381999668,
+      "f1_std": 0.06787298751865974
+    },
+    "ee": {
+      "accuracy_mean": 0.669047619047619,
+      "accuracy_std": 0.04902681019517619,
+      "precision_mean": 0.6124881811675018,
+      "precision_std": 0.061843813988713685,
+      "recall_mean": 0.6449130328121926,
+      "recall_std": 0.07702425935365928,
+      "f1_mean": 0.6160935547829867,
+      "f1_std": 0.06609215153794168
+    },
+    "toxic": {
+      "accuracy_mean": 0.9899397022800075,
+      "accuracy_std": 0.008218758400859814,
+      "precision_mean": 0.9715506987808504,
+      "precision_std": 0.04842922980294785,
+      "recall_mean": 0.9149122807017542,
+      "recall_std": 0.10407221017518521,
+      "f1_mean": 0.9297563400218267,
+      "f1_std": 0.06515845039804279
+    },
+    "biodist": {
+      "kl_divergence_mean": 0.16664848301105092,
+      "kl_divergence_std": 0.027183955851732975,
+      "js_divergence_mean": 0.03852358352105063,
+      "js_divergence_std": 0.004676087139485391
+    }
+  }
+}

models/final/best_params.json

@@ -1,11 +1,16 @@
 {
-  "dropout": 0.14731648771233646,
-  "lr": 0.00010399668603456206,
-  "weight_decay": 0.0007696017715340258,
-  "backbone_lr_ratio": 0.7472270478607838,
+  "dropout": 0.19323529834582587,
+  "lr": 0.000980484886752915,
+  "weight_decay": 0.00014367509113739864,
+  "backbone_lr_ratio": 0.13093310334961422,
+  "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
+  "head_hidden_dim": 128,
+  "set_transformer_block": "sab"
 }

models/final/class_weights.json

@@ -1,9 +1,7 @@
 {
   "pdi": [
-    0.011065198108553886,
-    0.03264812380075455,
-    0.8369068503379822,
-    3.119379997253418
+    0.524036169052124,
+    1.475963830947876
   ],
   "ee": [
     1.600011944770813,

models/final/epoch_mean.json

@@ -1 +1 @@
-{"epoch_mean": 23}
+{"epoch_mean": 11}

models/final/history.json

@@ -1,211 +1,126 @@
 {
   "train": [
     {
-      "loss": 17.60348994391305,
-      "loss_size": 12.354392801012311,
-      "loss_pdi": 1.4643298983573914,
-      "loss_ee": 1.090087456362588,
-      "loss_delivery": 0.8421656531946999,
-      "loss_biodist": 1.23681207214083,
-      "loss_toxic": 0.6157020671027047
+      "loss": 5.694288951046062,
+      "loss_size": 1.015234741259296,
+      "loss_pdi": 0.6592788083373375,
+      "loss_ee": 1.027098409409793,
+      "loss_delivery": 1.0114885227016683,
+      "loss_biodist": 0.7630386647750746,
+      "loss_toxic": 0.5493527664890829,
+      "loss_moe_lb": 2.045894391131851,
+      "loss_aux_moe": 1.0363417644545716,
+      "loss_aux_llm": 1.1829107824080396
     },
     {
-      "loss": 6.5880933318819315,
-      "loss_size": 1.750367718083518,
-      "loss_pdi": 1.3177664024489266,
-      "loss_ee": 1.068630073751722,
-      "loss_delivery": 0.8820391488926751,
-      "loss_biodist": 1.0014605011258806,
-      "loss_toxic": 0.567829430103302
+      "loss": 4.587394336484513,
+      "loss_size": 0.9309988603839334,
+      "loss_pdi": 0.6420835014784111,
+      "loss_ee": 0.9480358497151788,
+      "loss_delivery": 0.809673407519201,
+      "loss_biodist": 0.5112692536610477,
+      "loss_toxic": 0.30279909285171974,
+      "loss_moe_lb": 2.0554285611746446,
+      "loss_aux_moe": 0.9779460283242306,
+      "loss_aux_llm": 0.9288384830754883
     },
     {
-      "loss": 4.678432379450117,
-      "loss_size": 0.3203257258449282,
-      "loss_pdi": 1.180948099919728,
-      "loss_ee": 1.0613888587270464,
-      "loss_delivery": 0.8234607481530735,
-      "loss_biodist": 0.7979316924299512,
-      "loss_toxic": 0.49437726182597025
+      "loss": 4.160386994199933,
+      "loss_size": 0.9387357568122306,
+      "loss_pdi": 0.5914091456611201,
+      "loss_ee": 0.9026568987459507,
+      "loss_delivery": 0.84944020264413,
+      "loss_biodist": 0.3876211567307418,
+      "loss_toxic": 0.18268741205563102,
+      "loss_moe_lb": 2.1000781711542382,
+      "loss_aux_moe": 0.9373761130461715,
+      "loss_aux_llm": 0.9303386138295228
     },
     {
-      "loss": 4.174277833529881,
-      "loss_size": 0.4864327609539032,
-      "loss_pdi": 1.014147652047021,
-      "loss_ee": 1.0116309651306696,
-      "loss_delivery": 0.8047538192144462,
-      "loss_biodist": 0.5864214556557792,
-      "loss_toxic": 0.2708912321499416
+      "loss": 3.809376889804624,
+      "loss_size": 0.888136670766574,
+      "loss_pdi": 0.5791616782827197,
+      "loss_ee": 0.8553679860987753,
+      "loss_delivery": 0.8118948299648627,
+      "loss_biodist": 0.30093744889182866,
+      "loss_toxic": 0.15547859341890183,
+      "loss_moe_lb": 2.102432388179707,
+      "loss_aux_moe": 1.1255526001442153,
+      "loss_aux_llm": 0.9655606655298539
     },
     {
-      "loss": 3.4547910520008633,
-      "loss_size": 0.31596518414361136,
-      "loss_pdi": 0.9084383249282837,
-      "loss_ee": 0.9215172486645835,
-      "loss_delivery": 0.7124462143651077,
-      "loss_biodist": 0.42033374096666065,
-      "loss_toxic": 0.17609038363609994
+      "loss": 3.1629226522625618,
+      "loss_size": 0.7107668774870207,
+      "loss_pdi": 0.5377992931401955,
+      "loss_ee": 0.7858308424364846,
+      "loss_delivery": 0.7225946192190332,
+      "loss_biodist": 0.23169103110173964,
+      "loss_toxic": 0.10101612598441963,
+      "loss_moe_lb": 2.1294895545491634,
+      "loss_aux_moe": 0.8120556049189478,
+      "loss_aux_llm": 0.8958881971127582
     },
     {
-      "loss": 3.2108109337942943,
-      "loss_size": 0.2688417855118002,
-      "loss_pdi": 0.8587345012596675,
-      "loss_ee": 0.875971223626818,
-      "loss_delivery": 0.7385377415588924,
-      "loss_biodist": 0.3365256999220167,
-      "loss_toxic": 0.1321999922926937
+      "loss": 2.988117564399287,
+      "loss_size": 0.7456601148225227,
+      "loss_pdi": 0.5200111323370123,
+      "loss_ee": 0.7696504339856921,
+      "loss_delivery": 0.702117268874679,
+      "loss_biodist": 0.19998157122787438,
+      "loss_toxic": 0.08327638518992143,
+      "loss_moe_lb": 2.1596880116552675,
+      "loss_aux_moe": 0.763107772424536,
+      "loss_aux_llm": 0.8947458216604197
     },
     {
-      "loss": 2.9460756949016025,
-      "loss_size": 0.2779192647763661,
-      "loss_pdi": 0.7961623072624207,
-      "loss_ee": 0.8504888585635594,
-      "loss_delivery": 0.6400965334448431,
-      "loss_biodist": 0.2687057022537504,
-      "loss_toxic": 0.112703027070633
+      "loss": 2.7413068245042047,
+      "loss_size": 0.6423791460692883,
+      "loss_pdi": 0.5119175066081982,
+      "loss_ee": 0.7589052460103665,
+      "loss_delivery": 0.6498366547924168,
+      "loss_biodist": 0.1679054581612911,
+      "loss_toxic": 0.09372632653026332,
+      "loss_moe_lb": 2.148664339533392
     },
     {
-      "loss": 2.7530925273895264,
-      "loss_size": 0.25332520263535635,
-      "loss_pdi": 0.7469035514763424,
-      "loss_ee": 0.8471435691629138,
-      "loss_delivery": 0.5530903546937874,
-      "loss_biodist": 0.2454697404588972,
-      "loss_toxic": 0.10716012333120618
+      "loss": 2.7408997472727075,
+      "loss_size": 0.6384303912801562,
+      "loss_pdi": 0.49059138151834597,
+      "loss_ee": 0.7721629592607606,
+      "loss_delivery": 0.7107182373556326,
+      "loss_biodist": 0.15433137111787526,
+      "loss_toxic": 0.0748163205291417,
+      "loss_moe_lb": 2.154493790752483
     },
     {
-      "loss": 2.742221474647522,
-      "loss_size": 0.2413133829832077,
-      "loss_pdi": 0.6884610418762479,
-      "loss_ee": 0.8217804133892059,
-      "loss_delivery": 0.6628774159720966,
-      "loss_biodist": 0.23072052534137452,
-      "loss_toxic": 0.09706869375492845
+      "loss": 2.610751993251297,
+      "loss_size": 0.6624566971693399,
+      "loss_pdi": 0.4767126595636584,
+      "loss_ee": 0.7076389159796372,
+      "loss_delivery": 0.636365042331646,
+      "loss_biodist": 0.15519743205382014,
+      "loss_toxic": 0.06698143800538262,
+      "loss_moe_lb": 2.1544579132547916
     },
     {
-      "loss": 2.550677682672228,
-      "loss_size": 0.2765064158343843,
-      "loss_pdi": 0.6737975222723824,
-      "loss_ee": 0.7778623870440892,
-      "loss_delivery": 0.5299507762704577,
-      "loss_biodist": 0.1921106672712735,
-      "loss_toxic": 0.10044991970062256
+      "loss": 2.581533803130096,
+      "loss_size": 0.6987560360499148,
+      "loss_pdi": 0.4562466161431007,
+      "loss_ee": 0.7065791735109294,
+      "loss_delivery": 0.6151787998541346,
+      "loss_biodist": 0.1382373322223155,
+      "loss_toxic": 0.06480917761778845,
+      "loss_moe_lb": 2.137455944745046
     },
     {
-      "loss": 2.481816521712712,
-      "loss_size": 0.23023176033582007,
-      "loss_pdi": 0.6636314690113068,
-      "loss_ee": 0.7728129327297211,
-      "loss_delivery": 0.5704377527747836,
-      "loss_biodist": 0.1672141200729779,
-      "loss_toxic": 0.07748848732028689
-    },
-    {
-      "loss": 2.360851458140782,
-      "loss_size": 0.21546032758695738,
-      "loss_pdi": 0.6131412557193211,
-      "loss_ee": 0.830133033650262,
-      "loss_delivery": 0.4567379740016934,
-      "loss_biodist": 0.17073182123047964,
-      "loss_toxic": 0.07464704542819943
-    },
-    {
-      "loss": 2.2426107866423473,
-      "loss_size": 0.19084072799887508,
-      "loss_pdi": 0.6170341287340436,
-      "loss_ee": 0.7295470791203635,
-      "loss_delivery": 0.4835586505276816,
-      "loss_biodist": 0.1602931246161461,
-      "loss_toxic": 0.061337063022490056
-    },
-    {
-      "loss": 2.2947227358818054,
-      "loss_size": 0.21879314631223679,
-      "loss_pdi": 0.6161722391843796,
-      "loss_ee": 0.762788040297372,
-      "loss_delivery": 0.45067436354500906,
-      "loss_biodist": 0.18453349598816463,
-      "loss_toxic": 0.061761434456067424
-    },
-    {
-      "loss": 2.3243510808263506,
-      "loss_size": 0.23437534804855073,
-      "loss_pdi": 0.594099121434348,
-      "loss_ee": 0.77819305232593,
-      "loss_delivery": 0.4997542394059045,
-      "loss_biodist": 0.1680830376488822,
-      "loss_toxic": 0.049846263602375984
-    },
-    {
-      "loss": 2.2657968997955322,
-      "loss_size": 0.2230603982295309,
-      "loss_pdi": 0.6490842210394996,
-      "loss_ee": 0.7079165577888489,
-      "loss_delivery": 0.4766663594969681,
-      "loss_biodist": 0.16171239422900335,
-      "loss_toxic": 0.0473570313770324
-    },
-    {
-      "loss": 2.2227368525096347,
-      "loss_size": 0.22425250709056854,
-      "loss_pdi": 0.6083124216113772,
-      "loss_ee": 0.7293048288140979,
-      "loss_delivery": 0.4598711388451712,
-      "loss_biodist": 0.1544684906091009,
-      "loss_toxic": 0.04652748556275453
-    },
-    {
-      "loss": 2.1243858337402344,
-      "loss_size": 0.21263282213892257,
-      "loss_pdi": 0.573635390826634,
-      "loss_ee": 0.6900361818926675,
-      "loss_delivery": 0.4512601649122579,
-      "loss_biodist": 0.15764176366584642,
-      "loss_toxic": 0.03917955420911312
-    },
-    {
-      "loss": 2.289846863065447,
-      "loss_size": 0.21625602032457078,
-      "loss_pdi": 0.596784600189754,
-      "loss_ee": 0.7933975628444127,
-      "loss_delivery": 0.4503451883792877,
-      "loss_biodist": 0.18038264129843032,
-      "loss_toxic": 0.05268087850085327
-    },
-    {
-      "loss": 2.1171253493853976,
-      "loss_size": 0.21727498407874787,
-      "loss_pdi": 0.5871099403926304,
-      "loss_ee": 0.6942037897450584,
-      "loss_delivery": 0.4114475510349231,
-      "loss_biodist": 0.16847609941448485,
-      "loss_toxic": 0.038613016584089825
-    },
-    {
-      "loss": 2.2243686744144986,
-      "loss_size": 0.23041224852204323,
-      "loss_pdi": 0.5950902572699955,
-      "loss_ee": 0.7303137523787362,
-      "loss_delivery": 0.4817630084497588,
-      "loss_biodist": 0.13934307438986643,
-      "loss_toxic": 0.047446319966443946
-    },
-    {
-      "loss": 2.142884841987065,
-      "loss_size": 0.18598268926143646,
-      "loss_pdi": 0.5736084473984582,
-      "loss_ee": 0.7191481994731086,
-      "loss_delivery": 0.4928105877978461,
-      "loss_biodist": 0.13822122237512044,
-      "loss_toxic": 0.03311372648126313
-    },
-    {
-      "loss": 2.0640590446335927,
-      "loss_size": 0.1981512185718332,
-      "loss_pdi": 0.5239957081420081,
-      "loss_ee": 0.6807411738804409,
-      "loss_delivery": 0.45748772472143173,
-      "loss_biodist": 0.14945154903190477,
-      "loss_toxic": 0.0542316823931677
+      "loss": 2.3695818658144967,
+      "loss_size": 0.5367685263308714,
+      "loss_pdi": 0.47306264318385216,
+      "loss_ee": 0.6626323349070999,
+      "loss_delivery": 0.5844551013817765,
+      "loss_biodist": 0.13707212931564394,
+      "loss_toxic": 0.06755097333027735,
+      "loss_moe_lb": 2.1183437378901355
     }
   ],
   "val": []

models/final/model.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:241b22c5170dd7d5a82f409228589d61d16ca14c3ae1991a03bda8ca547ce743
-size 28583340
+oid sha256:5b1c826c11338f313eea643f839fbc27c4c1e0f0f19d58d672d405cc731caa2a
+size 58261734

models/final_backup_backbone/best_params.json

@@ -0,0 +1,11 @@
+{
+  "dropout": 0.14731648771233646,
+  "lr": 0.00010399668603456206,
+  "weight_decay": 0.0007696017715340258,
+  "backbone_lr_ratio": 0.7472270478607838,
+  "d_model": 256,
+  "num_heads": 8,
+  "n_attn_layers": 4,
+  "fusion_strategy": "attention",
+  "head_hidden_dim": 128
+}

models/final_backup_backbone/class_weights.json

@@ -0,0 +1,17 @@
+{
+  "pdi": [
+    0.011065198108553886,
+    0.03264812380075455,
+    0.8369068503379822,
+    3.119379997253418
+  ],
+  "ee": [
+    1.600011944770813,
+    0.9947696924209595,
+    0.40521833300590515
+  ],
+  "toxic": [
+    0.09425132721662521,
+    1.9057486057281494
+  ]
+}

models/final_backup_backbone/epoch_mean.json

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

models/final_backup_backbone/history.json

@@ -0,0 +1,212 @@
+{
+  "train": [
+    {
+      "loss": 17.60348994391305,
+      "loss_size": 12.354392801012311,
+      "loss_pdi": 1.4643298983573914,
+      "loss_ee": 1.090087456362588,
+      "loss_delivery": 0.8421656531946999,
+      "loss_biodist": 1.23681207214083,
+      "loss_toxic": 0.6157020671027047
+    },
+    {
+      "loss": 6.5880933318819315,
+      "loss_size": 1.750367718083518,
+      "loss_pdi": 1.3177664024489266,
+      "loss_ee": 1.068630073751722,
+      "loss_delivery": 0.8820391488926751,
+      "loss_biodist": 1.0014605011258806,
+      "loss_toxic": 0.567829430103302
+    },
+    {
+      "loss": 4.678432379450117,
+      "loss_size": 0.3203257258449282,
+      "loss_pdi": 1.180948099919728,
+      "loss_ee": 1.0613888587270464,
+      "loss_delivery": 0.8234607481530735,
+      "loss_biodist": 0.7979316924299512,
+      "loss_toxic": 0.49437726182597025
+    },
+    {
+      "loss": 4.174277833529881,
+      "loss_size": 0.4864327609539032,
+      "loss_pdi": 1.014147652047021,
+      "loss_ee": 1.0116309651306696,
+      "loss_delivery": 0.8047538192144462,
+      "loss_biodist": 0.5864214556557792,
+      "loss_toxic": 0.2708912321499416
+    },
+    {
+      "loss": 3.4547910520008633,
+      "loss_size": 0.31596518414361136,
+      "loss_pdi": 0.9084383249282837,
+      "loss_ee": 0.9215172486645835,
+      "loss_delivery": 0.7124462143651077,
+      "loss_biodist": 0.42033374096666065,
+      "loss_toxic": 0.17609038363609994
+    },
+    {
+      "loss": 3.2108109337942943,
+      "loss_size": 0.2688417855118002,
+      "loss_pdi": 0.8587345012596675,
+      "loss_ee": 0.875971223626818,
+      "loss_delivery": 0.7385377415588924,
+      "loss_biodist": 0.3365256999220167,
+      "loss_toxic": 0.1321999922926937
+    },
+    {
+      "loss": 2.9460756949016025,
+      "loss_size": 0.2779192647763661,
+      "loss_pdi": 0.7961623072624207,
+      "loss_ee": 0.8504888585635594,
+      "loss_delivery": 0.6400965334448431,
+      "loss_biodist": 0.2687057022537504,
+      "loss_toxic": 0.112703027070633
+    },
+    {
+      "loss": 2.7530925273895264,
+      "loss_size": 0.25332520263535635,
+      "loss_pdi": 0.7469035514763424,
+      "loss_ee": 0.8471435691629138,
+      "loss_delivery": 0.5530903546937874,
+      "loss_biodist": 0.2454697404588972,
+      "loss_toxic": 0.10716012333120618
+    },
+    {
+      "loss": 2.742221474647522,
+      "loss_size": 0.2413133829832077,
+      "loss_pdi": 0.6884610418762479,
+      "loss_ee": 0.8217804133892059,
+      "loss_delivery": 0.6628774159720966,
+      "loss_biodist": 0.23072052534137452,
+      "loss_toxic": 0.09706869375492845
+    },
+    {
+      "loss": 2.550677682672228,
+      "loss_size": 0.2765064158343843,
+      "loss_pdi": 0.6737975222723824,
+      "loss_ee": 0.7778623870440892,
+      "loss_delivery": 0.5299507762704577,
+      "loss_biodist": 0.1921106672712735,
+      "loss_toxic": 0.10044991970062256
+    },
+    {
+      "loss": 2.481816521712712,
+      "loss_size": 0.23023176033582007,
+      "loss_pdi": 0.6636314690113068,
+      "loss_ee": 0.7728129327297211,
+      "loss_delivery": 0.5704377527747836,
+      "loss_biodist": 0.1672141200729779,
+      "loss_toxic": 0.07748848732028689
+    },
+    {
+      "loss": 2.360851458140782,
+      "loss_size": 0.21546032758695738,
+      "loss_pdi": 0.6131412557193211,
+      "loss_ee": 0.830133033650262,
+      "loss_delivery": 0.4567379740016934,
+      "loss_biodist": 0.17073182123047964,
+      "loss_toxic": 0.07464704542819943
+    },
+    {
+      "loss": 2.2426107866423473,
+      "loss_size": 0.19084072799887508,
+      "loss_pdi": 0.6170341287340436,
+      "loss_ee": 0.7295470791203635,
+      "loss_delivery": 0.4835586505276816,
+      "loss_biodist": 0.1602931246161461,
+      "loss_toxic": 0.061337063022490056
+    },
+    {
+      "loss": 2.2947227358818054,
+      "loss_size": 0.21879314631223679,
+      "loss_pdi": 0.6161722391843796,
+      "loss_ee": 0.762788040297372,
+      "loss_delivery": 0.45067436354500906,
+      "loss_biodist": 0.18453349598816463,
+      "loss_toxic": 0.061761434456067424
+    },
+    {
+      "loss": 2.3243510808263506,
+      "loss_size": 0.23437534804855073,
+      "loss_pdi": 0.594099121434348,
+      "loss_ee": 0.77819305232593,
+      "loss_delivery": 0.4997542394059045,
+      "loss_biodist": 0.1680830376488822,
+      "loss_toxic": 0.049846263602375984
+    },
+    {
+      "loss": 2.2657968997955322,
+      "loss_size": 0.2230603982295309,
+      "loss_pdi": 0.6490842210394996,
+      "loss_ee": 0.7079165577888489,
+      "loss_delivery": 0.4766663594969681,
+      "loss_biodist": 0.16171239422900335,
+      "loss_toxic": 0.0473570313770324
+    },
+    {
+      "loss": 2.2227368525096347,
+      "loss_size": 0.22425250709056854,
+      "loss_pdi": 0.6083124216113772,
+      "loss_ee": 0.7293048288140979,
+      "loss_delivery": 0.4598711388451712,
+      "loss_biodist": 0.1544684906091009,
+      "loss_toxic": 0.04652748556275453
+    },
+    {
+      "loss": 2.1243858337402344,
+      "loss_size": 0.21263282213892257,
+      "loss_pdi": 0.573635390826634,
+      "loss_ee": 0.6900361818926675,
+      "loss_delivery": 0.4512601649122579,
+      "loss_biodist": 0.15764176366584642,
+      "loss_toxic": 0.03917955420911312
+    },
+    {
+      "loss": 2.289846863065447,
+      "loss_size": 0.21625602032457078,
+      "loss_pdi": 0.596784600189754,
+      "loss_ee": 0.7933975628444127,
+      "loss_delivery": 0.4503451883792877,
+      "loss_biodist": 0.18038264129843032,
+      "loss_toxic": 0.05268087850085327
+    },
+    {
+      "loss": 2.1171253493853976,
+      "loss_size": 0.21727498407874787,
+      "loss_pdi": 0.5871099403926304,
+      "loss_ee": 0.6942037897450584,
+      "loss_delivery": 0.4114475510349231,
+      "loss_biodist": 0.16847609941448485,
+      "loss_toxic": 0.038613016584089825
+    },
+    {
+      "loss": 2.2243686744144986,
+      "loss_size": 0.23041224852204323,
+      "loss_pdi": 0.5950902572699955,
+      "loss_ee": 0.7303137523787362,
+      "loss_delivery": 0.4817630084497588,
+      "loss_biodist": 0.13934307438986643,
+      "loss_toxic": 0.047446319966443946
+    },
+    {
+      "loss": 2.142884841987065,
+      "loss_size": 0.18598268926143646,
+      "loss_pdi": 0.5736084473984582,
+      "loss_ee": 0.7191481994731086,
+      "loss_delivery": 0.4928105877978461,
+      "loss_biodist": 0.13822122237512044,
+      "loss_toxic": 0.03311372648126313
+    },
+    {
+      "loss": 2.0640590446335927,
+      "loss_size": 0.1981512185718332,
+      "loss_pdi": 0.5239957081420081,
+      "loss_ee": 0.6807411738804409,
+      "loss_delivery": 0.45748772472143173,
+      "loss_biodist": 0.14945154903190477,
+      "loss_toxic": 0.0542316823931677
+    }
+  ],
+  "val": []
+}

models/final_backup_backbone/model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:241b22c5170dd7d5a82f409228589d61d16ca14c3ae1991a03bda8ca547ce743
+size 28583340

models/final_backup_backbone/optuna_trials.json

@@ -0,0 +1,962 @@
+[
+  {
+    "number": 0,
+    "value": 2.8156970342000327,
+    "params": {
+      "dropout": 0.249816047538945,
+      "lr": 0.0007969454818643932,
+      "weight_decay": 0.008471801418819975,
+      "backbone_lr_ratio": 0.15751320499779725
+    },
+    "user_attrs": {
+      "epoch_mean": 14,
+      "fold_best_epochs": [
+        10,
+        12,
+        21
+      ],
+      "fold_val_losses": [
+        3.733495855331421,
+        2.5024956703186034,
+        2.2110995769500734
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 1,
+    "value": 3.3739805857340492,
+    "params": {
+      "dropout": 0.1624074561769746,
+      "lr": 2.0511104188433963e-05,
+      "weight_decay": 1.7073967431528103e-05,
+      "backbone_lr_ratio": 0.5399484409787431
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        30,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        3.7005380153656007,
+        3.236070442199707,
+        3.1853332996368406
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 2,
+    "value": 2.6023060957590736,
+    "params": {
+      "dropout": 0.34044600469728353,
+      "lr": 0.0002607024758370766,
+      "weight_decay": 1.2087541473056957e-05,
+      "backbone_lr_ratio": 0.8706020878304853
+    },
+    "user_attrs": {
+      "epoch_mean": 16,
+      "fold_best_epochs": [
+        14,
+        19,
+        16
+      ],
+      "fold_val_losses": [
+        3.283898639678955,
+        2.366805338859558,
+        2.1562143087387087
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 3,
+    "value": 10.838356399536133,
+    "params": {
+      "dropout": 0.4329770563201687,
+      "lr": 2.6587543983272695e-05,
+      "weight_decay": 5.3370327626039544e-05,
+      "backbone_lr_ratio": 0.023270677083837805
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        30,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        9.760257911682128,
+        10.611498069763183,
+        12.143313217163087
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 4,
+    "value": 3.0307216803232833,
+    "params": {
+      "dropout": 0.2216968971838151,
+      "lr": 0.00011207606211860574,
+      "weight_decay": 0.0005342937261279777,
+      "backbone_lr_ratio": 0.038234752246751866
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        29,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        3.6398903369903564,
+        2.7796840190887453,
+        2.6725906848907472
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 5,
+    "value": 12.962510363260904,
+    "params": {
+      "dropout": 0.34474115788895177,
+      "lr": 1.9010245319870364e-05,
+      "weight_decay": 0.00014742753159914678,
+      "backbone_lr_ratio": 0.05404103854647329
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        30,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        11.658656311035156,
+        13.798492240905762,
+        13.430382537841798
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 6,
+    "value": 2.8127863725026447,
+    "params": {
+      "dropout": 0.28242799368681437,
+      "lr": 0.00037183641805732076,
+      "weight_decay": 6.290644294586152e-05,
+      "backbone_lr_ratio": 0.10677482709481352
+    },
+    "user_attrs": {
+      "epoch_mean": 21,
+      "fold_best_epochs": [
+        15,
+        20,
+        28
+      ],
+      "fold_val_losses": [
+        3.527592992782593,
+        2.6037728786468506,
+        2.306993246078491
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 7,
+    "value": 18.49700838724772,
+    "params": {
+      "dropout": 0.33696582754481696,
+      "lr": 1.2385137298860926e-05,
+      "weight_decay": 0.0026926469100861782,
+      "backbone_lr_ratio": 0.021930485556643693
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        30,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        18.425870513916017,
+        18.542006301879884,
+        18.523148345947266
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 8,
+    "value": 2.6576756556828816,
+    "params": {
+      "dropout": 0.12602063719411183,
+      "lr": 0.000790261954970823,
+      "weight_decay": 0.07286653737491042,
+      "backbone_lr_ratio": 0.4138040112561014
+    },
+    "user_attrs": {
+      "epoch_mean": 10,
+      "fold_best_epochs": [
+        9,
+        11,
+        10
+      ],
+      "fold_val_losses": [
+        3.56644024848938,
+        2.2665807485580443,
+        2.1400059700012206
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 9,
+    "value": 11.92358964284261,
+    "params": {
+      "dropout": 0.2218455076693483,
+      "lr": 1.5679933916722995e-05,
+      "weight_decay": 0.005456725485601475,
+      "backbone_lr_ratio": 0.07591104805282695
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        30,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        12.178006362915038,
+        12.020879745483398,
+        11.571882820129394
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 10,
+    "value": 2.744228076934814,
+    "params": {
+      "dropout": 0.48159581757804304,
+      "lr": 0.00013112992007873722,
+      "weight_decay": 1.0719864040714885e-05,
+      "backbone_lr_ratio": 0.8131735182005935
+    },
+    "user_attrs": {
+      "epoch_mean": 25,
+      "fold_best_epochs": [
+        21,
+        29,
+        24
+      ],
+      "fold_val_losses": [
+        3.2794203758239746,
+        2.4811975240707396,
+        2.472066330909729
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 11,
+    "value": 2.600485173861186,
+    "params": {
+      "dropout": 0.1398149921035961,
+      "lr": 0.0003676766469931975,
+      "weight_decay": 0.0814481577934799,
+      "backbone_lr_ratio": 0.416600060135216
+    },
+    "user_attrs": {
+      "epoch_mean": 14,
+      "fold_best_epochs": [
+        8,
+        14,
+        21
+      ],
+      "fold_val_losses": [
+        3.2930724143981935,
+        2.312326765060425,
+        2.196056342124939
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 12,
+    "value": 2.856875479221344,
+    "params": {
+      "dropout": 0.39786435978029433,
+      "lr": 0.00027104165213579705,
+      "weight_decay": 0.09324179801968951,
+      "backbone_lr_ratio": 0.24792013326598586
+    },
+    "user_attrs": {
+      "epoch_mean": 17,
+      "fold_best_epochs": [
+        14,
+        12,
+        25
+      ],
+      "fold_val_losses": [
+        3.5214386940002442,
+        2.671292519569397,
+        2.3778952240943907
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 13,
+    "value": 2.5929951747258504,
+    "params": {
+      "dropout": 0.1708198460450869,
+      "lr": 0.0002771027430797957,
+      "weight_decay": 0.022371906875975713,
+      "backbone_lr_ratio": 0.8529032090320552
+    },
+    "user_attrs": {
+      "epoch_mean": 11,
+      "fold_best_epochs": [
+        9,
+        12,
+        13
+      ],
+      "fold_val_losses": [
+        3.211256170272827,
+        2.29467031955719,
+        2.273059034347534
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 14,
+    "value": 2.7000243584314982,
+    "params": {
+      "dropout": 0.10319674017011438,
+      "lr": 5.318587054778435e-05,
+      "weight_decay": 0.026386070022612545,
+      "backbone_lr_ratio": 0.3335324680299515
+    },
+    "user_attrs": {
+      "epoch_mean": 28,
+      "fold_best_epochs": [
+        26,
+        29,
+        30
+      ],
+      "fold_val_losses": [
+        3.095314884185791,
+        2.559102272987366,
+        2.445655918121338
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 15,
+    "value": 2.811073144276937,
+    "params": {
+      "dropout": 0.17491537723907288,
+      "lr": 0.0004715550190685932,
+      "weight_decay": 0.020503728356529433,
+      "backbone_lr_ratio": 0.202005130829578
+    },
+    "user_attrs": {
+      "epoch_mean": 15,
+      "fold_best_epochs": [
+        11,
+        13,
+        22
+      ],
+      "fold_val_losses": [
+        3.665318727493286,
+        2.5160216808319094,
+        2.251879024505615
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 16,
+    "value": 2.51602889696757,
+    "params": {
+      "dropout": 0.157877413230548,
+      "lr": 0.00016221844139316182,
+      "weight_decay": 0.0011695139073230128,
+      "backbone_lr_ratio": 0.9818226878189111
+    },
+    "user_attrs": {
+      "epoch_mean": 19,
+      "fold_best_epochs": [
+        12,
+        21,
+        24
+      ],
+      "fold_val_losses": [
+        2.967289590835571,
+        2.41019024848938,
+        2.170606851577759
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 17,
+    "value": 2.617725872993469,
+    "params": {
+      "dropout": 0.1975499843454273,
+      "lr": 6.0680145665816345e-05,
+      "weight_decay": 0.0006054421043842105,
+      "backbone_lr_ratio": 0.9370662382902515
+    },
+    "user_attrs": {
+      "epoch_mean": 28,
+      "fold_best_epochs": [
+        23,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        3.112912178039551,
+        2.4265938997268677,
+        2.3136715412139894
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 18,
+    "value": 2.6030575354894006,
+    "params": {
+      "dropout": 0.2840632229983807,
+      "lr": 0.00016641263642332344,
+      "weight_decay": 0.0015861751549934148,
+      "backbone_lr_ratio": 0.5813736146009224
+    },
+    "user_attrs": {
+      "epoch_mean": 23,
+      "fold_best_epochs": [
+        21,
+        21,
+        27
+      ],
+      "fold_val_losses": [
+        3.163059639930725,
+        2.3479116916656495,
+        2.2982012748718263
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 19,
+    "value": 2.865672874450684,
+    "params": {
+      "dropout": 0.10525612628712336,
+      "lr": 6.627976747186786e-05,
+      "weight_decay": 0.00020126720225868128,
+      "backbone_lr_ratio": 0.12579626206936234
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        30,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        3.413493585586548,
+        2.58580904006958,
+        2.5977159976959228
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 20,
+    "value": 2.541643071174622,
+    "params": {
+      "dropout": 0.24563678706364706,
+      "lr": 0.0002197507105689095,
+      "weight_decay": 0.01585496331337175,
+      "backbone_lr_ratio": 0.2693624761288439
+    },
+    "user_attrs": {
+      "epoch_mean": 19,
+      "fold_best_epochs": [
+        15,
+        18,
+        24
+      ],
+      "fold_val_losses": [
+        3.2464155673980715,
+        2.2392319440841675,
+        2.139281702041626
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 21,
+    "value": 2.5363871733347576,
+    "params": {
+      "dropout": 0.24822650356184767,
+      "lr": 0.00018488380020633278,
+      "weight_decay": 0.016124874730349143,
+      "backbone_lr_ratio": 0.30071816195793755
+    },
+    "user_attrs": {
+      "epoch_mean": 21,
+      "fold_best_epochs": [
+        13,
+        26,
+        24
+      ],
+      "fold_val_losses": [
+        3.1824454784393312,
+        2.337935042381287,
+        2.088780999183655
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 22,
+    "value": 2.608898858229319,
+    "params": {
+      "dropout": 0.23999254948655954,
+      "lr": 0.00019074146892949267,
+      "weight_decay": 0.007473456802129268,
+      "backbone_lr_ratio": 0.24792013326598586
+    },
+    "user_attrs": {
+      "epoch_mean": 18,
+      "fold_best_epochs": [
+        12,
+        18,
+        25
+      ],
+      "fold_val_losses": [
+        3.299156427383423,
+        2.3785044670104982,
+        2.1490356802940367
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 23,
+    "value": 2.6372965017954506,
+    "params": {
+      "dropout": 0.2732819324439319,
+      "lr": 8.559450244862353e-05,
+      "weight_decay": 0.003126023459511234,
+      "backbone_lr_ratio": 0.3168235172426432
+    },
+    "user_attrs": {
+      "epoch_mean": 27,
+      "fold_best_epochs": [
+        28,
+        22,
+        30
+      ],
+      "fold_val_losses": [
+        3.0834681034088134,
+        2.466027092933655,
+        2.3623943090438844
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 24,
+    "value": 2.482991655667623,
+    "params": {
+      "dropout": 0.30169410062526664,
+      "lr": 0.000178917590650789,
+      "weight_decay": 0.013098468879191921,
+      "backbone_lr_ratio": 0.5322698244647086
+    },
+    "user_attrs": {
+      "epoch_mean": 19,
+      "fold_best_epochs": [
+        18,
+        16,
+        23
+      ],
+      "fold_val_losses": [
+        2.952860403060913,
+        2.3297120332717896,
+        2.166402530670166
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 25,
+    "value": 2.767317620913188,
+    "params": {
+      "dropout": 0.3194613602400577,
+      "lr": 4.342839523653342e-05,
+      "weight_decay": 0.001138022773453874,
+      "backbone_lr_ratio": 0.575108040368093
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        30,
+        30,
+        29
+      ],
+      "fold_val_losses": [
+        3.0871256828308105,
+        2.6297523975372314,
+        2.585074782371521
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 26,
+    "value": 2.6014270941416426,
+    "params": {
+      "dropout": 0.380465038157504,
+      "lr": 0.000131818792969444,
+      "weight_decay": 0.031283173085439854,
+      "backbone_lr_ratio": 0.5275684824553009
+    },
+    "user_attrs": {
+      "epoch_mean": 23,
+      "fold_best_epochs": [
+        21,
+        25,
+        24
+      ],
+      "fold_val_losses": [
+        3.164607048034668,
+        2.430464816093445,
+        2.209209418296814
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 27,
+    "value": 2.627247397104899,
+    "params": {
+      "dropout": 0.20197960767997714,
+      "lr": 8.186282230228678e-05,
+      "weight_decay": 0.003540606605897514,
+      "backbone_lr_ratio": 0.17090703553802636
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        29,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        3.2599289894104,
+        2.3099464178085327,
+        2.311866784095764
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 28,
+    "value": 2.6080079317092895,
+    "params": {
+      "dropout": 0.29875866523156275,
+      "lr": 0.0004703885391544067,
+      "weight_decay": 0.045327015648549004,
+      "backbone_lr_ratio": 0.6534099986720939
+    },
+    "user_attrs": {
+      "epoch_mean": 15,
+      "fold_best_epochs": [
+        9,
+        12,
+        25
+      ],
+      "fold_val_losses": [
+        3.350850296020508,
+        2.2914549112319946,
+        2.181718587875366
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 29,
+    "value": 2.6760944962501525,
+    "params": {
+      "dropout": 0.26572299607512456,
+      "lr": 0.0007379244149640388,
+      "weight_decay": 0.008589710887330617,
+      "backbone_lr_ratio": 0.40188133372791734
+    },
+    "user_attrs": {
+      "epoch_mean": 13,
+      "fold_best_epochs": [
+        9,
+        13,
+        17
+      ],
+      "fold_val_losses": [
+        3.524716091156006,
+        2.295975375175476,
+        2.207592022418976
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 30,
+    "value": 3.544327608744304,
+    "params": {
+      "dropout": 0.3647346234383415,
+      "lr": 3.7911482246002866e-05,
+      "weight_decay": 0.011264417139383486,
+      "backbone_lr_ratio": 0.18158532922746584
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        30,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        3.795345687866211,
+        3.331618309020996,
+        3.5060188293457033
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 31,
+    "value": 2.553473687171936,
+    "params": {
+      "dropout": 0.24271079071613583,
+      "lr": 0.0002079970663621355,
+      "weight_decay": 0.01383692673296454,
+      "backbone_lr_ratio": 0.2928581085972686
+    },
+    "user_attrs": {
+      "epoch_mean": 20,
+      "fold_best_epochs": [
+        14,
+        19,
+        28
+      ],
+      "fold_val_losses": [
+        3.171020269393921,
+        2.3204617738723754,
+        2.1689390182495116
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 32,
+    "value": 2.7147754907608026,
+    "params": {
+      "dropout": 0.3081417266050286,
+      "lr": 0.00015961988076370855,
+      "weight_decay": 0.005647416385662339,
+      "backbone_lr_ratio": 0.13126369611144165
+    },
+    "user_attrs": {
+      "epoch_mean": 25,
+      "fold_best_epochs": [
+        19,
+        28,
+        28
+      ],
+      "fold_val_losses": [
+        3.3805705070495606,
+        2.4436901807785034,
+        2.3200657844543455
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 33,
+    "value": 2.8405439217885338,
+    "params": {
+      "dropout": 0.19697712607188347,
+      "lr": 0.00022483991847689944,
+      "weight_decay": 0.04396257638256477,
+      "backbone_lr_ratio": 0.01101100938924494
+    },
+    "user_attrs": {
+      "epoch_mean": 26,
+      "fold_best_epochs": [
+        19,
+        30,
+        30
+      ],
+      "fold_val_losses": [
+        3.8356207847595214,
+        2.3643120765686034,
+        2.3216989040374756
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 34,
+    "value": 2.503487364451091,
+    "params": {
+      "dropout": 0.2561341328718051,
+      "lr": 0.00010300416486373038,
+      "weight_decay": 0.012808904827686988,
+      "backbone_lr_ratio": 0.6727535050162337
+    },
+    "user_attrs": {
+      "epoch_mean": 24,
+      "fold_best_epochs": [
+        17,
+        27,
+        28
+      ],
+      "fold_val_losses": [
+        3.041924476623535,
+        2.3080700874328612,
+        2.160467529296875
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 35,
+    "value": 2.4608302156130475,
+    "params": {
+      "dropout": 0.15463920314801388,
+      "lr": 0.00011604934996458467,
+      "weight_decay": 0.001746884129583056,
+      "backbone_lr_ratio": 0.7073818888939638
+    },
+    "user_attrs": {
+      "epoch_mean": 18,
+      "fold_best_epochs": [
+        16,
+        16,
+        23
+      ],
+      "fold_val_losses": [
+        2.9413320541381838,
+        2.2545334100723267,
+        2.1866251826286316
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 36,
+    "value": 2.452152466773987,
+    "params": {
+      "dropout": 0.14731648771233646,
+      "lr": 0.00010399668603456206,
+      "weight_decay": 0.0007696017715340258,
+      "backbone_lr_ratio": 0.7472270478607838
+    },
+    "user_attrs": {
+      "epoch_mean": 23,
+      "fold_best_epochs": [
+        19,
+        21,
+        29
+      ],
+      "fold_val_losses": [
+        2.921252131462097,
+        2.245358777046204,
+        2.1898464918136598
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 37,
+    "value": 2.5273226141929626,
+    "params": {
+      "dropout": 0.14194609604277394,
+      "lr": 0.00010021053483316553,
+      "weight_decay": 0.00038614379642781024,
+      "backbone_lr_ratio": 0.7094861329105961
+    },
+    "user_attrs": {
+      "epoch_mean": 21,
+      "fold_best_epochs": [
+        17,
+        19,
+        28
+      ],
+      "fold_val_losses": [
+        3.076446294784546,
+        2.3063220262527464,
+        2.1991995215415954
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 38,
+    "value": 2.8850494384765626,
+    "params": {
+      "dropout": 0.2174541133163097,
+      "lr": 3.3739662416512915e-05,
+      "weight_decay": 0.002017506096408397,
+      "backbone_lr_ratio": 0.4899201254858191
+    },
+    "user_attrs": {
+      "epoch_mean": 30,
+      "fold_best_epochs": [
+        30,
+        30,
+        29
+      ],
+      "fold_val_losses": [
+        3.2505751609802247,
+        2.640340280532837,
+        2.764232873916626
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  },
+  {
+    "number": 39,
+    "value": 2.597004270553589,
+    "params": {
+      "dropout": 0.4276269504938004,
+      "lr": 0.0001196820298534061,
+      "weight_decay": 0.00035992145065850994,
+      "backbone_lr_ratio": 0.7607765936389007
+    },
+    "user_attrs": {
+      "epoch_mean": 25,
+      "fold_best_epochs": [
+        24,
+        24,
+        27
+      ],
+      "fold_val_losses": [
+        3.053041934967041,
+        2.4371867895126345,
+        2.3007840871810914
+      ]
+    },
+    "state": "TrialState.COMPLETE"
+  }
+]

models/final_backup_backbone/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/pretrain/mpnn/pretrain_delivery.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ea823b87a22a3a87cb487a9e876a6605c27a138abe53eaf228bb348dad43667
+size 15970234

models/pretrain/mpnn/pretrain_history.json

@@ -0,0 +1,254 @@
+{
+  "train": [
+    {
+      "loss": 0.8207380050244176,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.7037654354365026,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.641996113549974,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.6070571800584409,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.5664057664468255,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.5478297611675892,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.5225721440771472,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.5069099566408475,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.4883425224024211,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.47361071388254816,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.4710996970063914,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.4520228576202772,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.45232121680418574,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.4500566631222651,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.4274507362961827,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.41204158604984553,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.39526375005000886,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.3738111386011387,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.3710037660054526,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.36222530734255337,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.3630701918634412,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.36140536376143495,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.3487493085334346,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.3467851376116652,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.34048351907452196,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.3541595635627183,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.33575822587381465,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.3219511436570785,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.30643386410068923,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.3101392985258709,
+      "n_samples": 8236
+    },
+    {
+      "loss": 0.31629752242895165,
+      "n_samples": 8236
+    }
+  ],
+  "val": [
+    {
+      "loss": 0.7836992233950629,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.7153084917606317,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.723381241791186,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6935155311345726,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6741305548682337,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6921254230824593,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6984183775181619,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6875423658664813,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6951277488855417,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6681769074567246,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.7054754480863044,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.7176714997016088,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.676661756540427,
+      "n_samples": 1454
+    },
+    {
+      "loss": 1.1059426443120963,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.7007260356185525,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.7131651905576661,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6576622851613135,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6736111434322604,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.7036862504531461,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6835714435970931,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6450633911679831,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6777869927014084,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.7330788161600472,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6992498341911925,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6878705129662767,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6829031571397427,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6679887015685091,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.730571337382436,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.902666473815005,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.6748524632053822,
+      "n_samples": 1454
+    },
+    {
+      "loss": 0.7610602847811608,
+      "n_samples": 1454
+    }
+  ]
+}

results/pretrained_encoders/baseline_pretrain_s42_summary.json

@@ -0,0 +1,932 @@
+{
+  "fold_results": [
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.3980773018180746,
+        "lr": 0.0009896030589860336,
+        "weight_decay": 0.0002988554096805342,
+        "backbone_lr_ratio": 0.2882797738700643,
+        "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.601088294900837,
+          "rmse": 1.265341177272295,
+          "mae": 0.5760945507650634,
+          "r2": 0.02009455288902051
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.569556123704086,
+          "rmse": 0.7546894220168228,
+          "mae": 0.5891120558411911,
+          "r2": 0.2760597975244864
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7738095238095238,
+          "precision": 0.7194960212201591,
+          "recall": 0.7359230220955096,
+          "f1": 0.7262905162064826
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.684408546477512,
+          "recall": 0.7380952380952381,
+          "f1": 0.6942196418797404
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2787090798576488,
+          "js_divergence": 0.05423684052771057
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.23085562232445592,
+        "lr": 0.0005227270589225132,
+        "weight_decay": 0.00461701039547356,
+        "backbone_lr_ratio": 0.12087164274488635,
+        "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": 84,
+          "mse": 0.3552515818625956,
+          "rmse": 0.5960298498083763,
+          "mae": 0.4518711648970133,
+          "r2": -0.08481203101542345
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.078371115810122,
+          "rmse": 1.038446491548853,
+          "mae": 0.6491757012354057,
+          "r2": 0.18375080395351728
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.5583333333333333,
+          "recall": 0.5634920634920635,
+          "f1": 0.5598915989159892
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5714285714285714,
+          "precision": 0.5153903903903904,
+          "recall": 0.5209076637648066,
+          "f1": 0.507505844715147
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.5193434798072222,
+          "js_divergence": 0.12637397289897973
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.318342405407684,
+        "lr": 0.0008937362022799393,
+        "weight_decay": 0.06633867694372529,
+        "backbone_lr_ratio": 0.20906844321970927,
+        "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.3646369784573562,
+          "rmse": 0.6038517851736105,
+          "mae": 0.47954003302203047,
+          "r2": 0.2798484825035997
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6872133280793995,
+          "rmse": 0.828983309891942,
+          "mae": 0.5440382711666947,
+          "r2": 0.08519994983143697
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7380952380952381,
+          "precision": 0.6785714285714286,
+          "recall": 0.7052785923753666,
+          "f1": 0.6867796610169492
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.6053846153846153,
+          "recall": 0.6720354808590102,
+          "f1": 0.6203502415458937
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.35680852482645714,
+          "js_divergence": 0.09025645010360116
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.12612237693012526,
+        "lr": 0.0009464809649969671,
+        "weight_decay": 0.07051438027639334,
+        "backbone_lr_ratio": 0.9343703462702573,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 19,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.6485083926359867,
+          "rmse": 1.2839425192102591,
+          "mae": 0.5286768473437932,
+          "r2": 0.06759108835091887
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.6919949158910226,
+          "rmse": 0.8318623178693831,
+          "mae": 0.5334780438945202,
+          "r2": 0.3272987252110757
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6375,
+          "recall": 0.6451612903225806,
+          "f1": 0.6407697790449038
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5952380952380952,
+          "precision": 0.5380952380952381,
+          "recall": 0.5259426847662142,
+          "f1": 0.5062561094819159
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.2642989361691452,
+          "js_divergence": 0.06629497269911663
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.20496607857127885,
+        "lr": 0.0009573095972231165,
+        "weight_decay": 0.00038069869878401156,
+        "backbone_lr_ratio": 0.04121720375781922,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 17,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.7197053265910517,
+          "rmse": 0.8483544816826583,
+          "mae": 0.511052241144214,
+          "r2": 0.036944874168868824
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.0001767893259463,
+          "rmse": 1.0000883907565103,
+          "mae": 0.60194506123662,
+          "r2": 0.03284555574221559
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6836007130124777,
+          "recall": 0.7265395894428153,
+          "f1": 0.6891391794046662
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6314814814814814,
+          "recall": 0.6444444444444445,
+          "f1": 0.6326984475152934
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2786833611146747,
+          "js_divergence": 0.07839598379605058
+        }
+      }
+    },
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.3980773018180746,
+        "lr": 0.0009896030589860336,
+        "weight_decay": 0.0002988554096805342,
+        "backbone_lr_ratio": 0.2882797738700643,
+        "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.5048982722984092,
+          "rmse": 1.2267429528219875,
+          "mae": 0.5805474416723093,
+          "r2": 0.07896521442969773
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.6736236117326778,
+          "rmse": 0.8207457655892462,
+          "mae": 0.6456835756142592,
+          "r2": 0.1437837404002562
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7857142857142857,
+          "precision": 0.7305773342836779,
+          "recall": 0.7305773342836779,
+          "f1": 0.7305773342836779
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6646103896103895,
+          "recall": 0.7014285714285715,
+          "f1": 0.6618818596621038
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9827586206896551,
+          "precision": 0.9912280701754386,
+          "recall": 0.75,
+          "f1": 0.8289085545722714
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2890762442571175,
+          "js_divergence": 0.05687567953491692
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.23085562232445592,
+        "lr": 0.0005227270589225132,
+        "weight_decay": 0.00461701039547356,
+        "backbone_lr_ratio": 0.12087164274488635,
+        "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": 84,
+          "mse": 0.34313442749289724,
+          "rmse": 0.5857767727495665,
+          "mae": 0.44384193185362075,
+          "r2": -0.04781054949350838
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.099845523641604,
+          "rmse": 1.0487352018701404,
+          "mae": 0.6761384357499783,
+          "r2": 0.1674962252922061
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.6215277777777778,
+          "recall": 0.6587301587301587,
+          "f1": 0.6149834044570887
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5595238095238095,
+          "precision": 0.5017255158384191,
+          "recall": 0.5057561486132914,
+          "f1": 0.4936173861292259
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.5199198163542676,
+          "js_divergence": 0.12515294319531267
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.318342405407684,
+        "lr": 0.0008937362022799393,
+        "weight_decay": 0.06633867694372529,
+        "backbone_lr_ratio": 0.20906844321970927,
+        "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.4877212975355261,
+          "rmse": 0.6983704586646876,
+          "mae": 0.5448984429905457,
+          "r2": 0.036759151467687046
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.7067592342133922,
+          "rmse": 0.8406897371880974,
+          "mae": 0.5448016864635671,
+          "r2": 0.05918096099439174
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7380952380952381,
+          "precision": 0.6851851851851851,
+          "recall": 0.7199413489736071,
+          "f1": 0.6936339522546419
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6240783117213323,
+          "recall": 0.6851073762838468,
+          "f1": 0.6414141414141414
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.38840109926844824,
+          "js_divergence": 0.09646626485112685
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.12612237693012526,
+        "lr": 0.0009464809649969671,
+        "weight_decay": 0.07051438027639334,
+        "backbone_lr_ratio": 0.9343703462702573,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 19,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.9244709959119324,
+          "rmse": 1.3872530396117113,
+          "mae": 0.5739306664004563,
+          "r2": -0.08849546342266934
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.7217060909432494,
+          "rmse": 0.849532866311392,
+          "mae": 0.5202912685256136,
+          "r2": 0.2984159330486872
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6267796610169492,
+          "recall": 0.6370967741935484,
+          "f1": 0.6307367680675224
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6190476190476191,
+          "precision": 0.5577586206896551,
+          "recall": 0.5861990950226245,
+          "f1": 0.5614861329147044
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.2617822353923735,
+          "js_divergence": 0.06083979574849321
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.20496607857127885,
+        "lr": 0.0009573095972231165,
+        "weight_decay": 0.00038069869878401156,
+        "backbone_lr_ratio": 0.04121720375781922,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 17,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.7148724493029326,
+          "rmse": 0.8455013005920999,
+          "mae": 0.5005925307438398,
+          "r2": 0.04341186430062394
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.0380505423911823,
+          "rmse": 1.018847654161888,
+          "mae": 0.6400383858588236,
+          "r2": -0.0037777382480910227
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7738095238095238,
+          "precision": 0.7213166144200627,
+          "recall": 0.7587976539589443,
+          "f1": 0.7325289089994973
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6308992939427722,
+          "recall": 0.6444444444444445,
+          "f1": 0.633663610371685
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2701739757789761,
+          "js_divergence": 0.07372463442305974
+        }
+      }
+    },
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.3980773018180746,
+        "lr": 0.0009896030589860336,
+        "weight_decay": 0.0002988554096805342,
+        "backbone_lr_ratio": 0.2882797738700643,
+        "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.4541105583695868,
+          "rmse": 1.2058650664023678,
+          "mae": 0.5784477213150766,
+          "r2": 0.11004854548874443
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.59667434347555,
+          "rmse": 0.7724469842491134,
+          "mae": 0.6142617367079546,
+          "r2": 0.24159090377534476
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.75,
+          "precision": 0.6916445623342176,
+          "recall": 0.7059871703492516,
+          "f1": 0.6974789915966386
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6695636813283872,
+          "recall": 0.7142857142857144,
+          "f1": 0.6750809541465826
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2587242230390428,
+          "js_divergence": 0.052354750762398025
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.23085562232445592,
+        "lr": 0.0005227270589225132,
+        "weight_decay": 0.00461701039547356,
+        "backbone_lr_ratio": 0.12087164274488635,
+        "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": 84,
+          "mse": 0.36298737292148275,
+          "rmse": 0.6024843341710079,
+          "mae": 0.45925771394034937,
+          "r2": -0.1084343866601285
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.0462031939942293,
+          "rmse": 1.022840747132333,
+          "mae": 0.634948962077987,
+          "r2": 0.20809960181702725
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6514423076923077,
+          "recall": 0.6904761904761905,
+          "f1": 0.6554552912223134
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5476190476190477,
+          "precision": 0.48089218250508575,
+          "recall": 0.48011512297226583,
+          "f1": 0.47062888038209943
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.5244687341232338,
+          "js_divergence": 0.12337098779147845
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.318342405407684,
+        "lr": 0.0008937362022799393,
+        "weight_decay": 0.06633867694372529,
+        "backbone_lr_ratio": 0.20906844321970927,
+        "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.4579348623566749,
+          "rmse": 0.6767088460753818,
+          "mae": 0.5252955685635763,
+          "r2": 0.09558682875266244
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6918557253424359,
+          "rmse": 0.8317786516510483,
+          "mae": 0.5500254434921469,
+          "r2": 0.07902011443593149
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6428571428571428,
+          "recall": 0.6796187683284457,
+          "f1": 0.6415362731152205
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6222222222222222,
+          "recall": 0.6851073762838468,
+          "f1": 0.636502849002849
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.3340915841461506,
+          "js_divergence": 0.08509576023385822
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.12612237693012526,
+        "lr": 0.0009464809649969671,
+        "weight_decay": 0.07051438027639334,
+        "backbone_lr_ratio": 0.9343703462702573,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 19,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.8202166309492056,
+          "rmse": 1.3491540427057267,
+          "mae": 0.5770047424306408,
+          "r2": -0.02952839998289769
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.8240920088267009,
+          "rmse": 0.9077951359347002,
+          "mae": 0.5719114731271135,
+          "r2": 0.1988846563023151
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6375,
+          "recall": 0.6451612903225806,
+          "f1": 0.6407697790449038
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5833333333333334,
+          "precision": 0.475235446249939,
+          "recall": 0.4788838612368025,
+          "f1": 0.4725999853769101
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.26692129352633476,
+          "js_divergence": 0.061695485072966134
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.20496607857127885,
+        "lr": 0.0009573095972231165,
+        "weight_decay": 0.00038069869878401156,
+        "backbone_lr_ratio": 0.04121720375781922,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 17,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.713518768290232,
+          "rmse": 0.844700401497615,
+          "mae": 0.5241211098341626,
+          "r2": 0.045223257644336656
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.0365573602046312,
+          "rmse": 1.018114610544722,
+          "mae": 0.6499280054764501,
+          "r2": -0.0023338557232996937
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.75,
+          "precision": 0.6949843260188088,
+          "recall": 0.7280058651026393,
+          "f1": 0.7043740573152337
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.639751552795031,
+          "recall": 0.6722222222222222,
+          "f1": 0.6522525093953665
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.3119448408030975,
+          "js_divergence": 0.08658354575321366
+        }
+      }
+    }
+  ],
+  "summary_stats": {
+    "size": {
+      "mse_mean": 0.9648704139917805,
+      "mse_std": 0.5891656318820191,
+      "rmse_mean": 0.9346718018959568,
+      "rmse_std": 0.3020911066753635,
+      "mae_mean": 0.5236781804611128,
+      "mae_std": 0.047009951044540414,
+      "r2_mean": 0.030359535294768855,
+      "r2_std": 0.09435084837045356
+    },
+    "delivery": {
+      "mse_mean": 0.8308453271717486,
+      "mse_std": 0.18767747504329296,
+      "rmse_mean": 0.9057064857810794,
+      "rmse_std": 0.10266980464496812,
+      "mae_mean": 0.5977185404312217,
+      "mae_std": 0.049065196777405395,
+      "r2_mean": 0.1530343582905001,
+      "r2_std": 0.10421669155896922
+    },
+    "pdi": {
+      "accuracy_mean": 0.7238095238095238,
+      "accuracy_std": 0.039411774660111545,
+      "precision_mean": 0.665421093848235,
+      "precision_std": 0.0446473083216482,
+      "recall_mean": 0.6887191408298253,
+      "recall_std": 0.04935509276909138,
+      "f1_mean": 0.669663032996382,
+      "f1_std": 0.04712559651378041
+    },
+    "ee": {
+      "accuracy_mean": 0.6460317460317461,
+      "accuracy_std": 0.058558307129195915,
+      "precision_mean": 0.5894331659154981,
+      "precision_std": 0.06893389274036124,
+      "recall_mean": 0.6169983629815563,
+      "recall_std": 0.08823505436531638,
+      "f1_mean": 0.5906772395955772,
+      "f1_std": 0.07683138902833272
+    },
+    "toxic": {
+      "accuracy_mean": 0.9511890595408049,
+      "accuracy_std": 0.016071213375949578,
+      "precision_mean": 0.7617961570593149,
+      "precision_std": 0.0689235340767993,
+      "recall_mean": 0.958294443004062,
+      "recall_std": 0.05613748087593979,
+      "f1_mean": 0.8151453855878633,
+      "f1_std": 0.03136647604932657
+    },
+    "biodist": {
+      "kl_divergence_mean": 0.34155649523094606,
+      "kl_divergence_std": 0.09686930050549591,
+      "js_divergence_mean": 0.08251453782615217,
+      "js_divergence_std": 0.024921020774773404
+    }
+  }
+}

results/pretrained_encoders/graphmvp_pretrain_s42_summary.json

@@ -0,0 +1,932 @@
+{
+  "fold_results": [
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.3980773018180746,
+        "lr": 0.0009896030589860336,
+        "weight_decay": 0.0002988554096805342,
+        "backbone_lr_ratio": 0.2882797738700643,
+        "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": 83,
+          "mse": 1.5221819405605161,
+          "rmse": 1.2337673770044806,
+          "mae": 0.5776655403784958,
+          "r2": 0.06838718401748567
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.614112021357094,
+          "rmse": 0.7836529980527696,
+          "mae": 0.6216259215926302,
+          "r2": 0.21942656293011076
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7976190476190477,
+          "precision": 0.7458333333333333,
+          "recall": 0.7523164647184604,
+          "f1": 0.7489010022859153
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6691176470588235,
+          "recall": 0.7114285714285714,
+          "f1": 0.6657848324514991
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2940147973370504,
+          "js_divergence": 0.06236386129092423
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.2902385468029013,
+        "lr": 0.00018739040729381568,
+        "weight_decay": 0.08181191083032673,
+        "backbone_lr_ratio": 0.6566372495858538,
+        "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.2934246330317906,
+          "rmse": 0.5416868403716215,
+          "mae": 0.4069202210133274,
+          "r2": 0.10398548983157396
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.156314335236939,
+          "rmse": 1.0753205732417375,
+          "mae": 0.6882066883329974,
+          "r2": 0.12475340569084314
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.5583333333333333,
+          "recall": 0.5634920634920635,
+          "f1": 0.5598915989159892
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5595238095238095,
+          "precision": 0.48693693693693696,
+          "recall": 0.49312592169735026,
+          "f1": 0.47630082959336323
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.5092761936607958,
+          "js_divergence": 0.13355010688479413
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.18117503720006686,
+        "lr": 8.135394049241399e-05,
+        "weight_decay": 0.015076469897649124,
+        "backbone_lr_ratio": 0.9586175198423679,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 18,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.34434236619012065,
+          "rmse": 0.586806924115693,
+          "mae": 0.46669395072967745,
+          "r2": 0.3199299791282213
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6575735692519624,
+          "rmse": 0.810909100486585,
+          "mae": 0.55640976827126,
+          "r2": 0.12465560610938209
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6190476190476191,
+          "precision": 0.6190476190476191,
+          "recall": 0.653958944281525,
+          "f1": 0.5961538461538461
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.6166666666666667,
+          "recall": 0.6785714285714285,
+          "f1": 0.6334814122048165
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.41784407477442437,
+          "js_divergence": 0.11469543309672324
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.1799936812783694,
+        "lr": 0.00033865324369584493,
+        "weight_decay": 0.002833833442807743,
+        "backbone_lr_ratio": 0.3278334190266754,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 8,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.6764751937886317,
+          "rmse": 1.2947877022078298,
+          "mae": 0.622077302143247,
+          "r2": 0.05177285245868457
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.725000410122951,
+          "rmse": 0.8514695591287753,
+          "mae": 0.6181600616487154,
+          "r2": 0.2952134634049711
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5714285714285714,
+          "recall": 0.5821114369501467,
+          "f1": 0.5728813559322035
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.5873015873015873,
+          "recall": 0.6022121669180492,
+          "f1": 0.5872162485065711
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.4920884182489342,
+          "js_divergence": 0.12950821301606347
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.19398015625011103,
+        "lr": 0.0008450080181209037,
+        "weight_decay": 0.002901987198512681,
+        "backbone_lr_ratio": 0.5310897731873246,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 16,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.7301598129307739,
+          "rmse": 0.8544938928575054,
+          "mae": 0.4922567930327551,
+          "r2": 0.022955472832782875
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.125780851035229,
+          "rmse": 1.0610282046370063,
+          "mae": 0.6240321383650961,
+          "r2": -0.08861149844600491
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.75,
+          "precision": 0.683728813559322,
+          "recall": 0.6986803519061584,
+          "f1": 0.6898188851767189
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.6060606060606061,
+          "recall": 0.6418803418803419,
+          "f1": 0.6178929765886287
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.26640631927709985,
+          "js_divergence": 0.07248344781159008
+        }
+      }
+    },
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.3980773018180746,
+        "lr": 0.0009896030589860336,
+        "weight_decay": 0.0002988554096805342,
+        "backbone_lr_ratio": 0.2882797738700643,
+        "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": 83,
+          "mse": 1.5188104198487684,
+          "rmse": 1.2324002677088188,
+          "mae": 0.5873989220259779,
+          "r2": 0.07045063768272763
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.6666016049923045,
+          "rmse": 0.8164567379796094,
+          "mae": 0.6511167181206161,
+          "r2": 0.15270913470266467
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7857142857142857,
+          "precision": 0.7305773342836779,
+          "recall": 0.7305773342836779,
+          "f1": 0.7305773342836779
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.691027926322044,
+          "recall": 0.741904761904762,
+          "f1": 0.6917852553884204
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9827586206896551,
+          "precision": 0.9912280701754386,
+          "recall": 0.75,
+          "f1": 0.8289085545722714
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2601734547154161,
+          "js_divergence": 0.056514013464432016
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.2902385468029013,
+        "lr": 0.00018739040729381568,
+        "weight_decay": 0.08181191083032673,
+        "backbone_lr_ratio": 0.6566372495858538,
+        "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.3231063887830616,
+          "rmse": 0.5684244794016717,
+          "mae": 0.42929340717160985,
+          "r2": 0.01334796030441776
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.1342861224694418,
+          "rmse": 1.0650286956084525,
+          "mae": 0.6788244451106084,
+          "r2": 0.14142717476551148
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.5595238095238095,
+          "precision": 0.5477272727272727,
+          "recall": 0.5634920634920635,
+          "f1": 0.5238241152137275
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5476190476190477,
+          "precision": 0.473011973011973,
+          "recall": 0.4779744065458351,
+          "f1": 0.46420797118471535
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.5813128628991394,
+          "js_divergence": 0.14714358123767265
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.18117503720006686,
+        "lr": 8.135394049241399e-05,
+        "weight_decay": 0.015076469897649124,
+        "backbone_lr_ratio": 0.9586175198423679,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 18,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.43267246099702694,
+          "rmse": 0.6577784284977936,
+          "mae": 0.5142740875349513,
+          "r2": 0.14547962007547455
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6198350807939489,
+          "rmse": 0.7872960566355892,
+          "mae": 0.5207243816383804,
+          "r2": 0.1748920752290377
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6418269230769231,
+          "recall": 0.6730205278592376,
+          "f1": 0.6452241715399609
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6190476190476191,
+          "recall": 0.6757703081232492,
+          "f1": 0.63584229390681
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.42003529077921203,
+          "js_divergence": 0.11409932566058836
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.1799936812783694,
+        "lr": 0.00033865324369584493,
+        "weight_decay": 0.002833833442807743,
+        "backbone_lr_ratio": 0.3278334190266754,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 8,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.709633019803758,
+          "rmse": 1.3075293571479603,
+          "mae": 0.629169082886214,
+          "r2": 0.03301853333873328
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.7678785210635485,
+          "rmse": 0.8762867801487983,
+          "mae": 0.6255631820472368,
+          "r2": 0.2535308451835053
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.5888594164456233,
+          "recall": 0.5982404692082112,
+          "f1": 0.5916666666666667
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.5874977416440831,
+          "recall": 0.6237556561085973,
+          "f1": 0.5958755389330425
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.46770485035756304,
+          "js_divergence": 0.12378056507326861
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.19398015625011103,
+        "lr": 0.0008450080181209037,
+        "weight_decay": 0.002901987198512681,
+        "backbone_lr_ratio": 0.5310897731873246,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 16,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.6133402094103233,
+          "rmse": 0.7831603982648275,
+          "mae": 0.4741146452669498,
+          "r2": 0.17927461319654014
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.0688240926391441,
+          "rmse": 1.0338394907523818,
+          "mae": 0.6684364553155571,
+          "r2": -0.0335352533249651
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7857142857142857,
+          "precision": 0.721875,
+          "recall": 0.7082111436950147,
+          "f1": 0.7142857142857142
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5752393980848153,
+          "recall": 0.6188034188034188,
+          "f1": 0.5883517571357862
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.28068032170399054,
+          "js_divergence": 0.07478332635258625
+        }
+      }
+    },
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.3980773018180746,
+        "lr": 0.0009896030589860336,
+        "weight_decay": 0.0002988554096805342,
+        "backbone_lr_ratio": 0.2882797738700643,
+        "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": 83,
+          "mse": 1.3818212591739194,
+          "rmse": 1.1755089362373727,
+          "mae": 0.5521946805697607,
+          "r2": 0.1542913759905169
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.6972128277816507,
+          "rmse": 0.8349927112146852,
+          "mae": 0.6792561926333041,
+          "r2": 0.11380042333630946
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7857142857142857,
+          "precision": 0.7305773342836779,
+          "recall": 0.7305773342836779,
+          "f1": 0.7305773342836779
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6891715244656421,
+          "recall": 0.7380952380952381,
+          "f1": 0.6965863304906686
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.30701952754304446,
+          "js_divergence": 0.057564049906523536
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.2902385468029013,
+        "lr": 0.00018739040729381568,
+        "weight_decay": 0.08181191083032673,
+        "backbone_lr_ratio": 0.6566372495858538,
+        "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.34204974602902377,
+          "rmse": 0.5848501910994163,
+          "mae": 0.44300510672231513,
+          "r2": -0.04449831793169534
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.1297399798287129,
+          "rmse": 1.0628922710363045,
+          "mae": 0.6765872548166357,
+          "r2": 0.14486827701797667
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5374198146828225,
+          "recall": 0.5396825396825397,
+          "f1": 0.5381231671554252
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5714285714285714,
+          "precision": 0.5202020202020202,
+          "recall": 0.5292564578278864,
+          "f1": 0.5119174942704355
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.5145521130605141,
+          "js_divergence": 0.1322571633039113
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.18117503720006686,
+        "lr": 8.135394049241399e-05,
+        "weight_decay": 0.015076469897649124,
+        "backbone_lr_ratio": 0.9586175198423679,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 18,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.4742411460721128,
+          "rmse": 0.6886516870465887,
+          "mae": 0.5345059656759813,
+          "r2": 0.06338220975850639
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6525032683185743,
+          "rmse": 0.807776744106052,
+          "mae": 0.5494959851183618,
+          "r2": 0.13140505545605896
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.5952380952380952,
+          "precision": 0.6319444444444444,
+          "recall": 0.6671554252199414,
+          "f1": 0.583673469387755
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.6091142490372272,
+          "recall": 0.669234360410831,
+          "f1": 0.6249493532803491
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.4235028564925827,
+          "js_divergence": 0.11443098458109038
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.1799936812783694,
+        "lr": 0.00033865324369584493,
+        "weight_decay": 0.002833833442807743,
+        "backbone_lr_ratio": 0.3278334190266754,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 8,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.7043255314620636,
+          "rmse": 1.305498192822213,
+          "mae": 0.6349209440782307,
+          "r2": 0.036020489197966854
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.7865588440055382,
+          "rmse": 0.886881527604188,
+          "mae": 0.6117917100851566,
+          "r2": 0.23537135185779967
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.616710875331565,
+          "recall": 0.6290322580645161,
+          "f1": 0.6208333333333333
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6092796092796093,
+          "recall": 0.636827551533434,
+          "f1": 0.6157297695238403
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.475069604298859,
+          "js_divergence": 0.12514035771620471
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.19398015625011103,
+        "lr": 0.0008450080181209037,
+        "weight_decay": 0.002901987198512681,
+        "backbone_lr_ratio": 0.5310897731873246,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 16,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.8580511263094861,
+          "rmse": 0.9263104913091971,
+          "mae": 0.5776896519043172,
+          "r2": -0.14817899060384732
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.043085209273457,
+          "rmse": 1.0213154308407648,
+          "mae": 0.59220993094917,
+          "r2": -0.008646178010456307
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6333333333333333,
+          "recall": 0.6583577712609971,
+          "f1": 0.6379310344827587
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6652777777777779,
+          "recall": 0.7064102564102565,
+          "f1": 0.6809523809523809
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.24694979259516875,
+          "js_divergence": 0.06657100087180373
+        }
+      }
+    }
+  ],
+  "summary_stats": {
+    "size": {
+      "mse_mean": 0.9283090169594252,
+      "mse_std": 0.5614727660610882,
+      "rmse_mean": 0.9161103444061993,
+      "rmse_std": 0.29841389684694647,
+      "mae_mean": 0.529478686742254,
+      "mae_std": 0.07272420920077967,
+      "r2_mean": 0.07130794061853929,
+      "r2_std": 0.10242823171449583
+    },
+    "delivery": {
+      "mse_mean": 0.8563537825446997,
+      "mse_std": 0.21315743069248957,
+      "rmse_mean": 0.91834312543158,
+      "rmse_std": 0.1140161677888585,
+      "mae_mean": 0.624162722269715,
+      "mae_std": 0.05009549555609884,
+      "r2_mean": 0.13208402972684966,
+      "r2_std": 0.10281086532328124
+    },
+    "pdi": {
+      "accuracy_mean": 0.6904761904761904,
+      "accuracy_std": 0.07299861769883757,
+      "precision_mean": 0.6372815612874345,
+      "precision_std": 0.06842608292617637,
+      "recall_mean": 0.6499270752265488,
+      "recall_std": 0.06582138762012522,
+      "f1_mean": 0.6322908686064913,
+      "f1_std": 0.07211332053746232
+    },
+    "ee": {
+      "accuracy_mean": 0.6555555555555556,
+      "accuracy_std": 0.05314498108820064,
+      "precision_mean": 0.6003302188598288,
+      "precision_std": 0.06446111150537735,
+      "recall_mean": 0.6363500564172834,
+      "recall_std": 0.07965961750314356,
+      "f1_mean": 0.6057916296274218,
+      "f1_std": 0.07002027044312158
+    },
+    "toxic": {
+      "accuracy_mean": 0.9500961633659415,
+      "accuracy_std": 0.01660762512329349,
+      "precision_mean": 0.7594152046783624,
+      "precision_std": 0.0698998626999633,
+      "recall_mean": 0.9577197303603838,
+      "recall_std": 0.05601633670665784,
+      "f1_mean": 0.8126176429273773,
+      "f1_std": 0.03238005389957549
+    },
+    "biodist": {
+      "kl_divergence_mean": 0.39710869851625297,
+      "kl_divergence_std": 0.10730729317255493,
+      "js_divergence_mean": 0.10165902868454511,
+      "js_divergence_std": 0.03125967849638632
+    }
+  }
+}

results/pretrained_encoders/grover_pretrain_s42_summary.json

@@ -0,0 +1,932 @@
+{
+  "fold_results": [
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.25888363054578495,
+        "lr": 0.0006536264082965735,
+        "weight_decay": 0.00932351457657158,
+        "backbone_lr_ratio": 0.1667894332719967,
+        "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.5789235054424788,
+          "rmse": 1.2565522294924627,
+          "mae": 0.5897082271031946,
+          "r2": 0.03365994962165875
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.5843462843532253,
+          "rmse": 0.7644254602989263,
+          "mae": 0.6106570679694414,
+          "r2": 0.2572606108431995
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.7065217391304348,
+          "recall": 0.7573057733428368,
+          "f1": 0.7039987743220468
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6388888888888888,
+          "recall": 0.6842857142857143,
+          "f1": 0.6433139534883722
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.295739533929628,
+          "js_divergence": 0.06360146142338208
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.37015922545878527,
+        "lr": 0.0009352630241255291,
+        "weight_decay": 0.005516105285933501,
+        "backbone_lr_ratio": 0.12042458689005697,
+        "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.3504606454005027,
+          "rmse": 0.5919971667166176,
+          "mae": 0.4625389823867452,
+          "r2": -0.07018221434674188
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.0871146054096508,
+          "rmse": 1.0426478817940652,
+          "mae": 0.6687804595490948,
+          "r2": 0.17713261263549906
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6067796610169491,
+          "recall": 0.6190476190476191,
+          "f1": 0.6108339272986457
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5357142857142857,
+          "precision": 0.4496353496353496,
+          "recall": 0.4523333809048095,
+          "f1": 0.4379382564939481
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.382227541927671,
+          "js_divergence": 0.09520406464421978
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.4042053577703755,
+        "lr": 0.0004438352448528134,
+        "weight_decay": 0.09116356780872863,
+        "backbone_lr_ratio": 0.14631991134515426,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 17,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.3317227949570641,
+          "rmse": 0.5759538132151432,
+          "mae": 0.44960591762459706,
+          "r2": 0.3448534068400446
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6524805809132939,
+          "rmse": 0.8077627008678315,
+          "mae": 0.5519438893922294,
+          "r2": 0.13143525632475672
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6586452762923352,
+          "recall": 0.6957478005865103,
+          "f1": 0.662107803700724
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.640510613250929,
+          "recall": 0.7009803921568626,
+          "f1": 0.6591013004136229
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.3461542146527445,
+          "js_divergence": 0.09035975387137388
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.12612237693012526,
+        "lr": 0.0009464809649969671,
+        "weight_decay": 0.07051438027639334,
+        "backbone_lr_ratio": 0.9343703462702573,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 16,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.630017453835127,
+          "rmse": 1.27672136891145,
+          "mae": 0.57301952361118,
+          "r2": 0.07804970427286173
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.7426488745509607,
+          "rmse": 0.8617707784271643,
+          "mae": 0.5635692574304797,
+          "r2": 0.27805705915089884
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6489665003563792,
+          "recall": 0.6532258064516129,
+          "f1": 0.6509485094850949
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6190476190476191,
+          "precision": 0.5598290598290598,
+          "recall": 0.5861990950226245,
+          "f1": 0.5606349206349206
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.32329647019984065,
+          "js_divergence": 0.07342759749750183
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.20496607857127885,
+        "lr": 0.0009573095972231165,
+        "weight_decay": 0.00038069869878401156,
+        "backbone_lr_ratio": 0.04121720375781922,
+        "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.714982360155395,
+          "rmse": 0.845566295541275,
+          "mae": 0.5158508041640744,
+          "r2": 0.04326478992735272
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.9337723075176421,
+          "rmse": 0.966318947096476,
+          "mae": 0.615490146495145,
+          "r2": 0.0970575934388902
+        },
+        "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.5733695652173912,
+          "recall": 0.6085470085470085,
+          "f1": 0.5861678004535147
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.3041751519174415,
+          "js_divergence": 0.0828665483142382
+        }
+      }
+    },
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.25888363054578495,
+        "lr": 0.0006536264082965735,
+        "weight_decay": 0.00932351457657158,
+        "backbone_lr_ratio": 0.1667894332719967,
+        "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.516128831525372,
+          "rmse": 1.2313118335845603,
+          "mae": 0.5887543337801134,
+          "r2": 0.07209183574368039
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.6296859698518799,
+          "rmse": 0.793527548262743,
+          "mae": 0.6343367575699913,
+          "r2": 0.199631134600178
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7976190476190477,
+          "precision": 0.7537037037037038,
+          "recall": 0.7929436920883821,
+          "f1": 0.765709598031173
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6626672044094911,
+          "recall": 0.7147619047619048,
+          "f1": 0.6696677261461316
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2648199540119534,
+          "js_divergence": 0.05676284260509181
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.37015922545878527,
+        "lr": 0.0009352630241255291,
+        "weight_decay": 0.005516105285933501,
+        "backbone_lr_ratio": 0.12042458689005697,
+        "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.3644874868613839,
+          "rmse": 0.603727990788388,
+          "mae": 0.4844940258633523,
+          "r2": -0.11301520130806475
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.0796036294761662,
+          "rmse": 1.039039763183376,
+          "mae": 0.6630113559973533,
+          "r2": 0.1828178799589153
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5498305084745763,
+          "recall": 0.5555555555555556,
+          "f1": 0.5509622238061297
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5238095238095238,
+          "precision": 0.4709239314502473,
+          "recall": 0.48534798534798534,
+          "f1": 0.4678875968992248
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.360009988398715,
+          "js_divergence": 0.08968488957961811
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.4042053577703755,
+        "lr": 0.0004438352448528134,
+        "weight_decay": 0.09116356780872863,
+        "backbone_lr_ratio": 0.14631991134515426,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 17,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.3763557616627376,
+          "rmse": 0.613478411733239,
+          "mae": 0.48970496371787575,
+          "r2": 0.2567040950520848
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.703379623118135,
+          "rmse": 0.8386773057130704,
+          "mae": 0.5631674273948496,
+          "r2": 0.06367980912955817
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6752941176470588,
+          "recall": 0.718475073313783,
+          "f1": 0.6785714285714286
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6362280701754386,
+          "recall": 0.7009803921568626,
+          "f1": 0.6528571942316682
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.33457114910094327,
+          "js_divergence": 0.08453543607257773
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.12612237693012526,
+        "lr": 0.0009464809649969671,
+        "weight_decay": 0.07051438027639334,
+        "backbone_lr_ratio": 0.9343703462702573,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 16,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.9634922624183433,
+          "rmse": 1.4012466815012778,
+          "mae": 0.6159229520596115,
+          "r2": -0.1105661891750771
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.8685113223217268,
+          "rmse": 0.9319395486412876,
+          "mae": 0.6189951948916256,
+          "r2": 0.1557038060862732
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6503347534996957,
+          "recall": 0.6810850439882699,
+          "f1": 0.6554552912223134
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5833333333333334,
+          "precision": 0.5217673814165041,
+          "recall": 0.5463298139768727,
+          "f1": 0.5207287433261615
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9,
+          "precision": 0.6666666666666666,
+          "recall": 0.9473684210526316,
+          "f1": 0.7222222222222222
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.24978094522784985,
+          "js_divergence": 0.06060532361779249
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.20496607857127885,
+        "lr": 0.0009573095972231165,
+        "weight_decay": 0.00038069869878401156,
+        "backbone_lr_ratio": 0.04121720375781922,
+        "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.666445125466178,
+          "rmse": 0.8163609039304723,
+          "mae": 0.49950868456203135,
+          "r2": 0.10821363903831627
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.9994847526266784,
+          "rmse": 0.9997423431198053,
+          "mae": 0.6419767608288033,
+          "r2": 0.033514743806195346
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6503347534996957,
+          "recall": 0.6810850439882699,
+          "f1": 0.6554552912223134
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.5829995828118482,
+          "recall": 0.614957264957265,
+          "f1": 0.5934111227214677
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2719915000176464,
+          "js_divergence": 0.0728753918681956
+        }
+      }
+    },
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.25888363054578495,
+        "lr": 0.0006536264082965735,
+        "weight_decay": 0.00932351457657158,
+        "backbone_lr_ratio": 0.1667894332719967,
+        "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.515637174855938,
+          "rmse": 1.2311121698918983,
+          "mae": 0.5759322074836636,
+          "r2": 0.07239274172745469
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.6078859434204266,
+          "rmse": 0.7796704069159138,
+          "mae": 0.6284612801311344,
+          "r2": 0.22734028369354697
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.75,
+          "precision": 0.7147058823529412,
+          "recall": 0.7601568068424804,
+          "f1": 0.7211948790896159
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6494435612082671,
+          "recall": 0.690952380952381,
+          "f1": 0.6540847855642071
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.28451853940500493,
+          "js_divergence": 0.06143098168643295
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.37015922545878527,
+        "lr": 0.0009352630241255291,
+        "weight_decay": 0.005516105285933501,
+        "backbone_lr_ratio": 0.12042458689005697,
+        "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.36607084927706934,
+          "rmse": 0.6050378907779821,
+          "mae": 0.4742527622064309,
+          "r2": -0.11785022720432559
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.072597205810797,
+          "rmse": 1.0356626892047416,
+          "mae": 0.6899142471737549,
+          "r2": 0.18812123758800225
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6309523809523809,
+          "precision": 0.5535714285714286,
+          "recall": 0.5634920634920635,
+          "f1": 0.553421368547419
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5595238095238095,
+          "precision": 0.49943718228031947,
+          "recall": 0.516245659102802,
+          "f1": 0.5015974929297562
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.33867682492719764,
+          "js_divergence": 0.08667274774942817
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.4042053577703755,
+        "lr": 0.0004438352448528134,
+        "weight_decay": 0.09116356780872863,
+        "backbone_lr_ratio": 0.14631991134515426,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 17,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.5053228762985142,
+          "rmse": 0.7108606588484935,
+          "mae": 0.544970038779346,
+          "r2": 0.001996347897613271
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.7184112972452166,
+          "rmse": 0.8475914683650471,
+          "mae": 0.5859764796448872,
+          "r2": 0.043670045518013456
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6836007130124777,
+          "recall": 0.7265395894428153,
+          "f1": 0.6891391794046662
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.6104363472784525,
+          "recall": 0.6666666666666666,
+          "f1": 0.6224764228882279
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.34415499439775205,
+          "js_divergence": 0.08771435148935566
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.12612237693012526,
+        "lr": 0.0009464809649969671,
+        "weight_decay": 0.07051438027639334,
+        "backbone_lr_ratio": 0.9343703462702573,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 16,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.7695798467912622,
+          "rmse": 1.3302555569480858,
+          "mae": 0.6023058171933854,
+          "r2": -0.0008878489144126256
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.7584981089453127,
+          "rmse": 0.8709179691252861,
+          "mae": 0.581826904741258,
+          "r2": 0.2626497202576923
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6489665003563792,
+          "recall": 0.6532258064516129,
+          "f1": 0.6509485094850949
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.571467539552646,
+          "recall": 0.583107088989442,
+          "f1": 0.5753294324722896
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.26330683934899884,
+          "js_divergence": 0.05955508350882244
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.20496607857127885,
+        "lr": 0.0009573095972231165,
+        "weight_decay": 0.00038069869878401156,
+        "backbone_lr_ratio": 0.04121720375781922,
+        "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.6649265530357125,
+          "rmse": 0.8154302870483242,
+          "mae": 0.5136514454137623,
+          "r2": 0.1102456775811359
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.9676062147836739,
+          "rmse": 0.9836697691723956,
+          "mae": 0.6175258674755179,
+          "r2": 0.06434076364622787
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7738095238095238,
+          "precision": 0.7122033898305085,
+          "recall": 0.7294721407624634,
+          "f1": 0.719359943731317
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5648809523809524,
+          "recall": 0.5871794871794872,
+          "f1": 0.573676603432701
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.29835153058548175,
+          "js_divergence": 0.08270536087983824
+        }
+      }
+    }
+  ],
+  "summary_stats": {
+    "size": {
+      "mse_mean": 0.9543035685322051,
+      "mse_std": 0.5981213334036609,
+      "rmse_mean": 0.9270408839286446,
+      "rmse_std": 0.3080564364803995,
+      "mae_mean": 0.5320147123966243,
+      "mae_std": 0.05351966804020229,
+      "r2_mean": 0.04726470045023875,
+      "r2_std": 0.126102268122003
+    },
+    "delivery": {
+      "mse_mean": 0.8270684480229857,
+      "mse_std": 0.17693567699093907,
+      "rmse_mean": 0.9042243053458754,
+      "rmse_std": 0.09719492602371275,
+      "mae_mean": 0.615708873112371,
+      "mae_std": 0.039477393320796444,
+      "r2_mean": 0.1574941704451898,
+      "r2_std": 0.07929093899582297
+    },
+    "pdi": {
+      "accuracy_mean": 0.7142857142857143,
+      "accuracy_std": 0.04123930494211614,
+      "precision_mean": 0.657051911797391,
+      "precision_std": 0.05438419992076548,
+      "recall_mean": 0.6835853401829536,
+      "recall_std": 0.06618289811754123,
+      "f1_mean": 0.6610780721243127,
+      "f1_std": 0.05631636693524698
+    },
+    "ee": {
+      "accuracy_mean": 0.6365079365079365,
+      "accuracy_std": 0.05715387900758358,
+      "precision_mean": 0.5754990153190521,
+      "precision_std": 0.06434539420274109,
+      "recall_mean": 0.6092582823339125,
+      "recall_std": 0.0807804696823919,
+      "f1_mean": 0.5812582234730809,
+      "f1_std": 0.0702133273457957
+    },
+    "toxic": {
+      "accuracy_mean": 0.9478174120312263,
+      "accuracy_std": 0.018643838747495024,
+      "precision_mean": 0.7425396825396825,
+      "precision_std": 0.03650310527039388,
+      "recall_mean": 0.9726010428369776,
+      "recall_std": 0.009861789747008767,
+      "f1_mean": 0.8108612462594762,
+      "f1_std": 0.038209835226095154
+    },
+    "biodist": {
+      "kl_divergence_mean": 0.3107850118699246,
+      "kl_divergence_std": 0.03838394769553977,
+      "js_divergence_mean": 0.07653345565385793,
+      "js_divergence_std": 0.012738301434340309
+    }
+  }
+}

results/pretrained_encoders/moe_pretrain_s42_summary.json

@@ -0,0 +1,977 @@
+{
+  "fold_results": [
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.4884866103742062,
+        "lr": 0.0002617457281978971,
+        "weight_decay": 9.524902838811199e-05,
+        "backbone_lr_ratio": 0.9539957173265954,
+        "moe_n_experts": 4,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 2,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 23,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.5279509180840005,
+          "rmse": 1.2361031179007682,
+          "mae": 0.578313840710255,
+          "r2": 0.0648564277702961
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.6655882063296256,
+          "rmse": 0.8158358942395374,
+          "mae": 0.674296624894286,
+          "r2": 0.15399722555537465
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7619047619047619,
+          "precision": 0.7122257053291536,
+          "recall": 0.7412687099073414,
+          "f1": 0.7214854111405835
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6748583260680036,
+          "recall": 0.7142857142857144,
+          "f1": 0.6773495605138607
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2500928345559933,
+          "js_divergence": 0.05156067014803717
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.10335173439505978,
+        "lr": 0.000196757629557381,
+        "weight_decay": 0.0003725444839719432,
+        "backbone_lr_ratio": 0.05035847475715459,
+        "moe_n_experts": 4,
+        "moe_top_k": 1,
+        "moe_expert_hidden_mult": 1,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 19,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.34213962908738765,
+          "rmse": 0.584927028856923,
+          "mae": 0.45376784687063526,
+          "r2": -0.044772788836473065
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.129722640308642,
+          "rmse": 1.0628841142423016,
+          "mae": 0.6647483138825561,
+          "r2": 0.1448814017846829
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6309523809523809,
+          "precision": 0.5535714285714286,
+          "recall": 0.5634920634920635,
+          "f1": 0.553421368547419
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5238095238095238,
+          "precision": 0.46498599439775906,
+          "recall": 0.46436896436896435,
+          "f1": 0.4483749534478767
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.49848379935596204,
+          "js_divergence": 0.1305599268889464
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.2514825943283291,
+        "lr": 0.00029260650717186565,
+        "weight_decay": 0.0001495891705646636,
+        "backbone_lr_ratio": 0.2434609849039314,
+        "moe_n_experts": 2,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 1,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 20,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.3851340701995523,
+          "rmse": 0.6205917097412375,
+          "mae": 0.4943395415542736,
+          "r2": 0.2393670925336251
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.7088597182801071,
+          "rmse": 0.8419380727108776,
+          "mae": 0.5629216438857838,
+          "r2": 0.0563848526375017
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6505882352941177,
+          "recall": 0.6876832844574781,
+          "f1": 0.6517857142857143
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.6139316239316239,
+          "recall": 0.6785714285714285,
+          "f1": 0.6277576489533011
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.2665180519269797,
+          "js_divergence": 0.06450481128744694
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.40463841601213524,
+        "lr": 0.0002572978619380859,
+        "weight_decay": 2.0354315009822035e-05,
+        "backbone_lr_ratio": 0.8823918523568081,
+        "moe_n_experts": 4,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 2,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 19,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.820237956048593,
+          "rmse": 1.3491619458199202,
+          "mae": 0.6053206579400653,
+          "r2": -0.029540461621649028
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.6693395316760985,
+          "rmse": 0.8181317324710603,
+          "mae": 0.5448598891843931,
+          "r2": 0.34932245037470455
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6159874608150471,
+          "recall": 0.6356304985337243,
+          "f1": 0.6199095022624435
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6309523809523809,
+          "precision": 0.5653410228683567,
+          "recall": 0.5826043237807944,
+          "f1": 0.5650366426763943
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.29215591110500194,
+          "js_divergence": 0.07269076599614506
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.3730299609537786,
+        "lr": 0.0007646220605180808,
+        "weight_decay": 0.004588258639584098,
+        "backbone_lr_ratio": 0.17129190112240483,
+        "moe_n_experts": 2,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 2,
+        "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.7676733172863143,
+          "rmse": 0.8761696852130381,
+          "mae": 0.5270063721386361,
+          "r2": -0.02724225576902284
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.9040350302083842,
+          "rmse": 0.9508075673912068,
+          "mae": 0.6287906430019387,
+          "r2": 0.12581304969093698
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6836007130124777,
+          "recall": 0.7265395894428153,
+          "f1": 0.6891391794046662
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6645240097721175,
+          "recall": 0.7038461538461539,
+          "f1": 0.6728729539789446
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.20431460217894573,
+          "js_divergence": 0.05043386132176378
+        }
+      }
+    },
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.4884866103742062,
+        "lr": 0.0002617457281978971,
+        "weight_decay": 9.524902838811199e-05,
+        "backbone_lr_ratio": 0.9539957173265954,
+        "moe_n_experts": 4,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 2,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 23,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.5550563846692986,
+          "rmse": 1.2470189993217018,
+          "mae": 0.5954494658094572,
+          "r2": 0.04826721502168674
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.704182837608261,
+          "rmse": 0.8391560269748772,
+          "mae": 0.6780449980548744,
+          "r2": 0.10494112024899149
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6737967914438503,
+          "recall": 0.7084818246614397,
+          "f1": 0.6785714285714286
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6931286549707601,
+          "recall": 0.7380952380952381,
+          "f1": 0.7007070707070707
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2453488369022721,
+          "js_divergence": 0.052406074597242516
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.10335173439505978,
+        "lr": 0.000196757629557381,
+        "weight_decay": 0.0003725444839719432,
+        "backbone_lr_ratio": 0.05035847475715459,
+        "moe_n_experts": 4,
+        "moe_top_k": 1,
+        "moe_expert_hidden_mult": 1,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 19,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.3727706424788934,
+          "rmse": 0.610549459486202,
+          "mae": 0.48057969708350445,
+          "r2": -0.13830901371429016
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.0803452340552668,
+          "rmse": 1.0393965720817377,
+          "mae": 0.6485996645432516,
+          "r2": 0.18225653875401804
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6309523809523809,
+          "precision": 0.5535714285714286,
+          "recall": 0.5634920634920635,
+          "f1": 0.553421368547419
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5476190476190477,
+          "precision": 0.49737237237237236,
+          "recall": 0.5073022215879358,
+          "f1": 0.4920019687461548
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9508196721311475,
+          "precision": 0.75,
+          "recall": 0.9741379310344828,
+          "f1": 0.8200589970501475
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.5055220562600806,
+          "js_divergence": 0.1300376337593102
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.2514825943283291,
+        "lr": 0.00029260650717186565,
+        "weight_decay": 0.0001495891705646636,
+        "backbone_lr_ratio": 0.2434609849039314,
+        "moe_n_experts": 2,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 1,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 20,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.501474748290053,
+          "rmse": 0.708148817897801,
+          "mae": 0.5538181106981245,
+          "r2": 0.009596331959947801
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.739867919930154,
+          "rmse": 0.8601557532971305,
+          "mae": 0.5739390922787909,
+          "r2": 0.015107561778814449
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7380952380952381,
+          "precision": 0.6923076923076923,
+          "recall": 0.7346041055718475,
+          "f1": 0.6998050682261208
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6228947368421053,
+          "recall": 0.6851073762838468,
+          "f1": 0.63836444060848
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.2914659417557568,
+          "js_divergence": 0.07149076417555128
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.40463841601213524,
+        "lr": 0.0002572978619380859,
+        "weight_decay": 2.0354315009822035e-05,
+        "backbone_lr_ratio": 0.8823918523568081,
+        "moe_n_experts": 4,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 2,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 19,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.7508482794444888,
+          "rmse": 1.3231962361813492,
+          "mae": 0.5995979247704507,
+          "r2": 0.009706868347147068
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.6800398997076931,
+          "rmse": 0.824645317520019,
+          "mae": 0.5383539610084588,
+          "r2": 0.33892042132757505
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6423197492163009,
+          "recall": 0.6664222873900294,
+          "f1": 0.6480643539467068
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6666666666666666,
+          "precision": 0.5898730898730898,
+          "recall": 0.6022121669180492,
+          "f1": 0.5907255864702673
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9166666666666666,
+          "precision": 0.6875,
+          "recall": 0.956140350877193,
+          "f1": 0.749791492910759
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.22697074487882685,
+          "js_divergence": 0.05572372821276913
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.3730299609537786,
+        "lr": 0.0007646220605180808,
+        "weight_decay": 0.004588258639584098,
+        "backbone_lr_ratio": 0.17129190112240483,
+        "moe_n_experts": 2,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 2,
+        "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.7388933186856751,
+          "rmse": 0.8595890405802502,
+          "mae": 0.4949821270059883,
+          "r2": 0.01126895729235744
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.0735857604041814,
+          "rmse": 1.0361398363175607,
+          "mae": 0.6398222603160759,
+          "r2": -0.038139707447659266
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6345029239766082,
+          "recall": 0.6517595307917888,
+          "f1": 0.6398559423769508
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.662962962962963,
+          "recall": 0.7064102564102565,
+          "f1": 0.6789544933874829
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.2849032950499658,
+          "js_divergence": 0.0742269697587313
+        }
+      }
+    },
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.4884866103742062,
+        "lr": 0.0002617457281978971,
+        "weight_decay": 9.524902838811199e-05,
+        "backbone_lr_ratio": 0.9539957173265954,
+        "moe_n_experts": 4,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 2,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 23,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.575393168842647,
+          "rmse": 1.255146672242988,
+          "mae": 0.5922032207340361,
+          "r2": 0.035820602519645584
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.6805529919450708,
+          "rmse": 0.8249563576002495,
+          "mae": 0.6910147559026192,
+          "r2": 0.13497608000435835
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7976190476190477,
+          "precision": 0.75,
+          "recall": 0.7794012829650748,
+          "f1": 0.7606837606837606
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7380952380952381,
+          "precision": 0.7064443475733798,
+          "recall": 0.7547619047619047,
+          "f1": 0.7126436781609194
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.23590056396282402,
+          "js_divergence": 0.047894235451583785
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.10335173439505978,
+        "lr": 0.000196757629557381,
+        "weight_decay": 0.0003725444839719432,
+        "backbone_lr_ratio": 0.05035847475715459,
+        "moe_n_experts": 4,
+        "moe_top_k": 1,
+        "moe_expert_hidden_mult": 1,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 19,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.3579244858231685,
+          "rmse": 0.5982679047242703,
+          "mae": 0.46379138462777647,
+          "r2": -0.09297412943304506
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.1096558552074713,
+          "rmse": 1.0534020387332992,
+          "mae": 0.6546663998885721,
+          "r2": 0.1600705115131673
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.5696286472148542,
+          "recall": 0.5793650793650793,
+          "f1": 0.571654650958326
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5476190476190477,
+          "precision": 0.49253034547152197,
+          "recall": 0.49681271109842545,
+          "f1": 0.4801587301587302
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.502867985249319,
+          "js_divergence": 0.13123665503825419
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.2514825943283291,
+        "lr": 0.00029260650717186565,
+        "weight_decay": 0.0001495891705646636,
+        "backbone_lr_ratio": 0.2434609849039314,
+        "moe_n_experts": 2,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 1,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 20,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.618457534400238,
+          "rmse": 0.7864207108159335,
+          "mae": 0.5679304331563235,
+          "r2": -0.22144258048008303
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.7362245815365097,
+          "rmse": 0.8580353031994137,
+          "mae": 0.5702473803190514,
+          "r2": 0.019957476658378215
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6523864289821737,
+          "recall": 0.6942815249266863,
+          "f1": 0.6473332296687918
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6228947368421053,
+          "recall": 0.6851073762838468,
+          "f1": 0.63836444060848
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.26356343521196823,
+          "js_divergence": 0.06378606399643424
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.40463841601213524,
+        "lr": 0.0002572978619380859,
+        "weight_decay": 2.0354315009822035e-05,
+        "backbone_lr_ratio": 0.8823918523568081,
+        "moe_n_experts": 4,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 2,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 19,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.8408361971730496,
+          "rmse": 1.3567741879815702,
+          "mae": 0.6257544057668153,
+          "r2": -0.04119098379947639
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.737776343129472,
+          "rmse": 0.8589390799873249,
+          "mae": 0.5497218715026975,
+          "r2": 0.28279373860245827
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6333333333333333,
+          "recall": 0.6583577712609971,
+          "f1": 0.6379310344827587
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5640697674418605,
+          "recall": 0.5891402714932127,
+          "f1": 0.5710279611967066
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.24488500401793606,
+          "js_divergence": 0.060842241588791174
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.3730299609537786,
+        "lr": 0.0007646220605180808,
+        "weight_decay": 0.004588258639584098,
+        "backbone_lr_ratio": 0.17129190112240483,
+        "moe_n_experts": 2,
+        "moe_top_k": 2,
+        "moe_expert_hidden_mult": 2,
+        "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.712297885673984,
+          "rmse": 0.8439774201209319,
+          "mae": 0.5084991393599748,
+          "r2": 0.046856950237922024
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.1055967310671657,
+          "rmse": 1.0514735997956228,
+          "mae": 0.6599975297677105,
+          "r2": -0.06909378763839769
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7619047619047619,
+          "precision": 0.711111111111111,
+          "recall": 0.750733137829912,
+          "f1": 0.7214854111405835
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7023809523809523,
+          "precision": 0.6536363636363637,
+          "recall": 0.7038461538461539,
+          "f1": 0.6717592592592593
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.1782104997564988,
+          "js_divergence": 0.04297224929274232
+        }
+      }
+    }
+  ],
+  "summary_stats": {
+    "size": {
+      "mse_mean": 0.9911392357458229,
+      "mse_std": 0.5817120308358389,
+      "rmse_mean": 0.9504028624589923,
+      "rmse_std": 0.29643487442535593,
+      "mae_mean": 0.542756944548421,
+      "mae_std": 0.054981488862431024,
+      "r2_mean": -0.008648784531427449,
+      "r2_std": 0.09954199193808574
+    },
+    "delivery": {
+      "mse_mean": 0.8483582187596068,
+      "mse_std": 0.18619352158793398,
+      "rmse_mean": 0.9157264844374813,
+      "rmse_std": 0.09901123400593626,
+      "mae_mean": 0.6186683352287373,
+      "mae_std": 0.05324906806818642,
+      "r2_mean": 0.13081259558966032,
+      "r2_std": 0.1207826593867956
+    },
+    "pdi": {
+      "accuracy_mean": 0.703968253968254,
+      "accuracy_std": 0.04618202071738462,
+      "precision_mean": 0.6485954432786384,
+      "precision_std": 0.056485137303304574,
+      "recall_mean": 0.6761008502725561,
+      "recall_std": 0.06599285309996766,
+      "f1_mean": 0.6529698282829115,
+      "f1_std": 0.05951390665637187
+    },
+    "ee": {
+      "accuracy_mean": 0.6595238095238095,
+      "accuracy_std": 0.0668930397006523,
+      "precision_mean": 0.6059632236682923,
+      "precision_std": 0.07313293968503662,
+      "recall_mean": 0.640831484108795,
+      "recall_std": 0.0904387508717925,
+      "f1_mean": 0.6110732925915952,
+      "f1_std": 0.08105368908168387
+    },
+    "toxic": {
+      "accuracy_mean": 0.947835626967474,
+      "accuracy_std": 0.01646827276703399,
+      "precision_mean": 0.741547619047619,
+      "precision_std": 0.034629149636327315,
+      "recall_mean": 0.9726111255149369,
+      "recall_std": 0.008721274087923442,
+      "f1_mean": 0.810171454978226,
+      "f1_std": 0.03512408564460891
+    },
+    "biodist": {
+      "kl_divergence_mean": 0.299413570811222,
+      "kl_divergence_std": 0.1056856517385418,
+      "js_divergence_mean": 0.0733577767675833,
+      "js_divergence_std": 0.029986885602095156
+    }
+  }
+}

results/pretrained_encoders/molformer_pretrain_s42_summary.json

@@ -0,0 +1,932 @@
+{
+  "fold_results": [
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.3980773018180746,
+        "lr": 0.0009896030589860336,
+        "weight_decay": 0.0002988554096805342,
+        "backbone_lr_ratio": 0.2882797738700643,
+        "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.6246690728469169,
+          "rmse": 1.274625071480597,
+          "mae": 0.5944420359916536,
+          "r2": 0.005662536347478486
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.7102819707127421,
+          "rmse": 0.842782279543621,
+          "mae": 0.6625874312594533,
+          "r2": 0.09718875402761229
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.8333333333333334,
+          "precision": 0.7891525423728813,
+          "recall": 0.8039914468995011,
+          "f1": 0.7958333333333334
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6829362569400521,
+          "recall": 0.7314285714285714,
+          "f1": 0.687769981711114
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9655172413793104,
+          "precision": 0.75,
+          "recall": 0.9821428571428572,
+          "f1": 0.8242424242424242
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.32097759110945434,
+          "js_divergence": 0.06712981080576393
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.396946460618591,
+        "lr": 0.0004854064403598687,
+        "weight_decay": 0.00010992238518795333,
+        "backbone_lr_ratio": 0.23864875345272699,
+        "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.38773907280523073,
+          "rmse": 0.6226869781882634,
+          "mae": 0.478045186826161,
+          "r2": -0.18401727831452308
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.1492300175260897,
+          "rmse": 1.0720214631834988,
+          "mae": 0.7018961062746458,
+          "r2": 0.13011572349706002
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.5952380952380952,
+          "precision": 0.5664174813110983,
+          "recall": 0.5873015873015872,
+          "f1": 0.5523510971786834
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5476190476190477,
+          "precision": 0.4826358616681197,
+          "recall": 0.4884639170353456,
+          "f1": 0.47545555874980455
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.4828981973271936,
+          "js_divergence": 0.12323903696473625
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.10827786083228613,
+        "lr": 0.0004876914007718381,
+        "weight_decay": 0.0852885971769344,
+        "backbone_lr_ratio": 0.21714091995810716,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 8,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.4123269005771989,
+          "rmse": 0.6421268570751413,
+          "mae": 0.49745388741471935,
+          "r2": 0.1856617383911774
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6503294595263629,
+          "rmse": 0.8064300710702466,
+          "mae": 0.5487457474150385,
+          "r2": 0.1342987717315124
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.5595238095238095,
+          "precision": 0.5652297220646625,
+          "recall": 0.5843108504398826,
+          "f1": 0.5357729648991785
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.6053846153846153,
+          "recall": 0.6720354808590102,
+          "f1": 0.6203502415458937
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.3376426852856137,
+          "js_divergence": 0.09061480035897176
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.1463802845991778,
+        "lr": 0.00017979738268552602,
+        "weight_decay": 0.026590603512062855,
+        "backbone_lr_ratio": 0.19962036354162335,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 18,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.7510092489615507,
+          "rmse": 1.3232570608016987,
+          "mae": 0.6278676227834749,
+          "r2": 0.009615822761402493
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.7222390124885935,
+          "rmse": 0.84984646406783,
+          "mae": 0.57879668119032,
+          "r2": 0.29789786998972123
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7380952380952381,
+          "precision": 0.6785714285714286,
+          "recall": 0.7052785923753666,
+          "f1": 0.6867796610169492
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5476190476190477,
+          "precision": 0.5483829895594602,
+          "recall": 0.575062845651081,
+          "f1": 0.5199383832295225
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.4613132992323466,
+          "js_divergence": 0.12165500439647682
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.35317424690388177,
+        "lr": 0.0009765192272158632,
+        "weight_decay": 0.00012175015526770499,
+        "backbone_lr_ratio": 0.023808704109897987,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 18,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.6604880002050739,
+          "rmse": 0.8127041283302762,
+          "mae": 0.4948491998892721,
+          "r2": 0.11618501260741076
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.936642398150463,
+          "rmse": 0.9678028715345202,
+          "mae": 0.6353725126848139,
+          "r2": 0.09428226317670119
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6071428571428571,
+          "precision": 0.6374353076480737,
+          "recall": 0.6752199413489737,
+          "f1": 0.5942028985507246
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.6475761475761476,
+          "recall": 0.7401709401709402,
+          "f1": 0.6653096672637463
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.40727903738537147,
+          "js_divergence": 0.10249950448751813
+        }
+      }
+    },
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.3980773018180746,
+        "lr": 0.0009896030589860336,
+        "weight_decay": 0.0002988554096805342,
+        "backbone_lr_ratio": 0.2882797738700643,
+        "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.545312966192772,
+          "rmse": 1.2431061765564404,
+          "mae": 0.5886967062770602,
+          "r2": 0.054230426962613154
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.6668586486255093,
+          "rmse": 0.8166141369248449,
+          "mae": 0.6276618008790859,
+          "r2": 0.15238241673384167
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.8452380952380952,
+          "precision": 0.8041666666666667,
+          "recall": 0.8121881682109764,
+          "f1": 0.8079831193951117
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6914328121224673,
+          "recall": 0.7452380952380953,
+          "f1": 0.701447533036203
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9827586206896551,
+          "precision": 0.9912280701754386,
+          "recall": 0.75,
+          "f1": 0.8289085545722714
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.27259534788302564,
+          "js_divergence": 0.06020453275369037
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.396946460618591,
+        "lr": 0.0004854064403598687,
+        "weight_decay": 0.00010992238518795333,
+        "backbone_lr_ratio": 0.23864875345272699,
+        "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.4005803081504934,
+          "rmse": 0.6329141396354591,
+          "mae": 0.4951437964946741,
+          "r2": -0.2232298457086057
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.1667453321847503,
+          "rmse": 1.0801598641797197,
+          "mae": 0.6814669590809794,
+          "r2": 0.11685789296077709
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.5595238095238095,
+          "precision": 0.5477272727272727,
+          "recall": 0.5634920634920635,
+          "f1": 0.5238241152137275
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5238095238095238,
+          "precision": 0.456992337164751,
+          "recall": 0.46650968079539507,
+          "f1": 0.45625740933077347
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.5730307313158203,
+          "js_divergence": 0.14005679605749516
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.10827786083228613,
+        "lr": 0.0004876914007718381,
+        "weight_decay": 0.0852885971769344,
+        "backbone_lr_ratio": 0.21714091995810716,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 8,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.45499801818836005,
+          "rmse": 0.6745354091434785,
+          "mae": 0.5309686183796397,
+          "r2": 0.1013870435125862
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6030715996767626,
+          "rmse": 0.7765768472448574,
+          "mae": 0.5282929584461574,
+          "r2": 0.1972071741079553
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6336898395721925,
+          "recall": 0.6649560117302054,
+          "f1": 0.635076427996782
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6547619047619048,
+          "precision": 0.5992402659069326,
+          "recall": 0.6507936507936507,
+          "f1": 0.6101851851851852
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.358983754945942,
+          "js_divergence": 0.09298507446003766
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.1463802845991778,
+        "lr": 0.00017979738268552602,
+        "weight_decay": 0.026590603512062855,
+        "backbone_lr_ratio": 0.19962036354162335,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 18,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.6469751588239248,
+          "rmse": 1.2833452999188975,
+          "mae": 0.6211532356262387,
+          "r2": 0.06845829708118345
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.7509799772948998,
+          "rmse": 0.866591009239595,
+          "mae": 0.6103691533050042,
+          "r2": 0.2699582374579258
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.6545454545454545,
+          "recall": 0.6994134897360704,
+          "f1": 0.625668449197861
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6190476190476191,
+          "precision": 0.5681642512077295,
+          "recall": 0.6041478129713425,
+          "f1": 0.563984674329502
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.4523145833037008,
+          "js_divergence": 0.1183277351587518
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.35317424690388177,
+        "lr": 0.0009765192272158632,
+        "weight_decay": 0.00012175015526770499,
+        "backbone_lr_ratio": 0.023808704109897987,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 18,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.6556117434601818,
+          "rmse": 0.8096985509806608,
+          "mae": 0.5088372581716006,
+          "r2": 0.12271004984074696
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.0003346592105318,
+          "rmse": 1.0001673156080095,
+          "mae": 0.6612379994114925,
+          "r2": 0.03269289816995369
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.7142857142857143,
+          "precision": 0.6592592592592592,
+          "recall": 0.6891495601173021,
+          "f1": 0.6657824933687002
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.6072310405643738,
+          "recall": 0.6940170940170941,
+          "f1": 0.6169253712734788
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.4083994319724002,
+          "js_divergence": 0.1006941777950375
+        }
+      }
+    },
+    {
+      "fold": 0,
+      "best_params": {
+        "dropout": 0.3980773018180746,
+        "lr": 0.0009896030589860336,
+        "weight_decay": 0.0002988554096805342,
+        "backbone_lr_ratio": 0.2882797738700643,
+        "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.444963857475926,
+          "rmse": 1.2020664946149717,
+          "mae": 0.575902431103659,
+          "r2": 0.11564655158081194
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 0.6487752159789035,
+          "rmse": 0.8054658378720376,
+          "mae": 0.6530212075288954,
+          "r2": 0.1753675508528415
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.8333333333333334,
+          "precision": 0.7936507936507937,
+          "recall": 0.7769066286528867,
+          "f1": 0.7844574780058651
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.7261904761904762,
+          "precision": 0.6868692000270947,
+          "recall": 0.7452380952380953,
+          "f1": 0.6987547892720306
+        },
+        "toxic": {
+          "n_samples": 58,
+          "accuracy": 0.9482758620689655,
+          "precision": 0.7,
+          "recall": 0.9732142857142857,
+          "f1": 0.7719528178243775
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.32634100712155006,
+          "js_divergence": 0.06706463335858984
+        }
+      }
+    },
+    {
+      "fold": 1,
+      "best_params": {
+        "dropout": 0.396946460618591,
+        "lr": 0.0004854064403598687,
+        "weight_decay": 0.00010992238518795333,
+        "backbone_lr_ratio": 0.23864875345272699,
+        "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.45034502529074416,
+          "rmse": 0.6710775106429541,
+          "mae": 0.5225315168943434,
+          "r2": -0.37519359936953656
+        },
+        "delivery": {
+          "n_samples": 61,
+          "mse": 1.1274144896048421,
+          "rmse": 1.0617977630438116,
+          "mae": 0.6428290372561725,
+          "r2": 0.14662850547534068
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6904761904761905,
+          "precision": 0.5972915181753385,
+          "recall": 0.6031746031746031,
+          "f1": 0.5997067448680352
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.5238095238095238,
+          "precision": 0.4141642671054436,
+          "recall": 0.42669235526378385,
+          "f1": 0.40981240981240985
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9016393442622951,
+          "precision": 0.6666666666666666,
+          "recall": 0.9482758620689655,
+          "f1": 0.7227272727272727
+        },
+        "biodist": {
+          "n_samples": 61,
+          "kl_divergence": 0.5404966442274723,
+          "js_divergence": 0.13248730719985696
+        }
+      }
+    },
+    {
+      "fold": 2,
+      "best_params": {
+        "dropout": 0.10827786083228613,
+        "lr": 0.0004876914007718381,
+        "weight_decay": 0.0852885971769344,
+        "backbone_lr_ratio": 0.21714091995810716,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 8,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.5404522607814335,
+          "rmse": 0.7351545829153441,
+          "mae": 0.5627435291618375,
+          "r2": -0.06738355880059732
+        },
+        "delivery": {
+          "n_samples": 60,
+          "mse": 0.6347980231433975,
+          "rmse": 0.7967421308951834,
+          "mae": 0.5446047547040507,
+          "r2": 0.1549738055263823
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.5952380952380952,
+          "precision": 0.6068181818181818,
+          "recall": 0.6378299120234604,
+          "f1": 0.5757575757575757
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6190476190476191,
+          "precision": 0.5553467000835423,
+          "recall": 0.588702147525677,
+          "f1": 0.5573950928686081
+        },
+        "toxic": {
+          "n_samples": 61,
+          "accuracy": 0.9344262295081968,
+          "precision": 0.7142857142857143,
+          "recall": 0.9655172413793103,
+          "f1": 0.7821428571428571
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.34103334231152865,
+          "js_divergence": 0.0899550896776649
+        }
+      }
+    },
+    {
+      "fold": 3,
+      "best_params": {
+        "dropout": 0.1463802845991778,
+        "lr": 0.00017979738268552602,
+        "weight_decay": 0.026590603512062855,
+        "backbone_lr_ratio": 0.19962036354162335,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 18,
+      "test_metrics": {
+        "size": {
+          "n_samples": 83,
+          "mse": 1.633019754709887,
+          "rmse": 1.2778966134667886,
+          "mae": 0.6318143501271865,
+          "r2": 0.07635158001484588
+        },
+        "delivery": {
+          "n_samples": 59,
+          "mse": 0.778873610785293,
+          "rmse": 0.8825381639256702,
+          "mae": 0.591008120058578,
+          "r2": 0.2428423116373991
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.5983050847457627,
+          "recall": 0.6063049853372434,
+          "f1": 0.6011957095129242
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6428571428571429,
+          "precision": 0.5736842105263158,
+          "recall": 0.6070889894419307,
+          "f1": 0.5785510785510786
+        },
+        "toxic": {
+          "n_samples": 60,
+          "accuracy": 0.9333333333333333,
+          "precision": 0.7142857142857143,
+          "recall": 0.9649122807017544,
+          "f1": 0.7818181818181817
+        },
+        "biodist": {
+          "n_samples": 60,
+          "kl_divergence": 0.46990571185549446,
+          "js_divergence": 0.12454201663905487
+        }
+      }
+    },
+    {
+      "fold": 4,
+      "best_params": {
+        "dropout": 0.35317424690388177,
+        "lr": 0.0009765192272158632,
+        "weight_decay": 0.00012175015526770499,
+        "backbone_lr_ratio": 0.023808704109897987,
+        "d_model": 256,
+        "num_heads": 8,
+        "n_attn_layers": 4,
+        "fusion_strategy": "attention",
+        "head_hidden_dim": 128,
+        "set_transformer_block": "sab"
+      },
+      "epoch_mean": 18,
+      "test_metrics": {
+        "size": {
+          "n_samples": 84,
+          "mse": 0.7197383052375195,
+          "rmse": 0.8483739182916454,
+          "mae": 0.5446056428176927,
+          "r2": 0.0369007446434223
+        },
+        "delivery": {
+          "n_samples": 58,
+          "mse": 1.0250908773516882,
+          "rmse": 1.0124677166960379,
+          "mae": 0.6531032593599682,
+          "r2": 0.008754043905628395
+        },
+        "pdi": {
+          "n_samples": 84,
+          "accuracy": 0.6190476190476191,
+          "precision": 0.6079545454545454,
+          "recall": 0.6392961876832844,
+          "f1": 0.5909920876445527
+        },
+        "ee": {
+          "n_samples": 84,
+          "accuracy": 0.6785714285714286,
+          "precision": 0.6275303643724697,
+          "recall": 0.7021367521367522,
+          "f1": 0.6415056360708534
+        },
+        "toxic": {
+          "n_samples": 59,
+          "accuracy": 0.9661016949152542,
+          "precision": 0.8,
+          "recall": 0.9821428571428572,
+          "f1": 0.865909090909091
+        },
+        "biodist": {
+          "n_samples": 58,
+          "kl_divergence": 0.40077254800841094,
+          "js_divergence": 0.10315238600509395
+        }
+      }
+    }
+  ],
+  "summary_stats": {
+    "size": {
+      "mse_mean": 0.9552153129138143,
+      "mse_std": 0.544124280885305,
+      "rmse_mean": 0.9369045861361744,
+      "rmse_std": 0.27825367812630597,
+      "mae_mean": 0.5516703345306142,
+      "mae_std": 0.050933508755883206,
+      "r2_mean": 0.0028657014366944233,
+      "r2_std": 0.1484374631238745
+    },
+    "delivery": {
+      "mse_mean": 0.8381110194840553,
+      "mse_std": 0.19956001015608432,
+      "rmse_mean": 0.9092002623352989,
+      "rmse_std": 0.1070789542976535,
+      "mae_mean": 0.6213995819236436,
+      "mae_std": 0.05084304982304164,
+      "r2_mean": 0.15009654795004349,
+      "r2_std": 0.07708969134269554
+    },
+    "pdi": {
+      "accuracy_mean": 0.6793650793650793,
+      "accuracy_std": 0.09408695848584789,
+      "precision_mean": 0.6493476732389074,
+      "precision_std": 0.08116705378595848,
+      "recall_mean": 0.6699209352348937,
+      "recall_std": 0.07666166785670459,
+      "f1_mean": 0.6383589437293336,
+      "f1_std": 0.08952118234954987
+    },
+    "ee": {
+      "accuracy_mean": 0.6341269841269842,
+      "accuracy_std": 0.06796728274922742,
+      "precision_mean": 0.583171421347301,
+      "precision_std": 0.08031795761416705,
+      "recall_mean": 0.6291817619044511,
+      "recall_std": 0.10148518503800956,
+      "f1_mean": 0.5869095341486802,
+      "f1_std": 0.08729125759225413
+    },
+    "toxic": {
+      "accuracy_mean": 0.9445751533791317,
+      "accuracy_std": 0.020063303116255492,
+      "precision_mean": 0.7481453634085212,
+      "precision_std": 0.07540464740954501,
+      "recall_mean": 0.9548256416904332,
+      "recall_std": 0.0554955125846467,
+      "f1_mean": 0.8001151448841631,
+      "f1_std": 0.039832204457600114
+    },
+    "biodist": {
+      "kl_divergence_mean": 0.41026559421902165,
+      "kl_divergence_std": 0.08303708399525898,
+      "js_divergence_mean": 0.10230719374124933,
+      "js_divergence_std": 0.023848212584556397
+    }
+  }
+}

scripts/check_prompt_len.py

@@ -0,0 +1,38 @@
+import inspect
+
+import pandas as pd
+from transformers import AutoTokenizer
+from lnp_ml.dataset import process_dataframe, SMILES_COL
+from lnp_ml.modeling.layers.llm_prompt import LLMPromptEncoder
+
+# 从源头读，避免脚本和模型的阈值各说各话
+MAXLEN = inspect.signature(LLMPromptEncoder.__init__).parameters["max_length"].default
+
+tok = AutoTokenizer.from_pretrained("models/qwen2.5-7b-instruct", trust_remote_code=True)
+df = process_dataframe(pd.read_csv("data/interim/internal.csv"))
+smis = sorted(set(df[SMILES_COL].dropna()), key=len)
+
+enc = LLMPromptEncoder.__new__(LLMPromptEncoder)   # 不加载 7B 权重，只借用格式化方法
+enc._is_biot5 = False
+# _build_rag_prompt 现在还要读分位边界。这里必须给非空值：留空会让 _qbin 返回空串，
+# 测出来的 token 数比真实 prompt 少约 8/邻居，等于白测。
+# 具体数值不影响长度（Q1/5 和 Q5/5 token 数相同），只决定落在哪个桶。
+enc._rag_pool_qedges = {
+    "delivery": [-0.80, -0.30, 0.20, 0.70],
+    "size": [-0.90, -0.20, 0.40, 1.10],
+}
+
+def nb_of(s):
+    return {"smiles": s, "sim": 0.812, "delivery": 1.234,
+            "extra": {"size": -0.456, "pdi": 1, "ee": 2, "toxic": 0,
+                      "biodist": [0.12, 0.34, 0.21, 0.05, 0.18, 0.07, 0.03]}}
+
+worst = smis[-1]
+print(f"分子数={len(smis)}  SMILES 长度 {len(smis[0])}~{len(worst)} 字符")
+print(f"分子数={len(smis)}  SMILES 长度 {len(smis[0])}~{len(worst)} 字符  max_length={MAXLEN}")
+for k in (2, 4, 8):
+    # 最坏情况：目标分子和全部邻居都取最长的那条
+    p = LLMPromptEncoder._build_rag_prompt(enc, worst, [nb_of(worst)] * k)
+    n = len(tok(p)["input_ids"])
+    status = "OK" if n <= MAXLEN else f"超出 {n - MAXLEN} tokens，会被截断"
+    print(f"最坏 rag_top_k={k}: {n:5d} tokens  余量 {MAXLEN - n:5d}  {status}")

scripts/fixed_hparams.json

@@ -0,0 +1,8 @@
+{
+    "dropout": 0.33, "lr": 0.0005, "weight_decay": 0.0001, "backbone_lr_ratio": 0.3,
+    "moe_n_experts": 4, "moe_top_k": 2, "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"
+  }