File size: 2,722 Bytes
7affe31 9f14291 7affe31 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | 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()
|