vllm_wrapper

Module Contents

Classes

VllmModel

A wrapper of vLLM model.

AsyncVllmModel

A asynchronous wrapper of vLLM model.

Functions

asdict_completion(output)

lagent.llms.vllm_wrapper.asdict_completion(output)
class lagent.llms.vllm_wrapper.VllmModel(path, tp=1, vllm_cfg=dict(), **kwargs)

Bases: lagent.llms.base_llm.BaseLLM

A wrapper of vLLM model.

参数:

  • path (str) –

    The path to the model. It could be one of the following options:

      1. A local directory path of a huggingface model.

      1. 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.

model
generate(inputs, do_preprocess=None, skip_special_tokens=False, return_dict=False, **kwargs)

Return the chat completions in non-stream mode.

参数:

  • 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.

  • return_dict (bool) –

返回:

(a list of/batched) text/chat completion

class lagent.llms.vllm_wrapper.AsyncVllmModel(path, tp=1, vllm_cfg=dict(), **kwargs)

Bases: lagent.llms.base_llm.AsyncBaseLLM

A asynchronous wrapper of vLLM model.

参数:

  • path (str) –

    The path to the model. It could be one of the following options:

      1. A local directory path of a huggingface model.

      1. 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.

model
async generate(inputs, session_ids=None, do_preprocess=None, skip_special_tokens=False, return_dict=False, **kwargs)

Return the chat completions in non-stream mode.

参数:

  • 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.

  • session_ids (Union[int, List[int]]) –

  • return_dict (bool) –

返回:

(a list of/batched) text/chat completion