Get GPT to answer questions with justifications, then extract those answers and justifications in a standardized format.

qa_gpt.py

# Setup OpenAI GPT API

import openai
import textwrap

# Set the GPT-3 API key
openai.api_key = "YOUR_API_KEY"

def get_gpt_response(prompt, max_tokens=2048, temperature=0.7, n=1):
    response = openai.Completion.create(
        engine="text-davinci-003",
        prompt=prompt,
        max_tokens=max_tokens,
        n = n,
        temperature=temperature
    )
    return response["choices"][0]["text"]

def format_prompt(context, options):
    alphabet = list("ABCDEFGHIJKLMNOPQRSTUVWXYZ"[:len(options)])
    any_answer = f"{alphabet[0]}/{alphabet[1]}/.../{alphabet[-1]}" if len(alphabet) > 2 else f"{alphabet[0]}/{alphabet[1]}"
    prompt = textwrap.dedent(f"""
    I want you to act as an expert question answerer. Consider the CONTEXT and answer exactly in the ANSWER_FORMAT.
    Refer to the answer choices solely by the letters. Answer exactly in the following format.

    CONTEXT:
    {context}
    """)

    for i, option in enumerate(options):
        prompt += f"{alphabet[i]} = {option}\n"

    prompt += textwrap.dedent(f"""
    ANSWER FORMAT:
    Answer = {any_answer}
    """)

    for i, option in enumerate(options):
        prompt += f"{alphabet[i]} => justification for why {alphabet[i]} was selected/not selected\n"

    prompt += "\nSOLUTION:"

    return prompt.strip()

def gpt_multiple_choice_justification(context, options):
    alphabet = list("ABCDEFGHIJKLMNOPQRSTUVWXYZ"[:len(options)])
    prompt = format_prompt(context, options)
    response = get_gpt_response(prompt, max_tokens=2048, temperature=0.7, n=1)
    lines = response.split("\n")
    for line in lines:
        if "Answer = " in line:
            answer = line.split("Answer = ")[1].strip()
            if answer in alphabet:
                answer = options[alphabet.index(answer)]
            else:
                answer = None

    js = []
    for option in options:
        alpha = alphabet[options.index(option)]
        justification = "Not sure about this option (please use best judgement)."
        for line in lines:
            if f"{alpha} => " in line:
                justification = line.split(f"{alpha} => ")[1]
        js.append(justification)

    return answer, js

if __name__ == "__main__":
    answer = gpt_multiple_choice_justification("What is the best machine learning model to use for an image classification problem?", ["Fully Connected", "Convolutional Neural Network", "Transformer", "Recurrent"])
    print(answer)