Skills › Scientific Agent Skills › pytdc › discover_metadata.py
discover_metadata.py
174 lines · skills/pytdc/scripts/discover_metadata.py at 5f2bc03
1#!/usr/bin/env python3
2"""Discover PyTDC registries without constructing loaders or downloading data."""
4from __future__ import annotations
6import argparse
7import sys
8from typing import Any, Iterable
10from _common import (
11 CliError,
12 bounded_int,
13 canonical_name,
14 emit_json,
15 load_pytdc_metadata,
16)
19def _window(values: Iterable[Any], offset: int, limit: int) -> dict[str, Any]:
20 items = list(values)
21 selected = items[offset : offset + limit]
22 return {
23 "items": selected,
24 "limit": limit,
25 "offset": offset,
26 "returned": len(selected),
27 "total": len(items),
28 "truncated": offset + len(selected) < len(items),
29 }
32def collect_metadata(
33 metadata: Any,
34 package_version: str,
35 *,
36 kind: str,
37 task: str | None,
38 offset: int,
39 limit: int,
40) -> dict[str, Any]:
41 """Collect a bounded view of the installed metadata registries."""
43 dataset_registry = dict(metadata.dataset_names)
44 if task is not None:
45 task = canonical_name(task, dataset_registry, "task")
47 sections: dict[str, Any] = {}
48 requested = (
49 ["tasks", "datasets", "benchmarks", "evaluators", "oracles"]
50 if kind == "all"
51 else [kind]
52 )
54 if "tasks" in requested:
55 sections["tasks"] = _window(
56 (
57 {"task": name, "dataset_count": len(names)}
58 for name, names in dataset_registry.items()
59 ),
60 offset,
61 limit,
62 )
64 if "datasets" in requested:
65 if task is not None:
66 values = (
67 {"task": task, "dataset": name}
68 for name in dataset_registry[task]
69 )
70 else:
71 values = (
72 {"task": task_name, "dataset": dataset_name}
73 for task_name, names in dataset_registry.items()
74 for dataset_name in names
75 )
76 sections["datasets"] = _window(values, offset, limit)
78 if "benchmarks" in requested:
79 values = (
80 {
81 "group": group_name,
82 "task": task_name,
83 "benchmark": benchmark_name,
84 }
85 for group_name, task_map in metadata.benchmark_names.items()
86 for task_name, names in task_map.items()
87 for benchmark_name in names
88 )
89 sections["benchmarks"] = _window(values, offset, limit)
91 if "evaluators" in requested:
92 sections["evaluators"] = _window(
93 ({"name": name} for name in dict.fromkeys(metadata.evaluator_name)),
94 offset,
95 limit,
96 )
98 if "oracles" in requested:
99 sections["oracles"] = _window(
100 ({"name": name} for name in dict.fromkeys(metadata.oracle_names)),
101 offset,
102 limit,
103 )
105 return {
106 "download_performed": False,
107 "kind": kind,
108 "package": "PyTDC",
109 "package_version": package_version,
110 "source": "installed tdc.metadata registry",
111 "task_filter": task,
112 **sections,
113 }
116def build_parser() -> argparse.ArgumentParser:
117 parser = argparse.ArgumentParser(
118 description=(
119 "Read task, dataset, benchmark, evaluator, and oracle names from the "
120 "installed PyTDC metadata module. This command never instantiates a "
121 "loader or downloads a dataset/model."
122 )
123 )
124 parser.add_argument(
125 "--kind",
126 choices=("tasks", "datasets", "benchmarks", "evaluators", "oracles", "all"),
127 default="tasks",
128 help="registry to show (default: tasks)",
129 )
130 parser.add_argument("--task", help="exact task filter for --kind datasets")
131 parser.add_argument(
132 "--offset",
133 type=bounded_int(0, 1_000_000),
134 default=0,
135 help="skip this many registry entries (default: 0)",
136 )
137 parser.add_argument(
138 "--limit",
139 type=bounded_int(1, 500),
140 default=50,
141 help="maximum entries returned per section (default: 50; max: 500)",
142 )
143 parser.add_argument("--output", help="write JSON to a relative workspace path")
144 parser.add_argument(
145 "--force", action="store_true", help="replace an existing --output file"
146 )
147 return parser
150def main(argv: list[str] | None = None) -> int:
151 parser = build_parser()
152 args = parser.parse_args(argv)
153 if args.task and args.kind not in {"datasets", "all"}:
154 parser.error("--task is only valid with --kind datasets or --kind all")
156 try:
157 metadata, package_version = load_pytdc_metadata()
158 result = collect_metadata(
159 metadata,
160 package_version,
161 kind=args.kind,
162 task=args.task,
163 offset=args.offset,
164 limit=args.limit,
165 )
166 emit_json(result, args.output, force=args.force)
167 except (CliError, OSError, TypeError, ValueError) as exc:
168 print(f"error: {exc}", file=sys.stderr)
169 return 2
170 return 0
173if __name__ == "__main__":
174 raise SystemExit(main())