|
import time |
|
from concurrent.futures import ThreadPoolExecutor |
|
from typing import Dict, List, Optional, Union |
|
|
|
import requests |
|
|
|
from opencompass.utils.prompt import PromptList |
|
|
|
from .base_api import BaseAPIModel |
|
|
|
PromptType = Union[PromptList, str] |
|
|
|
|
|
class ERNIEBot(BaseAPIModel): |
|
"""Model wrapper around ERNIE-Bot. |
|
|
|
Documentation: https://cloud.baidu.com/doc/WENXINWORKSHOP/s/jlil56u11 |
|
|
|
Args: |
|
path (str): The name of ENRIE-bot model. |
|
e.g. `erniebot` |
|
model_type (str): The type of the model |
|
e.g. `chat` |
|
secretkey (str): secretkey in order to obtain access_token |
|
key (str): Authorization key. |
|
query_per_second (int): The maximum queries allowed per second |
|
between two consecutive calls of the API. Defaults to 1. |
|
max_seq_len (int): Unused here. |
|
meta_template (Dict, optional): The model's meta prompt |
|
template if needed, in case the requirement of injecting or |
|
wrapping of any meta instructions. |
|
retry (int): Number of retires if the API call fails. Defaults to 2. |
|
""" |
|
|
|
def __init__(self, |
|
path: str, |
|
key: str, |
|
secretkey: str, |
|
url: str, |
|
query_per_second: int = 2, |
|
max_seq_len: int = 2048, |
|
meta_template: Optional[Dict] = None, |
|
retry: int = 2, |
|
generation_kwargs: Dict = { |
|
'temperature': 0.8, |
|
}): |
|
super().__init__(path=path, |
|
max_seq_len=max_seq_len, |
|
query_per_second=query_per_second, |
|
meta_template=meta_template, |
|
retry=retry, |
|
generation_kwargs=generation_kwargs) |
|
self.headers = {'Content_Type': 'application/json'} |
|
self.secretkey = secretkey |
|
self.key = key |
|
self.url = url |
|
access_token, _ = self._generate_access_token() |
|
self.access_token = access_token |
|
print(access_token) |
|
|
|
def _generate_access_token(self): |
|
try: |
|
BAIDU_APIKEY = self.key |
|
BAIDU_SECRETKEY = self.secretkey |
|
url = f'https://aip.baidubce.com/oauth/2.0/token?' \ |
|
f'client_id={BAIDU_APIKEY}&client_secret={BAIDU_SECRETKEY}' \ |
|
f'&grant_type=client_credentials' |
|
headers = { |
|
'Content-Type': 'application/json', |
|
'Accept': 'application/json' |
|
} |
|
response = requests.request('POST', url, headers=headers) |
|
resp_dict = response.json() |
|
if response.status_code == 200: |
|
access_token = resp_dict.get('access_token') |
|
refresh_token = resp_dict.get('refresh_token') |
|
if 'error' in resp_dict: |
|
raise ValueError(f'Failed to obtain certificate.' |
|
f'{resp_dict.get("error")}') |
|
else: |
|
return access_token, refresh_token |
|
else: |
|
error = resp_dict.get('error') |
|
raise ValueError( |
|
f'Failed to requests obtain certificate {error}.') |
|
except Exception as ex: |
|
raise ex |
|
|
|
def generate( |
|
self, |
|
inputs: List[str or PromptList], |
|
max_out_len: int = 512, |
|
) -> List[str]: |
|
"""Generate results given a list of inputs. |
|
|
|
Args: |
|
inputs (List[str or PromptList]): A list of strings or PromptDicts. |
|
The PromptDict should be organized in OpenCompass' |
|
API format. |
|
max_out_len (int): The maximum length of the output. |
|
|
|
Returns: |
|
List[str]: A list of generated strings. |
|
""" |
|
with ThreadPoolExecutor() as executor: |
|
results = list( |
|
executor.map(self._generate, inputs, |
|
[max_out_len] * len(inputs))) |
|
self.flush() |
|
return results |
|
|
|
def _generate( |
|
self, |
|
input: str or PromptList, |
|
max_out_len: int = 512, |
|
) -> str: |
|
"""Generate results given an input. |
|
|
|
Args: |
|
inputs (str or PromptList): A string or PromptDict. |
|
The PromptDict should be organized in OpenCompass' |
|
API format. |
|
max_out_len (int): The maximum length of the output. |
|
|
|
Returns: |
|
str: The generated string. |
|
""" |
|
assert isinstance(input, (str, PromptList)) |
|
""" |
|
{ |
|
"messages": [ |
|
{"role":"user","content":"请介绍一下你自己"}, |
|
{"role":"assistant","content":"我是百度公司开发的人工智能语言模型"}, |
|
{"role":"user","content": "我在上海，周末可以去哪里玩？"}, |
|
{"role":"assistant","content": "上海是一个充满活力和文化氛围的城市"}, |
|
{"role":"user","content": "周末这里的天气怎么样？"} |
|
] |
|
} |
|
|
|
""" |
|
|
|
if isinstance(input, str): |
|
messages = [{'role': 'user', 'content': input}] |
|
else: |
|
messages = [] |
|
for item in input: |
|
msg = {'content': item['prompt']} |
|
if item['role'] == 'HUMAN': |
|
msg['role'] = 'user' |
|
elif item['role'] == 'BOT': |
|
msg['role'] = 'assistant' |
|
|
|
messages.append(msg) |
|
data = {'messages': messages} |
|
data.update(self.generation_kwargs) |
|
|
|
max_num_retries = 0 |
|
while max_num_retries < self.retry: |
|
self.acquire() |
|
try: |
|
raw_response = requests.request('POST', |
|
url=self.url + |
|
self.access_token, |
|
headers=self.headers, |
|
json=data) |
|
response = raw_response.json() |
|
except Exception as err: |
|
print('Request Error:{}'.format(err)) |
|
time.sleep(3) |
|
continue |
|
|
|
self.release() |
|
|
|
if response is None: |
|
print('Connection error, reconnect.') |
|
|
|
|
|
|
|
self.wait() |
|
continue |
|
if raw_response.status_code == 200: |
|
try: |
|
msg = response['result'] |
|
return msg |
|
except KeyError: |
|
print(response) |
|
self.logger.error(str(response['error_code'])) |
|
if response['error_code'] == 336007: |
|
|
|
return '' |
|
|
|
time.sleep(1) |
|
continue |
|
|
|
if (response['error_code'] == 110 or response['error_code'] == 100 |
|
or response['error_code'] == 111 |
|
or response['error_code'] == 200 |
|
or response['error_code'] == 1000 |
|
or response['error_code'] == 1001 |
|
or response['error_code'] == 1002 |
|
or response['error_code'] == 21002 |
|
or response['error_code'] == 216100 |
|
or response['error_code'] == 336001 |
|
or response['error_code'] == 336003 |
|
or response['error_code'] == 336000 |
|
or response['error_code'] == 336007): |
|
print(response['error_msg']) |
|
return '' |
|
print(response) |
|
max_num_retries += 1 |
|
|
|
raise RuntimeError(response['error_msg']) |
|
|
