Function run_infer

light-eval/src/eval_bbh.py:137–187

(model, max_seq_len, tasks, data_path, infer_path, mode, overwrite = False)

    return prompt

def run_infer(model, max_seq_len, tasks, data_path, infer_path, mode, overwrite = False):

    for task in tasks:

        task_infer_path = os.path.join(infer_path, f'{task}_infer.jsonl')
        if not overwrite and os.path.exists(task_infer_path):
            print(f"{task_infer_path} existed, skip!")
            continue

        print(f'Testing {task} ...')
        task_data = json.load(open(os.path.join(data_path, f'bbh/{task}.json')))
        with open(os.path.join(data_path, f'cot-prompts/{task}.txt'), 'r') as f:
            task_prompt = f.readlines()[2:]
            task_prompt = ''.join(task_prompt)

        test_set = []
        answer_set = []
        for item in task_data['examples']:
            full_prompt = resize_prompt(
                model.tokenizer,
                max_seq_len,
                task_prompt + 
                '\n\nQ: ' + 
                item['input'] + 
                "\nA: Let's think step by step."
            )
            
            test_set.append(full_prompt)
            if mode == 'multiple_choice':
                answer_set.append(item['target'][1])
            elif mode == 'free_form':
                answer_set.append(item['target'])
            
        batch_prompt = batch_data(test_set, batch_size=8)
        res_completions = []
        for batch_input in tqdm(batch_prompt):
            
            outputs = model.generate(prompts=batch_input, images=None, max_gen_len=1024)
            
            for output in outputs:
                res_completions.append(output)

        torch.distributed.barrier()
        if torch.distributed.get_rank() == 0:
            with jsonlines.open(task_infer_path, mode='w') as writer:
                for (prompt, completion, prompt_answer) in zip(task_data['examples'], res_completions, answer_set):
                    record = {'prompt': prompt,
                        'completion': completion,
                        'target_ans': prompt_answer
                    }
                    writer.write(record)

def run_eval(tasks, infer_path, mode):