|
"""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)) |
|
|
