(subset: str)
| 37 | |
| 38 | |
| 39 | def _process_subset(subset: str) -> None: |
| 40 | print(f"Downloading {REPO_ID} config={subset} ...") |
| 41 | ds = load_dataset(REPO_ID, name=subset, split="train") |
| 42 | |
| 43 | for row in ds: |
| 44 | name = row["name"] |
| 45 | problem_dir = OUTPUT_DIR / subset / name |
| 46 | problem_dir.mkdir(parents=True, exist_ok=True) |
| 47 | |
| 48 | # definition.json |
| 49 | definition = _build_definition(row) |
| 50 | (problem_dir / "definition.json").write_text( |
| 51 | json.dumps(definition, indent=4) + "\n" |
| 52 | ) |
| 53 | |
| 54 | # reference.py |
| 55 | (problem_dir / "reference.py").write_text(row["reference"]) |
| 56 | |
| 57 | # workload.jsonl |
| 58 | workloads = json.loads(row["workloads"]) |
| 59 | with open(problem_dir / "workload.jsonl", "w") as f: |
| 60 | for workload in workloads: |
| 61 | f.write(json.dumps(workload) + "\n") |
| 62 | |
| 63 | print(f" -> {len(ds)} problems written to {OUTPUT_DIR / subset}") |
| 64 | |
| 65 | |
| 66 | def main() -> None: |
no test coverage detected