| | |
| | import json |
| | 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, float] |
| |
|
| |
|
| | class Gemini(BaseAPIModel): |
| | """Model wrapper around Gemini models. |
| | |
| | Documentation: |
| | |
| | Args: |
| | path (str): The name of Gemini model. |
| | e.g. `gemini-pro` |
| | 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, |
| | key: str, |
| | path: str, |
| | query_per_second: int = 2, |
| | max_seq_len: int = 2048, |
| | meta_template: Optional[Dict] = None, |
| | retry: int = 2, |
| | temperature: float = 1.0, |
| | top_p: float = 0.8, |
| | top_k: float = 10.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 = f'https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent?key={key}' |
| | self.temperature = temperature |
| | self.top_p = top_p |
| | self.top_k = top_k |
| | self.headers = { |
| | 'content-type': 'application/json', |
| | } |
| |
|
| | 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', 'parts': [{'text': input}]}] |
| | else: |
| | messages = [] |
| | system_prompt = None |
| | for item in input: |
| | if item['role'] == 'SYSTEM': |
| | system_prompt = item['prompt'] |
| | for item in input: |
| | if system_prompt is not None: |
| | msg = { |
| | 'parts': [{ |
| | 'text': system_prompt + '\n' + item['prompt'] |
| | }] |
| | } |
| | else: |
| | msg = {'parts': [{'text': item['prompt']}]} |
| | if item['role'] == 'HUMAN': |
| | msg['role'] = 'user' |
| | messages.append(msg) |
| | elif item['role'] == 'BOT': |
| | msg['role'] = 'model' |
| | messages.append(msg) |
| | elif item['role'] == 'SYSTEM': |
| | pass |
| |
|
| | |
| | |
| | assert msg['role'] in ['user', 'system'] |
| |
|
| | data = { |
| | 'model': |
| | self.path, |
| | 'contents': |
| | messages, |
| | 'safetySettings': [ |
| | { |
| | 'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', |
| | 'threshold': 'BLOCK_NONE' |
| | }, |
| | { |
| | 'category': 'HARM_CATEGORY_HATE_SPEECH', |
| | 'threshold': 'BLOCK_NONE' |
| | }, |
| | { |
| | 'category': 'HARM_CATEGORY_HARASSMENT', |
| | 'threshold': 'BLOCK_NONE' |
| | }, |
| | { |
| | 'category': 'HARM_CATEGORY_DANGEROUS_CONTENT', |
| | 'threshold': 'BLOCK_NONE' |
| | }, |
| | ], |
| | 'generationConfig': { |
| | 'candidate_count': 1, |
| | 'temperature': self.temperature, |
| | 'maxOutputTokens': 2048, |
| | 'topP': self.top_p, |
| | 'topK': self.top_k |
| | } |
| | } |
| |
|
| | for _ in range(self.retry): |
| | self.wait() |
| | raw_response = requests.post(self.url, |
| | headers=self.headers, |
| | data=json.dumps(data)) |
| | try: |
| | response = raw_response.json() |
| | except requests.JSONDecodeError: |
| | self.logger.error('JsonDecode error, got', |
| | str(raw_response.content)) |
| | time.sleep(1) |
| | continue |
| | if raw_response.status_code == 200 and response['msg'] == 'ok': |
| | body = response['body'] |
| | if 'candidates' not in body: |
| | self.logger.error(response) |
| | else: |
| | if 'content' not in body['candidates'][0]: |
| | return "Due to Google's restrictive policies, I am unable to respond to this question." |
| | else: |
| | return body['candidates'][0]['content']['parts'][0][ |
| | 'text'].strip() |
| | self.logger.error(response['msg']) |
| | self.logger.error(response) |
| | time.sleep(1) |
| |
|
| | raise RuntimeError('API call failed.') |
| |
|
| |
|
| | class GeminiAllesAPIN(Gemini): |
| | """Model wrapper around Gemini models. |
| | |
| | Documentation: |
| | |
| | Args: |
| | path (str): The name of Gemini model. |
| | e.g. `gemini-pro` |
| | 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, |
| | url: str, |
| | query_per_second: int = 2, |
| | max_seq_len: int = 2048, |
| | meta_template: Optional[Dict] = None, |
| | retry: int = 2, |
| | temperature: float = 1.0, |
| | top_p: float = 0.8, |
| | top_k: float = 10.0, |
| | ): |
| | super().__init__(key=key, |
| | path=path, |
| | max_seq_len=max_seq_len, |
| | query_per_second=query_per_second, |
| | meta_template=meta_template, |
| | retry=retry) |
| | |
| | self.url = url |
| | self.headers = { |
| | 'alles-apin-token': key, |
| | 'content-type': 'application/json', |
| | } |
| |
|
| | 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. |
| | """ |
| | return super().generate(inputs, max_out_len) |
| |
|
