run_audio.py

import argparse
import torch
from almeval.datasets import build_dataset
from almeval.models import build_model
from tqdm import tqdm
from almeval.utils import *
from loguru import logger
import sys
import os
from datetime import datetime
import warnings
import time
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=UserWarning)

def setup_logging(rank, model, data, work_dir):
    # Create log directory
    log_dir = os.path.join(work_dir, model, data, 'logs')
    os.makedirs(log_dir, exist_ok=True)
    
    log_file = os.path.join(log_dir, f'rank{rank}.log')
    
    # Remove default console output
    logger.remove()
    
    # Add file output
    log_file_handle = open(log_file, 'w', encoding='utf-8')
    logger.add(log_file_handle, 
               format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}",
               level="INFO")
    # Print rank=0 output to console as well
    if rank == 0:
        logger.add(sys.stdout, 
                format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}",
                level="INFO")
    
    return log_file


def merge_one_dataset(args, dataset, result_file, eval_file):
    model_data_dir = osp.join(args.work_dir, args.model, dataset.DATASET_NAME)
    os.makedirs(model_data_dir, exist_ok=True)

    if args.reeval:
        perf = dataset.evaluate(result_file, method=args.eval_method)
        with open(eval_file, 'w') as f:
            json.dump(perf, f, indent=4)
        return

    tmp_files = [osp.join(model_data_dir, f'{rank}_{args.world_size}_{dataset.DATASET_NAME}.pkl') for rank in range(args.world_size)]

    # Merge if all tmp_files exist
    if all(osp.exists(tmpfile) for tmpfile in tmp_files):
        data_all = {}
        for tmpfile in tmp_files:
            data_all.update(load(tmpfile))
        raw_data = dataset.data
        for x in raw_data:
            idx = int(x['index'])
            if idx not in data_all:
                logger.warning(f'index {idx} not found in data_all, details: {x}')
                x['prediction'] = 'null'
                x['real_prompt'] = ''
                continue
            x['prediction'] = str(data_all[idx]['prediction'])
            x['real_prompt'] = str(data_all[idx]['prompt'])

        dump(raw_data, result_file)

        for tmpfile in tmp_files:
            os.remove(tmpfile)

        logger.info(f'model {args.model}, data {dataset.DATASET_NAME}, all {args.world_size} result merged to {result_file}.')

    if args.skip_eval:
        logger.info(f'skip eval for {dataset.DATASET_NAME}')
        return
    perf = dataset.evaluate(result_file, method=args.eval_method)
    with open(eval_file, 'w') as f:
        json.dump(perf, f, indent=4)
    logger.info(f'model {args.model}, data {dataset.DATASET_NAME} evaluated.')


def do_reeval(dataset_name, result_file='auto', method='default'):
    datasets = []
    for dataset_name in args.data:
        d = build_dataset(dataset_name)
        if isinstance(d, list):
            datasets.extend(d)
        else:
            datasets.append(d)

    for dataset in datasets:
        if result_file == 'auto':
            benchmark_dir = osp.join(args.work_dir, args.model, dataset.DATASET_NAME)
            pred_result_file = osp.join(benchmark_dir, f'{args.model}_{dataset.DATASET_NAME}.jsonl')
        else:
            pred_result_file = result_file
        logger.info(f'evaluating {pred_result_file} with method {method}')
        perf = dataset.evaluate(pred_result_file, method=method)
        with open(pred_result_file.replace('.jsonl', f'_{method}_performance.json'), 'w') as f:
            json.dump(perf, f, indent=4)


