| import os | |
| import random | |
| import json | |
| from tqdm import tqdm | |
| import argparse | |
| import pathlib | |
| import openai | |
| openai.organization = os.getenv('OPENAI_ORG') | |
| openai.api_key = os.getenv('OPENAI_API_KEY') | |
| from load_aokvqa import load_aokvqa | |
| random.seed(0) | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--aokvqa-dir', type=pathlib.Path, required=True, dest='aokvqa_dir') | |
| parser.add_argument('--split', type=str, choices=['train', 'val', 'test'], required=True) | |
| parser.add_argument('--n', type=int, default=10, dest='num_examples') | |
| parser.add_argument('--train-context', type=argparse.FileType('r'), dest='train_context_file') | |
| parser.add_argument('--prefix', type=str, default='', dest='prompt_prefix') | |
| parser.add_argument('--include-choices', action='store_true', dest='include_choices') | |
| parser.add_argument('--context', type=argparse.FileType('r'), dest='context_file') | |
| parser.add_argument('--out', type=argparse.FileType('w'), required=True, dest='output_file') | |
| args = parser.parse_args() | |
| train_set = load_aokvqa(args.aokvqa_dir, 'train') | |
| eval_set = load_aokvqa(args.aokvqa_dir, args.split) | |
| train_context = {} | |
| context = {} | |
| if args.context_file is not None: | |
| train_context = json.load(args.train_context_file) | |
| context = json.load(args.context_file) | |
| predictions = {} | |
| for d in tqdm(eval_set): | |
| q = d['question_id'] | |
| prompt = args.prompt_prefix | |
| for e in random.sample(train_set, args.num_examples): | |
| prompt += prompt_element(e, | |
| context=train_context.get(q, None), | |
| include_choices=args.include_choices, | |
| answer=True | |
| ) | |
| prompt += '\n\n' | |
| prompt += prompt_element(d, | |
| context=context.get(q, None), | |
| include_choices=args.include_choices, | |
| answer=False | |
| ) | |
| response = openai.Completion.create( | |
| engine="text-curie-001", | |
| prompt=prompt, | |
| temperature=0.0, | |
| max_tokens=10, | |
| ) | |
| predictions[q] = response.choices[0].text.strip() | |
| json.dump(predictions, args.output_file) | |
| def prompt_element(d, context=None, include_choices=False, answer=False): | |
| return (f"Context: {context}\n" if context is not None else '') + \ | |
| f"Q: {d['question']}\n" + \ | |
| (f"Choices: {', '.join(d['choices'])}.\n" if include_choices else '') + \ | |
| f"A:" + (f" {d['choices'][d['correct_choice_idx']]}" if answer else '') | |
| if __name__ == '__main__': | |
| main() | |