Buckets:
| """Re-split an exported dataset by domain: held-out domains -> test_ood, the rest -> train/validation. | |
| python -m training.split_by_domain --source datasets/policy_30k --out datasets/policy_30k_domains \ | |
| --test-domains insurance_claims travel_booking --val-cases-per-domain 1 --seed 7 | |
| Validation takes whole cases (originals and their edited variants share a case_id) from every | |
| training domain, so no case or conversation is split across train and validation. | |
| """ | |
| import argparse | |
| import collections | |
| import json | |
| import random | |
| from pathlib import Path | |
| from training.common import read_rows | |
| def write_rows(path, rows): | |
| with path.open("w", encoding="utf-8") as handle: | |
| for row in rows: | |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--source", required=True) | |
| parser.add_argument("--out", required=True) | |
| parser.add_argument("--test-domains", nargs="+", required=True) | |
| parser.add_argument("--val-cases-per-domain", type=int, default=1) | |
| parser.add_argument("--seed", type=int, default=7) | |
| args = parser.parse_args() | |
| source, out = Path(args.source), Path(args.out) | |
| rows = [r for name in ("train", "validation", "test_ood") if (source / f"{name}.jsonl").exists() | |
| for r in read_rows(source / f"{name}.jsonl")] | |
| domains = sorted({r["domain"] for r in rows}) | |
| missing = set(args.test_domains) - set(domains) | |
| if missing: | |
| raise SystemExit(f"Unknown test domains: {sorted(missing)}; available: {domains}") | |
| rng = random.Random(args.seed) | |
| case_sizes = collections.Counter(r["case_id"] for r in rows) | |
| cases_by_domain = collections.defaultdict(set) | |
| for r in rows: | |
| if r["domain"] not in args.test_domains: | |
| cases_by_domain[r["domain"]].add(r["case_id"]) | |
| val_cases = set() | |
| for domain in sorted(cases_by_domain): | |
| # Only typical-sized cases (0.5x-1.5x the domain median), so no domain gets a near-empty validation slice. | |
| cases = sorted(cases_by_domain[domain]) | |
| median = sorted(case_sizes[c] for c in cases)[len(cases) // 2] | |
| typical = [c for c in cases if 0.5 * median <= case_sizes[c] <= 1.5 * median] | |
| val_cases.update(rng.sample(typical, args.val_cases_per_domain)) | |
| splits = {"train": [], "validation": [], "test_ood": []} | |
| for r in rows: | |
| if r["domain"] in args.test_domains: | |
| splits["test_ood"].append({**r, "held_out": True}) | |
| else: | |
| splits["validation" if r["case_id"] in val_cases else "train"].append({**r, "held_out": False}) | |
| out.mkdir(parents=True, exist_ok=True) | |
| summary = {} | |
| for name, split_rows in splits.items(): | |
| split_rows.sort(key=lambda r: r["pair_id"]) | |
| rng.shuffle(split_rows) | |
| write_rows(out / f"{name}.jsonl", split_rows) | |
| summary[name] = { | |
| "pairs": len(split_rows), | |
| "labels": dict(collections.Counter(r["label"] for r in split_rows)), | |
| "domains": dict(collections.Counter(r["domain"] for r in split_rows)), | |
| } | |
| manifest = { | |
| "selection_policy": "Domain split of the source export: held-out domains form test_ood; " | |
| "validation is whole cases sampled per training domain.", | |
| "source": str(source), | |
| "test_domains": sorted(args.test_domains), | |
| "train_domains": sorted(cases_by_domain), | |
| "val_cases_per_domain": args.val_cases_per_domain, | |
| "validation_cases": sorted(val_cases), | |
| "seed": args.seed, | |
| "splits": summary, | |
| } | |
| (out / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8") | |
| print(json.dumps({k: {"pairs": v["pairs"], "labels": v["labels"]} for k, v in summary.items()}, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 3.86 kB
- Xet hash:
- 237e6ee41efddfe36e673aedf1c22f64016822bf05c031a682426c2958d37cc4
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.