def main()

in scripts/write_out.py [0:0]

• 47 lines of code
• 16 McCabe index (conditional complexity)

def main():
    args = parse_args()
    np.random.seed(args.seed)

    if args.include_path is not None:
        eval_logger.info(f"Including path: {args.include_path}")

    task_manager = TaskManager(args.verbosity, include_path=args.include_path)

    if args.tasks == "all_tasks":
        task_names = task_manager.all_tasks
    else:
        task_names = args.tasks.split(",")
    task_dict = tasks.get_task_dict(task_names, task_manager)

    os.makedirs(args.output_base_path, exist_ok=True)
    for task_name, task in task_dict.items():
        if isinstance(task, tuple):
            _, task = task
        rnd = random.Random()
        rnd.seed(args.seed)

        iters = []

        for set in args.sets.split(","):
            docs = None
            if set == "train" and task.has_training_docs():
                docs = task.training_docs()
            if set == "val" and task.has_validation_docs():
                docs = task.validation_docs()
            if set == "test" and task.has_test_docs():
                docs = task.test_docs()
            if docs is not None:
                iters.append(docs)

        if len(iters) == 0:
            raise ValueError(
                f"Passed --sets '{args.sets}' but this task has no splits which match. Please specify a different --sets value."
            )

        docs = join_iters(iters)

        with open(
            os.path.join(args.output_base_path, task_name), "w", encoding="utf8"
        ) as f:
            for i, doc in (
                zip(range(args.num_examples), docs)
                if args.num_examples > 0
                else enumerate(docs)
            ):
                f.write(EXAMPLE_DIVIDER.format(i=i))
                ctx = task.fewshot_context(
                    doc=doc,
                    num_fewshot=args.num_fewshot,
                )
                f.write(ctx + "\n")