File size: 2,305 Bytes
80a598c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import asyncio
import json
from tenacity import retry, stop_after_attempt, wait_exponential
from openai import AsyncOpenAI
from tqdm.asyncio import tqdm
base_url="http://47.88.8.18:8088/v1"
api_key="eyJhbGciOiJIUzI1NiIsInR5cCI6Imp3dCJ9.eyJ1c2VybmFtZSI6IjQzNzkzNyIsInBhc3N3b3JkIjoiNDM3OTM3MTIzIiwiZXhwIjoyMDMxODEwOTAzfQ.M710cSMdw1OZ2TBVPRhlnoavZ8CQG5tXgj3WGl3FoIg"
client = AsyncOpenAI(base_url=base_url, api_key=api_key)
# 设置并发限制
MAX_RETRIES = 10
BASE_DELAY = 1
MAX_DELAY = 60
MAX_CONCURRENT = 64
model = "gpt-4"
# model = "gpt-4o-2024-08-06"
@retry(stop=stop_after_attempt(10), wait=wait_exponential(multiplier=1, min=4, max=60))
async def get_chat_completion(message: str, semaphore, retry_count=0) -> str:
try:
async with semaphore: # 使用传入的信号量限制并发
response = await client.chat.completions.create(
model=model,
messages=[{"role": "system", "content": "you are a helpful assistant"}, {"role": "user", "content": message}],
timeout=80
)
response_result = response.choices[0].message.content
# message[model] = response_result
temp = {}
temp["prompt"] = message
temp["label"] = response_result
return temp
except Exception as e:
print(f"Error in get_chat_completion for message {type(e).__name__} - {str(e)}")
raise
async def request_model(prompts):
semaphore = asyncio.Semaphore(MAX_CONCURRENT)
async def wrapped_get_chat_completion(prompt):
try:
return await get_chat_completion(prompt, semaphore)
except Exception as e:
print(f"Task failed after all retries with error: {e}")
return None
tasks = [wrapped_get_chat_completion(prompt) for prompt in prompts]
results = []
for future in tqdm.as_completed(tasks, total=len(tasks), desc="Processing prompts"):
result = await future
results.append(result)
return results
if __name__ == "__main__":
prompts = ["测试测试测试"+str(i) for i in range(10)]
results = asyncio.run(request_model(prompts))
with open(f'{model}_result.json', 'w', encoding="utf-8") as f:
f.write(json.dumps(results,ensure_ascii=False)) |