| | from concurrent.futures import ThreadPoolExecutor |
| | from typing import Dict, List, Optional, Union |
| |
|
| | from opencompass.registry import MODELS |
| | from opencompass.utils import PromptList |
| |
|
| | from ..base_api import BaseAPIModel |
| |
|
| | PromptType = Union[PromptList, str] |
| |
|
| |
|
| | @MODELS.register_module() |
| | class Claude(BaseAPIModel): |
| | """Model wrapper around Claude API. |
| | |
| | Args: |
| | key (str): Authorization key. |
| | path (str): The model to be used. Defaults to claude-2. |
| | 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 = 'claude-2', |
| | 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) |
| | try: |
| | from anthropic import AI_PROMPT, HUMAN_PROMPT, Anthropic |
| | except ImportError: |
| | raise ImportError('Import anthropic failed. Please install it ' |
| | 'with "pip install anthropic" and try again.') |
| |
|
| | self.anthropic = Anthropic(api_key=key) |
| | self.model = path |
| | self.human_prompt = HUMAN_PROMPT |
| | self.ai_prompt = AI_PROMPT |
| |
|
| | 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))) |
| | 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 = f'{self.human_prompt} {input}{self.ai_prompt}' |
| | else: |
| | messages = '' |
| | for item in input: |
| | if item['role'] == 'HUMAN' or item['role'] == 'SYSTEM': |
| | messages += f'{self.human_prompt} {item["prompt"]}' |
| | elif item['role'] == 'BOT': |
| | messages += f'{self.ai_prompt} {item["prompt"]}' |
| | if not messages.endswith(self.ai_prompt): |
| | messages += self.ai_prompt |
| |
|
| | num_retries = 0 |
| | while num_retries < self.retry: |
| | self.wait() |
| | try: |
| | completion = self.anthropic.completions.create( |
| | model=self.model, |
| | max_tokens_to_sample=max_out_len, |
| | prompt=messages) |
| | return completion.completion |
| | except Exception as e: |
| | self.logger.error(e) |
| | num_retries += 1 |
| | raise RuntimeError('Calling Claude API failed after retrying for ' |
| | f'{self.retry} times. Check the logs for details.') |
| |
|
