import base64 import hashlib import hmac import random import string import time from concurrent.futures import ThreadPoolExecutor from datetime import datetime from typing import Dict, List, Optional, Union import requests from opencompass.utils.prompt import PromptList from .base_api import BaseAPIModel PromptType = Union[PromptList, str] def generate_random_string(length=16): """生成随机串. :param length: 随机串长度，默认为 16 :return: 随机串 """ letters = string.ascii_letters + string.digits rand_str = ''.join(random.choice(letters) for i in range(length)) return rand_str def get_current_time(format='%Y-%m-%d %H:%M:%S'): """获取当前时间. :param format: 时间格式，默认为 '%H:%M:%S' :return: 当前时间字符串 """ now = datetime.now() time_str = now.strftime(format) return time_str def get_current_timestamp(): """ 获取当前时间时间戳 :return: """ timestamp_str = int(round(time.time() * 1000)) return str(timestamp_str) def encode_base64_string(s): """对字符串进行 Base64 编码. :param s: 字符串 :return: 编码后的字符串 """ encoded = base64.b64encode(s).decode() return encoded def get_current_time_gmt_format(): """ 获取当前时间的GMT 时间 :return: """ GMT_FORMAT = '%a, %d %b %Y %H:%M:%SGMT+00:00' now = datetime.now() time_str = now.strftime(GMT_FORMAT) return time_str class Yayi(BaseAPIModel): """Model wrapper around Yayi. Args: path (str): The name of Yayi model. url (str): The base URL for the API. url_path (str): The specific path for the API endpoint. x_tilake_app_key (str): The application key for authentication. x_tilake_app_secret (str): The application secret for authentication. x_tilake_ca_sginature_method (str): The signature method for authentication. query_per_second (int): The maximum queries allowed per second between two consecutive calls of the API. Defaults to 2. max_seq_len (int): The maximum sequence length. Defaults to 8192. 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 retries if the API call fails. Defaults to 2. temperature (float): The temperature for the model's response. Defaults to 0.0. """ def __init__( self, path: str, url: str, url_path: str, x_tilake_app_key: str, x_tilake_app_secret: str, x_tilake_ca_sginature_method: str, query_per_second: int = 2, max_seq_len: int = 8192, meta_template: Optional[Dict] = None, retry: int = 2, temperature: float = 0.0, ): super().__init__( path=path, max_seq_len=max_seq_len, query_per_second=query_per_second, meta_template=meta_template, retry=retry, ) self.url = url self.url_path = url_path self.X_TILAKE_APP_KEY = x_tilake_app_key self.X_TILAKE_APP_SECRET = x_tilake_app_secret self.X_TILAKE_CA_SGINATURE_METHOD = x_tilake_ca_sginature_method self.temperature = temperature self.model = path def generate_signature(self, method, accept, content_type, date, url_path): """ 生成签名. """ string_to_sign = (method + '\n' + accept + '\n' + content_type + '\n' + date + '\n' + url_path) string_to_sign = string_to_sign.encode('utf-8') secret_key = self.X_TILAKE_APP_SECRET.encode('utf-8') signature = hmac.new(secret_key, string_to_sign, hashlib.sha256).digest() return encode_base64_string(signature) def generate_header(self, content_type, accept, date, signature): """ 生成请求头参数. """ headers = { 'x-tilake-app-key': self.X_TILAKE_APP_KEY, 'x-tilake-ca-signature-method': self.X_TILAKE_CA_SGINATURE_METHOD, 'x-tilake-ca-timestamp': get_current_timestamp(), 'x-tilake-ca-nonce': generate_random_string(), 'x-tilake-ca-signature': signature, 'Date': date, 'Content-Type': content_type, 'Accept': accept, } return headers def generate( self, inputs: List[PromptType], max_out_len: int = 8192, ) -> List[str]: """Generate results given a list of inputs. Args: inputs (List[PromptType]): 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: PromptType, max_out_len: int = 8192, ) -> str: """Generate results given an input. Args: inputs (PromptType): 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 = [] msg_buffer, last_role = [], None for item in input: item['role'] = 'assistant' if item['role'] == 'BOT' else 'user' if item['role'] != last_role and last_role is not None: messages.append({ 'content': '\n'.join(msg_buffer), 'role': last_role }) msg_buffer = [] msg_buffer.append(item['prompt']) last_role = item['role'] messages.append({ 'content': '\n'.join(msg_buffer), 'role': last_role }) date = get_current_time_gmt_format() content_type = 'application/json' accept = '*/*' method = 'POST' data = { 'id': "001", 'messages': messages, 'max_new_tokens': 8192, 'temperature': 0.0, 'presence_penalty': 0.0, 'frequency_penalty': 0.0, 'top_p': 1.0, 'top_k': -1, } for _ in range(self.retry): signature_str = self.generate_signature(method=method, accept=accept, content_type=content_type, date=date, url_path=self.url_path) headers = self.generate_header(content_type=content_type, accept=accept, date=date, signature=signature_str) try: print(data) response = requests.post(self.url, json=data, headers=headers) except Exception as e: print(e) continue try: response = response.json() except Exception as e: print(e) continue print(response) try: return response['choices'][0]['message']['content'] except Exception as e: print(e) continue raise RuntimeError(f'Failed to respond in {self.retry} retrys')