defect-analysis/train_ml_models.py

"""训练和验证结构化 ML 模型。"""

import argparse

import pandas as pd

from defect_analysis.ml.datasets import build_supervised_dataset
from defect_analysis.ml.features import build_feature_frame
from defect_analysis.ml.model_registry import detect_optional_model_backends
from defect_analysis.ml.predict import predict_key_factors
from defect_analysis.ml.tabular_models import train_tabular_model
from defect_analysis.schemas import normalize_defect_schema


def load_defect_csv(csv_path):
    return normalize_defect_schema(pd.read_csv(csv_path, parse_dates=["timestamp"], encoding="utf-8-sig"))


def main():
    parser = argparse.ArgumentParser(description="训练/运行不良分析 ML 模型")
    parser.add_argument("--csv", default="defect_data.csv")
    parser.add_argument(
        "--model",
        default="random_forest",
        choices=["random_forest", "logistic_regression", "isolation_forest", "xgboost", "lightgbm"],
    )
    parser.add_argument("--target-defect-type")
    parser.add_argument("--target-severity")
    parser.add_argument("--top-n", type=int, default=10)
    parser.add_argument("--show-backends", action="store_true")
    args = parser.parse_args()

    if args.show_backends:
        print(detect_optional_model_backends())

    df = load_defect_csv(args.csv)
    if args.model == "isolation_forest":
        X = build_feature_frame(df)
        result = train_tabular_model("isolation_forest", X)
        scores = pd.Series(result["anomaly_scores"])
        print(f"IsolationForest 完成: 样本数={len(scores)}, 最高异常分={scores.max():.4f}, 平均异常分={scores.mean():.4f}")
        return

    X, y = build_supervised_dataset(
        df,
        target_defect_type=args.target_defect_type,
        target_severity=args.target_severity,
    )
    result = train_tabular_model(args.model, X, y)
    print(f"{args.model} 训练完成: {result['metrics']}")

    predictions = predict_key_factors(
        df,
        target_defect_type=args.target_defect_type,
        target_severity=args.target_severity,
        model_name=args.model,
        top_n=args.top_n,
    )
    if predictions.empty:
        print("未找到关键因子候选。")
    else:
        columns = ["维度", "因子值", "目标数", "异常倍数", "关键因子得分", "ml_probability", "model_name"]
        print(predictions[columns].to_string(index=False))


if __name__ == "__main__":
    main()