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.
140 lines
4.1 KiB
140 lines
4.1 KiB
import time
|
|
|
|
from tqdm import tqdm
|
|
import random
|
|
import requests
|
|
import json
|
|
import threading
|
|
from threading import Thread
|
|
import redis
|
|
|
|
lock = threading.RLock()
|
|
|
|
pool = redis.ConnectionPool(host='104.244.90.248', port=63179, max_connections=50, db=10, password='Zhicheng123*')
|
|
redis_ = redis.Redis(connection_pool=pool, decode_responses=True)
|
|
|
|
|
|
with open("api_key.txt", "r",) as f:
|
|
a = f.read()
|
|
a = a.split("\n")
|
|
|
|
redis_key_name_openaikey_list = "openaikey_list"
|
|
redis_zirenwu = "redis_zirenwu"
|
|
|
|
api_key_list = []
|
|
for i in a:
|
|
api_key_list.append(str(i.split("----")[-1]))
|
|
|
|
for i in api_key_list:
|
|
redis_.rpush(redis_key_name_openaikey_list, i)
|
|
|
|
lock = threading.RLock()
|
|
|
|
prompt_dict = {
|
|
"mulu_prompt": "为论文题目《{}》生成目录，要求只有一级标题和二级标题，一级标题使用中文数字 例如一、xxx；二级标题使用阿拉伯数字 例如1.1 xxx；一级标题不少于7个；每个一级标题至少包含3个二级标题",
|
|
}
|
|
|
|
with open("./data/题目4_new.txt", encoding="utf-8") as f:
|
|
text = f.read()
|
|
|
|
text_list = text.split("\n")
|
|
|
|
title_list = []
|
|
for i in text_list:
|
|
title_list.append(i.split("@@@@@")[0])
|
|
|
|
random.shuffle(title_list)
|
|
|
|
print(len(title_list))
|
|
|
|
zirenwu_list = []
|
|
|
|
for title in title_list:
|
|
for prompt in prompt_dict:
|
|
zirenwu_list.append((prompt, str(prompt_dict[prompt]).format(title)))
|
|
|
|
for i in zirenwu_list:
|
|
redis_.rpush(redis_zirenwu, str(i))
|
|
|
|
|
|
def request_api_chatgpt(api_key, task_type, prompt):
|
|
t1 = time.time()
|
|
global api_key_list
|
|
global zirenwu_list
|
|
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=1200)
|
|
|
|
res = response.json()
|
|
text = res["choices"][0]["message"]["content"]
|
|
|
|
# api_key_list.append(api_key)
|
|
t2 = time.time()
|
|
t_n = t2 - t1
|
|
lock.acquire()
|
|
with open("/home/majiahui/mulu_ner/data/paper_prompt_title_4/title_{}_data.txt".format(task_type), mode="a") as f:
|
|
f.write(prompt)
|
|
f.write("**************")
|
|
f.write(text)
|
|
f.write("\n")
|
|
f.write("=================================================================================================")
|
|
lock.release()
|
|
|
|
if t_n > 20:
|
|
redis_.rpush(redis_key_name_openaikey_list, api_key)
|
|
else:
|
|
time.sleep(20 - t_n)
|
|
redis_.rpush(redis_key_name_openaikey_list, api_key)
|
|
|
|
except:
|
|
print("task_type_bad", task_type)
|
|
print("api_key_bad", api_key)
|
|
time.sleep(20)
|
|
lock.acquire()
|
|
redis_.rpush(redis_key_name_openaikey_list, api_key)
|
|
redis_.rpush(redis_zirenwu, str((task_type, prompt)))
|
|
lock.release()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
while True:
|
|
if redis_.llen(redis_zirenwu) == 0:
|
|
time.sleep(1)
|
|
continue
|
|
elif redis_.llen(redis_zirenwu) != 0 and redis_.llen(redis_key_name_openaikey_list) != 0:
|
|
lock.acquire()
|
|
api_key = redis_.lpop(redis_key_name_openaikey_list)
|
|
api_key = api_key.decode('UTF-8')
|
|
dan_zirenwu = redis_.lpop(redis_zirenwu)
|
|
dan_zirenwu = dan_zirenwu.decode('UTF-8')
|
|
lock.release()
|
|
# dan_zirenwu = zirenwu_list.pop(0)
|
|
dan_zirenwu = eval(dan_zirenwu)
|
|
task_type, prompt = dan_zirenwu[0], dan_zirenwu[1]
|
|
t = Thread(target=request_api_chatgpt, args=(api_key, task_type, prompt))
|
|
t.start()
|
|
elif redis_.llen(redis_key_name_openaikey_list) == 0:
|
|
time.sleep(1)
|
|
continue
|
|
else:
|
|
time.sleep(1)
|
|
continue
|
|
|
|
|
|
|
|
|
|
|