OpenCompass/opencompass/models/baidu_api.py

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.')
                # if connect error, frequent requests will casuse
                # continuous unstable network, therefore wait here
                # to slow down the request
                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:
                        # exceed max length
                        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'])