|
"""Generate answers with GPT-3.5""" |
|
|
|
import argparse |
|
import json |
|
import os |
|
import time |
|
import concurrent.futures |
|
|
|
import openai |
|
import tqdm |
|
import shortuuid |
|
|
|
MODEL = "gpt-3.5-turbo" |
|
MODEL_ID = "gpt-3.5-turbo:20230327" |
|
|
|
|
|
def get_answer(question_id: int, question: str, max_tokens: int): |
|
ans = { |
|
"answer_id": shortuuid.uuid(), |
|
"question_id": question_id, |
|
"model_id": MODEL_ID, |
|
} |
|
for _ in range(3): |
|
try: |
|
response = openai.ChatCompletion.create( |
|
model=MODEL, |
|
messages=[ |
|
{"role": "system", "content": "You are a helpful assistant."}, |
|
{ |
|
"role": "user", |
|
"content": question, |
|
}, |
|
], |
|
max_tokens=max_tokens, |
|
) |
|
ans["text"] = response["choices"][0]["message"]["content"] |
|
return ans |
|
except Exception as e: |
|
print("[ERROR]", e) |
|
ans["text"] = "#ERROR#" |
|
time.sleep(1) |
|
return ans |
|
|
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser(description="ChatGPT answer generation.") |
|
parser.add_argument("-q", "--question") |
|
parser.add_argument("-o", "--output") |
|
parser.add_argument( |
|
"--max-tokens", |
|
type=int, |
|
default=1024, |
|
help="maximum number of tokens produced in the output", |
|
) |
|
args = parser.parse_args() |
|
|
|
questions_dict = {} |
|
with open(os.path.expanduser(args.question)) as f: |
|
for line in f: |
|
if not line: |
|
continue |
|
q = json.loads(line) |
|
questions_dict[q["question_id"]] = q["text"] |
|
|
|
answers = [] |
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor: |
|
futures = [] |
|
for qid, question in questions_dict.items(): |
|
future = executor.submit(get_answer, qid, question, args.max_tokens) |
|
futures.append(future) |
|
|
|
for future in tqdm.tqdm( |
|
concurrent.futures.as_completed(futures), total=len(futures) |
|
): |
|
answers.append(future.result()) |
|
|
|
answers.sort(key=lambda x: x["question_id"]) |
|
|
|
with open(os.path.expanduser(args.output), "w") as f: |
|
table = [json.dumps(ans) for ans in answers] |
|
f.write("\n".join(table)) |
|
|
