Function main

omnivoice/eval/wer/seedtts.py:282–409

()

def main():
    parser = get_parser()
    args = parser.parse_args()

    logging.basicConfig(
        format="%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s",
        level=logging.INFO,
        force=True,
    )

    logging.info(f"Calculating WER for {args.wav_path}")

    # 1. Prepare Data
    logging.info("Reading test list...")
    data_list = []
    samples = read_test_list(args.test_list)
    for s in samples:
        wav_path = str(Path(args.wav_path) / f"{s['id']}.{args.extension}")
        if not os.path.exists(wav_path):
            logging.warning(f"File missing: {wav_path}")
            continue
        data_list.append({"wav_path": wav_path, "truth_text": s["text"]})
    total_files = len(data_list)
    logging.info(f"Total files: {total_files}.")

    # 2. Worker config
    num_gpus = torch.cuda.device_count()
    assert num_gpus > 0, "No GPU found. GPU is required."
    total_workers = num_gpus * args.nj_per_gpu

    mp.set_start_method("spawn", force=True)
    manager = mp.Manager()
    rank_queue = manager.Queue()

    for _ in range(args.nj_per_gpu):
        for rank in range(num_gpus):
            rank_queue.put(rank)

    # 3. Scheduling: Split data into chunks for better load balancing
    chunk_size = max(1, args.batch_size)
    tasks = []
    for i in range(0, total_files, chunk_size):
        tasks.append(data_list[i : i + chunk_size])

    logging.info(
        f"Split data into {len(tasks)} chunks (size ~{chunk_size}). "
        f"Spawning {total_workers} workers."
    )

    # 4. Execution
    results = []

    with ProcessPoolExecutor(
        max_workers=total_workers,
        initializer=process_init,
        initargs=(rank_queue, args.model_dir, args.lang),
    ) as executor: