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