Bladeren bron

优化:model_bundle.py stratify 关闭时增加 UserWarning 提示

leod 4 dagen geleden
bovenliggende
commit
014a58e7cd
1 gewijzigde bestanden met toevoegingen van 8 en 1 verwijderingen
  1. 8 1
      defect_analysis/ml/model_bundle.py

+ 8 - 1
defect_analysis/ml/model_bundle.py

@@ -1,5 +1,6 @@
 """可持久化的 ML 模型包。"""
 
+import warnings
 from datetime import datetime, timezone
 
 import joblib
@@ -49,7 +50,13 @@ def create_model_bundle(
     if y.nunique() < 2:
         raise ValueError("目标标签只有一个类别,无法训练监督模型")
 
-    stratify = y if y.value_counts().min() >= 2 else None
+    min_count = int(y.value_counts().min())
+    if min_count < 2:
+        warnings.warn(
+            f"最小类别仅 {min_count} 个样本,已关闭分层抽样。验证集可能不包含少数类别。",
+            UserWarning,
+        )
+    stratify = y if min_count >= 2 else None
     X_train, X_valid, y_train, y_valid = train_test_split(
         X,
         y,