| from __future__ import annotations |
|
|
| import argparse |
| import shutil |
| import tempfile |
| from pathlib import Path |
|
|
| from timeout_risk_model import load_dataset_records, save_model, train_timeout_model |
|
|
| DEFAULT_DATASET_REPO_ID = "stack-all/FJNU-HuaxiangChaPandaTimeoutRecord" |
|
|
|
|
| def default_dataset_dir() -> Path: |
| return Path(__file__).resolve().parent.parent / "FJNU-HuaxiangChaPandaTimeoutRecord" |
|
|
|
|
| def download_dataset_snapshot(repo_id: str) -> Path: |
| try: |
| from huggingface_hub import snapshot_download |
| except ImportError as exc: |
| raise RuntimeError("Install huggingface_hub or use --dataset-dir with a local dataset checkout.") from exc |
|
|
| target_dir = Path(tempfile.mkdtemp(prefix="cha_panda_dataset_")) |
| snapshot_download( |
| repo_id=repo_id, |
| repo_type="dataset", |
| local_dir=str(target_dir), |
| allow_patterns=["records.csv", "timeout_records.csv", "tmp/*.csv"], |
| ) |
| return target_dir |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Train the Huaxiang ChaPanda timeout risk model.") |
| parser.add_argument( |
| "--dataset-dir", |
| type=Path, |
| default=default_dataset_dir(), |
| help="Local checkout of the Dataset repository.", |
| ) |
| parser.add_argument( |
| "--dataset-repo-id", |
| default=DEFAULT_DATASET_REPO_ID, |
| help="Hugging Face Dataset repo id used with --download-from-hub.", |
| ) |
| parser.add_argument( |
| "--download-from-hub", |
| action="store_true", |
| help="Download timeout_records.csv and tmp/*.csv from the Dataset repository before training.", |
| ) |
| parser.add_argument( |
| "--output", |
| type=Path, |
| default=Path(__file__).resolve().parent / "timeout_model.json", |
| help="Output JSON model path.", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| downloaded_dir: Path | None = None |
| dataset_dir = args.dataset_dir |
|
|
| if args.download_from_hub: |
| downloaded_dir = download_dataset_snapshot(args.dataset_repo_id) |
| dataset_dir = downloaded_dir |
|
|
| try: |
| records = load_dataset_records(dataset_dir) |
| model = train_timeout_model(records) |
| save_model(model, args.output) |
| finally: |
| if downloaded_dir: |
| shutil.rmtree(downloaded_dir, ignore_errors=True) |
|
|
| metadata = model["metadata"] |
| print(f"Saved model: {args.output}") |
| print(f"Mode: {model['mode']}") |
| print(f"Training records: {metadata['training_records']}") |
| print(f"Timeout / non-timeout: {metadata['timeout_records']} / {metadata['non_timeout_records']}") |
| print(f"Calibration: {metadata['calibration']}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
