14B / eval /evaluate_chatml_mmlu.py

Upload 3 files

50e6685 8 months ago

No virus

13.8 kB

	import os
	import pandas as pd
	import numpy as np
	import argparse
	import datasets
	import torch
	import re
	from thefuzz import process
	from typing import List
	from tqdm import tqdm
	from transformers.trainer_utils import set_seed

	from typing import Tuple, List, Union, Iterable

	import numpy as np
	import torch
	import torch.nn.functional as F
	from transformers import PreTrainedTokenizer
	from transformers import logging
	from transformers.generation import LogitsProcessor
	from typing import TYPE_CHECKING, Optional, Tuple, Union, Callable, List
	HistoryType = List[Tuple[str, str]]
	TokensType = List[int]
	BatchTokensType = List[List[int]]

	def make_context(
	tokenizer: PreTrainedTokenizer,
	query: str,
	history: List[Tuple[str, str]] = None,
	system: str = "",
	max_window_size: int = 6144,
	chat_format: str = "chatml",
	):
	if history is None:
	history = []

	im_start, im_end = "<\|im_start\|>", "<\|im_end\|>"
	im_start_tokens = [tokenizer.im_start_id]
	im_end_tokens = [tokenizer.im_end_id]
	nl_tokens = tokenizer.encode("\n")

	def _tokenize_str(role, content):
	return f"{role}\n{content}", tokenizer.encode(
	role
	) + nl_tokens + tokenizer.encode(content)

	system_text, system_tokens_part = _tokenize_str("system", system)
	system_tokens = im_start_tokens + system_tokens_part + im_end_tokens

	raw_text = ""
	context_tokens = []

	for turn_query, turn_response in reversed(history):
	query_text, query_tokens_part = _tokenize_str("user", turn_query)
	query_tokens = im_start_tokens + query_tokens_part + im_end_tokens
	response_text, response_tokens_part = _tokenize_str(
	"assistant", turn_response
	)
	response_tokens = im_start_tokens + response_tokens_part + im_end_tokens

	next_context_tokens = nl_tokens + query_tokens + nl_tokens + response_tokens
	prev_chat = (
	f"\n{im_start}{query_text}{im_end}\n{im_start}{response_text}{im_end}"
	)

	current_context_size = (
	len(system_tokens) + len(next_context_tokens) + len(context_tokens)
	)
	if current_context_size < max_window_size:
	context_tokens = next_context_tokens + context_tokens
	raw_text = prev_chat + raw_text
	else:
	break

	context_tokens = system_tokens + context_tokens
	raw_text = f"{im_start}{system_text}{im_end}" + raw_text
	context_tokens += (
	nl_tokens
	+ im_start_tokens
	+ _tokenize_str("user", query)[1]
	+ im_end_tokens
	+ nl_tokens
	+ im_start_tokens
	+ tokenizer.encode("assistant")
	+ nl_tokens
	)
	raw_text += f"\n{im_start}user\n{query}{im_end}\n{im_start}assistant\n"

	return raw_text, context_tokens

	def chat(
	model,
	tokenizer: PreTrainedTokenizer,
	query: str,
	history: Optional[HistoryType],
	system: str = "You are a helpful assistant.",
	append_history: bool = True
	) -> Tuple[str, HistoryType]:


	if history is None:
	history = []

	raw_text, context_tokens = make_context(
	tokenizer,
	query,
	history=history,
	system=system,
	max_window_size=6144,
	chat_format = "chatml",
	)

	stop_words_ids = [[tokenizer.im_end_id], [tokenizer.im_start_id]]
	input_ids = torch.tensor([context_tokens]).cuda()
	outputs = model.generate(
	input_ids,
	# stop_words_ids = stop_words_ids,
	return_dict_in_generate = False,
	)

	response = decode_tokens(
	outputs[0],
	tokenizer,
	raw_text_len=len(raw_text),
	context_length=len(context_tokens),
	chat_format='chatml',
	verbose=False,
	)

	if append_history:
	history.append((query, response))

	return response, history

	def decode_tokens(
	tokens: Union[torch.LongTensor, TokensType],
	tokenizer: PreTrainedTokenizer,
	raw_text_len: int,
	context_length: int,
	chat_format: str = "chatml",
	verbose: bool = False,
	return_end_reason: bool = False,
	) -> str:
	if torch.is_tensor(tokens):
	tokens = tokens.cpu().numpy().tolist()


	return _decode_chatml(
	tokens,
	stop_words=[],
	eod_token_ids=[tokenizer.im_start_id, tokenizer.im_end_id],
	tokenizer=tokenizer,
	raw_text_len=raw_text_len,
	context_length=context_length,
	verbose=verbose,
	return_end_reason=return_end_reason,
	)


	def _decode_chatml(
	tokens: List[int],
	*,
	stop_words: List[str],
	eod_token_ids: List[int],
	tokenizer: PreTrainedTokenizer,
	raw_text_len: int,
	context_length: int,
	verbose: bool = False,
	return_end_reason: bool = False,
	chat_format = "chatml",
	):
	end_reason = f"Gen length {len(tokens)}"
	eod_token_idx = context_length
	for eod_token_idx in range(context_length, len(tokens)):
	if tokens[eod_token_idx] in eod_token_ids:
	end_reason = f"Gen {tokenizer.decode([tokens[eod_token_idx]])!r}"
	break

	trim_decode_tokens = tokenizer.decode(tokens[:eod_token_idx])[raw_text_len:]
	if verbose:
	print("\nRaw Generate w/o EOD:", tokenizer.decode(tokens)[raw_text_len:])
	print("\nRaw Generate:", trim_decode_tokens)
	print("\nEnd Reason:", end_reason)
	for stop_word in stop_words:
	trim_decode_tokens = trim_decode_tokens.replace(stop_word, "").strip()
	trim_decode_tokens = trim_decode_tokens.strip()
	if verbose:
	print("\nGenerate:", trim_decode_tokens)

	if return_end_reason:
	return trim_decode_tokens, end_reason
	else:
	return trim_decode_tokens



	def load_models_tokenizer(args):
	from transformers import AutoModelForCausalLM, AutoTokenizer
	from transformers.generation import GenerationConfig

	tokenizer = AutoTokenizer.from_pretrained(args.checkpoint_path, trust_remote_code=True)
	model = AutoModelForCausalLM.from_pretrained(args.checkpoint_path, device_map="auto", trust_remote_code=True).eval()
	model.generation_config = GenerationConfig.from_pretrained(args.checkpoint_path, trust_remote_code=True)
	model.generation_config.do_sample = False # use greedy decoding
	return model, tokenizer


	def format_example(line):
	example = 'The following is a multiple-choice question. Please choose the most suitable one among A, B, C and D as the answer to this question.\n\n' + line['question'] + "\n"
	for choice in choices:
	example += f'{choice}. {line[f"{choice}"]}\n'
	return example


	def process_before_extraction(gen, choice_dict):
	# replace the choice by letter in the generated sentence
	# from longest one to shortest one
	for key, val in sorted(choice_dict.items(), key=lambda x: len(x[1]), reverse=True):
	pattern = re.compile(re.escape(val.rstrip(".")), re.IGNORECASE)
	gen = pattern.sub(key, gen)
	return gen

	def extract_choice(gen, choice_list):
	# answer is A \| choice is A \| choose A
	res = re.search(r"(?:(?:[Cc]hoose)\|(?:(?:[Aa]nswer\|[Cc]hoice)(?![^ABCD]{0,20}?(?:n't\|not))[^ABCD]{0,10}?\b(?:\|is\|:\|be))\b)[^ABCD]{0,20}?\b(A\|B\|C\|D)\b", gen)

	# A is correct \| A is right
	if res is None:
	res = re.search(r"\b(A\|B\|C\|D)\b(?![^ABCD]{0,8}?(?:n't\|not)[^ABCD]{0,5}?(?:correct\|right))[^ABCD]{0,10}?\b(?:correct\|right)\b", gen)

	# straight answer: A
	if res is None:
	res = re.search(r"^(A\|B\|C\|D)(?:\.\|,\|:\|$)", gen)

	# simply extract the first appearred letter
	if res is None:
	res = re.search(r"(?<![a-zA-Z])(A\|B\|C\|D)(?![a-zA-Z=])", gen)

	if res is None:
	return choices[choice_list.index(process.extractOne(gen, choice_list)[0])]
	else:
	return res.group(1)

	def extract_answer(response, row):
	gen = process_before_extraction(response, {choice: row[choice] for choice in choices})
	pred = extract_choice(gen, [row[choice] for choice in choices])
	return pred

	@torch.no_grad()
	def eval_subject(
	model,
	tokenizer,
	subject_name,
	test_df,
	save_result_dir=None,
	overwrite=False,
	**kwargs
	):
	result_path = os.path.join(save_result_dir, f'{subject_name}_result.csv')
	if not overwrite and os.path.exists(result_path):
	print(f"{result_path} existed, skip!")
	score = []
	for (_, datarow), (_, resultrow) in zip(test_df.iterrows(), pd.read_csv(result_path).iterrows()):
	# pred = extract_answer(resultrow['model_response'], datarow)
	pred = resultrow['model_output']
	correct = 1 if pred == datarow['answer'] else 0
	score.append(correct)
	return score

	result = []
	score = []

	for _, row in tqdm(test_df.iterrows(), total=len(test_df)):
	question = format_example(row)

	response, history = chat(
	model,
	tokenizer,
	question,
	history=None,
	)
	print(question)
	print(response)
	pred = extract_answer(response, row)
	print(pred)
	print("======================")

	if 'answer' in row:
	correct = 1 if pred == row['answer'] else 0
	score.append(correct)
	if args.debug: print(f'{question} pred: {pred} ref: {row["answer"]}')
	result.append(pred)

	if save_result_dir:
	test_df['model_output'] = result
	test_df['model_response'] = response
	if score:
	test_df["correctness"] = score
	os.makedirs(save_result_dir, exist_ok=True)
	test_df.to_csv(os.path.join(
	save_result_dir, f'{subject_name}_result.csv'), encoding="utf-8", index=False)

	return score


	def cal_mmlu(res):
	acc_sum_dict = dict()
	acc_norm_sum_dict = dict()
	cnt_dict = dict()
	acc_sum = 0.
	cnt = 0
	hard_cnt = 0
	hard_acc_sum = 0.

	for class_ in TASK_NAME_MAPPING.keys():
	acc_sum_dict[class_] = 0.
	acc_norm_sum_dict[class_] = 0.
	cnt_dict[class_] = 0.

	for tt in TASK_NAME_MAPPING[class_]:
	acc_sum += sum(res[tt])
	cnt += len(res[tt])

	acc_sum_dict[class_] += sum(res[tt])
	cnt_dict[class_] += len(res[tt])

	print('\n\n\n')
	for k in TASK_NAME_MAPPING.keys():
	if k in cnt_dict:
	print('%s ACC: %.2f ' % (
	k, acc_sum_dict[k] * 100 / cnt_dict[k]))
	print('AVERAGE ACC:%.2f ' % (acc_sum *100 / cnt))


	def main(args):
	print("loading model weights")
	if args.checkpoint_path is not None:
	model, tokenizer = load_models_tokenizer(args)
	else:
	model, tokenizer = None, None
	print("model loaded")

	dev_result = {}
	for subject_name in tqdm(SUBJECTS):
	# val_file_path = os.path.join(args.eval_data_path, 'val', f'{subject_name}_val.csv')
	# dev_file_path = os.path.join(args.eval_data_path, 'dev', f'{subject_name}_dev.csv')
	test_file_path = os.path.join(args.eval_data_path, 'test', f'{subject_name}_test.csv')
	# val_df = pd.read_csv(val_file_path, names=['question','A','B','C','D','answer'])
	# dev_df = pd.read_csv(dev_file_path, names=['question','A','B','C','D','answer'])
	test_df = pd.read_csv(test_file_path, names=['question','A','B','C','D','answer'])

	score = eval_subject(model, tokenizer, subject_name, test_df, save_result_dir=f"outs_chat/mmlu_eval_result", overwrite=args.overwrite)
	dev_result[subject_name] = score
	cal_mmlu(dev_result)


	TASK_NAME_MAPPING = {'stem': ['abstract_algebra', 'anatomy', 'astronomy', 'college_biology', 'college_chemistry', 'college_computer_science', 'college_mathematics', 'college_physics', 'computer_security', 'conceptual_physics', 'electrical_engineering', 'elementary_mathematics', 'high_school_biology', 'high_school_chemistry', 'high_school_computer_science', 'high_school_mathematics', 'high_school_physics', 'high_school_statistics', 'machine_learning'],
	'Humanities': ['formal_logic', 'high_school_european_history', 'high_school_us_history', 'high_school_world_history', 'international_law', 'jurisprudence', 'logical_fallacies', 'moral_disputes', 'moral_scenarios', 'philosophy', 'prehistory', 'professional_law', 'world_religions'],
	'other': ['business_ethics', 'college_medicine', 'human_aging', 'management', 'marketing', 'medical_genetics', 'miscellaneous', 'nutrition', 'professional_accounting', 'professional_medicine', 'virology', 'global_facts', 'clinical_knowledge'],
	'social': ['econometrics', 'high_school_geography', 'high_school_government_and_politics', 'high_school_macroeconomics', 'high_school_microeconomics', 'high_school_psychology', 'human_sexuality', 'professional_psychology', 'public_relations', 'security_studies', 'sociology', 'us_foreign_policy']}
	SUBJECTS = [v for vl in TASK_NAME_MAPPING.values() for v in vl]
	choices = ["A", "B", "C", "D"]

	if __name__ == '__main__':
	parser = argparse.ArgumentParser(description='Test HF checkpoint.')
	parser.add_argument('-c', '--checkpoint-path', type=str, help='Checkpoint path', default="Qwen/Qwen-7B-Chat")
	parser.add_argument('-s', '--seed', type=int, default=1234, help='Random seed')

	"""Provide extra arguments required for tasks."""
	group = parser.add_argument_group(title='Evaluation options')
	group.add_argument('-d', '--eval_data_path', type=str,
	help='Path to eval data')
	group.add_argument("--debug", action='store_true', default=False,
	help='Print infos.')
	group.add_argument("--overwrite", action='store_true', default=False,
	help='Overwrite existed results')

	args = parser.parse_args()
	set_seed(args.seed)

	main(args)