| import threading |
| from concurrent.futures import ThreadPoolExecutor |
| from typing import Dict, List, Optional, Union |
|
|
| from opencompass.models.base import BaseModel, LMTemplateParser |
| from opencompass.utils.logging import get_logger |
| from opencompass.utils.prompt import PromptList |
|
|
| PromptType = Union[PromptList, str] |
|
|
|
|
| def valid_str(string, coding='utf-8'): |
| """decode text according to its encoding type.""" |
| invalid_chars = [b'\xef\xbf\xbd'] |
| bstr = bytes(string, coding) |
| for invalid_char in invalid_chars: |
| bstr = bstr.replace(invalid_char, b'') |
| ret = bstr.decode(encoding=coding, errors='ignore') |
| return ret |
|
|
|
|
| class TurboMindAPIModel(BaseModel): |
| """Model wrapper for TurboMind Triton Inference Server gRPC API. |
| |
| Args: |
| path (str): The name of OpenAI's model. |
| tis_addr (str): The address (ip:port format) of turbomind's |
| triton inference server |
| max_seq_len (int): The maximum allowed sequence length of a model. |
| Note that the length of prompt + generated tokens shall not exceed |
| this value. Defaults to 2048. |
| meta_template (Dict, optional): The model's meta prompt |
| template if needed, in case the requirement of injecting or |
| wrapping of any meta instructions. |
| """ |
|
|
| is_api: bool = True |
|
|
| def __init__( |
| self, |
| path: str, |
| api_addr: str = 'http://0.0.0.0:23333', |
| max_seq_len: int = 2048, |
| meta_template: Optional[Dict] = None, |
| ): |
| super().__init__(path=path, |
| max_seq_len=max_seq_len, |
| meta_template=meta_template) |
| from lmdeploy.serve.openai.api_client import APIClient |
| self.chatbot = APIClient(api_addr) |
| self.model_name = self.chatbot.available_models[0] |
| self.logger = get_logger() |
| self.template_parser = LMTemplateParser(meta_template) |
| self.eos_token_id = None |
| if meta_template and 'eos_token_id' in meta_template: |
| self.eos_token_id = meta_template['eos_token_id'] |
| self.api_addr = api_addr |
|
|
| def generate( |
| self, |
| inputs: List[str or PromptList], |
| max_out_len: int = 512, |
| temperature: float = 1.0, |
| ) -> 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. |
| temperature (float): What sampling temperature to use, |
| between 0 and 2. Higher values like 0.8 will make the output |
| more random, while lower values like 0.2 will make it more |
| focused and deterministic. Defaults to 0.7. |
| |
| Returns: |
| List[str]: A list of generated strings. |
| """ |
|
|
| with ThreadPoolExecutor() as executor: |
| results = list( |
| executor.map(self._generate, inputs, |
| [max_out_len] * len(inputs), |
| [temperature] * len(inputs))) |
| return results |
|
|
| def get_token_len(self, prompt: str) -> int: |
| input_ids, length = self.chatbot.encode(prompt) |
| return length |
|
|
| def wait(self): |
| """Wait till the next query can be sent. |
| |
| Applicable in both single-thread and multi-thread environments. |
| """ |
| return self.token_bucket.get_token() |
|
|
| def _generate(self, prompt: str or PromptList, max_out_len: int, |
| temperature: float) -> str: |
| """Generate results given a list of inputs. |
| |
| Args: |
| prompt (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. |
| temperature (float): What sampling temperature to use, |
| between 0 and 2. Higher values like 0.8 will make the output |
| more random, while lower values like 0.2 will make it more |
| focused and deterministic. |
| |
| Returns: |
| str: The generated string. |
| """ |
| assert type( |
| prompt) is str, 'We only support string for TurboMind RPC API' |
|
|
| response = '' |
| for output in self.chatbot.completions_v1( |
| session_id=threading.currentThread().ident, |
| prompt=prompt, |
| model=self.model_name, |
| max_tokens=max_out_len, |
| temperature=temperature, |
| top_p=0.8, |
| top_k=1): |
| response += output['choices'][0]['text'] |
| response = valid_str(response) |
| return response |
|
|
