| """Fixed random subset of an items file. | |
| python finetune/make_subset.py <items.pt> <out.pt> [n=100000] [seed=0] [exclude_src=a,b]""" | |
| import random, sys, torch | |
| items = torch.load(sys.argv[1], weights_only=False) | |
| excl = set(sys.argv[5].split(",")) if len(sys.argv) > 5 and sys.argv[5] else set() | |
| pool = [i for i in range(len(items)) if items[i].get("src") not in excl] | |
| n = min(int(sys.argv[3]) if len(sys.argv) > 3 else 100000, len(pool)) | |
| idx = sorted(random.Random(int(sys.argv[4]) if len(sys.argv) > 4 else 0).sample(pool, n)) | |
| sub = [items[i] for i in idx] | |
| torch.save(sub, sys.argv[2]) | |
| import collections | |
| print(len(items), "->", len(sub), "src", dict(collections.Counter(i.get("src") for i in sub).most_common(8)), | |
| "goal_done", dict(collections.Counter(i["label"] for i in sub if i["qid"] == "goal_done"))) | |