import argparse
import json
import random

random.seed(42)

parser = argparse.ArgumentParser()
parser.add_argument(
    "--data_path",
    type=str,
    required=True,
)
parser.add_argument(
    "--output_path",
    type=str,
    required=True,
)
parser.add_argument(
    "--task_num",
    type=int,
    default=100,
)
parser.add_argument(
    "--passage_num",
    type=int,
    default=0,
)
parser.add_argument(
    "--collected_data",
    type=str,
    default=None,
)
args = parser.parse_args()


def filter_dataset(dataset):
    if not args.collected_data:
        return dataset
    known_passages = json.load(open(args.collected_data))
    known_passage_ids = []
    if isinstance(known_passages, list):
        for kp in known_passages:
            known_passage_ids.append(kp["passage_id"])
    elif isinstance(known_passages, dict):
        known_passage_ids = list(known_passages.keys())
    dataset = [d for d in dataset if not d["passage_id"] in known_passage_ids]
    return dataset


def read_dataset(data_path):
    new_dataset = []
    if data_path.endswith('jsonl'):
        raw_dataset = [json.loads(v) for v in open(data_path)]
    else:
        raw_dataset = json.load(open(data_path))
    if isinstance(raw_dataset, dict):
        for k, v in raw_dataset.items():
            v["passage_id"] = k
            new_dataset.append(v)
    elif isinstance(raw_dataset, list):
        for data in raw_dataset:
            if not "passage_id" in data and "id" in data:
                data["passage_id"] = data["id"]
            new_dataset.append(data)
    else:
        raise ValueError("Unrecognizable dataset data?")
    new_dataset = filter_dataset(new_dataset)
    random.shuffle(new_dataset)
    return new_dataset


new_dataset = []
passage_num = 0
if args.passage_num > 0:
    passage_num = args.passage_num
    dataset = read_dataset(args.data_path)
    if not len(dataset) >= args.task_num * args.passage_num:
        raise ValueError(
            """
            The number of passage is less than args.task_num * args.passage_num.
            Use --repeat to use
            """
        )
    for i in range(args.task_num):
        spos = i * args.passage_num
        new_dataset.append(dataset[spos : spos + args.passage_num])

elif args.passage_num == 0 and "," in args.data_path:
    data_paths = args.data_path.split(",")
    datasets = [read_dataset(data_path) for data_path in data_paths]
    for dataset in datasets:
        assert len(dataset) >= args.task_num
    passage_num = len(datasets)
    for i in range(args.task_num):
        data_item = []
        for d in datasets:
            data_item.append(d[i])
        new_dataset.append(data_item)

with open(args.output_path, "w") as f:
    json.dump(new_dataset, f)
print(f"wrote {len(new_dataset)} tasks with {passage_num} passages each")

import argparse
import json
import random

random.seed(42)

parser = argparse.ArgumentParser()
parser.add_argument(
    "--data_path",
    type=str,
    required=True,
)
parser.add_argument(
    "--output_path",
    type=str,
    required=True,
)
parser.add_argument(
    "--task_num",
    type=int,
    default=100,
)
parser.add_argument(
    "--passage_num",
    type=int,
    default=0,
)
parser.add_argument(
    "--collected_data",
    type=str,
    default=None,
)
args = parser.parse_args()


def filter_dataset(dataset):
    if not args.collected_data:
        return dataset
    known_passages = json.load(open(args.collected_data))
    known_passage_ids = []
    if isinstance(known_passages, list):
        for kp in known_passages:
            known_passage_ids.append(kp["passage_id"])
    elif isinstance(known_passages, dict):
        known_passage_ids = list(known_passages.keys())
    dataset = [d for d in dataset if not d["passage_id"] in known_passage_ids]
    return dataset


def read_dataset(data_path):
    new_dataset = []
    if data_path.endswith('jsonl'):
        raw_dataset = [json.loads(v) for v in open(data_path)]
    else:
        raw_dataset = json.load(open(data_path))
    if isinstance(raw_dataset, dict):
        for k, v in raw_dataset.items():
            v["passage_id"] = k
            new_dataset.append(v)
    elif isinstance(raw_dataset, list):
        for data in raw_dataset:
            if not "passage_id" in data and "id" in data:
                data["passage_id"] = data["id"]
            new_dataset.append(data)
    else:
        raise ValueError("Unrecognizable dataset data?")
    new_dataset = filter_dataset(new_dataset)
    random.shuffle(new_dataset)
    return new_dataset


new_dataset = []
passage_num = 0
if args.passage_num > 0:
    passage_num = args.passage_num
    dataset = read_dataset(args.data_path)
    if not len(dataset) >= args.task_num * args.passage_num:
        raise ValueError(
            """
            The number of passage is less than args.task_num * args.passage_num.
            Use --repeat to use
            """
        )
    for i in range(args.task_num):
        spos = i * args.passage_num
        new_dataset.append(dataset[spos : spos + args.passage_num])

elif args.passage_num == 0 and "," in args.data_path:
    data_paths = args.data_path.split(",")
    datasets = [read_dataset(data_path) for data_path in data_paths]
    for dataset in datasets:
        assert len(dataset) >= args.task_num
    passage_num = len(datasets)
    for i in range(args.task_num):
        data_item = []
        for d in datasets:
            data_item.append(d[i])
        new_dataset.append(data_item)

with open(args.output_path, "w") as f:
    json.dump(new_dataset, f)
print(f"wrote {len(new_dataset)} tasks with {passage_num} passages each")

The number of passage is less than args.task_num * args.passage_num.             Use --repeat to use