| from typing import Dict, List, Optional | |
| from opencompass.models.base import BaseModel | |
| from opencompass.utils import get_logger | |
| try: | |
| from vllm import LLM, SamplingParams | |
| except ImportError: | |
| LLM, SamplingParams = None, None | |
| DEFAULT_MODEL_KWARGS = dict(trust_remote_code=True) | |
| class VLLM(BaseModel): | |
| """Model Wrapper for VLLM.""" | |
| def __init__( | |
| self, | |
| path: str, | |
| max_seq_len: int = 2048, | |
| model_kwargs: dict = None, | |
| generation_kwargs: dict = dict(), | |
| meta_template: Optional[Dict] = None, | |
| mode: str = 'none', | |
| use_fastchat_template: bool = False, | |
| end_str: Optional[str] = None, | |
| ): | |
| super().__init__(path=path, | |
| max_seq_len=max_seq_len, | |
| meta_template=meta_template) | |
| assert LLM, ('Please install VLLM with `pip install vllm`. ' | |
| 'note: torch==2.1.2 is required.') | |
| self.logger = get_logger() | |
| self._load_model(path, model_kwargs) | |
| self.tokenizer = self.model.get_tokenizer() | |
| self.generation_kwargs = generation_kwargs | |
| self.generation_kwargs.pop('do_sample', None) | |
| assert mode in ['none', 'mid'] | |
| self.mode = mode | |
| self.use_fastchat_template = use_fastchat_template | |
| self.end_str = end_str | |
| def _load_model(self, | |
| path: str, | |
| add_model_kwargs: dict = None, | |
| num_retry: int = 3): | |
| model_kwargs = DEFAULT_MODEL_KWARGS.copy() | |
| if add_model_kwargs is not None: | |
| model_kwargs.update(add_model_kwargs) | |
| self.model = LLM(path, **model_kwargs) | |
| def generate(self, inputs: List[str], max_out_len: int, | |
| **kwargs) -> List[str]: | |
| """Generate results given a list of inputs. | |
| Args: | |
| inputs (List[str]): A list of strings. | |
| max_out_len (int): The maximum length of the output. | |
| Returns: | |
| List[str]: A list of generated strings. | |
| """ | |
| if self.mode == 'mid': | |
| input_ids = self.tokenizer(inputs, truncation=False)['input_ids'] | |
| inputs = [] | |
| for input_id in input_ids: | |
| if len(input_id) > self.max_seq_len - max_out_len: | |
| half = int((self.max_seq_len - max_out_len) / 2) | |
| inputs.append( | |
| self.tokenizer.decode(input_id[:half], | |
| skip_special_tokens=True) + | |
| self.tokenizer.decode(input_id[-half:], | |
| skip_special_tokens=True)) | |
| else: | |
| inputs.append( | |
| self.tokenizer.decode(input_id, | |
| skip_special_tokens=True)) | |
| generation_kwargs = kwargs.copy() | |
| generation_kwargs.update(self.generation_kwargs) | |
| generation_kwargs.update({'max_tokens': max_out_len}) | |
| sampling_kwargs = SamplingParams(**generation_kwargs) | |
| outputs = self.model.generate(inputs, sampling_kwargs) | |
| prompt_list, output_strs = [], [] | |
| for output in outputs: | |
| prompt = output.prompt | |
| generated_text = output.outputs[0].text | |
| if self.end_str: | |
| generated_text = generated_text.split(self.end_str)[0] | |
| prompt_list.append(prompt) | |
| output_strs.append(generated_text) | |
| return output_strs | |
| def prompts_preproccess(self, inputs: List[str]): | |
| if self.use_fastchat_template: | |
| try: | |
| from fastchat.model import get_conversation_template | |
| except ModuleNotFoundError: | |
| raise ModuleNotFoundError( | |
| 'Fastchat is not implemented. You can use ' | |
| "'pip install \"fschat[model_worker,webui]\"' " | |
| 'to implement fastchat.') | |
| conv = get_conversation_template('vicuna') | |
| conv.append_message(conv.roles[0], inputs[0]) | |
| conv.append_message(conv.roles[1], None) | |
| inputs = [conv.get_prompt()] | |
| return inputs | |
| def get_token_len(self, prompt: str) -> int: | |
| """Get lengths of the tokenized strings. | |
| Args: | |
| prompt (str): Input string. | |
| Returns: | |
| int: Length of the input tokens | |
| """ | |
| return len(self.model.get_tokenizer().encode(prompt)) | |