File size: 2,341 Bytes
d573b56 |
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 |
from openai import OpenAI
import time
import os
import json
# siliconflow
siliconflow_interlm = {
"api_key": "your-api",
"base_url": "https://api.siliconflow.cn/v1",
# "model_name": "Qwen/Qwen2-Math-72B-Instruct",
# "model_name": "internlm/internlm2_5-20b-chat",
"model_name": "deepseek-ai/DeepSeek-V2.5",
"temperature": 0.5,
}
client = OpenAI(
api_key=siliconflow_interlm["api_key"],
base_url=siliconflow_interlm["base_url"],
)
def response(messages, max_tokens=512, temperature=0.5, sleep_time=0.05):
try:
res = client.chat.completions.create(
model=siliconflow_interlm["model_name"],
messages=messages,
# max_tokens=max_tokens,
temperature=temperature,
).choices[0].message.content
except Exception as e:
print(f"Got exception {e}, sleep(0.05s) and retry again")
if sleep_time>0:
time.sleep(sleep_time)
res = client.chat.completions.create(
model=siliconflow_interlm["model_name"],
messages=messages,
# max_tokens=max_tokens,
temperature=temperature,
).choices[0].message.content
return res
def save_json(data_list, file_path, jsonl=False):
dirname, filename = os.path.split(file_path)
if not os.path.exists(dirname):
os.makedirs(dirname, exist_ok=True)
if filename.split(".")[-1] == "jsonl":
jsonl = True
with open(file_path, 'at', encoding='utf-8') as f:
if jsonl and isinstance(data_list, list):
for item in data_list:
f.write(json.dumps(item, ensure_ascii=False, indent=4) + '\n')
else:
f.write(json.dumps(data_list, ensure_ascii=False, indent=4))
def save_json_once(data, file_path):
dirname, filename = os.path.split(file_path)
if not os.path.exists(dirname):
os.makedirs(dirname, exist_ok=True)
with open(file_path, 'at', encoding='utf-8') as f:
## 不加 indent，单条数据就是1行
f.write(json.dumps(data, ensure_ascii=False) + '\n')
# file_names = [ "游戏概述", "章节与故事情节", "主要角色", "人物剧情梳理", "游戏世界观", "建筑与环境", "战斗系统", "游戏玩法", "艺术与音乐", "文化内涵", "市场影响", "彩蛋梗数据" ] |