|
from concurrent.futures import ThreadPoolExecutor |
|
from typing import Dict, List, Optional, Union |
|
|
|
from opencompass.utils.prompt import PromptList |
|
|
|
from .base_api import BaseAPIModel |
|
|
|
PromptType = Union[PromptList, str] |
|
|
|
|
|
class ZhiPuAI(BaseAPIModel): |
|
"""Model wrapper around ZhiPuAI. |
|
|
|
Args: |
|
path (str): The name of OpenAI's model. |
|
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, |
|
query_per_second: int = 2, |
|
max_seq_len: int = 2048, |
|
meta_template: Optional[Dict] = None, |
|
retry: int = 2, |
|
): |
|
super().__init__(path=path, |
|
max_seq_len=max_seq_len, |
|
query_per_second=query_per_second, |
|
meta_template=meta_template, |
|
retry=retry) |
|
import zhipuai |
|
self.zhipuai = zhipuai |
|
self.zhipuai.api_key = key |
|
self.model = path |
|
|
|
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)) |
|
|
|
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 = {'model': self.model, 'prompt': messages} |
|
|
|
max_num_retries = 0 |
|
while max_num_retries < self.retry: |
|
self.acquire() |
|
response = self.zhipuai.model_api.invoke(**data) |
|
self.release() |
|
|
|
if response is None: |
|
print('Connection error, reconnect.') |
|
|
|
|
|
|
|
self.wait() |
|
continue |
|
if response['code'] == 200 and response['success']: |
|
msg = response['data']['choices'][0]['content'] |
|
return msg |
|
|
|
|
|
if (response['code'] == 1301 or response['code'] == 1261 |
|
or response['code'] == 1234 or response['code'] == 1214): |
|
print(response['msg']) |
|
return '' |
|
print(response) |
|
max_num_retries += 1 |
|
|
|
raise RuntimeError(response['msg']) |
|
|
