You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
93 lines
2.5 KiB
93 lines
2.5 KiB
import time
|
|
|
|
from tqdm import tqdm
|
|
import random
|
|
import requests
|
|
import json
|
|
import threading
|
|
from threading import Thread
|
|
|
|
|
|
api_key_list = [
|
|
"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",
|
|
"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",
|
|
"sk-nYbapOeC5VmSReJB1JgEr3BlbkFJnOo2J9qFJRKXrOSNiYFO",
|
|
"sk-tOy3uBFkPsg9uVWTpDOor3BlbkFJkbXgo0sHAubK8VWyaeso",
|
|
"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",
|
|
"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",
|
|
"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",
|
|
"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",
|
|
"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",
|
|
"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"
|
|
]
|
|
|
|
lock = threading.RLock()
|
|
|
|
def request_api_chatgpt(api_key, prompt):
|
|
try:
|
|
OPENAI_API_KEY = api_key
|
|
url = "https://api.openai.com/v1/chat/completions"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {OPENAI_API_KEY}"
|
|
}
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
"temperature": 0.5
|
|
}
|
|
response = requests.post(url,
|
|
headers=headers,
|
|
data=json.dumps(data),
|
|
timeout=240)
|
|
|
|
res = response.json()
|
|
text = res["choices"][0]["message"]["content"]
|
|
|
|
lock.acquire()
|
|
api_key_list.append(api_key)
|
|
|
|
|
|
with open("/home/majiahui/mulu_ner/data/prompt_small_gen.txt", mode="a") as f:
|
|
f.write(prompt)
|
|
f.write("**************")
|
|
f.write(text)
|
|
f.write("\n")
|
|
lock.release()
|
|
|
|
except:
|
|
time.sleep(5)
|
|
lock.acquire()
|
|
api_key_list.append(api_key)
|
|
lock.release()
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
with open("./data/prompt_shuffle.txt", encoding="utf-8") as f:
|
|
text = f.read()
|
|
|
|
text_list = text.split("\n")
|
|
|
|
index = 0
|
|
while True:
|
|
if index == len(text_list):
|
|
break
|
|
|
|
if api_key_list == []:
|
|
time.sleep(1)
|
|
continue
|
|
else:
|
|
api_key = api_key_list.pop(0)
|
|
prompt = text_list[index]
|
|
|
|
t = Thread(target=request_api_chatgpt, args=(api_key, prompt))
|
|
t.start()
|
|
lock.acquire()
|
|
index += 1
|
|
print(index)
|
|
lock.release()
|
|
|
|
|
|
|