def main()

in src/scripts/gen_completions.py [0:0]


def main():
    parser = argparse.ArgumentParser(description="Generate completions for given messages.")

    parser.add_argument("-d", "--dataset", required=True, help="Dataset name", type=Dataset, choices=list(Dataset))
    parser.add_argument("-e", "--engine", required=True, help="OpenAI model compatible with ChatCompletion API",
                        type=Model, choices=list(Model))
    parser.add_argument("-s", "--size", type=int, help="Sample size", default=3)
    parser.add_argument("-i", "--indices_path", type=str, help="Path to file with indices", default=None)
    parser.add_argument("-en", "--entity_name", help="Entity name", type=Entity,
                        choices=[i for i in list(Entity) if i != Entity.OBJECT_NAME], default=None)

    args = parser.parse_args()

    get_completions(args.dataset, args.engine, args.size, args.indices_path, args.entity_name)