"""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")))