mulu_ner/generate_sencend.py

import time
import os
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_bad_list = "openaikey_bad_list"
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()

path_output = "chatgpt_data_v1"
file_small_title = r'data/title_mulu_to_/small_title_prompt_shuffle.txt'
file_references = r'data/title_mulu_to_/references_prompt.txt'
file_zhaiyao = r'data/title_mulu_to_/zhaiyao_prompt.txt'
file_task_book = r'data/jianjie_to_/task_book_prompt.txt'


path_list = [file_small_title, file_references, file_zhaiyao, file_task_book]
zirenwu_list = []
for path in path_list:
    with open(path, encoding="utf-8") as f:
        type_prompt = path.split("/")[-1].split(".")[0]
        texts = f.readlines()
        for i in texts:
            zirenwu_list.append([type_prompt, json.loads(i)])

import random
random.shuffle(zirenwu_list)
for i in zirenwu_list:
    redis_.rpush(redis_zirenwu, json.dumps(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=240)

        res = response.json()

        text = res["choices"][0]["message"]["content"]
        path_root = '/home/majiahui/mulu_ner/data/{}/paper_prompt_title_1_1'.format(path_output)
        if not os.path.exists(path_root):
            os.makedirs(path_root)

        lock.acquire()
        with open(path_root + "/sencend_{}_data.txt".format(task_type), mode="a") as f:
            f.write(prompt)
            f.write("*" * 20)
            f.write(text)
            f.write("@" * 20)
        lock.release()
        t2 = time.time()
        t_n = t2 - t1
        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:
        time.sleep(20)
        lock.acquire()
        redis_.rpush(redis_key_name_openaikey_list, api_key)
        redis_.rpush(redis_zirenwu, json.dumps([type_prompt, 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 = json.loads(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