def process_dataset(args, dataset, model):
    # Assign different subsets to each process
    model_data_dir = osp.join(args.work_dir, args.model, dataset.DATASET_NAME)
    result_file = osp.join(model_data_dir, f'{args.model}_{dataset.DATASET_NAME}.jsonl')
    eval_file = osp.join(model_data_dir, f'{args.model}_{dataset.DATASET_NAME}_{args.eval_method}_performance.json')
    os.makedirs(model_data_dir, exist_ok=True)
    rank = int(args.rank)
    
    if os.path.exists(result_file) and not args.force_reinfer:
        if args.reeval or not os.path.exists(eval_file):
            if rank==0:
                logger.info(f'file {result_file} exists, reevaluating...')
                merge_one_dataset(args, dataset, result_file, eval_file)
            return
        else:
            # Exit if not reeval and result file exists
            return 

    else:
        if args.debug:
            dataset.set_demo_mode()
        sample_indices = [i for i in range(len(dataset))]

        # Distribute data to each rank
        world_size = int(args.world_size)
        rank = int(args.rank)
        sample_indices_sub = sample_indices[rank::world_size]

        tmpl = osp.join(model_data_dir, f'{rank}_{args.world_size}_{dataset.DATASET_NAME}.pkl')
        out_file = tmpl.format(rank)
        res = load(out_file) if osp.exists(out_file) else {}
        
        processed_samples = 0
        for i in tqdm(sample_indices_sub, disable=args.rank != 0):
            msg = dataset[i]
            idx = int(msg['index'])
            if not args.force_reinfer:
                if idx in res:
                    continue

            if processed_samples==0:
                logger.info(f'Msg example: {msg}')

            real_prompt, response = model(msg)
            torch.cuda.empty_cache()
            if response is None:
                continue
            
            # we need response and prompt, because model may change prompt
            res[idx] = {
                'prompt': real_prompt,
                'prediction': response,
            }
            processed_samples += 1
            if processed_samples % 20 == 0:
                dump(res, out_file)
        dump(res, out_file)

        # Write a file to indicate this rank is done
        with open(osp.join(model_data_dir, f'{rank}_{args.world_size}_{dataset.DATASET_NAME}.done'), 'w') as f:
            f.write('done')

        # Rank 0 needs to wait for other ranks to finish, then merge results
        time_elapsed = 0
        if rank == 0:
            while True:
                all_success_files = [osp.join(model_data_dir, f'{rank}_{args.world_size}_{dataset.DATASET_NAME}.done') 
                                for rank in range(args.world_size)]
                if len(all_success_files) == args.world_size and \
                    all(osp.exists(success_file) for success_file in all_success_files):
                    # Delete all done files
                    for success_file in all_success_files:
                        os.remove(success_file)
                    break
                else:
                    time.sleep(10)
                    time_elapsed += 10
                    logger.info(f'waiting for other ranks to finish, time elapsed: {time_elapsed}s')
            merge_one_dataset(args, dataset, result_file, eval_file)

def main(args):
    datasets = []
    for dataset_name in args.data:
        d = build_dataset(dataset_name)
        if isinstance(d, list):
            datasets.extend(d)
        else:
            datasets.append(d)
            
    model = build_model(args.model)
    logger.info(f"Datasets: {datasets}")
    for dataset in datasets:
        setup_logging(args.rank, args.model, dataset.DATASET_NAME, args.work_dir)
        logger.info(f"Running {args.model} on dataset: {dataset.DATASET_NAME}")
        process_dataset(args, dataset, model)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, nargs='+', required=True, help='List of dataset names')
    parser.add_argument('--model', type=str, help='Model name')
    parser.add_argument('--work-dir', type=str, default='./eval_result', help='Working directory')
    parser.add_argument('--rank', type=int, default=0, help='Current GPU rank')
    parser.add_argument('--world-size', type=int, default=1, help='Total number of GPUs')
    parser.add_argument('--reeval', action='store_true', help='Whether to re-evaluate')
    parser.add_argument('--eval-file', type=str, default='auto', help='Evaluation file path, used when run_eval_only=True')
    parser.add_argument('--debug', action='store_true', help='Debug mode, only run 10 samples per dataset')
    parser.add_argument('--eval-method', type=str, default='default', help='Evaluation method')
    parser.add_argument('--force-reinfer', action='store_true', help='Whether to force re-inference')
    parser.add_argument('--skip-eval', action='store_true', help='Whether to skip evaluation')
    args = parser.parse_args()
    if args.reeval:
        do_reeval(args.data, args.eval_file, args.eval_method)
    else:
        main(args)