Spaces:
Sleeping
Sleeping
File size: 5,454 Bytes
c65e58e 499d83b c65e58e 499d83b c65e58e |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
import time
import requests
import json
from volcenginesdkarkruntime import Ark
from util.config_util import read_config as config
from util.config_util import load_json
from util import logger
import volcenginesdkcore
import volcenginesdkark
from volcenginesdkcore.rest import ApiException
from util.logger_util import log_decorate
class DouBaoService:
def __init__(self, model_name):
config = load_json('./conf/config.json')
self.conf = config[f"{model_name}ModelInfo"]
self.client = self.init_client()
self._complete_args = {}
def init_client(self):
base_url = self.conf["BASE_URL"]
ak = self.conf["ACCESS_KEY"]
sk = self.conf["SECRET_KEY"]
# api_key = self.conf["API_KEY"]
client = Ark(ak=ak, sk=sk, base_url=base_url)
# client = Ark(ak=api_key, base_url=base_url)
return client
def get_api_key(self):
configuration = volcenginesdkcore.Configuration()
configuration.ak = self.conf["ACCESS_KEY"]
configuration.sk = self.conf["SECRET_KEY"]
configuration.region = "cn-beijing"
endpoint_id = self.conf["ENDPOINT_ID"]
volcenginesdkcore.Configuration.set_default(configuration)
# use global default configuration
api_instance = volcenginesdkark.ARKApi()
get_api_key_request = volcenginesdkark.GetApiKeyRequest(
duration_seconds=30 * 24 * 3600,
resource_type="endpoint",
resource_ids=[
endpoint_id
],
)
try:
resp = api_instance.get_api_key(get_api_key_request)
return resp.api_key
except ApiException as e:
logger.error(f"Exception when calling api: {e}")
def set_complete_args(self, temperature=None, top_p=None, max_token=None):
if temperature is not None:
self._complete_args["temperature"] = temperature
if top_p is not None:
self._complete_args["top_p"] = top_p
if max_token is not None:
self._complete_args["max_tokens"] = max_token
def form_user_role(self, content):
return {"role": "user", "content": content}
def form_sys_role(self, content):
return {"role": "system", "content": content}
def form_assistant_role(self, content):
return {"role": "assistant", "content": content}
@property
def complete_args(self):
return {"temperature": 0.01, "top_p": 0.7}
@log_decorate
def chat_complete(self, messages):
endpoint_id = self.conf["ENDPOINT_ID"]
completion = self.client.chat.completions.create(
model=endpoint_id,
messages=messages,
**self.complete_args
)
logger.info(f"complete doubao task, id: {completion.id}")
return completion.choices[0].message.content
def prd_to_keypoint(self, prd_content):
role_desc = {"role": "system", "content": PRD2KP_SYS}
messages = [
role_desc,
{"role": "user", "content": prd_content}
]
return self.chat_complete(messages)
def prd_to_cases(self, prd_content, case_language="Chinese"):
role_desc = {"role": "system", "content": PRD_CASE_SYS[case_language]}
messages = [
role_desc,
{"role": "user", "content": prd_content}
]
return self.chat_complete(messages)
def keypoint_to_case(self, key_points):
role_desc = {"role": "system", "content": KP2CASE_SYS}
messages = [
role_desc,
{"role": "user", "content": key_points}
]
return self.chat_complete(messages)
def case_merge_together(self, case_suits):
role_desc = {"role": "system", "content": CASE_AGG_SYS}
content_case_suits = ""
for i, case_suit in enumerate(case_suits):
case_suit_expr = json.dumps(case_suit, indent=4, ensure_ascii=False)
content_case_suits += f"来自初级测试工程师{i + 1}的测试用例：\n```json\n{case_suit_expr}\n```\n"
messages = [
role_desc,
{"role": "user", "content": content_case_suits}
]
completion = self.chat_complete(messages)
return completion
def cycle_more_case(self, prd_content, case_language="Chinese"):
role_desc = {"role": "system", "content": PRD_CASE_SYS[case_language]}
messages = [
role_desc,
{"role": "user", "content": PRD_CASE_1[case_language] + prd_content + "\n" + PRD_CASE_2[case_language]}
]
result = []
for sys in MORE_CASE_PROMPT[case_language]:
if sys:
messages.append({"role": "user", "content": sys})
reply = self.chat_complete(messages)
result.append(reply)
messages.append({"role": "assistant", "content": reply})
time.sleep(10)
return result
if __name__ == "__main__":
cli = DouBaoService("DouBao128Pro")
# print(cli.get_api_key())
# prd_content = requests.get("https://tosv.byted.org/obj/music-qa-bucket/xmind-test/de3ebc67410c43603034e21bfefa76a0.md").text
# aa = cli.cycle_more_case(prd_content, "English")
# print(aa)
print(cli.chat_complete(messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Introduce LLM shortly."},
]))
|