# flake8: noqa import ast import json import os import pandas as pd import tiktoken from tqdm import tqdm from .constructions import ChatGPTSchema, ResultsForHumanSchema from .utils import extract_answer, read_jsonl, save_jsonl # define the datasets english_qa_datasets = [ 'lsat-ar', 'lsat-lr', 'lsat-rc', 'logiqa-en', 'sat-math', 'sat-en', 'aqua-rat', 'sat-en-without-passage', 'gaokao-english' ] chinese_qa_datasets = [ 'logiqa-zh', 'jec-qa-kd', 'jec-qa-ca', 'gaokao-chinese', 'gaokao-geography', 'gaokao-history', 'gaokao-biology', 'gaokao-chemistry', 'gaokao-physics', 'gaokao-mathqa' ] english_cloze_datasets = ['math'] chinese_cloze_datasets = ['gaokao-mathcloze'] multi_choice_datasets = ['jec-qa-kd', 'jec-qa-ca', 'gaokao-physics'] math_output_datasets = ['gaokao-mathcloze', 'math'] def convert_zero_shot(line, dataset_name): try: passage = line['passage'] if line['passage'] is not None else '' if dataset_name in english_qa_datasets: option_string = 'ABCDEFG' count = len(line['options']) if count == 1: count = 5 return passage + 'Q: ' + line['question'] + ' ' \ + 'Answer Choices: ' + ' '.join(line['options']) + '\n' + \ 'A: Among A through {}, the answer is'.format(option_string[count - 1]) elif dataset_name in chinese_qa_datasets: option_string = 'ABCDEFG' count = len(line['options']) if count == 1: count = 4 return passage + '问题：' + line['question'] + ' ' \ + '选项：' + ' '.join(line['options']) + '\n' + \ '答案：从A到{}, 我们应选择'.format(option_string[count - 1]) elif dataset_name in english_cloze_datasets: return passage + 'Q: ' + line['question'] + '\n' \ 'A: The answer is' elif dataset_name in chinese_cloze_datasets: return passage + '问题：' + line['question'] + '\n' \ '答案：' except NameError: print('Dataset not defined.') prefix = '该问题为单选题，所有选项中必有一个正确答案，且只有一个正确答案。\n' def convert_zero_shot_CoT_stage1(line, dataset_name): try: passage = line['passage'] if line['passage'] is not None else '' if dataset_name in english_qa_datasets: return passage + 'Q: ' + line['question'] + ' ' \ + 'Answer Choices: ' + ' '.join(line['options']) + '\n' + \ "Let's think step by step." elif dataset_name in chinese_qa_datasets: option_string = 'ABCDEFG' count = len(line['options']) if count == 1: count = 4 return passage + '问题：' + line['question'] + ' ' \ + '选项：' + ' '.join(line['options']) + '\n' + \ '从A到{}, 我们应选择什么？让我们逐步思考：'.format(option_string[count - 1]) elif dataset_name in english_cloze_datasets: return passage + 'Q: ' + line['question'] + '\n' \ "A: Let's think step by step." elif dataset_name in chinese_cloze_datasets: return passage + '问题：' + line['question'] + '\n' \ '答案：让我们逐步思考：' except NameError: print('Dataset not defined.') # process few-shot raw_prompts def combine_prompt(prompt_path, dataset_name, load_explanation=True, chat_mode=False): skip_passage = False if dataset_name == 'sat-en-without-passage': skip_passage = True dataset_name = 'sat-en' demostrations = [] # read the prompts by context and explanation context_row = [0, 1, 3, 5, 7, 9] explanation_row = [0, 2, 4, 6, 8, 10] raw_prompts_context = pd.read_csv(prompt_path, header=0, skiprows=lambda x: x not in context_row, keep_default_na=False) raw_prompts_explanation = pd.read_csv( prompt_path, header=0, skiprows=lambda x: x not in explanation_row, keep_default_na=False).replace(r'\n\n', '\n', regex=True) contexts = [] for line in list(raw_prompts_context[dataset_name]): if line: # print(line) contexts.append(ast.literal_eval(line)) explanations = [ exp for exp in raw_prompts_explanation[dataset_name] if exp ] for idx, (con, exp) in enumerate(zip(contexts, explanations)): passage = con['passage'] if con[ 'passage'] is not None and not skip_passage else '' question = con['question'] options = con['options'] if con['options'] is not None else '' label = con['label'] if con['label'] is not None else '' answer = con[ 'answer'] if 'answer' in con and con['answer'] is not None else '' if dataset_name in english_qa_datasets: question_input = 'Problem {}. '.format(idx + 1) + passage + ' ' + question + '\n' \ + 'Choose from the following options: ' + ' '.join(options) + '\n' question_output = (('Explanation for Problem {}: '.format(idx + 1) + exp + '\n') if load_explanation else '') \ + 'The answer is therefore {}'.format(label) elif dataset_name in chinese_qa_datasets: question_input = '问题 {}. '.format(idx + 1) + passage + ' ' + question + '\n' \ + '从以下选项中选择: ' + ' '.join(options) + '\n' question_output = (('问题 {}的解析: '.format(idx + 1) + exp + '\n') if load_explanation else '') \ + '答案是 {}'.format(label) elif dataset_name in english_cloze_datasets: question_input = 'Problem {}. '.format(idx + 1) + question + '\n' question_output = (('Explanation for Problem {}: '.format(idx + 1) + exp + '\n') if load_explanation else '') \ + 'The answer is therefore {}'.format(answer) elif dataset_name in chinese_cloze_datasets: question_input = '问题 {}. '.format(idx + 1) + question + '\n' question_output = (('问题 {}的解析: '.format(idx + 1) + exp + '\n') if load_explanation else '') \ + '答案是 {}'.format(answer) else: raise ValueError( f'During loading few-sot examples, found unknown dataset: {dataset_name}' ) if chat_mode: demostrations.append((question_input, question_output)) else: demostrations.append(question_input + question_output + '\n') return demostrations enc = None def _lazy_load_enc(): global enc if enc is None: enc = tiktoken.encoding_for_model('gpt-4') # cut prompt if reach max token length def concat_prompt(demos, dataset_name, max_tokens, end_of_example='\n', verbose=False): _lazy_load_enc() demostration_en = 'Here are the answers for the problems in the exam.\n' demostration_zh = '以下是考试中各个问题的答案。\n' for i in range(len(demos)): # print(len(enc.encode(demostration_en)), len(enc.encode(demostration_zh))) if dataset_name in english_qa_datasets: demostration_en = demostration_en + demos[i] + end_of_example elif dataset_name in chinese_qa_datasets: demostration_zh = demostration_zh + demos[i] + end_of_example elif dataset_name in english_cloze_datasets: demostration_en = demostration_en + demos[i] + end_of_example elif dataset_name in chinese_cloze_datasets: demostration_zh = demostration_zh + demos[i] + end_of_example # break if reach max token limit if len(enc.encode(demostration_en)) < max_tokens and len( enc.encode(demostration_zh)) < max_tokens: output = demostration_en if len(demostration_en) > len( demostration_zh) else demostration_zh prompt_num = i + 1 else: break if verbose: print('max_tokens set as ', max_tokens, 'actual_tokens is', len(enc.encode(output)), 'num_shot is', prompt_num) return output, prompt_num def concat_prompt_chat_mode(demos, dataset_name, max_tokens, end_of_example='\n', verbose=False): _lazy_load_enc() answers = [] sentences = '' for i in range(len(demos)): answers += [ { 'role': 'user', 'content': demos[i][0] }, { 'role': 'assistant', 'content': demos[i][1] }, ] sentences += json.dumps(answers[-1]) # break if reach max token limit if len(enc.encode(sentences)) > max_tokens: answers.pop() answers.pop() break if verbose: print('max_tokens set as ', max_tokens, 'actual_tokens is', len(enc.encode(sentences)), 'num_shot is', len(answers) // 2) return answers, len(answers) // 2 def convert_few_shot(line, dataset_name, demo, n_shot, chat_mode=False): passage = line['passage'] if line['passage'] is not None else '' question = line['question'] options = line['options'] if line['options'] is not None else '' if dataset_name in english_qa_datasets: question_input = 'Problem {}. '.format(n_shot + 1) + passage + ' ' + question + '\n' \ + 'Choose from the following options: ' + ' '.join(options) + '\n' # + "Explanation for Problem {}: ".format(n_shot + 1) if dataset_name in chinese_qa_datasets: question_input = '问题 {}. '.format(n_shot + 1) + passage + ' ' + question + '\n' \ + '从以下选项中选择: ' + ' '.join(options) + '\n' # + "问题 {}的解析: ".format(n_shot + 1) if dataset_name in english_cloze_datasets: question_input = 'Problem {}. '.format(n_shot + 1) + question + '\n' # + "Explanation for Problem {}: ".format(n_shot + 1) if dataset_name in chinese_cloze_datasets: question_input = '问题 {}. '.format(n_shot + 1) + question + '\n' # + "问题 {}的解析: ".format(n_shot + 1) if chat_mode: return demo + [ { 'role': 'user', 'content': question_input }, ] else: return demo + question_input def load_dataset(dataset_name, setting_name, parent_path, prompt_path=None, max_tokens=None, end_of_example='\n', chat_mode=False, verbose=False): test_path = os.path.join(parent_path, dataset_name + '.jsonl') loaded_jsonl = read_jsonl(test_path) processed = [] if setting_name == 'few-shot-CoT' or setting_name == 'few-shot': # process demo once if it is few-shot-CoT processed_demos = combine_prompt( prompt_path, dataset_name, load_explanation=setting_name == 'few-shot-CoT', chat_mode=chat_mode) if chat_mode: chosen_prompt, n_shot = concat_prompt_chat_mode(processed_demos, dataset_name, max_tokens, end_of_example, verbose=verbose) else: chosen_prompt, n_shot = concat_prompt(processed_demos, dataset_name, max_tokens, end_of_example, verbose=verbose) if verbose: loaded_jsonl = tqdm(loaded_jsonl) for meta_idx, line in enumerate(loaded_jsonl): if setting_name == 'zero-shot': ctxt = convert_zero_shot(line, dataset_name) elif setting_name == 'zero-shot-CoT': ctxt = convert_zero_shot_CoT_stage1(line, dataset_name) elif setting_name == 'few-shot-CoT' or setting_name == 'few-shot': ctxt = convert_few_shot(line, dataset_name, chosen_prompt, n_shot, chat_mode) try: new_instance = ChatGPTSchema(context=ctxt, metadata=meta_idx) processed.append(new_instance.to_dict()) except NameError: print('Dataset not defined.') return processed def generate_second_stage_input(dataset_name, input_list, output_list, with_format_prompt=False): try: english_format_prompt = 'Based on the previous results, your task is to extract the final answer and provide the output enclosed in brackets【】, such as 【0】 or 【A】.' chinese_format_prompt = '根据以上内容，你的任务是把最终的答案提取出来并填在【】中，例如【0】或者【A】。' if dataset_name in english_qa_datasets: prompt_suffix = 'Therefore, among A through E, the answer is' if with_format_prompt: prompt_suffix = english_format_prompt + prompt_suffix elif dataset_name in chinese_qa_datasets: prompt_suffix = '因此，从A到D, 我们应选择' if with_format_prompt: prompt_suffix = chinese_format_prompt + prompt_suffix elif dataset_name in english_cloze_datasets: prompt_suffix = 'Therefore, the answer is' if with_format_prompt: prompt_suffix = english_format_prompt + prompt_suffix elif dataset_name in chinese_cloze_datasets: prompt_suffix = '因此，答案是' if with_format_prompt: prompt_suffix = chinese_format_prompt + prompt_suffix except NameError: print('Dataset not defined.') processed = [] for i in range(len(input_list)): ctxt = '{0}\n{1}\n{2}'.format(input_list[i]['context'], extract_answer(output_list[i]), prompt_suffix) new_instance = ChatGPTSchema(context=ctxt, metadata=input_list[i]['metadata']) processed.append(new_instance.to_dict()) return processed def load_dataset_as_result_schema(dataset_name, parent_path): test_path = os.path.join(parent_path, dataset_name + '.jsonl') loaded_jsonl = read_jsonl(test_path) processed = [] for i, line in enumerate(loaded_jsonl): problem_input = convert_zero_shot(line, dataset_name) processed.append( ResultsForHumanSchema( index=i, problem_input=problem_input, label=line['label'] if line['label'] else line['answer'], )) return processed if __name__ == '__main__': # set variables parent_dir = '../../data/V1_1/' raw_prompt_path = '../data/few_shot_prompts.csv' # set dataset name to process setting_name = 'few-shot-CoT' # setting_name can be chosen from ["zero-shot", "zero-shot-CoT", "few-shot-CoT"] data_name = 'jec-qa-kd' save_dir = '../../experiment_input/{}/'.format(setting_name) if not os.path.exists(save_dir): os.makedirs(save_dir) processed_data = load_dataset(data_name, setting_name, parent_dir, prompt_path=raw_prompt_path, max_tokens=2048) save_jsonl(processed_data, os.path.join(save_dir, '{}.jsonl'.format(data_name)))