import argparse import os from tqdm import tqdm import json from copy import deepcopy from vllm import LLM, SamplingParams from pebble import ProcessPool from concurrent.futures import TimeoutError import random from eval.utils import generate_completions, load_hf_lm_and_tokenizer from transformers import AutoTokenizer from data_processing.answer_extraction import * from eval.eval_script import * from few_shot_prompts import * def evaluate(eval_fn, tasks, _timeout=15): with ProcessPool() as pool: timeout_cnt = 0 iterator = pool.map(eval_fn, tasks, timeout=_timeout).result() labels = [] while True: try: labels.append(int(next(iterator))) except StopIteration: break except TimeoutError as error: labels.append(0) timeout_cnt += 1 except Exception as error: print(error.traceback, flush=True) exit() return labels, timeout_cnt def infer(args, test_data): global tokenizer if tokenizer is None: tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name_or_path, trust_remote_code=True) if args.prompt_format == 'few_shot': assert args.few_shot_prompt is not None prompting = eval(args.few_shot_prompt)() prompts = [] for example in test_data: prompt = "" if args.prompt_format == 'few_shot': prompt = prompting.format_prompt(example['messages'][-2]['content'], example['messages'][-1]['content']) else: for mess in example['messages']: if args.prompt_format == 'sft': if mess['role'] == 'user': prompt += f"{tokenizer.eos_token}User: {mess['content'].strip()}\n\nAssistant:" elif mess['role'] == 'assistant': prompt += mess['content'].rstrip() else: raise NotImplementedError() prompt = prompt.lstrip() if args.prompt_format == 'sft' and prompt.startswith(tokenizer.eos_token): prompt = prompt[len(tokenizer.eos_token):].lstrip() example['prompt'] = prompt prompts.append(prompt.lstrip()) global model print("Loading model and tokenizer...") if args.use_vllm: if model is None: model = LLM(model=args.model_name_or_path, tokenizer=args.tokenizer_name_or_path, trust_remote_code=True, tensor_parallel_size=len(os.environ['CUDA_VISIBLE_DEVICES'].split(","))) eos_token = tokenizer.eos_token if tokenizer is not None and tokenizer.eos_token is not None else '' stop_words = [eos_token] if args.prompt_format == 'few_shot': stop_words.extend(prompting.stop_words()) outputs = model.generate(prompts, SamplingParams(temperature=args.temperature, top_p=1.0, max_tokens=1024, n=1, stop=stop_words)) outputs = sorted(outputs, key=lambda x: int(x.request_id)) # sort outputs by request_id outputs = [output.outputs[0].text for output in outputs] else: model, tokenizer = load_hf_lm_and_tokenizer( model_name_or_path=args.model_name_or_path, tokenizer_name_or_path=args.tokenizer_name_or_path, load_in_8bit=args.load_in_8bit, load_in_half=args.load_in_half, gptq_model=args.gptq ) stop_id_sequences = [] if tokenizer.eos_token_id is not None: stop_id_sequences = [[tokenizer.eos_token_id]] if args.prompt_format == 'few_shot': stop_id_sequences.extend([tokenizer.encode(word) for word in prompting.stop_words()]) outputs, finish_completion = generate_completions( model=model, tokenizer=tokenizer, prompts=prompts, max_new_tokens=512, batch_size=args.eval_batch_size, stop_id_sequences=stop_id_sequences if stop_id_sequences else None, end_of_generation_id_sequence=[tokenizer.eos_token_id] if tokenizer.eos_token_id is not None else None ) if args.complete_partial_output: model_outputs = [example['messages'][-1]['content'] + output for example, output in zip(test_data, outputs)] else: model_outputs = outputs predictions = [eval(args.answer_extraction_fn)(item['messages'][-2]['content'], output, task='cot') for item, output in tqdm(zip(test_data, model_outputs), desc="extract answer", total=len(model_outputs))] assert len(model_outputs) > 0, f"{len(model_outputs)}" results = [] for example, output, pred in zip(test_data, model_outputs, predictions): item = deepcopy(example) item.update({ 'model_output': output, 'prediction': pred, }) results.append(item) return results def main(args): random.seed(42) print("Loading data...") test_data = [] with open(os.path.join(args.data_dir, f"train.jsonl" if args.infer_train_set else f"test.jsonl")) as fin: for line in fin: example = json.loads(line) messages = example['messages'] assert messages[-1]['role'] == 'assistant' if not args.complete_partial_output: example['reference'] = example.get('reference', '') or [mess['content'] for mess in messages if mess['role'] == 'assistant'] for mess in messages: if mess['role'] == 'assistant': mess['content'] = '' example['messages'] = messages test_data.append(example) if args.max_num_examples and len(test_data) > args.max_num_examples: test_data = random.sample(test_data, args.max_num_examples) if args.n_subsets > 1: assert args.subset_id >= 0 and args.subset_id < args.n_subsets test_data = [item for i, item in enumerate(test_data) if i % args.n_subsets == args.subset_id] if not test_data: return if not os.path.exists(args.save_dir): os.makedirs(args.save_dir, exist_ok=True) results = infer(args, test_data) labels, eval_timeout_cnt = evaluate(eval(args.eval_fn), results) for item, label in zip(results, labels): item['accuracy'] = label print("Calculating accuracy...") acc = 0 for item in results: acc += item['accuracy'] print("output acc = {:.5f}".format(acc / len(results) * 100), flush=True) print(f"Timeout count >>> output eval = {eval_timeout_cnt}", flush=True) pred_fname = "predictions.json" if args.n_subsets > 1: pred_fname = f"predictions.{args.subset_id}.json" with open(os.path.join(args.save_dir, pred_fname), "w") as fout: json.dump(results, fout, ensure_ascii=True) metric_fname = "metrics.json" if args.n_subsets > 1: metric_fname = f"metrics.{args.subset_id}.json" with open(os.path.join(args.save_dir, metric_fname), "w") as fout: json.dump({ "n_samples": len(results), "accuracy": sum(item['accuracy'] for item in results) / len(results), }, fout, indent=4) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--data_dir", type=str, default="data/mgsm") parser.add_argument("--max_num_examples", type=int, default=None, help="maximum number of examples to evaluate.") parser.add_argument("--save_dir", type=str, default="results/mgsm") parser.add_argument("--model_name_or_path", type=str, default=None, help="if specified, we will load the model to generate the predictions.") parser.add_argument("--tokenizer_name_or_path", type=str, default=None, help="if specified, we will load the tokenizer from here.") parser.add_argument("--eval_batch_size", type=int, default=1, help="batch size for evaluation.") parser.add_argument("--load_in_8bit", action="store_true", help="load model in 8bit mode, which will reduce memory and speed up inference.") parser.add_argument("--gptq", action="store_true", help="If given, we're evaluating a 4-bit quantized GPTQ model.") parser.add_argument("--use_vllm", action="store_true") parser.add_argument("--load_in_half", action='store_true') parser.add_argument("--infer_train_set", action="store_true") parser.add_argument("--n_subsets", type=int, default=1) parser.add_argument("--subset_id", type=int, default=0) parser.add_argument("--temperature", type=float, default=0.0) parser.add_argument("--repeat_id_start", type=int, default=0) parser.add_argument("--n_repeat_sampling", type=int, default=1) parser.add_argument("--complete_partial_output", action='store_true') parser.add_argument("--prompt_format", type=str, choices=['sft', 'few_shot'], default='sft') parser.add_argument("--few_shot_prompt", type=str, default=None) parser.add_argument("--answer_extraction_fn", type=str, required=True) parser.add_argument("--eval_fn", type=str, required=True) parser.add_argument("--gpus", type=str, default=None) args, unparsed_args = parser.parse_known_args() if args.gpus is not None: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus print(unparsed_args, flush=True) if 'math6' in args.data_dir: args.multi_turn = True # model_name_or_path cannot be both None or both not None. model = None tokenizer = None pool = None if args.n_repeat_sampling > 1 or args.repeat_id_start != 0: assert args.temperature > 0 save_dir = args.save_dir for i in range(args.repeat_id_start, args.repeat_id_start + args.n_repeat_sampling): print(f"working on the {i} trials ...", flush=True) args.save_dir = os.path.join(save_dir, str(i)) os.makedirs(args.save_dir, exist_ok=True) main(args) else: main(args) if pool is not None: pool.close()