|
import asyncio |
|
from typing import List, Union |
|
|
|
from lagent.llms.base_llm import AsyncBaseLLM, BaseLLM |
|
from lagent.utils.util import filter_suffix |
|
|
|
|
|
def asdict_completion(output): |
|
return { |
|
key: getattr(output, key) |
|
for key in [ |
|
'text', 'token_ids', 'cumulative_logprob', 'logprobs', |
|
'finish_reason', 'stop_reason' |
|
] |
|
} |
|
|
|
|
|
class VllmModel(BaseLLM): |
|
""" |
|
A wrapper of vLLM model. |
|
|
|
Args: |
|
path (str): The path to the model. |
|
It could be one of the following options: |
|
- i) A local directory path of a huggingface model. |
|
- ii) The model_id of a model hosted inside a model repo |
|
on huggingface.co, such as "internlm/internlm-chat-7b", |
|
"Qwen/Qwen-7B-Chat ", "baichuan-inc/Baichuan2-7B-Chat" |
|
and so on. |
|
tp (int): tensor parallel |
|
vllm_cfg (dict): Other kwargs for vllm model initialization. |
|
""" |
|
|
|
def __init__(self, path: str, tp: int = 1, vllm_cfg=dict(), **kwargs): |
|
|
|
super().__init__(path=path, **kwargs) |
|
from vllm import LLM |
|
self.model = LLM( |
|
model=self.path, |
|
trust_remote_code=True, |
|
tensor_parallel_size=tp, |
|
**vllm_cfg) |
|
|
|
def generate(self, |
|
inputs: Union[str, List[str]], |
|
do_preprocess: bool = None, |
|
skip_special_tokens: bool = False, |
|
return_dict: bool = False, |
|
**kwargs): |
|
"""Return the chat completions in non-stream mode. |
|
|
|
Args: |
|
inputs (Union[str, List[str]]): input texts to be completed. |
|
do_preprocess (bool): whether pre-process the messages. Default to |
|
True, which means chat_template will be applied. |
|
skip_special_tokens (bool): Whether or not to remove special tokens |
|
in the decoding. Default to be False. |
|
Returns: |
|
(a list of/batched) text/chat completion |
|
""" |
|
from vllm import SamplingParams |
|
|
|
batched = True |
|
if isinstance(inputs, str): |
|
inputs = [inputs] |
|
batched = False |
|
prompt = inputs |
|
gen_params = self.update_gen_params(**kwargs) |
|
max_new_tokens = gen_params.pop('max_new_tokens') |
|
stop_words = gen_params.pop('stop_words') |
|
|
|
sampling_config = SamplingParams( |
|
skip_special_tokens=skip_special_tokens, |
|
max_tokens=max_new_tokens, |
|
stop=stop_words, |
|
**gen_params) |
|
response = self.model.generate(prompt, sampling_params=sampling_config) |
|
texts = [resp.outputs[0].text for resp in response] |
|
|
|
texts = filter_suffix(texts, self.gen_params.get('stop_words')) |
|
for resp, text in zip(response, texts): |
|
resp.outputs[0].text = text |
|
if batched: |
|
return [asdict_completion(resp.outputs[0]) |
|
for resp in response] if return_dict else texts |
|
return asdict_completion( |
|
response[0].outputs[0]) if return_dict else texts[0] |
|
|
|
|
|
class AsyncVllmModel(AsyncBaseLLM): |
|
""" |
|
A asynchronous wrapper of vLLM model. |
|
|
|
Args: |
|
path (str): The path to the model. |
|
It could be one of the following options: |
|
- i) A local directory path of a huggingface model. |
|
- ii) The model_id of a model hosted inside a model repo |
|
on huggingface.co, such as "internlm/internlm-chat-7b", |
|
"Qwen/Qwen-7B-Chat ", "baichuan-inc/Baichuan2-7B-Chat" |
|
and so on. |
|
tp (int): tensor parallel |
|
vllm_cfg (dict): Other kwargs for vllm model initialization. |
|
""" |
|
|
|
def __init__(self, path: str, tp: int = 1, vllm_cfg=dict(), **kwargs): |
|
super().__init__(path=path, **kwargs) |
|
from vllm import AsyncEngineArgs, AsyncLLMEngine |
|
|
|
engine_args = AsyncEngineArgs( |
|
model=self.path, |
|
trust_remote_code=True, |
|
tensor_parallel_size=tp, |
|
**vllm_cfg) |
|
self.model = AsyncLLMEngine.from_engine_args(engine_args) |
|
|
|
async def generate(self, |
|
inputs: Union[str, List[str]], |
|
session_ids: Union[int, List[int]] = None, |
|
do_preprocess: bool = None, |
|
skip_special_tokens: bool = False, |
|
return_dict: bool = False, |
|
**kwargs): |
|
"""Return the chat completions in non-stream mode. |
|
|
|
Args: |
|
inputs (Union[str, List[str]]): input texts to be completed. |
|
do_preprocess (bool): whether pre-process the messages. Default to |
|
True, which means chat_template will be applied. |
|
skip_special_tokens (bool): Whether or not to remove special tokens |
|
in the decoding. Default to be False. |
|
Returns: |
|
(a list of/batched) text/chat completion |
|
""" |
|
from vllm import SamplingParams |
|
|
|
batched = True |
|
if isinstance(inputs, str): |
|
inputs = [inputs] |
|
batched = False |
|
if session_ids is None: |
|
session_ids = list(range(len(inputs))) |
|
elif isinstance(session_ids, (int, str)): |
|
session_ids = [session_ids] |
|
assert len(inputs) == len(session_ids) |
|
|
|
prompt = inputs |
|
gen_params = self.update_gen_params(**kwargs) |
|
max_new_tokens = gen_params.pop('max_new_tokens') |
|
stop_words = gen_params.pop('stop_words') |
|
|
|
sampling_config = SamplingParams( |
|
skip_special_tokens=skip_special_tokens, |
|
max_tokens=max_new_tokens, |
|
stop=stop_words, |
|
**gen_params) |
|
|
|
async def _inner_generate(uid, text): |
|
resp, generator = '', self.model.generate( |
|
text, sampling_params=sampling_config, request_id=uid) |
|
async for out in generator: |
|
resp = out.outputs[0] |
|
return resp |
|
|
|
response = await asyncio.gather(*[ |
|
_inner_generate(sid, inp) for sid, inp in zip(session_ids, prompt) |
|
]) |
|
texts = [resp.text for resp in response] |
|
|
|
texts = filter_suffix(texts, self.gen_params.get('stop_words')) |
|
for resp, text in zip(response, texts): |
|
resp.text = text |
|
if batched: |
|
return [asdict_completion(resp) |
|
for resp in response] if return_dict else texts |
|
return asdict_completion(response[0]) if return_dict else texts[0] |
|
|
