()
| 34 | |
| 35 | |
| 36 | def main() -> int: |
| 37 | parser = argparse.ArgumentParser(description=__doc__) |
| 38 | parser.add_argument("dataset_local", help="Path to the local SWE-bench dataset dir.") |
| 39 | parser.add_argument("preds_path", type=Path, help="Predictions JSONL to skip-already-done.") |
| 40 | parser.add_argument("n", type=int, help="How many to pick.") |
| 41 | parser.add_argument( |
| 42 | "--random", |
| 43 | action="store_true", |
| 44 | help="Uniform random sample from unseen pool (instead of stratified round-robin).", |
| 45 | ) |
| 46 | parser.add_argument( |
| 47 | "--seed", |
| 48 | type=int, |
| 49 | default=None, |
| 50 | help="Random seed; only meaningful with --random. Omit for non-reproducible sampling.", |
| 51 | ) |
| 52 | args = parser.parse_args() |
| 53 | |
| 54 | from datasets import load_from_disk |
| 55 | |
| 56 | ds = load_from_disk(args.dataset_local)["test"] |
| 57 | all_rows = [(r["instance_id"], r["repo"]) for r in ds] |
| 58 | |
| 59 | seen: set[str] = set() |
| 60 | if args.preds_path.exists(): |
| 61 | with args.preds_path.open(encoding="utf-8") as f: |
| 62 | for line in f: |
| 63 | line = line.strip() |
| 64 | if not line: |
| 65 | continue |
| 66 | try: |
| 67 | rec = json.loads(line) |
| 68 | except json.JSONDecodeError: |
| 69 | continue |
| 70 | iid = rec.get("instance_id") |
| 71 | if iid: |
| 72 | seen.add(iid) |
| 73 | |
| 74 | unseen = [(iid, repo) for iid, repo in all_rows if iid not in seen] |
| 75 | if not unseen: |
| 76 | print( |
| 77 | f"All {len(all_rows)} instances already predicted in {args.preds_path}", |
| 78 | file=sys.stderr, |
| 79 | ) |
| 80 | return 1 |
| 81 | |
| 82 | if args.random: |
| 83 | rng = random.Random(args.seed) |
| 84 | sample = rng.sample(unseen, k=min(args.n, len(unseen))) |
| 85 | chosen = [iid for iid, _ in sample] |
| 86 | mode = f"random{f' (seed={args.seed})' if args.seed is not None else ''}" |
| 87 | else: |
| 88 | # Stratified round-robin: pick from each repo in turn so the batch |
| 89 | # spans the dataset breadth rather than clustering on one repo. |
| 90 | by_repo: dict[str, list[str]] = defaultdict(list) |
| 91 | for iid, repo in unseen: |
| 92 | by_repo[repo].append(iid) |
| 93 | chosen = [] |
no test coverage detected