Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """CLI for preparing, training, and querying the zero-shot recommender.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import sys | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parent.parent | |
| if str(ROOT) not in sys.path: | |
| sys.path.insert(0, str(ROOT)) | |
| from scripts.hp_search_head_sweep import _load_train_fixed_test_dataset | |
| from src.zero_shot_recommender.data import best_record_groups, load_trials | |
| from src.zero_shot_recommender.meta_features import extract_meta_features | |
| from src.zero_shot_recommender.training import ( | |
| load_recommender, | |
| recommend, | |
| save_recommender, | |
| train_recommender, | |
| ) | |
| DEFAULT_DATASETS = ("massbench_adenocarcinoma", "massbench_alzheimer", "massbench_benchmark") | |
| def extract_all(dataset_names) -> dict: | |
| output = {} | |
| for dataset in dataset_names: | |
| X, y, batches, *_ = _load_train_fixed_test_dataset(dataset) | |
| output[dataset] = extract_meta_features(X, y, batches) | |
| return output | |
| def parse_args(argv=None): | |
| parser = argparse.ArgumentParser() | |
| sub = parser.add_subparsers(dest="command", required=True) | |
| meta = sub.add_parser("extract-meta") | |
| meta.add_argument("--datasets", nargs="+", default=list(DEFAULT_DATASETS)) | |
| meta.add_argument("--output", type=Path, default=ROOT / "results" / "zero_shot_recommender" / "meta_features.json") | |
| train = sub.add_parser("train") | |
| train.add_argument("--trials", nargs="+", type=Path, default=[ROOT / "results" / "zero_shot_hpo" / "unified_trials.jsonl"]) | |
| train.add_argument("--meta", type=Path, default=ROOT / "results" / "zero_shot_recommender" / "meta_features.json") | |
| train.add_argument("--output-dir", type=Path, default=ROOT / "results" / "zero_shot_recommender" / "model") | |
| train.add_argument("--hidden-size", type=int, default=64) | |
| train.add_argument("--epochs", type=int, default=500) | |
| query = sub.add_parser("recommend") | |
| query.add_argument("--dataset", required=True) | |
| query.add_argument("--model-dir", type=Path, default=ROOT / "results" / "zero_shot_recommender" / "model") | |
| query.add_argument("--target", default="all", help="Scoped target name, 'global', or 'all'") | |
| return parser.parse_args(argv) | |
| def main(argv=None): | |
| args = parse_args(argv) | |
| if args.command == "extract-meta": | |
| values = extract_all(args.datasets) | |
| args.output.parent.mkdir(parents=True, exist_ok=True) | |
| args.output.write_text(json.dumps(values, indent=2)) | |
| print(f"Wrote {len(values)} dataset descriptors to {args.output}") | |
| return | |
| if args.command == "train": | |
| groups = best_record_groups(load_trials(args.trials)) | |
| meta = json.loads(args.meta.read_text()) | |
| manifest = {"targets": {}} | |
| for target, records in groups.items(): | |
| model, metadata = train_recommender(records, meta, hidden_size=args.hidden_size, epochs=args.epochs) | |
| save_recommender(model, metadata, args.output_dir / target) | |
| manifest["targets"][target] = { | |
| "training_rows": len(records), | |
| "model_family": records[0].model_family if target != "global" else None, | |
| "head_type": records[0].config.get("head_type") if target != "global" else None, | |
| } | |
| args.output_dir.mkdir(parents=True, exist_ok=True) | |
| (args.output_dir / "manifest.json").write_text(json.dumps(manifest, indent=2)) | |
| print(f"Trained {len(groups)} scoped recommenders in {args.output_dir}") | |
| return | |
| meta = extract_all([args.dataset])[args.dataset] | |
| manifest_path = args.model_dir / "manifest.json" | |
| manifest = json.loads(manifest_path.read_text()) | |
| targets = list(manifest["targets"]) if args.target == "all" else [args.target] | |
| unknown = sorted(set(targets) - set(manifest["targets"])) | |
| if unknown: | |
| raise ValueError(f"Unknown target(s) {unknown}; available={list(manifest['targets'])}") | |
| output = {} | |
| for target in targets: | |
| model, metadata = load_recommender(args.model_dir / target) | |
| output[target] = recommend(model, metadata, meta) | |
| print(json.dumps(output, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |