initial commit

models/common/vqa_tools/aokvqa/gpt3/query_gpt3.py

@@ -0,0 +1,79 @@
+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()