import yaml
from tqdm import tqdm


def main() -> None:
    subset = ["extended", "diamond", "main"]

    for task in tqdm(subset):
        file_name = f"gpqa_{task}_n_shot.yaml"
        try:
            with open(f"{file_name}", "w") as f:
                f.write("# Generated by _generate_configs.py\n")
                yaml.dump(
                    {
                        "include": "_gpqa_n_shot_yaml",
                        "task": f"gpqa_{task}_n_shot",
                        "dataset_name": f"gpqa_{task}",
                    },
                    f,
                )
        except FileExistsError:
            pass


if __name__ == "__main__":
    main()
