fix: 修正 size 预测的反归一化并将 MoE+LLM 模型接入应用

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

@@ -65,7 +65,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,7 +87,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=0, ge=0, le=1000, description="两阶段推理：粗筛后用完整模型重排的候选数，0 表示关闭")
+    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="组分范围配置（默认使用标准范围）")
@@ -127,7 +127,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: 有毒)
 
@@ -146,6 +146,9 @@ class HealthResponse(BaseModel):
     model_loaded: bool
     device: str
     available_organs: List[str]
+    use_moe: bool = False
+    use_llm: bool = False
+    use_rag: bool = False
 
 
 # ============ Global State ============
@@ -239,11 +242,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)),
     )
 
 
@@ -368,6 +376,8 @@ async def optimize_formulation(request: OptimizeRequest):
         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,
@@ -270,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()
         
@@ -352,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(
                 "自定义迭代步长",
@@ -512,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",
@@ -528,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)
@@ -550,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:
@@ -602,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,

app/optimize.py

@@ -146,6 +146,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 +171,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 +309,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 +607,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 +702,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,
@@ -817,7 +831,7 @@ def optimize(
     routes: Optional[List[str]] = None,
     scoring_weights: Optional[ScoringWeights] = None,
     batch_size: int = 256,
-    rerank_top_n: int = 0,
+    rerank_top_n: int = 200,
     rerank_batch_size: int = 16,
 ) -> List[Formulation]:
     """

app/size_zscore_stats.json

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

lnp_ml/modeling/models.py

@@ -459,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()

scripts/preprocess_internal.py

@@ -68,6 +68,12 @@ def main(
     df["size"] = pd.to_numeric(df["size"], errors="coerce")
     df["size"] = np.log(df["size"].replace(0, np.nan))  # 避免 log(0)
 
+    size_stats = {"mean": float(df["size"].mean()), "std": float(df["size"].std())}
+    size_stats_path = APP_DIR / "size_zscore_stats.json"
+    with open(size_stats_path, "w") as f:
+        json.dump(size_stats, f, indent=2)
+    logger.success(f"Saved size stats to {size_stats_path}: {size_stats}")
+
     # 保存
     output_path.parent.mkdir(parents=True, exist_ok=True)
     df.to_csv(output_path, index=False)