Compressor/models/vibration/train_model.py

#!/usr/bin/env python3
"""Train the vibration XGBoost assistant model without reading the test set."""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from models.common import modeling, paths
from models.vibration import features_vibration


def main() -> None:
    parser = argparse.ArgumentParser(description="训练振动模型，不读取测试集")
    parser.add_argument("--datasets-dir", default=str(paths.VIBRATION_DATAS))
    parser.add_argument("--models-dir", default=str(paths.VIBRATION_MODELS))
    parser.add_argument("--normal-ratio", type=int, default=None)
    parser.add_argument("--all-normal", action="store_true")
    parser.add_argument("--seed", type=int, default=20260921)
    parser.add_argument("--rounds", type=int, default=800)
    parser.add_argument("--early-stopping-rounds", type=int, default=50)
    args = parser.parse_args()
    if args.normal_ratio is not None and args.normal_ratio < 1:
        parser.error("--normal-ratio 必须大于等于 1")
    modeling.train(
        "vibration", features_vibration, Path(args.datasets_dir), Path(args.models_dir),
        args.seed, args.rounds, args.early_stopping_rounds, args.all_normal, args.normal_ratio,
    )


if __name__ == "__main__":
    main()