import os import json import itertools import argparse import numpy as np from functools import partial from models import gpt, gpt_usage from tasks import get_task def get_value(task, x, y, n_evaluate_sample, cache_value=True): value_prompt = task.value_prompt_wrap(x, y) if cache_value and value_prompt in task.value_cache: return task.value_cache[value_prompt] value_outputs = gpt(value_prompt, n=n_evaluate_sample, stop=None) value = task.value_outputs_unwrap(x, y, value_outputs) if cache_value: task.value_cache[value_prompt] = value return value def get_values(task, x, ys, n_evaluate_sample, cache_value=True): values = [] local_value_cache = {} for y in ys: # each partial output if y in local_value_cache: # avoid duplicate candidates value = 0 else: value = get_value(task, x, y, n_evaluate_sample, cache_value=cache_value) local_value_cache[y] = value values.append(value) return values def get_votes(task, x, ys, n_evaluate_sample): vote_prompt = task.vote_prompt_wrap(x, ys) vote_outputs = gpt(vote_prompt, n=n_evaluate_sample, stop=None) values = task.vote_outputs_unwrap(vote_outputs, len(ys)) return values def get_proposals(task, x, y): propose_prompt = task.propose_prompt_wrap(x, y) proposals = gpt(propose_prompt, n=1, stop=None)[0].split('\n') return [y + _ + '\n' for _ in proposals] def get_samples(task, x, y, n_generate_sample, prompt_sample, stop): if prompt_sample == 'standard': prompt = task.standard_prompt_wrap(x, y) elif prompt_sample == 'cot': prompt = task.cot_prompt_wrap(x, y) else: raise ValueError(f'prompt_sample {prompt_sample} not recognized') samples = gpt(prompt, n=n_generate_sample, stop=stop) return [y + _ for _ in samples] def solve(args, task, idx, to_print=True): print(gpt) x = task.get_input(idx) # input ys = [''] # current output candidates infos = [] for step in range(task.steps): # generation if args.method_generate == 'sample': new_ys = [get_samples(task, x, y, args.n_generate_sample, prompt_sample=args.prompt_sample, stop=task.stops[step]) for y in ys] elif args.method_generate == 'propose': new_ys = [get_proposals(task, x, y) for y in ys] new_ys = list(itertools.chain(*new_ys)) ids = list(range(len(new_ys))) # evaluation if args.method_evaluate == 'vote': values = get_votes(task, x, new_ys, args.n_evaluate_sample) elif args.method_evaluate == 'value': values = get_values(task, x, new_ys, args.n_evaluate_sample) # selection if args.method_select == 'sample': ps = np.array(values) / sum(values) select_ids = np.random.choice(ids, size=args.n_select_sample, p=ps).tolist() elif args.method_select == 'greedy': select_ids = sorted(ids, key=lambda x: values[x], reverse=True)[:args.n_select_sample] select_new_ys = [new_ys[select_id] for select_id in select_ids] # log if to_print: sorted_new_ys, sorted_values = zip(*sorted(zip(new_ys, values), key=lambda x: x[1], reverse=True)) print(f'-- new_ys --: {sorted_new_ys}\n-- sol values --: {sorted_values}\n-- choices --: {select_new_ys}\n') infos.append({'step': step, 'x': x, 'ys': ys, 'new_ys': new_ys, 'values': values, 'select_new_ys': select_new_ys}) ys = select_new_ys if to_print: print(ys) return ys, {'steps': infos} def naive_solve(args, task, idx, to_print=True): x = task.get_input(idx) # input ys = get_samples(task, x, '', args.n_generate_sample, args.prompt_sample, stop=None) return ys, {} def run(args): task = get_task(args.task, args.task_file_path) logs, cnt_avg, cnt_any = [], 0, 0 global gpt gpt = partial(gpt, model=args.backend, temperature=args.temperature) if args.naive_run: file = f'logs/{args.task}/{args.backend}_{args.temperature}_naive_{args.prompt_sample}_sample_{args.n_generate_sample}_start{args.task_start_index}_end{args.task_end_index}.json' else: file = f'logs/{args.task}/{args.backend}_{args.temperature}_{args.method_generate}{args.n_generate_sample}_{args.method_evaluate}{args.n_evaluate_sample}_{args.method_select}{args.n_select_sample}_start{args.task_start_index}_end{args.task_end_index}.json' os.makedirs(os.path.dirname(file), exist_ok=True) for i in range(args.task_start_index, args.task_end_index): # solve if args.naive_run: ys, info = naive_solve(args, task, i) else: ys, info = solve(args, task, i) # log infos = [task.test_output(i, y) for y in ys] info.update({'idx': i, 'ys': ys, 'infos': infos, 'usage_so_far': gpt_usage(args.backend)}) logs.append(info) with open(file, 'w') as f: json.dump(logs, f, indent=4) # log main metric accs = [info['r'] for info in infos] cnt_avg += sum(accs) / len(accs) cnt_any += any(accs) print(i, 'sum(accs)', sum(accs), 'cnt_avg', cnt_avg, 'cnt_any', cnt_any, '\n') n = args.task_end_index - args.task_start_index print(cnt_avg / n, cnt_any / n) print('usage_so_far', gpt_usage(args.backend)) def parse_args(): args = argparse.ArgumentParser() args.add_argument('--backend', type=str, choices=['gpt-4', 'gpt-3.5-turbo'], default='gpt-4') args.add_argument('--temperature', type=float, default=0.7) args.add_argument('--task', type=str, required=True, choices=['game24', 'text', 'crosswords']) args.add_argument('--task_file_path', type=str, required=True) args.add_argument('--task_start_index', type=int, default=900) args.add_argument('--task_end_index', type=int, default=1000) args.add_argument('--naive_run', action='store_true') args.add_argument('--prompt_sample', type=str, choices=['standard', 'cot']) # only used when method_generate = sample, or naive_run args.add_argument('--method_generate', type=str, choices=['sample', 'propose']) args.add_argument('--method_evaluate', type=str, choices=['value', 'vote']) args.add_argument('--method_select', type=str, choices=['sample', 'greedy']) args.add_argument('--n_generate_sample', type=int, default=1) # only thing needed if naive_run args.add_argument('--n_evaluate_sample', type=int, default=1) args.add_argument('--n_select_sample', type=int, default=1) args = args.parse_args() return args if __name__ == '__main__': args = parse_args() print(args) run(args)