84 lines
3.1 KiB
Python
84 lines
3.1 KiB
Python
import tiktoken
|
|
import random
|
|
import requests
|
|
import time
|
|
import json
|
|
import concurrent.futures
|
|
import urllib3
|
|
|
|
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
|
|
|
OPENAI_API_KEY = open('openai_api_key.txt', 'r').read().split('\n')
|
|
|
|
|
|
def get_tokens_num(text, model):
|
|
encoding = tiktoken.encoding_for_model(model)
|
|
tokens_num = len(encoding.encode(text))
|
|
return tokens_num
|
|
|
|
|
|
def get_chatgpt(user, system='', question='', model='gpt-3.5-turbo-16k-0613', temperature=1.0, max_tokens=2048,
|
|
top_p=0.8,
|
|
frequency_penalty=0.0, presence_penalty=0.0):
|
|
try_count = 0
|
|
while try_count < 20:
|
|
key_choice = random.choice(OPENAI_API_KEY)
|
|
url = r'https://api.openai.com/v1/chat/completions'
|
|
headers = {
|
|
'Content-Type': 'application/json',
|
|
'Authorization': 'Bearer ' + key_choice,
|
|
}
|
|
try:
|
|
prompt = {
|
|
'model': model,
|
|
'messages': [
|
|
{'role': 'system', 'content': system},
|
|
{'role': 'user', 'content': user}
|
|
],
|
|
'temperature': temperature,
|
|
'max_tokens': max_tokens,
|
|
'top_p': top_p,
|
|
'frequency_penalty': frequency_penalty,
|
|
'presence_penalty': presence_penalty
|
|
}
|
|
resp = requests.post(url, json=prompt, headers=headers, verify=False, timeout=60)
|
|
answer = json.loads(resp.text)
|
|
if 'choices' in answer:
|
|
result = answer['choices'][0]['message']['content']
|
|
else:
|
|
time.sleep(1)
|
|
continue
|
|
if result == '':
|
|
time.sleep(1)
|
|
continue
|
|
else:
|
|
if question:
|
|
return {'question': question, 'answer': result}
|
|
else:
|
|
return {'system': system, 'user': user, 'answer': result}
|
|
except Exception as e:
|
|
print(key_choice + ': query request failed \n' + str(e))
|
|
print('retrying...')
|
|
try_count += 1
|
|
time.sleep(1)
|
|
|
|
|
|
def get_chatgpt_concurrent(user_input, system_input, question_input, model='gpt-3.5-turbo-16k-0613', temperature=1.0,
|
|
max_tokens=2048,
|
|
top_p=0.8,
|
|
frequency_penalty=0.0, presence_penalty=0.0):
|
|
work_count = 0
|
|
result = []
|
|
print('Query size is ' + str(len(user_input)))
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor:
|
|
futures = [executor.submit(get_chatgpt, user, system, question, model, temperature, max_tokens, top_p,
|
|
frequency_penalty,
|
|
presence_penalty) for user, system, question in
|
|
zip(user_input, system_input, question_input)]
|
|
for future in concurrent.futures.as_completed(futures):
|
|
if future.result():
|
|
result.append(future.result())
|
|
work_count += 1
|
|
print(str(work_count) + '/' + str(len(user_input)) + ' finished')
|
|
return result
|