mirror of
https://github.com/deepseek-ai/DeepSeek-Math
synced 2024-11-28 23:10:48 +00:00
230 lines
9.7 KiB
Python
Executable File
230 lines
9.7 KiB
Python
Executable File
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 '</s>'
|
|
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()
|