| 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)) |
|
|
