| from concurrent.futures import ThreadPoolExecutor |
| from typing import Dict, List, Optional, Union |
|
|
| from opencompass.models.base import BaseModel |
| 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 TurboMindModel(BaseModel): |
| """Model wrapper for TurboMind Python API. |
| |
| Args: |
| path (str): path of the turbomind model |
| concurrency (int): the maximum allowed concurrency of turbomind. |
| 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. |
| engine_config (Dict, optional): The engine config to set |
| arguments like session_len, max_batch_size for TurboMind. |
| gen_config (Dict, optional): Generation config to set |
| arguments like top_k, top_p, temperature. |
| end_str (str, optional): Whether to trim generated strings with end_str |
| if the model has special ending strings that are not handled well. |
| Defaults to None. |
| """ |
|
|
| def __init__(self, |
| path: str, |
| concurrency: int = 8, |
| max_seq_len: int = 2048, |
| meta_template: Optional[Dict] = None, |
| engine_config: Optional[Dict] = None, |
| gen_config: Optional[Dict] = None, |
| end_str: Optional[str] = None): |
| super().__init__(path=path, |
| max_seq_len=max_seq_len, |
| meta_template=meta_template) |
| from lmdeploy.turbomind import TurboMind |
|
|
| if engine_config is not None: |
| from lmdeploy.messages import TurbomindEngineConfig |
| engine_config = TurbomindEngineConfig(**engine_config) |
| if gen_config is not None: |
| from lmdeploy.messages import EngineGenerationConfig |
| gen_config = EngineGenerationConfig(**gen_config) |
| self.logger = get_logger() |
| tm_model = TurboMind.from_pretrained(path, engine_config=engine_config) |
| self.tokenizer = tm_model.tokenizer |
| self.generators = [ |
| tm_model.create_instance() for i in range(concurrency) |
| ] |
| self.generator_ids = [i + 1 for i in range(concurrency)] |
| self.gen_config = gen_config |
| self.end_str = end_str |
|
|
| def generate( |
| self, |
| inputs: List[str], |
| max_out_len: int = 512, |
| ) -> List[str]: |
| """Generate results given a list of inputs. |
| |
| Args: |
| inputs (List[str]): A list of prompts |
| max_out_len (int): The maximum length of the output. |
| |
| Returns: |
| List[str]: A list of generated strings. |
| """ |
| assert isinstance( |
| inputs, List), f'List(str) is expected, but got {type(inputs)}' |
|
|
| |
| batch_size = len(self.generators) |
| batch_inputs = [ |
| inputs[i:i + batch_size] for i in range(0, len(inputs), batch_size) |
| ] |
|
|
| results = [] |
| for batch_input in batch_inputs: |
| with ThreadPoolExecutor() as executor: |
| _results = list( |
| executor.map( |
| self._generate, |
| self.generators[:len(batch_input)], |
| self.generator_ids[:len(batch_input)], |
| batch_input, |
| [max_out_len] * len(batch_input), |
| [self.gen_config] * len(batch_input), |
| [self.end_str] * len(batch_input), |
| )) |
| results += _results |
| return results |
|
|
| def get_token_len(self, prompt: str) -> int: |
| input_ids = self.tokenizer.encode(prompt) |
| return len(input_ids) |
|
|
| 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, |
| generator, |
| session_id, |
| prompt: str or PromptList, |
| max_out_len: int, |
| gen_config=None, |
| end_str: Optional[str] = None) -> 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. |
| gen_config (EngineGenerationConfig, optional): Generation |
| config to set arguments like top_k, top_p, temperature. |
| end_str (str, optional): Whether to trim generated strings |
| with end_str if the model has special ending strings |
| that are not handled well. |
| Defaults to None. |
| Returns: |
| str: The generated string. |
| """ |
| assert type( |
| prompt) is str, 'We only support string for TurboMind Python API' |
|
|
| input_ids = self.tokenizer.encode(prompt) |
|
|
| for outputs in generator.stream_infer(session_id=session_id, |
| input_ids=[input_ids], |
| gen_config=gen_config, |
| request_output_len=max_out_len, |
| sequence_start=True, |
| sequence_end=True, |
| step=0, |
| stream_output=False): |
| _, output_ids, _ = outputs |
| response = self.tokenizer.decode(output_ids) |
| response = valid_str(response) |
| |
| if end_str: |
| response = response.split(end_str)[0] |
| return response |
|
|
