Function main

light-eval/src/eval_bbh.py:210–255

(args, multiple_choice_tasks=MULTIPLE_CHOICE_TASKS, free_form_tasks=FREE_FORM_TASKS)

                

def main(args, multiple_choice_tasks=MULTIPLE_CHOICE_TASKS, free_form_tasks=FREE_FORM_TASKS):

    run_multiple_choice = args.task == 'all' or args.task == 'multiple_choice'
    run_free_form = args.task == 'all' or args.task == 'free_form'

    path_split = args.pretrained_path[0].split('/') if isinstance(args.pretrained_path,list) else args.pretrained_path.split('/')
    if path_split[-1] == '':
        path_split.pop(-1)
    model_name = path_split[-1] 
    infer_path = os.path.join('results', model_name, 'bbh/infer')
    os.makedirs(infer_path, exist_ok=True)
    eval_path = os.path.join('results', model_name, 'bbh/eval')
    os.makedirs(eval_path, exist_ok=True)

    model = load(args)
    

    if run_multiple_choice:
        run_infer(model, args.max_seq_len, multiple_choice_tasks, args.data_dir, infer_path, 'multiple_choice', args.overwrite)
    if run_free_form:
        run_infer(model, args.max_seq_len, free_form_tasks, args.data_dir, infer_path, 'free_form', args.overwrite)
    
    
    torch.distributed.barrier()
    if torch.distributed.get_rank() == 0:

        score = {}
        total_results = []

        if run_multiple_choice:    
            score['multiple_choice'], task_results = run_eval(multiple_choice_tasks, infer_path, mode='multiple_choice')
        total_results.extend(task_results)

        if run_free_form:
            score['free_form'], task_results = run_eval(free_form_tasks, infer_path, mode='free_form')
        total_results.extend(task_results)

        if args.task == 'all':
            score['TOTAL'] = '%.4f' %(sum(total_results) / len(total_results))
        
        result_path = os.path.join(eval_path, 'run_results.json')

        with open(result_path, 'w') as f:
            json.dump(score, f, ensure_ascii=False, indent=2) 
    
    return 

if __name__ == '__main__':