Function main

tasks/AutoTTS/code/main.py:242–281

()

    return optimized_data

async def main():
    # Main function
    parser = argparse.ArgumentParser(description='Run experiments on various datasets')
    parser.add_argument('--dataset', type=str, default='math',
                        choices=list(DATASET_CONFIGS.keys()),
                        help='Dataset to run experiment on')
    parser.add_argument('--start', type=int, default=0,
                        help='Start index of the dataset')
    parser.add_argument('--end', type=int, default=2,
                        help='End index of the dataset (-1 for all)')
    parser.add_argument('--model', type=str, default='gpt-4o-mini',
                        help='Model to use for the experiment')
    parser.add_argument('--mode', type=str, choices=['atom', 'plugin'], default='atom',
                        help='Mode: atom (standard experiment) or plugin (generate contracted dataset)')
    parser.add_argument('--max_concurrent', type=int, default=1000,
                        help='Maximum number of concurrent tasks (default: 50)')

    args = parser.parse_args()

    if args.mode == 'plugin':
        # Run plugin mode
        await optimize_dataset(
            dataset=args.dataset,
            model=args.model,
            start=args.start,
            end=args.end
        )
    elif args.mode == 'atom':
        # Run standard experiment
        runner = ExperimentRunner(
            dataset=args.dataset,
            model=args.model,
            start=args.start,
            end=args.end,
            mode=args.mode,
            max_concurrent=args.max_concurrent
        )
        await runner.run()
    else:
        raise ValueError(f"Invalid mode: {args.mode}")

if __name__ == "__main__":
    asyncio.run(main())