def main()

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


def main():
    parser = argparse.ArgumentParser(description="Generate embeddings 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 for calculating embeddings",
                        type=ModelType, choices=list(ModelType))
    parser.add_argument("-en", "--entity_name", help="Calculate embeddings for the entity_name",
                        type=Entity, choices=list(Entity))
    parser.add_argument("-i", "--indices_path", type=str, help="Path to file with indices", default=None)

    args = parser.parse_args()

    get_embeddings(args.dataset, args.engine, args.entity_name, args.indices_path)